fix: main.py — Повторный экзамен #2: Сравнительный обзор 3 сущностей (Tavily)
This commit is contained in:
@@ -1,207 +1,235 @@
|
||||
import os
|
||||
import argparse
|
||||
import asyncio
|
||||
from typing import TypedDict, List, Dict, Optional
|
||||
from typing import TypedDict, List, Dict, Any, Tuple, Annotated
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
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()
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
|
||||
from deepagents import create_deep_agent
|
||||
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
|
||||
|
||||
# LLM configuration
|
||||
from langgraph.graph import StateGraph, START, END
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
load_dotenv()
|
||||
|
||||
# -------------------- LLM --------------------
|
||||
llm = ChatOpenAI(
|
||||
model="openai/gpt-oss-20b:free",
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
api_key=OPENAI_API_KEY,
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
# Tavily client
|
||||
tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
|
||||
# -------------------- Tavily tool --------------------
|
||||
tavily = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
|
||||
|
||||
# State definition
|
||||
class CompareState(TypedDict):
|
||||
entities: List[str]
|
||||
criteria: List[str]
|
||||
findings: Dict[str, List[str]]
|
||||
final_table: Optional[str]
|
||||
verdict: Optional[str]
|
||||
|
||||
# Pydantic models for parsing
|
||||
class CriteriaOutput(BaseModel):
|
||||
criteria: List[str] = Field(description="List of comparison criteria")
|
||||
|
||||
class VerdictOutput(BaseModel):
|
||||
verdict: str = Field(description="Verdict text")
|
||||
|
||||
criteria_parser = PydanticOutputParser(pydantic_object=CriteriaOutput)
|
||||
verdict_parser = PydanticOutputParser(pydantic_object=VerdictOutput)
|
||||
|
||||
# Node: plan_criteria
|
||||
def plan_criteria(state: CompareState) -> CompareState:
|
||||
prompt = (
|
||||
f"Given the following entities: {', '.join(state['entities'])}. "
|
||||
"Generate 3 to 5 distinct criteria for comparing these entities. "
|
||||
"Return a JSON object with a field 'criteria' that is a list of strings."
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
parsed = criteria_parser.parse(response.content)
|
||||
state["criteria"] = parsed.criteria
|
||||
# Initialize findings dict
|
||||
state["findings"] = {entity: [] for entity in state["entities"]}
|
||||
return state
|
||||
|
||||
# Node: research_entity
|
||||
def research_entity(state: CompareState) -> CompareState:
|
||||
# Find first entity with missing notes
|
||||
for entity in state["entities"]:
|
||||
if len(state["findings"][entity]) < len(state["criteria"]):
|
||||
idx = len(state["findings"][entity])
|
||||
criterion = state["criteria"][idx]
|
||||
query = f"{entity} {criterion}"
|
||||
# Perform Tavily search
|
||||
results = tavily_client.search(query)
|
||||
# Take first result snippet
|
||||
if results and results[0].snippet:
|
||||
note = results[0].snippet.strip()
|
||||
else:
|
||||
note = f"No relevant info found for {entity} on {criterion}."
|
||||
state["findings"][entity].append(note)
|
||||
break
|
||||
return state
|
||||
|
||||
# Node: check_done
|
||||
def check_done(state: CompareState) -> str:
|
||||
for entity in state["entities"]:
|
||||
if len(state["findings"][entity]) < len(state["criteria"]):
|
||||
return "continue"
|
||||
return "done"
|
||||
|
||||
# Node: build_table
|
||||
def build_table(state: CompareState) -> CompareState:
|
||||
header = ["Criterion"] + state["entities"]
|
||||
rows = []
|
||||
for criterion in state["criteria"]:
|
||||
row = [criterion]
|
||||
for entity in state["entities"]:
|
||||
notes = state["findings"][entity]
|
||||
idx = state["criteria"].index(criterion)
|
||||
note = notes[idx] if idx < len(notes) else ""
|
||||
row.append(note)
|
||||
rows.append(row)
|
||||
# Build markdown table
|
||||
table_lines = ["| " + " | ".join(header) + " |"]
|
||||
table_lines.append("|" + "|".join(["---"] * len(header)) + "|")
|
||||
for row in rows:
|
||||
table_lines.append("| " + " | ".join(row) + " |")
|
||||
table_md = "\n".join(table_lines)
|
||||
state["final_table"] = table_md
|
||||
return state
|
||||
|
||||
# Node: verdict
|
||||
def verdict(state: CompareState) -> CompareState:
|
||||
prompt = (
|
||||
f"Here is a markdown table comparing the entities:\n\n{state['final_table']}\n\n"
|
||||
"Based on this table, provide a concise verdict recommending which entity "
|
||||
"is best suited for a typical vector database use case. "
|
||||
"Return a JSON object with a field 'verdict' that is a string."
|
||||
)
|
||||
response = llm.invoke(prompt)
|
||||
parsed = verdict_parser.parse(response.content)
|
||||
state["verdict"] = parsed.verdict
|
||||
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_done", check_done)
|
||||
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_edge("research_entity", "check_done")
|
||||
graph.add_conditional_edges(
|
||||
"check_done",
|
||||
lambda x: x,
|
||||
{
|
||||
"continue": "research_entity",
|
||||
"done": "build_table",
|
||||
},
|
||||
)
|
||||
graph.add_edge("build_table", "verdict")
|
||||
graph.add_edge("verdict", END)
|
||||
|
||||
compiled_graph = graph.compile()
|
||||
|
||||
# Tool: compare_entities
|
||||
@tool
|
||||
def compare_entities(query: str) -> str:
|
||||
def web_search(query: str) -> str:
|
||||
"""
|
||||
Compare three entities based on user query.
|
||||
Expected format: "Compare 3 vector DBs: Chroma, FAISS, Qdrant"
|
||||
Perform a web search using Tavily and return a concise summary of the top result.
|
||||
"""
|
||||
# Extract entities after colon
|
||||
if ":" in query:
|
||||
parts = query.split(":", 1)
|
||||
entities_part = parts[1]
|
||||
else:
|
||||
entities_part = query
|
||||
entities = [e.strip() for e in entities_part.split(",") if e.strip()]
|
||||
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 --------------------
|
||||
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)
|
||||
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'])}"),
|
||||
]
|
||||
)
|
||||
response = await invoke_llm([{"role": "system", "content": prompt.messages[0].content},
|
||||
{"role": "user", "content": prompt.messages[1].content}])
|
||||
try:
|
||||
import json
|
||||
|
||||
criteria = json.loads(response)
|
||||
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"]}
|
||||
return state
|
||||
|
||||
|
||||
# -------------------- Node: research entity --------------------
|
||||
async def research_entity(state: CompareState) -> CompareState:
|
||||
if not state["pairs"]:
|
||||
return state
|
||||
entity, criterion = state["pairs"].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]}...")
|
||||
return state
|
||||
|
||||
|
||||
# -------------------- 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)
|
||||
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 ""),
|
||||
]
|
||||
)
|
||||
response = await invoke_llm([{"role": "system", "content": prompt.messages[0].content},
|
||||
{"role": "user", "content": prompt.messages[1].content}])
|
||||
state["verdict"] = response.strip()
|
||||
return state
|
||||
|
||||
|
||||
# -------------------- 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",
|
||||
)
|
||||
|
||||
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",
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
if len(entities) != 3:
|
||||
return "Please provide exactly three entities separated by commas."
|
||||
print("Please provide exactly three entities.")
|
||||
return
|
||||
|
||||
initial_state: CompareState = {
|
||||
"entities": entities,
|
||||
"criteria": [],
|
||||
"pairs": [],
|
||||
"findings": {},
|
||||
"final_table": None,
|
||||
"verdict": None,
|
||||
"messages": [],
|
||||
}
|
||||
final_state = compiled_graph.invoke(initial_state)
|
||||
table = final_state["final_table"] or ""
|
||||
verdict_text = final_state["verdict"] or ""
|
||||
return f"{table}\n\nVerdict:\n{verdict_text}"
|
||||
|
||||
# DeepAgent setup
|
||||
backend = CompositeBackend([LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend()])
|
||||
# Run the graph
|
||||
async for event in app.astream(initial_state):
|
||||
# The graph updates state internally; we only need final output after END
|
||||
pass
|
||||
|
||||
agent = create_deep_agent(
|
||||
model=llm,
|
||||
tools=[compare_entities],
|
||||
backend=backend,
|
||||
system_prompt="You are a helpful assistant that can compare three entities.",
|
||||
)
|
||||
final = event # last state after END
|
||||
print("\n=== План критериев ===")
|
||||
print(", ".join(final["criteria"]))
|
||||
print("\n=== Итоговая таблица ===")
|
||||
print(final["final_table"])
|
||||
print("\n=== Вердикт ===")
|
||||
print(final["verdict"])
|
||||
|
||||
# CLI
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Compare three entities.")
|
||||
parser.add_argument(
|
||||
"--query",
|
||||
type=str,
|
||||
default="Compare 3 vector DBs: Chroma, FAISS, Qdrant",
|
||||
help="Comparison query in the format 'Compare 3 vector DBs: A, B, C'",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
async def run():
|
||||
result = await agent.ainvoke(
|
||||
{"messages": [{"role": "user", "content": args.query}]},
|
||||
{"configurable": {"thread_id": "session-1"}},
|
||||
)
|
||||
print(result["messages"][-1]["content"])
|
||||
asyncio.run(run())
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user