Reworked to produce research brief without tables. Updated README. Code compiles.: update src/main.py

src/main.py

@@ -1,7 +1,7 @@
 """
-LangGraph comparison agent demo.
+LangGraph research brief generator.
 Usage:
-    python -m src.main
+    python -m src.main "Topic"
 """
 import os
 from typing import TypedDict, List, Dict
@@ -13,117 +13,90 @@ 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
+class BriefState(TypedDict):
+    topic: str
+    sections: List[str]
+    content: Dict[str, str]
+    brief: 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"])
+async def plan_sections(state: BriefState) -> BriefState:
     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."
+        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."
     )
     response = await llm.ainvoke({"role": "user", "content": prompt})
-    # parse JSON
     import json
     try:
-        criteria = json.loads(response["content"].strip())
+        sections = json.loads(response["content"].strip())
     except Exception as e:
-        raise ValueError(f"Failed to parse criteria: {e}")
-    state["criteria"] = criteria
+        raise ValueError(f"Failed to parse sections: {e}")
+    state["sections"] = sections
     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}"
+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", "")
-            note = f"{crit}: {snippet}" if snippet else f"{crit}: no info"
-            state["findings"][ent].append(note)
+            content = snippet if snippet else "No relevant information found."
+            state["content"][sec] = content
             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."
+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."
     )
-    response = await llm.ainvoke({"role": "user", "content": prompt})
-    state["verdict"] = response["content"].strip()
+    summary_resp = await llm.ainvoke({"role": "user", "content": summary_prompt})
+    lines.append(summary_resp["content"].strip() + "\n")
+    state["brief"] = "\n".join(lines)
     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 = StateGraph(BriefState)
+builder.add_node("plan_sections", plan_sections)
+builder.add_node("fetch_section", fetch_section)
+builder.add_node("build_brief", build_brief)
 
-builder.set_entry_point("plan_criteria")
-# transition logic
+builder.set_entry_point("plan_sections")
+# after planning, fetch sections until all done
 builder.add_conditional_edges(
-    "plan_criteria",
-    lambda _: "research_entity",
+    "plan_sections",
+    lambda _: "fetch_section",
 )
 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",
+    "fetch_section",
+    lambda state: "build_brief" if all(sec in state["content"] for sec in state["sections"]) else "fetch_section",
 )
-builder.add_edge("build_table", "verdict")
-builder.add_edge("verdict", "END")
+builder.add_edge("build_brief", "END")
 
 app = builder.compile()
 
 # ---------- Demo ----------
 if __name__ == "__main__":
     import argparse
-    parser = argparse.ArgumentParser(description="Compare three entities.")
-    parser.add_argument("entities", nargs='*', help="Three entities to compare (default: Chroma, FAISS, Qdrant)")
+    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")
     args = parser.parse_args()
-    if not args.entities:
-        args.entities = ["Chroma", "FAISS", "Qdrant"]
-    init_state: CompareState = {
-        "entities": list(args.entities),
-        "criteria": [],
-        "findings": {},
-        "final_table": None,
-        "verdict": None,
+    init_state: BriefState = {
+        "topic": args.topic,
+        "sections": [],
+        "content": {},
+        "brief": 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"])
+    print(result["brief"])