diff --git a/src/main.py b/src/main.py index 5752bb9..35d1bc9 100644 --- a/src/main.py +++ b/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"])