feat: solution for 'Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily)'

src/nodes.py

@@ -1,117 +1,99 @@
 import os
-from typing import Dict, List, Tuple
-
-from langgraph.graph import StateGraph, END
-from langchain_openai import ChatOpenAI
-from tavily import TavilyClient
+from typing import Dict, List
 from dotenv import load_dotenv
-
-from .state import CompareState
+from langchain_openai import ChatOpenAI
+from langchain_tavily import TavilySearchTool
+from langgraph.prebuilt import create_chat_agent
+from langgraph import add_messages, StateGraph
+from src.state import CompareState
 
 load_dotenv()
 
-# Initialize LLM and Tavily client
-llm = ChatOpenAI(
-    temperature=0.2,
-    model="gpt-4o-mini",
-    openai_api_key=os.getenv("OPENAI_API_KEY"),
-)
+# LLM and Tavily tool
+llm = ChatOpenAI(temperature=0.7)
+tavily = TavilySearchTool(api_key=os.getenv("TAVILY_API_KEY"))
 
-tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+# Helper to format findings for LLM
+def format_findings(findings: Dict[str, List[str]]) -> str:
+    parts = []
+    for entity, notes in findings.items():
+        parts.append(f"**{entity}**:")
+        for note in notes:
+            parts.append(f"- {note}")
+    return "\n".join(parts)
 
+# Node: Generate comparison criteria
 def plan_criteria(state: CompareState) -> CompareState:
-    """
-    Generate 3–5 comparison criteria for the given entities.
-    """
-    entities = state.get("entities", [])
-    if not entities:
-        raise ValueError("No entities provided for criteria planning.")
-
+    entities = state["entities"]
     prompt = (
-        f"Given the following entities: {', '.join(entities)}.\n"
-        "Suggest 3 to 5 key criteria to compare them. "
-        "Return the criteria as a numbered list, one per line."
+        f"Generate 3-5 concise comparison criteria for the following entities: "
+        f"{', '.join(entities)}. Return a numbered list."
     )
     response = llm.invoke(prompt)
-    criteria_text = response.content.strip()
     # Parse numbered list
     criteria = []
-    for line in criteria_text.splitlines():
+    for line in response.splitlines():
         line = line.strip()
         if line:
-            # Remove leading numbers if present
+            # Remove leading numbers
             if line[0].isdigit() and (len(line) > 1 and line[1] in ". "):
                 line = line[2:].strip()
             criteria.append(line)
     state["criteria"] = criteria
+    # Initialize findings dict
+    state["findings"] = {entity: [] for entity in entities}
     return state
 
+# Node: Research one entity-criterion pair
 def research_entity(state: CompareState) -> CompareState:
-    """
-    For each entity–criterion pair, perform a Tavily web search
-    and store a short note in findings.
-    """
-    entities = state.get("entities", [])
-    criteria = state.get("criteria", [])
-    findings: Dict[str, List[str]] = {entity: [] for entity in entities}
+    entities = state["entities"]
+    criteria = state["criteria"]
+    findings = state["findings"]
 
+    # Find next entity needing research
     for entity in entities:
-        for criterion in criteria:
+        if len(findings[entity]) < len(criteria):
+            idx = len(findings[entity])
+            criterion = criteria[idx]
             query = f"{entity} {criterion}"
-            try:
-                result = tavily.search(query=query, max_results=1)
-                if result and result["results"]:
-                    snippet = result["results"][0]["content"][:200]
-                else:
-                    snippet = "No relevant information found."
-            except Exception as e:
-                snippet = f"Error during search: {e}"
-            findings[entity].append(snippet)
-
+            # Tavily search
+            results = tavily.invoke({"query": query, "max_results": 3})
+            # Take first result snippet
+            if results and "results" in results and len(results["results"]) > 0:
+                snippet = results["results"][0]["snippet"]
+                source = results["results"][0]["url"]
+                note = f"{criterion}: {snippet} (Source: {source})"
+            else:
+                note = f"{criterion}: No recent information found."
+            findings[entity].append(note)
+            break
     state["findings"] = findings
     return state
 
-def build_table(state: CompareState) -> CompareState:
-    """
-    Build a Markdown table from findings.
-    Rows: criteria, Columns: entities.
-    """
-    entities = state.get("entities", [])
-    criteria = state.get("criteria", [])
-    findings = state.get("findings", {})
-
-    header = "| Criterion | " + " | ".join(entities) + " |\n"
-    separator = "|---" * (len(entities) + 1) + "|\n"
-
-    rows = ""
-    for idx, criterion in enumerate(criteria):
-        row = f"| {criterion} | "
-        for entity in entities:
-            notes = findings.get(entity, [])
-            note = notes[idx] if idx < len(notes) else ""
-            # Escape pipe characters
-            note = note.replace("|", "\\|")
-            row += f"{note} | "
-        rows += row + "\n"
-
-    table = header + separator + rows
-    state["final_table"] = table
+# Node: Build cohesive research brief
+def build_brief(state: CompareState) -> CompareState:
+    findings_text = format_findings(state["findings"])
+    prompt = (
+        f"Using the following findings, write a cohesive research brief that summarizes "
+        f"the strengths and weaknesses of each entity. The brief should be clear, "
+        f"structured, and suitable for a technical audience.\n\n"
+        f"Findings:\n{findings_text}"
+    )
+    brief = llm.invoke(prompt)
+    state["final_brief"] = brief
     return state
 
+# Node: Generate verdict/recommendation
 def verdict(state: CompareState) -> CompareState:
-    """
-    Generate a verdict recommendation based on the table and criteria.
-    """
-    table = state.get("final_table", "")
-    criteria = state.get("criteria", [])
-    entities = state.get("entities", [])
-
+    brief = state["final_brief"]
+    criteria = state["criteria"]
     prompt = (
-        f"Here is a comparative table of the following entities: {', '.join(entities)}.\n\n"
-        f"{table}\n\n"
-        f"Based on the criteria: {', '.join(criteria)}.\n"
-        "Provide a concise recommendation (2–4 sentences) indicating which entity is best suited for which use case."
+        f"Based on the research brief below and the comparison criteria, provide a "
+        f"clear recommendation on which entity is best suited for a typical use case. "
+        f"Explain your reasoning in 2-4 sentences.\n\n"
+        f"Research Brief:\n{brief}\n\n"
+        f"Criteria:\n- " + "\n- ".join(criteria)
     )
-    response = llm.invoke(prompt)
-    state["verdict"] = response.content.strip()
+    recommendation = llm.invoke(prompt)
+    state["verdict"] = recommendation
     return state