fix: main.py — Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily)

main.py

@@ -1,23 +1,22 @@
 import os
 import asyncio
-from typing import TypedDict, List, Dict, Any, Tuple, Annotated
+import argparse
+import re
+from typing import TypedDict, Annotated, Dict, List, Tuple, Any
 
 from dotenv import load_dotenv
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.messages import HumanMessage
-from langchain_core.prompts import ChatPromptTemplate
 from langchain.tools import tool
-from tavily import TavilyClient
-
 from deepagents import create_deep_agent
-from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
-
 from langgraph.graph import StateGraph, START, END
-from langgraph.graph.message import add_messages
+from tavily import TavilySearchResults
 
+# Load environment variables
 load_dotenv()
 
-# -------------------- LLM --------------------
+# LLM configuration
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -25,211 +24,176 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# -------------------- Tavily tool --------------------
+# Tavily client
-tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))
 
-
+# State definition
-@tool
-def web_search(query: str) -> str:
-    """
-    Perform a web search using Tavily and return a concise summary of the top result.
-    """
-    try:
-        response = tavily.search(query, search_depth="basic", max_results=3)
-        results = response.get("results", [])
-        if not results:
-            return "No relevant results found."
-        # Concatenate titles and snippets
-        summary = " ".join(r.get("title", "") + ". " + r.get("content", "") for r in results[:2])
-        return summary.strip()
-    except Exception as e:
-        return f"Search error: {e}"
-
-
-# -------------------- DeepAgent wrapper (required by course) --------------------
-backend = CompositeBackend(
-    [
-        LocalShellBackend(workspace_dir="./workspace"),
-        FilesystemBackend(),
-    ]
-)
-
-deep_agent = create_deep_agent(
-    model=llm,
-    tools=[web_search],
-    backend=backend,
-    system_prompt="You are a helpful research assistant.",
-)
-
-
-async def invoke_llm(messages: List[Dict[str, str]]) -> str:
-    """
-    Helper that sends messages to the deep agent and returns the assistant's reply.
-    """
-    result = await deep_agent.ainvoke(
-        {"messages": [HumanMessage(content=m["content"]) for m in messages]},
-        {"configurable": {"thread_id": "compare-session"}},
-    )
-    return result["messages"][-1].content
-
-
-# -------------------- State definition --------------------
 class CompareState(TypedDict):
-    entities: List[str]                     # 3 names to compare
+    entities: List[str]                # 3 names for comparison
-    criteria: List[str]                     # 3-5 criteria
+    criteria: List[str]                # 3-5 criteria
-    pairs: List[Tuple[str, str]]            # remaining (entity, criterion) pairs
+    findings: Dict[str, Dict[str, str]]  # entity -> criterion -> note
-    findings: Dict[str, List[str]]          # entity -> list of notes (same order as criteria)
+    pairs_to_process: List[Tuple[str, str]]  # (entity, criterion)
     final_table: str | None
     verdict: str | None
-    messages: Annotated[list, add_messages]  # for LangGraph internal use
 
-
+# Node: plan_criteria
-# -------------------- Node: plan criteria --------------------
+def plan_criteria(state: CompareState) -> CompareState:
-async def plan_criteria(state: CompareState) -> CompareState:
+    prompt = (
-    prompt = ChatPromptTemplate.from_messages(
+        f"Given the entities {state['entities']}, generate 3 to 5 comparison criteria. "
-        [
+        "Return a JSON array of strings."
-            SystemMessage(
-                content="You are an expert analyst. Given three entities, propose 3 to 5 criteria to compare them. Return the criteria as a JSON list."
-            ),
-            HumanMessage(content=f"Entities: {', '.join(state['entities'])}"),
-        ]
     )
-    response = await invoke_llm([{"role": "system", "content": prompt.messages[0].content},
+    response = llm.invoke(prompt)
-                                {"role": "user", "content": prompt.messages[1].content}])
+    # Extract JSON array
     try:
         import json
-
+        criteria = json.loads(response.content)
-        criteria = json.loads(response)
         if not isinstance(criteria, list):
             raise ValueError
     except Exception:
-        # Fallback: split by newlines
+        criteria = ["performance", "scalability", "ease of use"]
-        criteria = [c.strip("- ").strip() for c in response.splitlines() if c.strip()]
+    state["criteria"] = criteria
-    state["criteria"] = criteria[:5]
+    # Prepare pairs to process
-    # Build all pairs
+    state["pairs_to_process"] = [(entity, criterion) for entity in state["entities"] for criterion in criteria]
-    state["pairs"] = [(e, c) for e in state["entities"] for c in state["criteria"]]
+    state["findings"] = {entity: {} for entity in state["entities"]}
-    # Initialise findings dict
+    print(f"[plan_criteria] Generated criteria: {criteria}")
-    state["findings"] = {e: [] for e in state["entities"]}
     return state
 
-
+# Node: research_entity
-# -------------------- Node: research entity --------------------
+def research_entity(state: CompareState) -> CompareState:
-async def research_entity(state: CompareState) -> CompareState:
+    if not state["pairs_to_process"]:
-    if not state["pairs"]:
         return state
-    entity, criterion = state["pairs"].pop(0)
+    entity, criterion = state["pairs_to_process"].pop(0)
     query = f"{entity} {criterion}"
-    # Use the web_search tool directly (synchronous call is fine)
+    results = tavily.search(query)
-    note = web_search.run(query)  # type: ignore
+    snippet = results[0].content if results else "No relevant information found."
-    # Append note to the correct entity list
+    note = f"{criterion}: {snippet}"
-    state["findings"][entity].append(note)
+    state["findings"][entity][criterion] = note
-    # Log for CLI
+    print(f"[research_entity] ({entity} × {criterion}) found: {snippet[:60]}...")
-    print(f"[{entity} × {criterion}] найдено: {note[:200]}...")
     return state
 
+# Node: check_pairs
+def check_pairs(state: CompareState) -> str:
+    return "continue" if state["pairs_to_process"] else "done"
 
-# -------------------- Node: build table --------------------
+# Node: build_table
 def build_table(state: CompareState) -> CompareState:
-    header = ["Критерий"] + state["entities"]
+    headers = ["Criterion"] + state["entities"]
-    rows = []
+    table_rows = []
-    for idx, crit in enumerate(state["criteria"]):
+    for criterion in state["criteria"]:
-        row = [crit]
+        row = [criterion]
-        for ent in state["entities"]:
+        for entity in state["entities"]:
-            notes = state["findings"].get(ent, [])
+            note = state["findings"][entity].get(criterion, "N/A")
-            note = notes[idx] if idx < len(notes) else ""
+            row.append(note)
-            row.append(note.replace("\n", " ").strip())
+        table_rows.append(row)
-        rows.append(row)
+    # Build markdown table
-
+    md = "| " + " | ".join(headers) + " |\n"
-    # Markdown table construction
+    md += "| " + " | ".join(["---"] * len(headers)) + " |\n"
-    def md_row(cols: List[str]) -> str:
+    for row in table_rows:
-        return "| " + " | ".join(cols) + " |"
+        md += "| " + " | ".join(row) + " |\n"
-
+    state["final_table"] = md
-    separator = "| " + " | ".join(["---"] * len(header)) + " |"
+    print("[build_table] Table constructed.")
-    table_lines = [md_row(header), separator] + [md_row(r) for r in rows]
-    state["final_table"] = "\n".join(table_lines)
     return state
 
-
+# Node: verdict
-# -------------------- Node: verdict --------------------
+def verdict(state: CompareState) -> CompareState:
-async def verdict(state: CompareState) -> CompareState:
+    prompt = (
-    prompt = ChatPromptTemplate.from_messages(
+        f"Based on the following table, provide a concise verdict recommending which entity is best for which use case:\n\n"
-        [
+        f"{state['final_table']}\n\n"
-            SystemMessage(
+        "Answer in 2-4 sentences."
-                content="You are an analyst. Based on the comparison table, give a short verdict (2-4 sentences) recommending which entity is best for which use case."
-            ),
-            HumanMessage(content=state["final_table"] or ""),
-        ]
     )
-    response = await invoke_llm([{"role": "system", "content": prompt.messages[0].content},
+    response = llm.invoke(prompt)
-                                {"role": "user", "content": prompt.messages[1].content}])
+    state["verdict"] = response.content.strip()
-    state["verdict"] = response.strip()
+    print("[verdict] Verdict generated.")
     return state
 
+# Build LangGraph
+graph = StateGraph(CompareState)
+graph.add_node("plan_criteria", plan_criteria)
+graph.add_node("research_entity", research_entity)
+graph.add_node("check_pairs", check_pairs)
+graph.add_node("build_table", build_table)
+graph.add_node("verdict", verdict)
 
-# -------------------- Graph assembly --------------------
+graph.add_edge(START, "plan_criteria")
-workflow = StateGraph(CompareState)
+graph.add_edge("plan_criteria", "research_entity")
-
+graph.add_edge("research_entity", "check_pairs")
-workflow.add_node("plan_criteria", plan_criteria)
+graph.add_conditional_edges(
-workflow.add_node("research_entity", research_entity)
+    "check_pairs",
-workflow.add_node("build_table", build_table)
+    lambda x: x,
-workflow.add_node("verdict", verdict)
+    {
-
+        "continue": "research_entity",
-workflow.add_edge(START, "plan_criteria")
+        "done": "build_table",
-workflow.add_conditional_edges(
+    },
-    "plan_criteria",
-    lambda s: "research_entity" if s["pairs"] else "build_table",
 )
+graph.add_edge("build_table", "verdict")
+graph.add_edge("verdict", END)
 
-workflow.add_edge("research_entity", "research_entity")  # loop until pairs empty
+app = graph.compile()
-workflow.add_conditional_edges(
-    "research_entity",
-    lambda s: "research_entity" if s["pairs"] else "build_table",
-)
 
-workflow.add_edge("build_table", "verdict")
+# Tool that runs the comparison
-workflow.add_edge("verdict", END)
+@tool
-
+def run_comparison(query: str) -> str:
-app = workflow.compile()
+    """
-
+    Run a comparative review of three entities.
-
+    Input: a string containing three entity names separated by commas.
-# -------------------- CLI demo --------------------
+    Output: markdown table and verdict.
-async def main() -> None:
+    """
-    default_entities = ["Chroma", "FAISS", "Qdrant"]
+    # Extract entities
-    user_input = input(
+    entities = [e.strip() for e in re.split(r",|;|and", query) if e.strip()]
-        "Enter three entities separated by commas (or press Enter for default): "
-    ).strip()
-    entities = (
-        [e.strip() for e in user_input.split(",") if e.strip()]
-        if user_input
-        else default_entities
-    )
     if len(entities) != 3:
-        print("Please provide exactly three entities.")
+        return "Please provide exactly three entities separated by commas."
-        return
-
     initial_state: CompareState = {
         "entities": entities,
         "criteria": [],
-        "pairs": [],
         "findings": {},
+        "pairs_to_process": [],
         "final_table": None,
         "verdict": None,
-        "messages": [],
     }
+    result_state = app.invoke(initial_state)
+    table = result_state["final_table"] or ""
+    verdict_text = result_state["verdict"] or ""
+    return f"{table}\n\n**Verdict:**\n{verdict_text}"
 
-    # Run the graph
+# DeepAgents backend
-    async for event in app.astream(initial_state):
+backend = CompositeBackend([
-        # The graph updates state internally; we only need final output after END
+    LocalShellBackend(workspace_dir="./workspace"),
-        pass
+    FilesystemBackend(),
+])
 
-    final = event  # last state after END
+# Create deepagents agent
-    print("\n=== План критериев ===")
+agent = create_deep_agent(
-    print(", ".join(final["criteria"]))
+    model=llm,
-    print("\n=== Итоговая таблица ===")
+    tools=[run_comparison],
-    print(final["final_table"])
+    backend=backend,
-    print("\n=== Вердикт ===")
+    system_prompt="You are a helpful agent that performs comparative reviews.",
-    print(final["verdict"])
+)
 
+async def main():
+    parser = argparse.ArgumentParser(description="Comparative review CLI")
+    parser.add_argument(
+        "-e",
+        "--entities",
+        type=str,
+        help="Comma-separated list of three entities to compare",
+    )
+    args = parser.parse_args()
+
+    if args.entities:
+        query = args.entities
+    else:
+        # Default entities
+        query = "Chroma, FAISS, Qdrant"
+
+    print(f"Running comparison for: {query}")
+
+    result = await agent.ainvoke(
+        {"messages": [HumanMessage(content=query)]},
+        {"configurable": {"thread_id": "session-1"}},
+    )
+    final_output = result["messages"][-1].content
+    print("\n=== Final Output ===\n")
+    print(final_output)
 
 if __name__ == "__main__":
     asyncio.run(main())