From c3a933ec425f442c45dda55f88072c4c0353c9f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= Date: Thu, 2 Jul 2026 13:06:51 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20main.py=20=E2=80=94=20=D0=9F=D0=BE=D0=B2?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BD=D1=8B=D0=B9=20=D1=8D=D0=BA=D0=B7=D0=B0?= =?UTF-8?q?=D0=BC=D0=B5=D0=BD=20#2:=20=D0=A1=D1=80=D0=B0=D0=B2=D0=BD=D0=B8?= =?UTF-8?q?=D1=82=D0=B5=D0=BB=D1=8C=D0=BD=D1=8B=D0=B9=20=D0=BE=D0=B1=D0=B7?= =?UTF-8?q?=D0=BE=D1=80=203=20=D1=81=D1=83=D1=89=D0=BD=D0=BE=D1=81=D1=82?= =?UTF-8?q?=D0=B5=D0=B9=20(Tavily)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 316 +++++++++++++++++++++++++------------------------------- 1 file changed, 140 insertions(+), 176 deletions(-) diff --git a/main.py b/main.py index a3089a9..77124f6 100644 --- a/main.py +++ b/main.py @@ -1,23 +1,22 @@ import os import asyncio -from typing import TypedDict, List, Dict, Any, Tuple, Annotated +import argparse +import re +from typing import TypedDict, Annotated, Dict, List, Tuple, Any from dotenv import load_dotenv -from langchain_openai import ChatOpenAI, OpenAIEmbeddings -from langchain_core.messages import HumanMessage, SystemMessage -from langchain_core.prompts import ChatPromptTemplate +from langchain_openai import ChatOpenAI +from langchain_core.messages import HumanMessage from langchain.tools import tool -from tavily import TavilyClient - from deepagents import create_deep_agent -from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend - +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend from langgraph.graph import StateGraph, START, END -from langgraph.graph.message import add_messages +from tavily import TavilySearchResults +# Load environment variables load_dotenv() -# -------------------- LLM -------------------- +# LLM configuration llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -25,211 +24,176 @@ llm = ChatOpenAI( temperature=0.0, ) -# -------------------- Tavily tool -------------------- -tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")) +# Tavily client +tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) - -@tool -def web_search(query: str) -> str: - """ - Perform a web search using Tavily and return a concise summary of the top result. - """ - try: - response = tavily.search(query, search_depth="basic", max_results=3) - results = response.get("results", []) - if not results: - return "No relevant results found." - # Concatenate titles and snippets - summary = " ".join(r.get("title", "") + ". " + r.get("content", "") for r in results[:2]) - return summary.strip() - except Exception as e: - return f"Search error: {e}" - - -# -------------------- DeepAgent wrapper (required by course) -------------------- -backend = CompositeBackend( - [ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), - ] -) - -deep_agent = create_deep_agent( - model=llm, - tools=[web_search], - backend=backend, - system_prompt="You are a helpful research assistant.", -) - - -async def invoke_llm(messages: List[Dict[str, str]]) -> str: - """ - Helper that sends messages to the deep agent and returns the assistant's reply. - """ - result = await deep_agent.ainvoke( - {"messages": [HumanMessage(content=m["content"]) for m in messages]}, - {"configurable": {"thread_id": "compare-session"}}, - ) - return result["messages"][-1].content - - -# -------------------- State definition -------------------- +# State definition class CompareState(TypedDict): - entities: List[str] # 3 names to compare - criteria: List[str] # 3-5 criteria - pairs: List[Tuple[str, str]] # remaining (entity, criterion) pairs - findings: Dict[str, List[str]] # entity -> list of notes (same order as criteria) + entities: List[str] # 3 names for comparison + criteria: List[str] # 3-5 criteria + findings: Dict[str, Dict[str, str]] # entity -> criterion -> note + pairs_to_process: List[Tuple[str, str]] # (entity, criterion) final_table: str | None verdict: str | None - messages: Annotated[list, add_messages] # for LangGraph internal use - -# -------------------- Node: plan criteria -------------------- -async def plan_criteria(state: CompareState) -> CompareState: - prompt = ChatPromptTemplate.from_messages( - [ - SystemMessage( - content="You are an expert analyst. Given three entities, propose 3 to 5 criteria to compare them. Return the criteria as a JSON list." - ), - HumanMessage(content=f"Entities: {', '.join(state['entities'])}"), - ] +# Node: plan_criteria +def plan_criteria(state: CompareState) -> CompareState: + prompt = ( + f"Given the entities {state['entities']}, generate 3 to 5 comparison criteria. " + "Return a JSON array of strings." ) - response = await invoke_llm([{"role": "system", "content": prompt.messages[0].content}, - {"role": "user", "content": prompt.messages[1].content}]) + response = llm.invoke(prompt) + # Extract JSON array try: import json - - criteria = json.loads(response) + criteria = json.loads(response.content) if not isinstance(criteria, list): raise ValueError except Exception: - # Fallback: split by newlines - criteria = [c.strip("- ").strip() for c in response.splitlines() if c.strip()] - state["criteria"] = criteria[:5] - # Build all pairs - state["pairs"] = [(e, c) for e in state["entities"] for c in state["criteria"]] - # Initialise findings dict - state["findings"] = {e: [] for e in state["entities"]} + criteria = ["performance", "scalability", "ease of use"] + state["criteria"] = criteria + # Prepare pairs to process + state["pairs_to_process"] = [(entity, criterion) for entity in state["entities"] for criterion in criteria] + state["findings"] = {entity: {} for entity in state["entities"]} + print(f"[plan_criteria] Generated criteria: {criteria}") return state - -# -------------------- Node: research entity -------------------- -async def research_entity(state: CompareState) -> CompareState: - if not state["pairs"]: +# Node: research_entity +def research_entity(state: CompareState) -> CompareState: + if not state["pairs_to_process"]: return state - entity, criterion = state["pairs"].pop(0) + entity, criterion = state["pairs_to_process"].pop(0) query = f"{entity} {criterion}" - # Use the web_search tool directly (synchronous call is fine) - note = web_search.run(query) # type: ignore - # Append note to the correct entity list - state["findings"][entity].append(note) - # Log for CLI - print(f"[{entity} × {criterion}] найдено: {note[:200]}...") + results = tavily.search(query) + snippet = results[0].content if results else "No relevant information found." + note = f"{criterion}: {snippet}" + state["findings"][entity][criterion] = note + print(f"[research_entity] ({entity} × {criterion}) found: {snippet[:60]}...") return state +# Node: check_pairs +def check_pairs(state: CompareState) -> str: + return "continue" if state["pairs_to_process"] else "done" -# -------------------- Node: build table -------------------- +# Node: build_table def build_table(state: CompareState) -> CompareState: - header = ["Критерий"] + state["entities"] - rows = [] - for idx, crit in enumerate(state["criteria"]): - row = [crit] - for ent in state["entities"]: - notes = state["findings"].get(ent, []) - note = notes[idx] if idx < len(notes) else "" - row.append(note.replace("\n", " ").strip()) - rows.append(row) - - # Markdown table construction - def md_row(cols: List[str]) -> str: - return "| " + " | ".join(cols) + " |" - - separator = "| " + " | ".join(["---"] * len(header)) + " |" - table_lines = [md_row(header), separator] + [md_row(r) for r in rows] - state["final_table"] = "\n".join(table_lines) + headers = ["Criterion"] + state["entities"] + table_rows = [] + for criterion in state["criteria"]: + row = [criterion] + for entity in state["entities"]: + note = state["findings"][entity].get(criterion, "N/A") + row.append(note) + table_rows.append(row) + # Build markdown table + md = "| " + " | ".join(headers) + " |\n" + md += "| " + " | ".join(["---"] * len(headers)) + " |\n" + for row in table_rows: + md += "| " + " | ".join(row) + " |\n" + state["final_table"] = md + print("[build_table] Table constructed.") return state - -# -------------------- Node: verdict -------------------- -async def verdict(state: CompareState) -> CompareState: - prompt = ChatPromptTemplate.from_messages( - [ - SystemMessage( - content="You are an analyst. Based on the comparison table, give a short verdict (2-4 sentences) recommending which entity is best for which use case." - ), - HumanMessage(content=state["final_table"] or ""), - ] +# Node: verdict +def verdict(state: CompareState) -> CompareState: + prompt = ( + f"Based on the following table, provide a concise verdict recommending which entity is best for which use case:\n\n" + f"{state['final_table']}\n\n" + "Answer in 2-4 sentences." ) - response = await invoke_llm([{"role": "system", "content": prompt.messages[0].content}, - {"role": "user", "content": prompt.messages[1].content}]) - state["verdict"] = response.strip() + response = llm.invoke(prompt) + state["verdict"] = response.content.strip() + print("[verdict] Verdict generated.") return state +# Build LangGraph +graph = StateGraph(CompareState) +graph.add_node("plan_criteria", plan_criteria) +graph.add_node("research_entity", research_entity) +graph.add_node("check_pairs", check_pairs) +graph.add_node("build_table", build_table) +graph.add_node("verdict", verdict) -# -------------------- Graph assembly -------------------- -workflow = StateGraph(CompareState) - -workflow.add_node("plan_criteria", plan_criteria) -workflow.add_node("research_entity", research_entity) -workflow.add_node("build_table", build_table) -workflow.add_node("verdict", verdict) - -workflow.add_edge(START, "plan_criteria") -workflow.add_conditional_edges( - "plan_criteria", - lambda s: "research_entity" if s["pairs"] else "build_table", +graph.add_edge(START, "plan_criteria") +graph.add_edge("plan_criteria", "research_entity") +graph.add_edge("research_entity", "check_pairs") +graph.add_conditional_edges( + "check_pairs", + lambda x: x, + { + "continue": "research_entity", + "done": "build_table", + }, ) +graph.add_edge("build_table", "verdict") +graph.add_edge("verdict", END) -workflow.add_edge("research_entity", "research_entity") # loop until pairs empty -workflow.add_conditional_edges( - "research_entity", - lambda s: "research_entity" if s["pairs"] else "build_table", -) +app = graph.compile() -workflow.add_edge("build_table", "verdict") -workflow.add_edge("verdict", END) - -app = workflow.compile() - - -# -------------------- CLI demo -------------------- -async def main() -> None: - default_entities = ["Chroma", "FAISS", "Qdrant"] - user_input = input( - "Enter three entities separated by commas (or press Enter for default): " - ).strip() - entities = ( - [e.strip() for e in user_input.split(",") if e.strip()] - if user_input - else default_entities - ) +# Tool that runs the comparison +@tool +def run_comparison(query: str) -> str: + """ + Run a comparative review of three entities. + Input: a string containing three entity names separated by commas. + Output: markdown table and verdict. + """ + # Extract entities + entities = [e.strip() for e in re.split(r",|;|and", query) if e.strip()] if len(entities) != 3: - print("Please provide exactly three entities.") - return - + return "Please provide exactly three entities separated by commas." initial_state: CompareState = { "entities": entities, "criteria": [], - "pairs": [], "findings": {}, + "pairs_to_process": [], "final_table": None, "verdict": None, - "messages": [], } + result_state = app.invoke(initial_state) + table = result_state["final_table"] or "" + verdict_text = result_state["verdict"] or "" + return f"{table}\n\n**Verdict:**\n{verdict_text}" - # Run the graph - async for event in app.astream(initial_state): - # The graph updates state internally; we only need final output after END - pass +# DeepAgents backend +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) - final = event # last state after END - print("\n=== План критериев ===") - print(", ".join(final["criteria"])) - print("\n=== Итоговая таблица ===") - print(final["final_table"]) - print("\n=== Вердикт ===") - print(final["verdict"]) +# Create deepagents agent +agent = create_deep_agent( + model=llm, + tools=[run_comparison], + backend=backend, + system_prompt="You are a helpful agent that performs comparative reviews.", +) +async def main(): + parser = argparse.ArgumentParser(description="Comparative review CLI") + parser.add_argument( + "-e", + "--entities", + type=str, + help="Comma-separated list of three entities to compare", + ) + args = parser.parse_args() + + if args.entities: + query = args.entities + else: + # Default entities + query = "Chroma, FAISS, Qdrant" + + print(f"Running comparison for: {query}") + + result = await agent.ainvoke( + {"messages": [HumanMessage(content=query)]}, + {"configurable": {"thread_id": "session-1"}}, + ) + final_output = result["messages"][-1].content + print("\n=== Final Output ===\n") + print(final_output) if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file