338 lines
13 KiB
Python
338 lines
13 KiB
Python
"""POST /api/chat — SSE streaming chat with Agent."""
|
||
|
||
import json
|
||
import os
|
||
import traceback
|
||
from pathlib import Path
|
||
from typing import AsyncGenerator
|
||
|
||
from fastapi import APIRouter
|
||
from pydantic import BaseModel
|
||
from sse_starlette.sse import EventSourceResponse
|
||
|
||
from graph.agent import agent_manager
|
||
from graph.session_manager import session_manager
|
||
|
||
router = APIRouter()
|
||
|
||
|
||
class ChatRequest(BaseModel):
|
||
message: str
|
||
session_id: str = "default"
|
||
stream: bool = True
|
||
|
||
|
||
async def _generate_title(session_id: str) -> str | None:
|
||
"""Generate a title for a session using DeepSeek. Returns title or None."""
|
||
try:
|
||
messages = session_manager.load_session(session_id)
|
||
first_user = ""
|
||
first_assistant = ""
|
||
for msg in messages:
|
||
if msg["role"] == "user" and not first_user:
|
||
first_user = msg["content"][:200]
|
||
elif msg["role"] == "assistant" and not first_assistant:
|
||
first_assistant = msg["content"][:200]
|
||
if first_user and first_assistant:
|
||
break
|
||
|
||
if not first_user:
|
||
return None
|
||
|
||
from langchain_deepseek import ChatDeepSeek
|
||
from langchain_core.messages import HumanMessage as HM
|
||
|
||
llm = ChatDeepSeek(
|
||
model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
|
||
api_key=os.getenv("DEEPSEEK_API_KEY"),
|
||
api_base=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
|
||
temperature=0.3,
|
||
)
|
||
|
||
prompt = (
|
||
f"根据以下对话内容，生成一个不超过10个字的中文标题，只输出标题文本，不要加引号或标点。\n\n"
|
||
f"用户: {first_user}\n"
|
||
f"助手: {first_assistant}"
|
||
)
|
||
|
||
result = await llm.ainvoke([HM(content=prompt)])
|
||
title = result.content.strip().strip('"\'""''')[:20]
|
||
session_manager.update_title(session_id, title)
|
||
return title
|
||
except Exception:
|
||
traceback.print_exc()
|
||
return None
|
||
|
||
|
||
async def event_generator(message: str, session_id: str) -> AsyncGenerator[dict, None]:
|
||
"""Generate SSE events from agent stream.
|
||
|
||
Tracks multiple response segments — each time the agent finishes
|
||
tool calls and starts generating new text, a new_response event is
|
||
emitted and a new segment begins. Each segment is saved as a
|
||
separate assistant message in the session history.
|
||
"""
|
||
try:
|
||
# Use merged history for agent context (combines consecutive assistant msgs)
|
||
history = session_manager.load_session_for_agent(session_id)
|
||
is_first_message = len(history) == 0
|
||
|
||
# Per-segment tracking
|
||
segments: list[dict] = [] # finalized segments
|
||
current_segment: dict = {"content": "", "tool_calls": []}
|
||
|
||
async for event in agent_manager.astream(message, history, session_id):
|
||
event_type = event.get("type", "unknown")
|
||
|
||
if event_type == "retrieval":
|
||
yield {
|
||
"event": "retrieval",
|
||
"data": json.dumps(
|
||
{"query": event["query"], "results": event["results"]},
|
||
ensure_ascii=False,
|
||
),
|
||
}
|
||
|
||
elif event_type == "token":
|
||
current_segment["content"] += event["content"]
|
||
yield {
|
||
"event": "token",
|
||
"data": json.dumps({"content": event["content"]}, ensure_ascii=False),
|
||
}
|
||
|
||
elif event_type == "new_response":
|
||
# Finalize current segment and start a new one
|
||
segments.append(current_segment)
|
||
current_segment = {"content": "", "tool_calls": []}
|
||
yield {
|
||
"event": "new_response",
|
||
"data": json.dumps({}, ensure_ascii=False),
|
||
}
|
||
|
||
elif event_type == "tool_start":
|
||
current_segment["tool_calls"].append({
|
||
"tool": event["tool"],
|
||
"input": event.get("input", ""),
|
||
})
|
||
yield {
|
||
"event": "tool_start",
|
||
"data": json.dumps(
|
||
{"tool": event["tool"], "input": event["input"]},
|
||
ensure_ascii=False,
|
||
),
|
||
}
|
||
|
||
elif event_type == "tool_end":
|
||
# Match with current segment's tool_calls to add output
|
||
for tc in reversed(current_segment["tool_calls"]):
|
||
if tc["tool"] == event["tool"] and "output" not in tc:
|
||
tc["output"] = event["output"]
|
||
break
|
||
yield {
|
||
"event": "tool_end",
|
||
"data": json.dumps(
|
||
{"tool": event["tool"], "output": event["output"]},
|
||
ensure_ascii=False,
|
||
),
|
||
}
|
||
|
||
elif event_type == "done":
|
||
# Finalize the last segment
|
||
segments.append(current_segment)
|
||
|
||
# Save user message once
|
||
session_manager.save_message(session_id, "user", message)
|
||
|
||
# Save each segment as a separate assistant message
|
||
for seg in segments:
|
||
tc = seg["tool_calls"] if seg["tool_calls"] else None
|
||
session_manager.save_message(
|
||
session_id, "assistant", seg["content"], tool_calls=tc
|
||
)
|
||
|
||
yield {
|
||
"event": "done",
|
||
"data": json.dumps(
|
||
{"content": event["content"], "session_id": session_id},
|
||
ensure_ascii=False,
|
||
),
|
||
}
|
||
|
||
# Auto-generate title on first conversation
|
||
if is_first_message:
|
||
title = await _generate_title(session_id)
|
||
if title:
|
||
yield {
|
||
"event": "title",
|
||
"data": json.dumps(
|
||
{"session_id": session_id, "title": title},
|
||
ensure_ascii=False,
|
||
),
|
||
}
|
||
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
yield {
|
||
"event": "error",
|
||
"data": json.dumps({"error": str(e)}, ensure_ascii=False),
|
||
}
|
||
|
||
|
||
@router.post("/chat")
|
||
async def chat(request: ChatRequest):
|
||
if request.stream:
|
||
return EventSourceResponse(
|
||
event_generator(request.message, request.session_id)
|
||
)
|
||
# Non-streaming fallback
|
||
result = await agent_manager.ainvoke(request.message, request.session_id)
|
||
return {"reply": result}
|
||
|
||
|
||
@router.get("/chat/raw-traces")
|
||
async def get_raw_traces(session_id: str = "default", limit: int = 20):
|
||
"""获取指定 session 的原始 LLM 调用报文
|
||
|
||
Args:
|
||
session_id: 会话 ID
|
||
limit: 返回最近 N 条记录 (默认 20)
|
||
|
||
Returns:
|
||
{"traces": [...]} 原始报文列表，按时间倒序
|
||
"""
|
||
try:
|
||
from graph.raw_trace import load_raw_traces
|
||
|
||
# 从 agent_manager 获取 raw_logs_dir
|
||
raw_logs_dir = agent_manager._raw_logs_dir
|
||
if not raw_logs_dir:
|
||
return {"traces": []}
|
||
|
||
traces = load_raw_traces(session_id, raw_logs_dir, limit=limit)
|
||
return {"traces": traces}
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
return {"traces": [], "error": str(e)}
|
||
|
||
|
||
@router.get("/chat/trajectory")
|
||
async def get_trajectory(session_id: str = "default"):
|
||
"""获取会话完整交互轨迹 - 包含完整的 human/assistant/tool 消息流
|
||
|
||
Args:
|
||
session_id: 会话 ID
|
||
|
||
Returns:
|
||
{
|
||
"session_id": "...",
|
||
"messages": [
|
||
{"role": "human", "content": "...", "timestamp": ...},
|
||
{"role": "assistant", "content": "...", "timestamp": ..., "tool_calls": [...]},
|
||
{"role": "tool", "name": "...", "content": "...", "timestamp": ...},
|
||
{"role": "assistant", "content": "...", "timestamp": ...}
|
||
],
|
||
"steps": [...] // 保留旧格式用于向后兼容
|
||
}
|
||
"""
|
||
try:
|
||
import json
|
||
from graph.raw_trace import load_raw_traces
|
||
|
||
raw_logs_dir = agent_manager._raw_logs_dir
|
||
if not raw_logs_dir:
|
||
# 无 raw_logs_dir，返回空
|
||
return {"session_id": session_id, "messages": [], "steps": []}
|
||
|
||
# 优先读取完整消息流日志文件
|
||
msgs_log_file = raw_logs_dir / f"{session_id}.msgs.jsonl"
|
||
messages = []
|
||
|
||
if msgs_log_file.exists():
|
||
# 从 msgs.jsonl 读取完整消息流（包含 tool 消息）
|
||
with open(msgs_log_file, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line:
|
||
try:
|
||
msg = json.loads(line)
|
||
messages.append(msg)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
else:
|
||
# 回退到旧逻辑：从 session 文件读取（无 tool 消息）
|
||
session_messages = session_manager.load_session(session_id)
|
||
for msg in session_messages:
|
||
messages.append({
|
||
"role": msg["role"],
|
||
"content": msg["content"],
|
||
"timestamp": None,
|
||
})
|
||
|
||
# 构建旧格式 steps（向后兼容）
|
||
steps = []
|
||
traces = load_raw_traces(session_id, raw_logs_dir, limit=1000)
|
||
traces.sort(key=lambda x: x.get("timestamp_start", 0))
|
||
|
||
# 按 user/assistant 消息构建 steps（忽略 tool 消息）
|
||
session_messages = session_manager.load_session(session_id)
|
||
assistant_count = sum(1 for m in session_messages if m["role"] == "assistant")
|
||
|
||
for msg_idx, msg in enumerate(session_messages):
|
||
step = {
|
||
"role": msg["role"],
|
||
"content": msg["content"],
|
||
"timestamp": None,
|
||
}
|
||
|
||
if msg["role"] == "assistant":
|
||
llm_calls = []
|
||
assistant_seq = sum(1 for m in session_messages[:msg_idx+1] if m["role"] == "assistant")
|
||
traces_per_assistant = len(traces) / max(assistant_count, 1)
|
||
start_trace = int((assistant_seq - 1) * traces_per_assistant)
|
||
end_trace = int(assistant_seq * traces_per_assistant)
|
||
|
||
for i in range(start_trace, min(end_trace, len(traces))):
|
||
trace = traces[i]
|
||
llm_call = {
|
||
"model": trace.get("request", {}).get("model", "unknown"),
|
||
"finish_reason": trace.get("finish_reason", "unknown"),
|
||
"request_messages": trace.get("request", {}).get("messages", []),
|
||
"response_text": "",
|
||
"tool_calls_list": [],
|
||
"usage": trace.get("usage", {}),
|
||
"duration_ms": trace.get("duration_ms", 0),
|
||
}
|
||
|
||
response_gens = trace.get("response", {}).get("generations", [])
|
||
if response_gens:
|
||
llm_call["response_text"] = response_gens[0].get("text", "")
|
||
msg_data = response_gens[0].get("message", {})
|
||
if isinstance(msg_data, dict) and "tool_calls" in msg_data:
|
||
tool_calls = msg_data["tool_calls"]
|
||
if tool_calls:
|
||
llm_call["tool_calls_list"] = [
|
||
{
|
||
"name": tc.get("name", "") if isinstance(tc, dict) else getattr(tc, "name", ""),
|
||
"args": tc.get("args", {}) if isinstance(tc, dict) else getattr(tc, "args", {}),
|
||
}
|
||
for tc in tool_calls
|
||
]
|
||
|
||
llm_calls.append(llm_call)
|
||
if step["timestamp"] is None:
|
||
step["timestamp"] = trace.get("timestamp_start")
|
||
|
||
step["llm_calls"] = llm_calls
|
||
|
||
steps.append(step)
|
||
|
||
return {
|
||
"session_id": session_id,
|
||
"messages": messages, # 新格式：完整消息流（含 tool）
|
||
"steps": steps, # 旧格式：向后兼容
|
||
}
|
||
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
return {"session_id": session_id, "messages": [], "steps": [], "error": str(e)}
|