235 lines
7.9 KiB
Python
235 lines
7.9 KiB
Python
import os
|
||
import asyncio
|
||
from typing import TypedDict, List, Dict, Any, Tuple, Annotated
|
||
|
||
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.tools import tool
|
||
from tavily import TavilyClient
|
||
|
||
from deepagents import create_deep_agent
|
||
from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
|
||
|
||
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=os.getenv("OPENAI_API_KEY"),
|
||
temperature=0.0,
|
||
)
|
||
|
||
# -------------------- Tavily tool --------------------
|
||
tavily = TavilyClient(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 --------------------
|
||
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:
|
||
print("Please provide exactly three entities.")
|
||
return
|
||
|
||
initial_state: CompareState = {
|
||
"entities": entities,
|
||
"criteria": [],
|
||
"pairs": [],
|
||
"findings": {},
|
||
"final_table": None,
|
||
"verdict": None,
|
||
"messages": [],
|
||
}
|
||
|
||
# Run the graph
|
||
async for event in app.astream(initial_state):
|
||
# The graph updates state internally; we only need final output after END
|
||
pass
|
||
|
||
final = event # last state after END
|
||
print("\n=== План критериев ===")
|
||
print(", ".join(final["criteria"]))
|
||
print("\n=== Итоговая таблица ===")
|
||
print(final["final_table"])
|
||
print("\n=== Вердикт ===")
|
||
print(final["verdict"])
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main()) |