Week 5-6：LangChain / LangGraph 入门

# ===== LangGraph 伪代码 =====
from langgraph.graph import StateGraph
from typing import TypedDict

class MyState(TypedDict):
    # ... 定义状态字段

graph = StateGraph(MyState)
graph.add_node("node_a", handler_a)
graph.add_node("node_b", handler_b)
graph.set_entry_point("node_a")
graph.add_edge("node_a", "node_b")        # 固定流转
graph.add_conditional_edges(...)          # 条件流转
app = graph.compile()
result = await app.ainvoke(initial_state)

# app/agent/trading_agent.py
"""
用 LangGraph 构建的交易 Agent
能力：
1. 对话理解
2. 工具调用（查询价格、下单）
3. 决策（要不要下单）
"""

from typing import TypedDict, Literal
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool

# ===== 第一步：定义 State =====
class TradingAgentState(TypedDict):
    """Agent 在各节点之间流转的状态"""
    messages: list              # 对话历史
    user_intent: str            # 用户意图分类
    tools_called: list[str]     # 已调用过的工具名
    final_answer: str           # 最终回答

# ===== 第二步：定义 Tools（Agent 可调用的函数）=====

@tool
def get_stock_price(symbol: str) -> dict:
    """获取交易标的（股票/加密货币）的当前价格"""
    mock_prices = {"BTC": 99999.99, "ETH": 3500.0, "AAPL": 210.5, "特斯拉": 245.0}
    price = mock_prices.get(symbol.upper(), mock_prices.get(symbol, 0))
    return {"symbol": symbol.upper(), "price": price, "currency": "USD"}

@tool
def place_order(symbol: str, side: str, quantity: int) -> dict:
    """
    下单交易。注意：这是 Mock 实现，不会真的下单。
    
    Args:
        symbol: 交易标的代码，如 BTC, ETH, AAPL
        side: buy 或 sell
        quantity: 交易数量
    """
    return {
        "order_id": f"ORD{123456:06d}",
        "symbol": symbol.upper(),
        "side": side,
        "quantity": quantity,
        "status": "pending",
        "timestamp": "2024-01-15T10:30:00Z",
    }

@tool
def get_market_news(symbol: str) -> list[str]:
    """获取某标的的近期市场新闻摘要"""
    news_map = {
        "BTC": ["SEC 批准比特币 ETF", "贝莱德宣布增持 BTC"],
        "ETH": ["以太坊坎昆升级完成", "Layer2 TVL 创新高"],
    }
    return news_map.get(symbol.upper(), [f"暂无 {symbol} 的相关新闻"])

TOOLS = [get_stock_price, place_order, get_market_news]

# ===== 第三步：初始化 LLM =====
# 硅基流动是 OpenAI 兼容接口，所以用 ChatOpenAI
llm = ChatOpenAI(
    model="Qwen/Qwen2.5-7B-Instruct",
    base_url="https://api.siliconflow.cn/v1",
    api_key="your-api-key-here",
    temperature=0.7,
)

# 让 LLM 知道有哪些工具可以用
llm_with_tools = llm.bind_tools(TOOLS)

# ===== 第四步：定义图节点 =====

def analyze_intent(state: TradingAgentState) -> TradingAgentState:
    """
    节点1：分析用户意图，决定要不要调工具、调哪些工具
    这是一个"路由器"节点
    """
    # 取最后一条用户消息
    last_msg = state["messages"][-1].content if state["messages"] else ""
    
    # 简单的意图分类（实际项目用 LLM 分类）
    intent_keywords = {
        "get_price": ["价格", "多少钱", "行情", "现在多少", "price"],
        "place_order": ["买", "卖", "下单", "交易", "买入", "卖出"],
        "get_news": ["新闻", "消息", "分析", "资讯"],
        "chat": ["你好", "help", "帮助", "能做什么"],
    }
    
    detected_intents = []
    for intent, keywords in intent_keywords.items():
        if any(kw.lower() in last_msg.lower() for kw in keywords):
            detected_intents.append(intent)
    
    # 用 LLM 做最终判断（比关键词匹配更准）
    analysis_prompt = f"""
    用户消息："{last_msg}"
    检测到的意图关键词：{detected_intents}
    
    请判断用户的主要意图是以下哪一种（只返回关键词）：
    - query_price: 查询价格
    - place_order: 下单交易
    - query_news: 查询新闻
    - general_chat: 闲聊或帮助
    
    只返回意图关键词，不要其他内容。
    """
    response = llm.invoke([{"role": "user", "content": analysis_prompt}])
    user_intent = response.content.strip()
    
    return {**state, "user_intent": user_intent, "tools_called": []}


def llm_with_tool_calling(state: TradingAgentState) -> TradingAgentState:
    """
    节点2：LLM + Function Calling 主逻辑
    调用 LLM，如果要调工具就让 LangGraph 自动路由到 ToolNode
    """
    messages = state["messages"]
    response = llm_with_tools.invoke(messages)
    return {**state, "messages": [*messages, response]}


def generate_final_response(state: TradingAgentState) -> TradingAgentState:
    """
    节点3：生成最终回答（对工具返回结果做总结）
    """
    messages = state["messages"]
    
    summary_prompt = f"""
    请根据以下对话历史和工具返回结果，给用户一个清晰、自然的回答。
    不要提到内部工具调用过程，用面向用户的语言。
    
    对话历史：
    {messages}
    """
    response = llm.invoke([{"role": "user", "content": summary_prompt}])
    return {**state, "final_answer": response.content}


def direct_chat_response(state: TradingAgentState) -> TradingAgentState:
    """
    简化模式：不需要工具的闲聊，直接让 LLM 回答
    """
    messages = state["messages"]
    response = llm.invoke(messages)
    return {**state, "final_answer": response.content}


# ===== 第五步：构建 StateGraph =====

def should_use_tools(state: TradingAgentState) -> str:
    """条件函数：根据意图决定走哪条路"""
    intent = state["user_intent"]
    if intent in ("query_price", "place_order", "query_news"):
        return "use_tools"
    return "direct_chat"


# 构建图
graph = StateGraph(TradingAgentState)

# 添加节点
graph.add_node("analyze_intent", analyze_intent)
graph.add_node("use_tools", llm_with_tool_calling)
graph.add_node("tools", ToolNode(TOOLS))              # LangGraph 内置的工具执行节点
graph.add_node("summarize", generate_final_response)
graph.add_node("direct_chat", direct_chat_response)

# 边定义
graph.add_edge(START, "analyze_intent")                          # 入口 → 意图分析
graph.add_conditional_edges(                                      # 意图分析 → 分支
    "analyze_intent",
    should_use_tools,
    {
        "use_tools": "use_tools",
        "direct_chat": "direct_chat",
    }
)
graph.add_edge("use_tools", "tools")                              # LLM 调用 → 执行工具
graph.add_edge("tools", "summarize")                              # 工具执行完 → 总结
graph.add_edge("direct_chat", END)                                # 闲聊直接结束

# 编译成可运行的应用
app = graph.compile()


# ===== 第六步：暴露给 FastAPI 用的接口 =====

async def run_agent(user_input: str, history: list = None) -> str:
    """运行 Agent，返回最终回答"""
    initial_state: TradingAgentState = {
        "messages": history or [],
        "user_intent": "",
        "tools_called": [],
        "final_answer": "",
    }
    initial_state["messages"].append({"role": "user", "content": user_input})
    
    result = await app.ainvoke(initial_state)
    return result["final_answer"]


# 测试
if __name__ == "__main__":
    import asyncio

    async def test():
        # 测试 1：查询价格（触发工具）
        r1 = await run_agent("BTC 现在多少钱？")
        print(f"Q: BTC 现在多少钱？\nA: {r1}\n")

        # 测试 2：下单（触发工具）
        r2 = await run_agent("帮我买 0.5 个 ETH")
        print(f"Q: 帮我买 0.5 个 ETH\nA: {r2}\n")

        # 测试 3：闲聊（不触发工具）
        r3 = await run_agent("你好，你能做什么？")
        print(f"Q: 你好，你能做什么？\nA: {r3}\n")

    asyncio.run(test())