From 08ab537f6b0ec944c214b33230e9fc4bb8c0db96 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, 18 Jun 2026 09:55:02 +0000 Subject: [PATCH] Solution published: update src/main.py --- src/main.py | 142 ++++++++++++++++++++++++++++++---------------------- 1 file changed, 83 insertions(+), 59 deletions(-) diff --git a/src/main.py b/src/main.py index 35d1bc9..4a4ec11 100644 --- a/src/main.py +++ b/src/main.py @@ -1,11 +1,12 @@ """ -LangGraph research brief generator. +LangGraph comparison agent. Usage: - python -m src.main "Topic" + python -m src.main "Entity1,Entity2,Entity3" """ import os +import json from typing import TypedDict, List, Dict -from langgraph.graph import StateGraph +from langgraph.graph import StateGraph, END from langchain_openai import ChatOpenAI from langchain_tavily import TavilySearch from dotenv import load_dotenv @@ -13,90 +14,113 @@ from dotenv import load_dotenv load_dotenv() # ---------- State ---------- -class BriefState(TypedDict): - topic: str - sections: List[str] - content: Dict[str, str] - brief: str | None +class CompareState(TypedDict): + entities: List[str] # 3 names for comparison + criteria: List[str] # 3–5 criteria + findings: Dict[str, Dict[str, str]] # entity -> criterion -> note + final_table: str | None + verdict: str | None -# ---------- Nodes ---------- +# ---------- LLM and tools ---------- llm = ChatOpenAI(temperature=0) search_tool = TavilySearch(api_key=os.getenv("TAVILY_API_KEY")) -async def plan_sections(state: BriefState) -> BriefState: +# ---------- Nodes ---------- +async def plan_criteria(state: CompareState) -> CompareState: prompt = ( - f"You are a helpful assistant. Given the topic: {state['topic']}. " - "Generate 5 concise section headings for a research brief. Return a JSON array of strings." + f"You are a helpful assistant. Given the entities: {', '.join(state['entities'])}. " + "Generate 3 to 5 concise criteria for comparing these entities. " + "Return a JSON array of strings." ) response = await llm.ainvoke({"role": "user", "content": prompt}) - import json try: - sections = json.loads(response["content"].strip()) + criteria = json.loads(response["content"].strip()) except Exception as e: - raise ValueError(f"Failed to parse sections: {e}") - state["sections"] = sections + raise ValueError(f"Failed to parse criteria: {e}") + state["criteria"] = criteria return state -async def fetch_section(state: BriefState) -> BriefState: - # find next section not yet fetched - for sec in state["sections"]: - if sec not in state["content"]: - query = f"{state['topic']} {sec}" - result = await search_tool.ainvoke({"query": query, "max_results": 1}) - snippet = result.get("results", [{}])[0].get("snippet", "") - content = snippet if snippet else "No relevant information found." - state["content"][sec] = content - 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 entity not in state["findings"] or criterion not in state["findings"][entity]: + query = f"{entity} {criterion}" + result = await search_tool.ainvoke({"query": query, "max_results": 1}) + snippet = result.get("results", [{}])[0].get("snippet", "") + note = snippet if snippet else "No relevant information found." + state.setdefault("findings", {}) + state["findings"][entity] = state["findings"].get(entity, {}) + state["findings"][entity][criterion] = note + return state return state -async def build_brief(state: BriefState) -> BriefState: - lines = [f"# Research Brief: {state['topic']}\n"] - for sec in state["sections"]: - lines.append(f"## {sec}\n") - lines.append(state["content"][sec] + "\n") - lines.append("---\n") - lines.append("**Summary**\n") - # simple summary using LLM - summary_prompt = ( - f"Based on the following sections:\n{''.join(lines)}\n" - "Provide a concise summary of the brief in 3-4 sentences." +async def build_table(state: CompareState) -> CompareState: + # Build markdown table + header = "| Criterion | " + " | ".join(state["entities"]) + " |" + separator = "|---|" + "---|" * len(state["entities"]) # simple separator + rows = [] + for criterion in state["criteria"]: + cells = [criterion] + for entity in state["entities"]: + note = state["findings"].get(entity, {}).get(criterion, "") + cells.append(note.replace("\n", " ")) + rows.append("| " + " | ".join(cells) + " |") + table = "\n".join([header, separator] + rows) + state["final_table"] = table + return state + +async def verdict(state: CompareState) -> CompareState: + prompt = ( + f"You have the following comparison table:\n{state['final_table']}\n" + "Based on this table, provide a concise verdict: which entity is best for a general-purpose use case, and why." ) - summary_resp = await llm.ainvoke({"role": "user", "content": summary_prompt}) - lines.append(summary_resp["content"].strip() + "\n") - state["brief"] = "\n".join(lines) + response = await llm.ainvoke({"role": "user", "content": prompt}) + state["verdict"] = response["content"].strip() return state # ---------- Graph ---------- -builder = StateGraph(BriefState) -builder.add_node("plan_sections", plan_sections) -builder.add_node("fetch_section", fetch_section) -builder.add_node("build_brief", build_brief) +builder = StateGraph(CompareState) +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) -builder.set_entry_point("plan_sections") -# after planning, fetch sections until all done +builder.set_entry_point("plan_criteria") +# After planning, loop research until all pairs processed builder.add_conditional_edges( - "plan_sections", - lambda _: "fetch_section", + "plan_criteria", + lambda _: "research_entity", ) builder.add_conditional_edges( - "fetch_section", - lambda state: "build_brief" if all(sec in state["content"] for sec in state["sections"]) else "fetch_section", + "research_entity", + lambda state: "build_table" if all( + criterion in state["findings"].get(entity, {}) for entity in state["entities"] for criterion in state["criteria"] + ) else "research_entity", ) -builder.add_edge("build_brief", "END") +builder.add_edge("build_table", "verdict") +builder.add_edge("verdict", END) app = builder.compile() # ---------- Demo ---------- if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Generate a research brief on a topic.") - parser.add_argument("topic", nargs='?', default="Artificial Intelligence", help="Topic for the research brief") + parser = argparse.ArgumentParser(description="Compare three entities and produce a table and verdict.") + parser.add_argument("entities", nargs='?', default="Chroma,FAISS,Qdrant", help="Comma-separated list of three entities to compare") args = parser.parse_args() - init_state: BriefState = { - "topic": args.topic, - "sections": [], - "content": {}, - "brief": None, + entities = [e.strip() for e in args.entities.split(',') if e.strip()] + if len(entities) != 3: + raise ValueError("Please provide exactly three entities separated by commas.") + init_state: CompareState = { + "entities": entities, + "criteria": [], + "findings": {}, + "final_table": None, + "verdict": None, } result = app.invoke(init_state) - print(result["brief"]) + print("## Comparison Table\n") + print(result["final_table"], "\n") + print("## Verdict\n") + print(result["verdict"], "\n")