Initial solution ready: add src/main.py

src/main.py

@@ -0,0 +1,127 @@
+"""
+LangGraph comparison agent demo.
+Usage:
+    python -m src.main
+"""
+import os
+from typing import TypedDict, List, Dict
+from langgraph.graph import StateGraph
+from langchain_openai import ChatOpenAI
+from langchain_tavily import TavilySearch
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# ---------- State ----------
+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
+
+# ---------- Nodes ----------
+llm = ChatOpenAI(temperature=0)
+search_tool = TavilySearch(api_key=os.getenv("TAVILY_API_KEY"))
+
+async def plan_criteria(state: CompareState) -> CompareState:
+    entities_str = ", ".join(state["entities"])
+    prompt = (
+        f"You are a helpful assistant. Given the following entities: {entities_str}. "
+        "Generate 3-5 concise criteria for comparing them. Return a JSON array of strings."
+    )
+    response = await llm.agenerate([{"role": "user", "content": prompt}])
+    # parse JSON
+    import json
+    try:
+        criteria = json.loads(response.generations[0][0].text.strip())
+    except Exception as e:
+        raise ValueError(f"Failed to parse criteria: {e}")
+    state["criteria"] = criteria
+    return state
+
+async def research_entity(state: CompareState) -> CompareState:
+    # find next unprocessed entity-criterion pair
+    for ent in state["entities"]:
+        if ent not in state["findings"]:
+            state["findings"][ent] = []
+        for crit in state["criteria"]:
+            if any(crit.lower() in note.lower() for note in state["findings"][ent]):
+                continue
+            # perform search
+            query = f"{ent} {crit}"
+            result = await search_tool.ainvoke({"query": query, "max_results": 1})
+            snippet = result.get("results", [{}])[0].get("snippet", "")
+            note = f"{crit}: {snippet}" if snippet else f"{crit}: no info"
+            state["findings"][ent].append(note)
+            return state
+    # all processed
+    return state
+
+async def build_table(state: CompareState) -> CompareState:
+    rows = []
+    for crit in state["criteria"]:
+        row = [crit]
+        for ent in state["entities"]:
+            notes = state["findings"][ent]
+            note = next((n for n in notes if n.startswith(crit)), "")
+            row.append(note)
+        rows.append(row)
+    # markdown table
+    header = " | ".join(["Criterion"] + state["entities"])
+    sep = " | ".join([":---:"] * (len(state["entities"]) + 1))
+    table_rows = [f"{row[0]} | {' | '.join(row[1:])}" for row in rows]
+    state["final_table"] = f"{header}\n{sep}\n" + "\n".join(table_rows)
+    return state
+
+async def verdict(state: CompareState) -> CompareState:
+    prompt = (
+        f"Based on the following table:\n{state['final_table']}\n"
+        "Provide a short recommendation (2-4 sentences) on which entity is best for each use case."
+    )
+    response = await llm.agenerate([{"role": "user", "content": prompt}])
+    state["verdict"] = response.generations[0][0].text.strip()
+    return state
+
+# ---------- Graph ----------
+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_criteria")
+# transition logic
+builder.add_conditional_edges(
+    "plan_criteria",
+    lambda _: "research_entity",
+)
+builder.add_conditional_edges(
+    "research_entity",
+    lambda state: "build_table" if all(ent in state["findings"] and len(state["findings"][ent]) == len(state["criteria"]) for ent in state["entities"]) else "research_entity",
+)
+builder.add_edge("build_table", "verdict")
+builder.add_edge("verdict", "END")
+
+app = builder.compile()
+
+# ---------- Demo ----------
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="Compare three entities.")
+    parser.add_argument("entities", nargs=3, help="Three entities to compare")
+    args = parser.parse_args()
+    init_state: CompareState = {
+        "entities": list(args.entities),
+        "criteria": [],
+        "findings": {},
+        "final_table": None,
+        "verdict": None,
+    }
+    result = app.invoke(init_state)
+    print("\n=== Criteria ===")
+    print(result["criteria"])
+    print("\n=== Table ===")
+    print(result["final_table"])
+    print("\n=== Verdict ===")
+    print(result["verdict"])