From a1833273e59a19c8e9b0baad638da5357ee3b582 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=98=D0=BB=D1=8C=D1=8F=205f1b81b8-4f5d-11e8-9c2d-fa7ae01?= =?UTF-8?q?bbebc?= Date: Wed, 1 Jul 2026 18:40:07 +0000 Subject: [PATCH] add: main.py --- main.py | 177 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..f7f0473 --- /dev/null +++ b/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())