From 54f0708b6d36d6efcc996c945f878e2543069514 Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Mon, 15 Jun 2026 12:10:07 +0000 Subject: [PATCH] add: main.py --- main.py | 156 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..bcd13d2 --- /dev/null +++ b/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())