FiDA_Python/src/server/agent/graph.py

from typing import Literal
from langchain_core.messages import AIMessage
from langchain_core.runnables import RunnableConfig
from langchain_qwq import ChatQwen
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.graph import StateGraph, END, START
from pydantic import BaseModel
from pymongo import MongoClient

from src.core.config import MONGO_URI, settings
from src.server.agent.state import AgentState
from src.server.agent.agents import designer_node, researcher_node, visualizer_node, suggester_node
from langgraph.checkpoint.mongodb import MongoDBSaver


# --- Supervisor (路由逻辑) ---
# 定义路由的输出结构，强制 LLM 选择一个
class RouteResponse(BaseModel):
    # 将 FINISH 替换或增加 Suggester
    next: Literal["Designer", "Researcher", "Visualizer", "Suggester", "FINISH"]


llm_supervisor = ChatQwen(
    model="qwen3.5-flash",
    max_tokens=3_000,
    timeout=None,
    max_retries=2,
    api_key=settings.QWEN_API_KEY)


def supervisor_node(state: AgentState, config: RunnableConfig):
    use_report = config["configurable"].get("use_report", False)
    messages = state["messages"]
    if not messages:
        return {"next": "Suggester"}

    last_message = messages[-1]

    # --- 拦截逻辑修改 ---
    # 如果专家已经回复完了（AIMessage 且无工具调用），则交给 Suggester 生成按钮
    if isinstance(last_message, AIMessage) and not last_message.tool_calls:
        should_go_to_suggester = state.get("require_suggestion", False)

        # 如果符合建议条件
        if should_go_to_suggester:
            return {"next": "Suggester"}
        else:
            return {"next": "FINISH"}

    system_prompt = f"""你是家具设计主管，负责分配任务。
         当前设定：
         - 是否需要市场研究报告：{'是' if use_report else '否'}
        
         严格遵守以下规则：
         - 如果 **不需要** 市场研究报告（use_report = False），**绝对不能** 选择 Researcher
         - 只有在 **明确需要** 市场报告、竞争分析、材质趋势、价格区间等外部资讯时，才选择 Researcher，且 **必须** use_report = True
         - 常见分配：
         - 纯设计、风格、尺寸、材质建议 → Designer
         - 需要生成图片、渲染 → Visualizer
         - 需要产生建议按钮 → Suggester
         - 需要市场报告 → Researcher（但只有 use_report=True 时才允许）
         - 对话已完整、无需继续 → FINISH
        
         用户最后说了什么？请根据实际需求决定下一步。
    """

    chain = llm_supervisor.with_structured_output(RouteResponse)
    decision = chain.invoke([{"role": "system", "content": system_prompt}] + messages)
    next_node = decision.next

    if next_node == "Researcher" and not use_report:
        print("警告：LLM 違規選擇 Researcher，已強制改為 FINISH 或 Suggester")
        next_node = "Suggester" if state.get("require_suggestion", False) else "FINISH"

    return {"next": next_node}


# --- 构建 Graph ---
workflow = StateGraph(AgentState)

workflow.add_node("Supervisor", supervisor_node)
workflow.add_node("Designer", designer_node)
workflow.add_node("Researcher", researcher_node)
workflow.add_node("Visualizer", visualizer_node)
workflow.add_node("Suggester", suggester_node)  # 新增节点
workflow.add_edge(START, "Supervisor")

# 修改条件边映射
workflow.add_conditional_edges(
    "Supervisor",
    lambda state: state["next"],
    {
        "Designer": "Designer",
        "Researcher": "Researcher",
        "Visualizer": "Visualizer",
        "Suggester": "Suggester",  # 原本的 FINISH 现在指向 Suggester
        "FINISH": END  # 直接结束，不给建议
    }
)

# 专家执行完依然回到 Supervisor
workflow.add_edge("Designer", "Supervisor")
workflow.add_edge("Researcher", "Supervisor")
workflow.add_edge("Visualizer", "Supervisor")
# 重点：Suggester 可以是整个流程的终点
workflow.add_edge("Suggester", END)

client = MongoClient(MONGO_URI)
checkpointer = MongoDBSaver(
    client=client["furniture_agent_db"],
    db_name="langgraph",
    collection_name="checkpoints",
    serde=JsonPlusSerializer(pickle_fallback=True),  # ← 關鍵這一行
)
app = workflow.compile(checkpointer=checkpointer)