add: main.py

main.py

@@ -0,0 +1,156 @@
+import os, asyncio
+from typing import TypedDict, Annotated, List, Dict
+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 langgraph.graph.message import add_messages
+from langchain_tavily import TavilySearchResults
+
+# ---------- 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,
+)
+
+# ---------- Backend ----------
+backend = CompositeBackend([
+    LocalShellBackend(workspace_dir="./workspace"),
+    FilesystemBackend(),
+])
+
+# ---------- Tavily tool ----------
+@tool
+def tavily_search(query: str) -> str:
+    """Search the web using Tavily and return a short summary."""
+    tavily = TavilySearchResults(max_results=3)
+    results = tavily.run(query)
+    # Return first 3 results as a concise note
+    notes = []
+    for r in results:
+        notes.append(f"{r['title']}: {r['url']} – {r.get('content', '')[:120]}...")
+    return "\n".join(notes) if notes else "No relevant info found."
+
+# ---------- State ----------
+class CompareState(TypedDict):
+    entities: List[str]
+    criteria: List[str]
+    findings: Dict[str, List[str]]
+    final_table: str | None
+    verdict: str | None
+
+# ---------- Nodes ----------
+async def plan_criteria(state: CompareState) -> CompareState:
+    entities = state["entities"]
+    prompt = (
+        f"You are an expert analyst. Given the entities: {', '.join(entities)}\n"
+        "Generate 3-5 concise criteria for comparing them."
+    )
+    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    criteria = [c.strip() for c in response.content.split("\n") if c.strip()]
+    state["criteria"] = criteria
+    state["findings"] = {e: [] for e in entities}
+    return state
+
+async def research_entity(state: CompareState) -> CompareState:
+    # Find next unprocessed entity-criterion pair
+    for entity in state["entities"]:
+        for criterion in state["criteria"]:
+            if len(state["findings"][entity]) < len(state["criteria"]):
+                # Build query
+                query = f"{entity} {criterion}"
+                note = tavily_search(query)
+                state["findings"][entity].append(f"{criterion}: {note}")
+                return state
+    return state
+
+async def build_table(state: CompareState) -> CompareState:
+    headers = " | ".join(state["entities"]) + ""
+    rows = []
+    for criterion in state["criteria"]:
+        row = []
+        for entity in state["entities"]:
+            # Find note for this criterion
+            note = next((n for n in state["findings"][entity] if n.startswith(criterion)), "N/A")
+            row.append(note)
+        rows.append(" | ".join(row))
+    table = "| " + headers + " |\n| " + " | ".join(["---"] * len(state["entities"])) + " |\n"
+    table += "| " + " | ".join(rows) + " |"
+    state["final_table"] = table
+    return state
+
+async def verdict(state: CompareState) -> CompareState:
+    prompt = (
+        f"Based on the following comparison table, provide a concise verdict on which entity is best for each use case:\n\n"
+        f"{state['final_table']}"
+    )
+    response = await llm.ainvoke([HumanMessage(content=prompt)])
+    state["verdict"] = response.content.strip()
+    return state
+
+# ---------- Graph ----------
+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)
+
+# Edge logic
+graph.set_entry_point("plan_criteria")
+graph.add_edge("plan_criteria", "research_entity")
+# research_entity loops until all findings filled
+graph.add_conditional_edges(
+    "research_entity",
+    lambda state: "done" if all(len(state["findings"][e]) == len(state["criteria"]) for e in state["entities"]) else "research_entity",
+    {"done": "build_table"},
+)
+graph.add_edge("build_table", "verdict")
+graph.add_edge("verdict", END)
+
+app = graph.compile()
+
+# ---------- DeepAgent ----------
+agent = create_deep_agent(
+    model=llm,
+    tools=[tavily_search],
+    backend=backend,
+    system_prompt="You are a comparison assistant.",
+)
+
+# ---------- CLI ----------
+async def main():
+    # Default entities
+    entities = ["Chroma", "FAISS", "Qdrant"]
+    # Optional custom input
+    user_input = input("Enter 3 entities separated by commas (or press Enter for default): ")
+    if user_input.strip():
+        entities = [e.strip() for e in user_input.split(",")[:3]]
+    # Prepare initial state
+    state: CompareState = {
+        "entities": entities,
+        "criteria": [],
+        "findings": {},
+        "final_table": None,
+        "verdict": None,
+    }
+    # Run graph
+    result = await app.ainvoke(state)
+    # Print outputs
+    print("\n=== Criteria ===")
+    print("\n".join(result["criteria"]))
+    print("\n=== Findings ===")
+    for e in result["entities"]:
+        print(f"\n{e}:")
+        for f in result["findings"][e]:
+            print(f"- {f}")
+    print("\n=== Final Table ===")
+    print(result["final_table"])
+    print("\n=== Verdict ===")
+    print(result["verdict"])
+
+if __name__ == "__main__":
+    asyncio.run(main())