# 安装: pip install langgraph langchain langchain-openai

from typing import TypedDict, Annotated
import operator
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain.tools import DuckDuckGoSearchRun

# 1. 定义状态
class ResearchState(TypedDict):
    topic: str                    # 研究主题
    search_results: Annotated[list, operator.add]  # 搜索结果（累加）
    analysis: str                 # 分析结果
    report: str                   # 最终报告

# 2. 定义节点
def search_node(state: ResearchState):
    """搜索节点：搜索相关信息"""
    search = DuckDuckGoSearchRun()
    result = search.run(state["topic"])
    return {"search_results": [result]}

def analyze_node(state: ResearchState):
    """分析节点：提炼关键信息"""
    llm = ChatOpenAI(model="gpt-4o")
    prompt = f"从以下搜索结果中提取3个关键发现（中文）：\n{state['search_results']}"
    analysis = llm.invoke(prompt).content
    return {"analysis": analysis}

def write_report_node(state: ResearchState):
    """报告节点：生成最终报告"""
    llm = ChatOpenAI(model="gpt-4o")
    prompt = f"""
主题：{state['topic']}
研究发现：{state['analysis']}

请生成一份结构化的研究报告（Markdown）：
1. 背景概述
2. 三个关键发现
3. 趋势判断
4. 建议
"""
    report = llm.invoke(prompt).content
    return {"report": report}

# 3. 构建工作流图
workflow = StateGraph(ResearchState)

workflow.add_node("search", search_node)
workflow.add_node("analyze", analyze_node)
workflow.add_node("write_report", write_report_node)

# 定义边
workflow.set_entry_point("search")
workflow.add_edge("search", "analyze")
workflow.add_edge("analyze", "write_report")
workflow.add_edge("write_report", END)

app = workflow.compile()

# 4. 执行
result = app.invoke({"topic": "2026年中国AI大模型市场格局"})
print(result["report"])