From be2c9466862a8572b667efacd3153ec2d9f79406 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A0=D0=B8=D0=BD=D0=B0=D1=80=20=D0=9C=D0=B8=D1=80=D0=B7?= =?UTF-8?q?=D0=B0=D0=B3=D0=B8=D1=82=D0=BE=D0=B2?= Date: Thu, 11 Jun 2026 15:33:18 +0000 Subject: [PATCH] =?UTF-8?q?=D0=A1=D0=B4=D0=B5=D0=BB=D0=B0=D0=BD=D0=BE=20?= =?UTF-8?q?=D1=80=D0=B5=D1=88=D0=B5=D0=BD=D0=B8=D0=B5:=20add=20compare=5Fa?= =?UTF-8?q?gent.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- compare_agent.py | 145 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 145 insertions(+) create mode 100644 compare_agent.py diff --git a/compare_agent.py b/compare_agent.py new file mode 100644 index 0000000..3cbe5a6 --- /dev/null +++ b/compare_agent.py @@ -0,0 +1,145 @@ +""" +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"])