add: main.py

main.py

@@ -0,0 +1,177 @@
+import os
+import sys
+import asyncio
+from typing import TypedDict, List, Dict
+from dotenv import load_dotenv
+
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
+from langchain.tools import tool
+from deepagents import create_deep_agent
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+
+from langgraph.graph import StateGraph, START, END
+from langchain_core.output_parsers import PydanticOutputParser
+from pydantic import BaseModel, Field
+
+from tavily import TavilyClient
+
+# ---------- Load environment variables ----------
+load_dotenv()
+
+# ---------- LLM ----------
+llm = ChatOpenAI(
+    model="openai/gpt-oss-20b:free",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+    temperature=0.0,
+)
+
+# ---------- State definition ----------
+class CompareState(TypedDict):
+    entities: List[str]                # 3 имени для сравнения
+    criteria: List[str]                 # 3–5 критериев
+    findings: Dict[str, List[str]]      # entity -> список заметок по критериям
+    final_table: str | None
+    verdict: str | None
+
+# ---------- Pydantic parser for criteria ----------
+class CriteriaOutput(BaseModel):
+    criteria: List[str] = Field(..., description="List of comparison criteria")
+
+criteria_parser = PydanticOutputParser(pydantic_object=CriteriaOutput)
+
+# ---------- Tavily client ----------
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+if not TAVILY_API_KEY:
+    raise RuntimeError("TAVILY_API_KEY not set in environment")
+client = TavilyClient(api_key=TAVILY_API_KEY)
+
+# ---------- Graph nodes ----------
+async def plan_criteria(state: CompareState) -> CompareState:
+    entities = ", ".join(state["entities"])
+    prompt = (
+        f"You are a comparison helper. Given the following entities: {entities}.\n"
+        "Generate 3-5 concise criteria for comparing these entities.\n"
+        "Return a JSON object with a field 'criteria' that is a list of strings."
+    )
+    result = await llm.ainvoke(prompt)
+    parsed = criteria_parser.parse(result.content)
+    state["criteria"] = parsed.criteria
+    # initialize findings dict
+    state["findings"] = {e: [] for e in state["entities"]}
+    return state
+
+async def research_entity(state: CompareState) -> CompareState:
+    # Find next unprocessed entity
+    for entity in state["entities"]:
+        processed = len(state["findings"][entity])
+        if processed < len(state["criteria"]):
+            criterion = state["criteria"][processed]
+            query = f"{entity} {criterion}"
+            # Perform Tavily search
+            results = client.search(query, max_results=3)
+            # Take first result content, truncate to 200 chars
+            content = results[0]["content"][:200] if results else "No data found."
+            note = f"{criterion}: {content}"
+            state["findings"][entity].append(note)
+            break
+    return state
+
+async def build_table(state: CompareState) -> CompareState:
+    headers = ["Criterion"] + state["entities"]
+    table = "| " + " | ".join(headers) + " |\n"
+    table += "| " + " | ".join(["---" for _ in headers]) + " |\n"
+    for idx, criterion in enumerate(state["criteria"]):
+        row = [criterion]
+        for entity in state["entities"]:
+            notes = state["findings"][entity]
+            note = notes[idx] if idx < len(notes) else ""
+            row.append(note)
+        table += "| " + " | ".join(row) + " |\n"
+    state["final_table"] = table
+    return state
+
+async def verdict(state: CompareState) -> CompareState:
+    prompt = (
+        "You are an expert advisor. Based on the following comparison table,\n"
+        "provide a concise verdict of 2–4 sentences recommending which entity\n"
+        "is best for each possible use case.\n"
+        f"Comparison Table:\n{state['final_table']}"
+    )
+    result = await llm.ainvoke(prompt)
+    state["verdict"] = result.content.strip()
+    return state
+
+# ---------- Graph construction ----------
+
+def create_compare_graph() -> StateGraph[CompareState]:
+    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.set_entry_point("plan_criteria")
+    graph.add_edge("plan_criteria", "research_entity")
+
+    def is_done(state: CompareState) -> str:
+        return "build_table" if all(
+            len(state["findings"][e]) >= len(state["criteria"]) for e in state["entities"]
+        ) else "research_entity"
+
+    graph.add_conditional_edges("research_entity", is_done)
+    graph.add_edge("build_table", "verdict")
+    graph.add_edge("verdict", END)
+
+    return graph
+
+# ---------- DeepAgent tool ----------
+@tool
+def compare_entities(entities: List[str]) -> str:
+    """Compare three entities and return a markdown table with verdict."""
+    if len(entities) != 3:
+        return "Error: Please provide exactly three entities to compare."
+    # Initialize state
+    state: CompareState = {
+        "entities": entities,
+        "criteria": [],
+        "findings": {},
+        "final_table": None,
+        "verdict": None,
+    }
+    graph = create_compare_graph()
+    # Run graph synchronously
+    final_state = graph.invoke(state)
+    table = final_state["final_table"] or "No table generated."
+    verdict = final_state["verdict"] or "No verdict generated."
+    return f"{table}\n\nVerdict:\n{verdict}"
+
+# ---------- DeepAgent setup ----------
+backend = CompositeBackend([
+    LocalShellBackend(workspace_dir="./workspace"),
+    FilesystemBackend(),
+])
+
+agent = create_deep_agent(
+    model=llm,
+    tools=[compare_entities],
+    backend=backend,
+    system_prompt="You are a helpful comparison assistant. Use the provided tool to compare entities.",
+)
+
+# ---------- CLI ----------
+async def main():
+    if len(sys.argv) >= 4:
+        entities = sys.argv[1:4]
+    else:
+        entities = ["Chroma", "FAISS", "Qdrant"]
+    user_msg = f"Compare these entities: {', '.join(entities)}"
+    result = await agent.ainvoke(
+        {"messages": [HumanMessage(content=user_msg)]},
+        {"configurable": {"thread_id": "session-1"}},
+    )
+    print(result["messages"][-1].content)
+
+if __name__ == "__main__":
+    asyncio.run(main())