Solution published: update src/main.py

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):
+class CompareState(TypedDict):
-    topic: str
+    entities: List[str]                # 3 names for comparison
-    sections: List[str]
+    criteria: List[str]                # 3–5 criteria
-    content: Dict[str, str]
+    findings: Dict[str, Dict[str, str]]  # entity -> criterion -> note
-    brief: str | None
+    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']}. "
+        f"You are a helpful assistant. Given the entities: {', '.join(state['entities'])}. "
-        "Generate 5 concise section headings for a research brief. Return a JSON array of strings."
+        "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}")
+        raise ValueError(f"Failed to parse criteria: {e}")
-    state["sections"] = sections
+    state["criteria"] = criteria
     return state
 
-async def fetch_section(state: BriefState) -> BriefState:
+async def research_entity(state: CompareState) -> CompareState:
-    # find next section not yet fetched
+    # Find next unprocessed entity-criterion pair
-    for sec in state["sections"]:
+    for entity in state["entities"]:
-        if sec not in state["content"]:
+        for criterion in state["criteria"]:
-            query = f"{state['topic']} {sec}"
+            if entity not in state["findings"] or criterion not in state["findings"][entity]:
-            result = await search_tool.ainvoke({"query": query, "max_results": 1})
+                query = f"{entity} {criterion}"
-            snippet = result.get("results", [{}])[0].get("snippet", "")
+                result = await search_tool.ainvoke({"query": query, "max_results": 1})
-            content = snippet if snippet else "No relevant information found."
+                snippet = result.get("results", [{}])[0].get("snippet", "")
-            state["content"][sec] = content
+                note = snippet if snippet else "No relevant information found."
-            return state
+                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:
+async def build_table(state: CompareState) -> CompareState:
-    lines = [f"# Research Brief: {state['topic']}\n"]
+    # Build markdown table
-    for sec in state["sections"]:
+    header = "| Criterion | " + " | ".join(state["entities"]) + " |"
-        lines.append(f"## {sec}\n")
+    separator = "|---|" + "---|" * len(state["entities"])  # simple separator
-        lines.append(state["content"][sec] + "\n")
+    rows = []
-    lines.append("---\n")
+    for criterion in state["criteria"]:
-    lines.append("**Summary**\n")
+        cells = [criterion]
-    # simple summary using LLM
+        for entity in state["entities"]:
-    summary_prompt = (
+            note = state["findings"].get(entity, {}).get(criterion, "")
-        f"Based on the following sections:\n{''.join(lines)}\n"
+            cells.append(note.replace("\n", " "))
-        "Provide a concise summary of the brief in 3-4 sentences."
+        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})
+    response = await llm.ainvoke({"role": "user", "content": prompt})
-    lines.append(summary_resp["content"].strip() + "\n")
+    state["verdict"] = response["content"].strip()
-    state["brief"] = "\n".join(lines)
     return state
 
 # ---------- Graph ----------
-builder = StateGraph(BriefState)
+builder = StateGraph(CompareState)
-builder.add_node("plan_sections", plan_sections)
+builder.add_node("plan_criteria", plan_criteria)
-builder.add_node("fetch_section", fetch_section)
+builder.add_node("research_entity", research_entity)
-builder.add_node("build_brief", build_brief)
+builder.add_node("build_table", build_table)
+builder.add_node("verdict", verdict)
 
-builder.set_entry_point("plan_sections")
+builder.set_entry_point("plan_criteria")
-# after planning, fetch sections until all done
+# After planning, loop research until all pairs processed
 builder.add_conditional_edges(
-    "plan_sections",
+    "plan_criteria",
-    lambda _: "fetch_section",
+    lambda _: "research_entity",
 )
 builder.add_conditional_edges(
-    "fetch_section",
+    "research_entity",
-    lambda state: "build_brief" if all(sec in state["content"] for sec in state["sections"]) else "fetch_section",
+    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 = argparse.ArgumentParser(description="Compare three entities and produce a table and verdict.")
-    parser.add_argument("topic", nargs='?', default="Artificial Intelligence", help="Topic for the research brief")
+    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 = {
+    entities = [e.strip() for e in args.entities.split(',') if e.strip()]
-        "topic": args.topic,
+    if len(entities) != 3:
-        "sections": [],
+        raise ValueError("Please provide exactly three entities separated by commas.")
-        "content": {},
+    init_state: CompareState = {
-        "brief": None,
+        "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")