From 473ea64565d3403d6e80a92ccc2480a6edf36ed6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=93=D0=BB=D0=B5=D0=B1=20=D0=9D=D0=B8=D0=BA=D0=B8=D1=88?= =?UTF-8?q?=D0=B8=D0=BD?= Date: Thu, 11 Jun 2026 15:46:33 +0000 Subject: [PATCH] =?UTF-8?q?=D0=A1=D0=BE=D0=B7=D0=B4=D0=B0=D0=BD=20main.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 148 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..d68801c --- /dev/null +++ b/main.py @@ -0,0 +1,148 @@ +import os +import asyncio +from typing import TypedDict, List, Dict + +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage +from langchain.tools import tool +from langgraph.graph import StateGraph, START, END +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, +) + +# ---------- Tavily ---------- +search = TavilySearchResults(tavily_api_key=os.getenv("TAVILY_API_KEY")) + +# ---------- State ---------- +class CompareState(TypedDict): + entities: List[str] + criteria: List[str] + findings: Dict[str, List[str]] + final_table: str | None + verdict: str | None + +# ---------- Tool ---------- +@tool +def tavily_search(query: str) -> str: + """Return a short summary of Tavily search results for the query.""" + results = search.run(query) + snippets = [] + for r in results[:3]: + snippets.append(f"{r['title']}: {r['content'][:200]}...") + return " | ".join(snippets) + +# ---------- Agent ---------- +from langchain.agents import create_openai_functions_agent +from langchain.schema import AgentAction, AgentFinish + +# Simple function calling agent using the tool +from langchain.agents import Tool, AgentExecutor + +agent_executor = AgentExecutor.from_agent_and_tools( + agent=llm, + tools=[Tool(name="tavily_search", func=tavily_search, description="Search the web with Tavily and return a short summary.")], + verbose=False, +) + +# ---------- LangGraph nodes ---------- +async def plan_criteria(state: CompareState) -> CompareState: + entities = state["entities"] + prompt = f"You are given three entities: {', '.join(entities)}. Generate 3-5 concise criteria to compare them. Return a JSON array of strings." + response = await agent_executor.ainvoke({"input": prompt}) + import json + try: + criteria = json.loads(response) + except Exception: + criteria = ["Performance", "Scalability", "Ease of Use"] + state["criteria"] = criteria + state["findings"] = {e: [] for e in entities} + return state + +async def research_entity(state: CompareState) -> CompareState: + for entity in state["entities"]: + for criterion in state["criteria"]: + if len(state["findings"][entity]) <= state["criteria"].index(criterion): + query = f"{entity} {criterion} comparison" + result = await agent_executor.ainvoke({"input": query}) + note = result + state["findings"][entity].append(note) + print(f"[{entity} × {criterion}] найдено: {note[:60]}...") + return state + return state + +async def build_table(state: CompareState) -> CompareState: + headers = " | ".join(state["entities"]) + "" + rows = [] + for criterion in state["criteria"]: + row = [criterion] + for entity in state["entities"]: + notes = state["findings"][entity] + idx = state["criteria"].index(criterion) + row.append(notes[idx] if idx < len(notes) else "N/A") + rows.append(" | ".join(row)) + table = "| " + headers + " |\n| " + " | ".join(["---"] * len(state["entities"])) + " |\n" + for r in rows: + table += "| " + r + " |\n" + state["final_table"] = table + return state + +async def verdict(state: CompareState) -> CompareState: + prompt = f"Based on the following table, provide a concise verdict on which entity is best for which use case.\n\n{state['final_table']}" + result = await agent_executor.ainvoke({"input": prompt}) + state["verdict"] = result + return state + +# ---------- Graph ---------- +from langgraph.graph import StateGraph + +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.add_edge(START, "plan_criteria") +graph.add_edge("plan_criteria", "research_entity") + +graph.add_conditional_edges( + "research_entity", + lambda state: "build_table" if all( + len(state["findings"][e]) == len(state["criteria"]) for e in state["entities"] + ) else "research_entity", +) + +graph.add_edge("build_table", "verdict") + +graph.add_edge("verdict", END) + +app = graph.compile() + +# ---------- CLI ---------- +async def main(): + default_entities = ["Chroma", "FAISS", "Qdrant"] + user_input = input("Введите 3 сущности через запятую (или нажмите Enter для по умолчанию): ") + if user_input.strip(): + entities = [e.strip() for e in user_input.split(",")[:3]] + else: + entities = default_entities + initial_state: CompareState = { + "entities": entities, + "criteria": [], + "findings": {}, + "final_table": None, + "verdict": None, + } + result = await app.ainvoke(initial_state) + print("\n--- Итоговая таблица ---") + print(result["final_table"]) + print("\n--- Вердикт ---") + print(result["verdict"]) + +if __name__ == "__main__": + asyncio.run(main())