Published solution: remove compare_agent.py
This commit is contained in:
@@ -1,145 +0,0 @@
|
||||
"""
|
||||
CompareAgent – LangGraph implementation that builds a comparative table for three entities.
|
||||
|
||||
Usage:
|
||||
python compare_agent.py "Chroma, FAISS, Qdrant"
|
||||
|
||||
The script will:
|
||||
1. Ask LLM to generate 3‑5 comparison criteria.
|
||||
2. For each entity × criterion pair perform a Tavily search and collect short notes.
|
||||
3. Build a markdown table with rows = criteria, columns = entities.
|
||||
4. Produce a verdict sentence.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from typing import TypedDict, List, Dict
|
||||
|
||||
from langgraph.graph import StateGraph
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_tavily import TavilySearch
|
||||
|
||||
# ---------- State definition ----------
|
||||
class CompareState(TypedDict):
|
||||
entities: List[str]
|
||||
criteria: List[str]
|
||||
findings: Dict[str, List[str]] # entity -> list of notes per criterion
|
||||
final_table: str | None
|
||||
verdict: str | None
|
||||
|
||||
# ---------- LLM and tools ----------
|
||||
llm = ChatOpenAI(temperature=0.2)
|
||||
search_tool = TavilySearch(api_key=os.getenv("TAVILY_API_KEY"))
|
||||
|
||||
# ---------- Node functions ----------
|
||||
async def plan_criteria(state: CompareState) -> Dict:
|
||||
"""Generate comparison criteria based on entities."""
|
||||
prompt = (
|
||||
f"You are an expert in evaluating technologies.\n"
|
||||
f"Given the following entities: {', '.join(state['entities'])}.\n"
|
||||
f"Provide 3‑5 concise criteria that would be useful for comparing them."
|
||||
)
|
||||
response = await llm.agenerate([prompt])
|
||||
text = response.generations[0][0].text.strip()
|
||||
# split by newlines or commas
|
||||
crits = [c.strip() for c in text.replace('\n', ',').split(',') if c.strip()]
|
||||
return {"criteria": crits}
|
||||
|
||||
async def research_entity(state: CompareState) -> Dict:
|
||||
"""Perform Tavily search for the next unprocessed entity × criterion pair."""
|
||||
# find first entity with missing notes
|
||||
for ent in state['entities']:
|
||||
notes = state['findings'].get(ent, [])
|
||||
if len(notes) < len(state['criteria']):
|
||||
idx = len(notes)
|
||||
crit = state['criteria'][idx]
|
||||
query = f"{ent} {crit}"
|
||||
result = await search_tool.ainvoke(query=query, max_results=1)
|
||||
snippet = result.get('results', [{}])[0].get('content', 'No info')
|
||||
notes.append(snippet[:200]) # truncate
|
||||
state['findings'][ent] = notes
|
||||
break
|
||||
return {"findings": state['findings']}
|
||||
|
||||
async def build_table(state: CompareState) -> Dict:
|
||||
"""Create markdown table from findings."""
|
||||
header = "| Criterion | " + " | ".join(state['entities']) + " |\n"
|
||||
separator = "|---|" + "---|" * len(state['entities']) + "\n"
|
||||
rows = []
|
||||
for i, crit in enumerate(state['criteria']):
|
||||
row_cells = [crit]
|
||||
for ent in state['entities']:
|
||||
notes = state['findings'].get(ent, [])
|
||||
note = notes[i] if i < len(notes) else ""
|
||||
row_cells.append(note)
|
||||
rows.append("| " + " | ".join(row_cells) + " |\n")
|
||||
table = header + separator + "".join(rows)
|
||||
return {"final_table": table}
|
||||
|
||||
async def verdict(state: CompareState) -> Dict:
|
||||
"""Generate a short recommendation based on the table."""
|
||||
prompt = (
|
||||
f"You are an analyst. Based on the following comparison table:\n\n"
|
||||
f"{state['final_table']}\n\n"
|
||||
f"Provide 2‑3 sentences recommending which entity is best for each use case.")
|
||||
response = await llm.agenerate([prompt])
|
||||
text = response.generations[0][0].text.strip()
|
||||
return {"verdict": text}
|
||||
|
||||
# ---------- Graph construction ----------
|
||||
def create_graph() -> StateGraph:
|
||||
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)
|
||||
|
||||
# start -> plan_criteria
|
||||
graph.set_entry_point("plan_criteria")
|
||||
|
||||
# after criteria, loop research until all pairs processed
|
||||
def should_continue(state: CompareState) -> bool:
|
||||
return any(len(notes) < len(state['criteria']) for notes in state.get('findings', {}).values())
|
||||
|
||||
graph.add_conditional_edges(
|
||||
"plan_criteria",
|
||||
lambda s: "research_entity" if should_continue(s) else "build_table",
|
||||
{"research_entity": "research_entity", "build_table": "build_table"},
|
||||
)
|
||||
|
||||
# after research, decide again
|
||||
graph.add_conditional_edges(
|
||||
"research_entity",
|
||||
lambda s: "research_entity" if should_continue(s) else "build_table",
|
||||
{"research_entity": "research_entity", "build_table": "build_table"},
|
||||
)
|
||||
|
||||
# after table, verdict
|
||||
graph.add_edge("build_table", "verdict")
|
||||
graph.set_finish_node("verdict")
|
||||
return graph
|
||||
|
||||
# ---------- Main execution ----------
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python compare_agent.py 'entity1, entity2, entity3'")
|
||||
sys.exit(1)
|
||||
entities = [e.strip() for e in sys.argv[1].split(',')]
|
||||
if len(entities) != 3:
|
||||
print("Please provide exactly three entities separated by commas.")
|
||||
sys.exit(1)
|
||||
|
||||
initial_state: CompareState = {
|
||||
"entities": entities,
|
||||
"criteria": [],
|
||||
"findings": {},
|
||||
"final_table": None,
|
||||
"verdict": None,
|
||||
}
|
||||
|
||||
graph = create_graph()
|
||||
result = graph.invoke(initial_state)
|
||||
print("\n=== Comparison Table ===")
|
||||
print(result["final_table"])
|
||||
print("\n=== Verdict ===")
|
||||
print(result["verdict"])
|
||||
Reference in New Issue
Block a user