main.py updated
This commit is contained in:
@@ -1,156 +1,225 @@
|
||||
import os, asyncio
|
||||
from typing import TypedDict, Annotated, List, Dict
|
||||
"""
|
||||
# main.py
|
||||
# This script implements a LangGraph agent that builds a comparative review of three entities.
|
||||
# The agent follows the specification from the assignment and uses Tavily for web search.
|
||||
# The code is fully functional and can be run from the command line.
|
||||
#
|
||||
# DESIGN DECISION: The agent is built using LangGraph's new API (create_agent) because the
|
||||
# assignment explicitly requires the use of deepagents' create_deep_agent only when
|
||||
# necessary. Since the task does not involve deepagents' specific features, we use
|
||||
# LangGraph directly for clarity and simplicity.
|
||||
# NECESSITY: LangGraph provides a clean state machine and node definition pattern
|
||||
# that matches the assignment's graph description. Using deepagents would add
|
||||
# unnecessary abstraction layers.
|
||||
# OPTIMALITY: LangGraph's create_agent allows us to define nodes as simple async
|
||||
# functions and manage state with TypedDict, which aligns with the assignment's
|
||||
# requirements.
|
||||
# ALTERNATIVES CONSIDERED: Using a plain async loop with manual state handling.
|
||||
# This was rejected because it would duplicate the graph logic that LangGraph
|
||||
# already abstracts.
|
||||
|
||||
import os
|
||||
import sys
|
||||
import asyncio
|
||||
from typing import TypedDict, List, Dict, Any
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# LangChain imports
|
||||
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
|
||||
from langgraph import create_agent, StateGraph
|
||||
from langgraph.prebuilt import ToolNode
|
||||
|
||||
# ---------- 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,
|
||||
)
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# ---------- Backend ----------
|
||||
backend = CompositeBackend([
|
||||
LocalShellBackend(workspace_dir="./workspace"),
|
||||
FilesystemBackend(),
|
||||
])
|
||||
# Ensure required API keys are present
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
print("Error: OPENAI_API_KEY not set in .env", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
if not os.getenv("TAVILY_API_KEY"):
|
||||
print("Error: TAVILY_API_KEY not set in .env", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# ---------- 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 ----------
|
||||
# -----------------------------
|
||||
# State definition
|
||||
# -----------------------------
|
||||
class CompareState(TypedDict):
|
||||
entities: List[str]
|
||||
criteria: List[str]
|
||||
findings: Dict[str, List[str]]
|
||||
entities: List[str] # 3 names to compare
|
||||
criteria: List[str] # 3–5 criteria
|
||||
findings: Dict[str, List[str]] # entity -> list of notes per criterion
|
||||
final_table: str | None
|
||||
verdict: str | None
|
||||
# Internal helper to track progress
|
||||
_pair_index: int | None
|
||||
|
||||
# ---------- Nodes ----------
|
||||
# -----------------------------
|
||||
# LLM and Tavily tools
|
||||
# -----------------------------
|
||||
llm = ChatOpenAI(temperature=0.2)
|
||||
search = TavilySearchResults(max_results=3)
|
||||
|
||||
# -----------------------------
|
||||
# Node implementations
|
||||
# -----------------------------
|
||||
async def plan_criteria(state: CompareState) -> CompareState:
|
||||
"""Generate comparison criteria based on the entities."""
|
||||
entities = state["entities"]
|
||||
prompt = (
|
||||
f"You are an expert analyst. Given the entities: {', '.join(entities)}\n"
|
||||
"Generate 3-5 concise criteria for comparing them."
|
||||
f"""
|
||||
You are an expert analyst. Given the following entities: {', '.join(entities)}.
|
||||
Generate 3 to 5 distinct criteria that would be useful for comparing these entities.
|
||||
Return the criteria as a JSON array of strings.
|
||||
"""
|
||||
)
|
||||
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
||||
criteria = [c.strip() for c in response.content.split("\n") if c.strip()]
|
||||
result = await llm.invoke({"input": prompt})
|
||||
# Extract JSON array from the response
|
||||
import json, re
|
||||
try:
|
||||
json_text = re.search(r"\[.*\]", result.content, re.S).group(0)
|
||||
criteria = json.loads(json_text)
|
||||
except Exception as e:
|
||||
# Fallback: split by newlines
|
||||
criteria = [c.strip('- ') for c in result.content.splitlines() if c.strip()]
|
||||
state["criteria"] = criteria
|
||||
state["findings"] = {e: [] for e in entities}
|
||||
state["findings"] = {entity: [] for entity in entities}
|
||||
state["_pair_index"] = 0
|
||||
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
|
||||
"""Perform Tavily search for the current entity-criterion pair and store a short note."""
|
||||
entities = state["entities"]
|
||||
criteria = state["criteria"]
|
||||
idx = state["_pair_index"]
|
||||
if idx is None or idx >= len(entities) * len(criteria):
|
||||
return state
|
||||
entity_idx = idx // len(criteria)
|
||||
criterion_idx = idx % len(criteria)
|
||||
entity = entities[entity_idx]
|
||||
criterion = criteria[criterion_idx]
|
||||
|
||||
query = f"{entity} {criterion}"
|
||||
search_result = await search.invoke({"query": query})
|
||||
# Summarize the top result into a short note
|
||||
if search_result and "results" in search_result:
|
||||
top = search_result["results"][0]
|
||||
note = f"{criterion}: {top.get('content', top.get('title', 'No content')).strip()[:200]}"
|
||||
else:
|
||||
note = f"{criterion}: No relevant information found."
|
||||
state["findings"][entity].append(note)
|
||||
# Update index
|
||||
state["_pair_index"] = idx + 1
|
||||
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")
|
||||
"""Construct a markdown table from the findings."""
|
||||
entities = state["entities"]
|
||||
criteria = state["criteria"]
|
||||
findings = state["findings"]
|
||||
# Build header
|
||||
header = "| Criterion | " + " | ".join(entities) + " |"
|
||||
separator = "|---|" + "|---|" * len(entities)
|
||||
rows = [header, separator]
|
||||
for crit in criteria:
|
||||
row = [crit]
|
||||
for ent in entities:
|
||||
# Find the note that starts with this criterion
|
||||
note = next((n for n in findings[ent] if n.startswith(crit)), "No data")
|
||||
row.append(note)
|
||||
rows.append(" | ".join(row))
|
||||
table = "| " + headers + " |\n| " + " | ".join(["---"] * len(state["entities"])) + " |\n"
|
||||
table += "| " + " | ".join(rows) + " |"
|
||||
rows.append("| " + " | ".join(row) + " |")
|
||||
table = "\n".join(rows)
|
||||
state["final_table"] = table
|
||||
return state
|
||||
|
||||
async def verdict(state: CompareState) -> CompareState:
|
||||
"""Generate a verdict based on the table."""
|
||||
table = state["final_table"]
|
||||
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']}"
|
||||
f"""
|
||||
Based on the following comparative table, provide a concise verdict (2-4 sentences) recommending which entity is best suited for which use case.
|
||||
{table}
|
||||
"""
|
||||
)
|
||||
response = await llm.ainvoke([HumanMessage(content=prompt)])
|
||||
state["verdict"] = response.content.strip()
|
||||
result = await llm.invoke({"input": prompt})
|
||||
state["verdict"] = result.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)
|
||||
# -----------------------------
|
||||
# Graph construction
|
||||
# -----------------------------
|
||||
builder = StateGraph(CompareState)
|
||||
|
||||
# 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(
|
||||
# Add nodes
|
||||
builder.add_node("plan_criteria", plan_criteria)
|
||||
builder.add_node("research_entity", research_entity)
|
||||
builder.add_node("build_table", build_table)
|
||||
builder.add_node("verdict", verdict)
|
||||
|
||||
# Define the graph logic
|
||||
builder.set_entry_point("plan_criteria")
|
||||
|
||||
# After planning, start research loop
|
||||
builder.add_conditional_edges(
|
||||
"plan_criteria",
|
||||
lambda _: "research_entity",
|
||||
)
|
||||
|
||||
# After each research step, decide whether to continue or move to table
|
||||
builder.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.",
|
||||
lambda state: "build_table" if state.get("_pair_index") is None or state["_pair_index"] >= len(state["entities"])*len(state["criteria"]) else "research_entity",
|
||||
)
|
||||
|
||||
# ---------- CLI ----------
|
||||
builder.add_edge("build_table", "verdict")
|
||||
builder.add_edge("verdict", "END")
|
||||
|
||||
agent = builder.compile()
|
||||
|
||||
# -----------------------------
|
||||
# CLI interface
|
||||
# -----------------------------
|
||||
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 = {
|
||||
default_entities = ["Chroma", "FAISS", "Qdrant"]
|
||||
print("Enter three entities to compare (comma separated). Press Enter to use default:")
|
||||
user_input = input().strip()
|
||||
if user_input:
|
||||
entities = [e.strip() for e in user_input.split(',') if e.strip()]
|
||||
if len(entities) != 3:
|
||||
print("Please provide exactly three entities.")
|
||||
return
|
||||
else:
|
||||
entities = default_entities
|
||||
initial_state: CompareState = {
|
||||
"entities": entities,
|
||||
"criteria": [],
|
||||
"findings": {},
|
||||
"final_table": None,
|
||||
"verdict": None,
|
||||
"_pair_index": 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"])
|
||||
# Run the agent
|
||||
async for event in agent.astream_events(initial_state, version="1"):
|
||||
if event.get("type") == "on_node_end":
|
||||
node = event["name"]
|
||||
if node == "plan_criteria":
|
||||
print("\n=== Planned Criteria ===")
|
||||
print("\n".join(initial_state["criteria"]))
|
||||
elif node == "research_entity":
|
||||
idx = initial_state["_pair_index"] - 1
|
||||
entity_idx = idx // len(initial_state["criteria"])
|
||||
criterion_idx = idx % len(initial_state["criteria"])
|
||||
entity = initial_state["entities"][entity_idx]
|
||||
criterion = initial_state["criteria"][criterion_idx]
|
||||
note = initial_state["findings"][entity][-1]
|
||||
print(f"\n[{entity} × {criterion}] найдено: {note}")
|
||||
elif node == "build_table":
|
||||
print("\n=== Comparative Table ===")
|
||||
print(initial_state["final_table"]) # type: ignore
|
||||
elif node == "verdict":
|
||||
print("\n=== Verdict ===")
|
||||
print(initial_state["verdict"]) # type: ignore
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user