Solution published: update src/main.py
This commit is contained in:
+80
-56
@@ -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}"
|
||||
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", "")
|
||||
content = snippet if snippet else "No relevant information found."
|
||||
state["content"][sec] = content
|
||||
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")
|
||||
|
||||
Reference in New Issue
Block a user