157 lines
5.3 KiB
Python
157 lines
5.3 KiB
Python
import os, asyncio
|
||
from typing import TypedDict, Annotated, List, Dict
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage
|
||
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 langgraph.graph.message import add_messages
|
||
from langchain_tavily import TavilySearchResults
|
||
|
||
# ---------- 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,
|
||
)
|
||
|
||
# ---------- Backend ----------
|
||
backend = CompositeBackend([
|
||
LocalShellBackend(workspace_dir="./workspace"),
|
||
FilesystemBackend(),
|
||
])
|
||
|
||
# ---------- Tavily tool ----------
|
||
@tool
|
||
def tavily_search(query: str) -> str:
|
||
"""Search the web using Tavily and return a short summary."""
|
||
tavily = TavilySearchResults(max_results=3)
|
||
results = tavily.run(query)
|
||
# Return first 3 results as a concise note
|
||
notes = []
|
||
for r in results:
|
||
notes.append(f"{r['title']}: {r['url']} – {r.get('content', '')[:120]}...")
|
||
return "\n".join(notes) if notes else "No relevant info found."
|
||
|
||
# ---------- State ----------
|
||
class CompareState(TypedDict):
|
||
entities: List[str]
|
||
criteria: List[str]
|
||
findings: Dict[str, List[str]]
|
||
final_table: str | None
|
||
verdict: str | None
|
||
|
||
# ---------- Nodes ----------
|
||
async def plan_criteria(state: CompareState) -> CompareState:
|
||
entities = state["entities"]
|
||
prompt = (
|
||
f"You are an expert analyst. Given the entities: {', '.join(entities)}\n"
|
||
"Generate 3-5 concise criteria for comparing them."
|
||
)
|
||
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
||
criteria = [c.strip() for c in response.content.split("\n") if c.strip()]
|
||
state["criteria"] = criteria
|
||
state["findings"] = {e: [] for e in entities}
|
||
return state
|
||
|
||
async def research_entity(state: CompareState) -> CompareState:
|
||
# Find next unprocessed entity-criterion pair
|
||
for entity in state["entities"]:
|
||
for criterion in state["criteria"]:
|
||
if len(state["findings"][entity]) < len(state["criteria"]):
|
||
# Build query
|
||
query = f"{entity} {criterion}"
|
||
note = tavily_search(query)
|
||
state["findings"][entity].append(f"{criterion}: {note}")
|
||
return state
|
||
return state
|
||
|
||
async def build_table(state: CompareState) -> CompareState:
|
||
headers = " | ".join(state["entities"]) + ""
|
||
rows = []
|
||
for criterion in state["criteria"]:
|
||
row = []
|
||
for entity in state["entities"]:
|
||
# Find note for this criterion
|
||
note = next((n for n in state["findings"][entity] if n.startswith(criterion)), "N/A")
|
||
row.append(note)
|
||
rows.append(" | ".join(row))
|
||
table = "| " + headers + " |\n| " + " | ".join(["---"] * len(state["entities"])) + " |\n"
|
||
table += "| " + " | ".join(rows) + " |"
|
||
state["final_table"] = table
|
||
return state
|
||
|
||
async def verdict(state: CompareState) -> CompareState:
|
||
prompt = (
|
||
f"Based on the following comparison table, provide a concise verdict on which entity is best for each use case:\n\n"
|
||
f"{state['final_table']}"
|
||
)
|
||
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
||
state["verdict"] = response.content.strip()
|
||
return state
|
||
|
||
# ---------- Graph ----------
|
||
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)
|
||
|
||
# Edge logic
|
||
graph.set_entry_point("plan_criteria")
|
||
graph.add_edge("plan_criteria", "research_entity")
|
||
# research_entity loops until all findings filled
|
||
graph.add_conditional_edges(
|
||
"research_entity",
|
||
lambda state: "done" if all(len(state["findings"][e]) == len(state["criteria"]) for e in state["entities"]) else "research_entity",
|
||
{"done": "build_table"},
|
||
)
|
||
graph.add_edge("build_table", "verdict")
|
||
graph.add_edge("verdict", END)
|
||
|
||
app = graph.compile()
|
||
|
||
# ---------- DeepAgent ----------
|
||
agent = create_deep_agent(
|
||
model=llm,
|
||
tools=[tavily_search],
|
||
backend=backend,
|
||
system_prompt="You are a comparison assistant.",
|
||
)
|
||
|
||
# ---------- CLI ----------
|
||
async def main():
|
||
# Default entities
|
||
entities = ["Chroma", "FAISS", "Qdrant"]
|
||
# Optional custom input
|
||
user_input = input("Enter 3 entities separated by commas (or press Enter for default): ")
|
||
if user_input.strip():
|
||
entities = [e.strip() for e in user_input.split(",")[:3]]
|
||
# Prepare initial state
|
||
state: CompareState = {
|
||
"entities": entities,
|
||
"criteria": [],
|
||
"findings": {},
|
||
"final_table": None,
|
||
"verdict": None,
|
||
}
|
||
# Run graph
|
||
result = await app.ainvoke(state)
|
||
# Print outputs
|
||
print("\n=== Criteria ===")
|
||
print("\n".join(result["criteria"]))
|
||
print("\n=== Findings ===")
|
||
for e in result["entities"]:
|
||
print(f"\n{e}:")
|
||
for f in result["findings"][e]:
|
||
print(f"- {f}")
|
||
print("\n=== Final Table ===")
|
||
print(result["final_table"])
|
||
print("\n=== Verdict ===")
|
||
print(result["verdict"])
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|