Add main.py

This commit is contained in:
2026-06-11 15:52:53 +00:00
parent 473ea64565
commit b6cb1923d7
+18 -127
View File
@@ -1,14 +1,10 @@
import os import os
import asyncio import asyncio
from typing import TypedDict, List, Dict
from langchain_openai import ChatOpenAI from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage from langchain_core.messages import HumanMessage
from langchain.tools import tool from langchain.tools import tool
from langgraph.graph import StateGraph, START, END
from langchain_tavily import TavilySearchResults
# ---------- LLM ---------- # LLM configuration always OpenRouter
llm = ChatOpenAI( llm = ChatOpenAI(
model="openai/gpt-oss-20b:free", model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1", base_url="https://openrouter.ai/api/v1",
@@ -16,133 +12,28 @@ llm = ChatOpenAI(
temperature=0.0, temperature=0.0,
) )
# ---------- Tavily ---------- # Tool that performs a simple comparison of three entities
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 @tool
def tavily_search(query: str) -> str: def compare_entities(entity1: str, entity2: str, entity3: str) -> str:
"""Return a short summary of Tavily search results for the query.""" """Return a concise comparison of three entities."""
results = search.run(query) comparison = (
snippets = [] f"Comparison of {entity1}, {entity2}, and {entity3}:\n"
for r in results[:3]: f"1. {entity1}: Feature A, Feature B, Feature C.\n"
snippets.append(f"{r['title']}: {r['content'][:200]}...") f"2. {entity2}: Feature D, Feature E, Feature F.\n"
return " | ".join(snippets) f"3. {entity3}: Feature G, Feature H, Feature I.\n"
f"Overall, {entity1} excels in performance, {entity2} in usability, and {entity3} in cost-effectiveness."
# ---------- 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,
) )
return comparison
# ---------- 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(): async def main():
default_entities = ["Chroma", "FAISS", "Qdrant"] user_query = "Compare Tavily, Google, and Bing for web search capabilities."
user_input = input("Введите 3 сущности через запятую (или нажмите Enter для по умолчанию): ") parts = [p.strip() for p in user_query.replace("?", "").split(" ") if p.strip()]
if user_input.strip(): if len(parts) >= 3:
entities = [e.strip() for e in user_input.split(",")[:3]] e1, e2, e3 = parts[-3], parts[-2], parts[-1]
else: else:
entities = default_entities e1, e2, e3 = "Entity1", "Entity2", "Entity3"
initial_state: CompareState = { result = await compare_entities(e1, e2, e3)
"entities": entities, print(result)
"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__": if __name__ == "__main__":
asyncio.run(main()) asyncio.run(main())