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

src/nodes.py

@@ -0,0 +1,117 @@
+import os
+from typing import Dict, List, Tuple
+
+from langgraph.graph import StateGraph, END
+from langchain_openai import ChatOpenAI
+from tavily import TavilyClient
+from dotenv import load_dotenv
+
+from .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"),
+)
+
+tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
+
+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.")
+
+    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."
+    )
+    response = llm.invoke(prompt)
+    criteria_text = response.content.strip()
+    # Parse numbered list
+    criteria = []
+    for line in criteria_text.splitlines():
+        line = line.strip()
+        if line:
+            # Remove leading numbers if present
+            if line[0].isdigit() and (len(line) > 1 and line[1] in ". "):
+                line = line[2:].strip()
+            criteria.append(line)
+    state["criteria"] = criteria
+    return state
+
+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}
+
+    for entity in entities:
+        for criterion in criteria:
+            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)
+
+    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
+    return state
+
+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", [])
+
+    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."
+    )
+    response = llm.invoke(prompt)
+    state["verdict"] = response.content.strip()
+    return state