fix: main.py — Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily)
This commit is contained in:
@@ -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
|
||||
entities: List[str] # 3 names for comparison
|
||||
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)
|
||||
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())
|
||||
Reference in New Issue
Block a user