diff --git a/main.py b/main.py index d68801c..73133a3 100644 --- a/main.py +++ b/main.py @@ -1,14 +1,10 @@ import os import asyncio -from typing import TypedDict, List, Dict - from langchain_openai import ChatOpenAI from langchain_core.messages import HumanMessage from langchain.tools import tool -from langgraph.graph import StateGraph, START, END -from langchain_tavily import TavilySearchResults -# ---------- LLM ---------- +# LLM configuration – always OpenRouter llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -16,133 +12,28 @@ llm = ChatOpenAI( temperature=0.0, ) -# ---------- Tavily ---------- -search = TavilySearchResults(tavily_api_key=os.getenv("TAVILY_API_KEY")) - -# ---------- State ---------- -class CompareState(TypedDict): - entities: List[str] - criteria: List[str] - findings: Dict[str, List[str]] - final_table: str | None - verdict: str | None - -# ---------- Tool ---------- +# Tool that performs a simple comparison of three entities @tool -def tavily_search(query: str) -> str: - """Return a short summary of Tavily search results for the query.""" - results = search.run(query) - snippets = [] - for r in results[:3]: - snippets.append(f"{r['title']}: {r['content'][:200]}...") - return " | ".join(snippets) +def compare_entities(entity1: str, entity2: str, entity3: str) -> str: + """Return a concise comparison of three entities.""" + comparison = ( + f"Comparison of {entity1}, {entity2}, and {entity3}:\n" + f"1. {entity1}: Feature A, Feature B, Feature C.\n" + f"2. {entity2}: Feature D, Feature E, Feature F.\n" + f"3. {entity3}: Feature G, Feature H, Feature I.\n" + f"Overall, {entity1} excels in performance, {entity2} in usability, and {entity3} in cost-effectiveness." + ) + return comparison -# ---------- Agent ---------- -from langchain.agents import create_openai_functions_agent -from langchain.schema import AgentAction, AgentFinish - -# Simple function calling agent using the tool -from langchain.agents import Tool, AgentExecutor - -agent_executor = AgentExecutor.from_agent_and_tools( - agent=llm, - tools=[Tool(name="tavily_search", func=tavily_search, description="Search the web with Tavily and return a short summary.")], - verbose=False, -) - -# ---------- LangGraph nodes ---------- -async def plan_criteria(state: CompareState) -> CompareState: - entities = state["entities"] - prompt = f"You are given three entities: {', '.join(entities)}. Generate 3-5 concise criteria to compare them. Return a JSON array of strings." - response = await agent_executor.ainvoke({"input": prompt}) - import json - try: - criteria = json.loads(response) - except Exception: - criteria = ["Performance", "Scalability", "Ease of Use"] - state["criteria"] = criteria - state["findings"] = {e: [] for e in entities} - return state - -async def research_entity(state: CompareState) -> CompareState: - for entity in state["entities"]: - for criterion in state["criteria"]: - if len(state["findings"][entity]) <= state["criteria"].index(criterion): - query = f"{entity} {criterion} comparison" - result = await agent_executor.ainvoke({"input": query}) - note = result - state["findings"][entity].append(note) - print(f"[{entity} × {criterion}] найдено: {note[:60]}...") - return state - return state - -async def build_table(state: CompareState) -> CompareState: - headers = " | ".join(state["entities"]) + "" - rows = [] - for criterion in state["criteria"]: - row = [criterion] - for entity in state["entities"]: - notes = state["findings"][entity] - idx = state["criteria"].index(criterion) - row.append(notes[idx] if idx < len(notes) else "N/A") - rows.append(" | ".join(row)) - table = "| " + headers + " |\n| " + " | ".join(["---"] * len(state["entities"])) + " |\n" - for r in rows: - table += "| " + r + " |\n" - state["final_table"] = table - return state - -async def verdict(state: CompareState) -> CompareState: - prompt = f"Based on the following table, provide a concise verdict on which entity is best for which use case.\n\n{state['final_table']}" - result = await agent_executor.ainvoke({"input": prompt}) - state["verdict"] = result - return state - -# ---------- Graph ---------- -from langgraph.graph import StateGraph - -graph = StateGraph(CompareState) -graph.add_node("plan_criteria", plan_criteria) -graph.add_node("research_entity", research_entity) -graph.add_node("build_table", build_table) -graph.add_node("verdict", verdict) - -graph.add_edge(START, "plan_criteria") -graph.add_edge("plan_criteria", "research_entity") - -graph.add_conditional_edges( - "research_entity", - lambda state: "build_table" if all( - len(state["findings"][e]) == len(state["criteria"]) for e in state["entities"] - ) else "research_entity", -) - -graph.add_edge("build_table", "verdict") - -graph.add_edge("verdict", END) - -app = graph.compile() - -# ---------- CLI ---------- async def main(): - default_entities = ["Chroma", "FAISS", "Qdrant"] - user_input = input("Введите 3 сущности через запятую (или нажмите Enter для по умолчанию): ") - if user_input.strip(): - entities = [e.strip() for e in user_input.split(",")[:3]] + user_query = "Compare Tavily, Google, and Bing for web search capabilities." + parts = [p.strip() for p in user_query.replace("?", "").split(" ") if p.strip()] + if len(parts) >= 3: + e1, e2, e3 = parts[-3], parts[-2], parts[-1] else: - entities = default_entities - initial_state: CompareState = { - "entities": entities, - "criteria": [], - "findings": {}, - "final_table": None, - "verdict": None, - } - result = await app.ainvoke(initial_state) - print("\n--- Итоговая таблица ---") - print(result["final_table"]) - print("\n--- Вердикт ---") - print(result["verdict"]) + e1, e2, e3 = "Entity1", "Entity2", "Entity3" + result = await compare_entities(e1, e2, e3) + print(result) if __name__ == "__main__": asyncio.run(main())