gongzhiyongandClaude Opus 4.6
3759ecd1a6
- Add web_search tool with Jina Search/Reader/Rerank - Flash mode: top 3, 8s timeout, no Reader/Rerank - Pro mode: top 10, 20s timeout, concurrent Reader + Rerank top 5 - Add Redis async cache (TTL=300s) for search results - Register "search" in ALL_TOOLS mapping Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
141 lines
4.6 KiB
Python
141 lines
4.6 KiB
Python
"""SSE streaming chat endpoint."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
|
|
from langchain_core.messages import HumanMessage
|
|
from litestar import post
|
|
from litestar.response import Stream
|
|
|
|
from app.graph.builder import get_chat_graph
|
|
from app.schemas import ChatRequest
|
|
from app.store.postgres import Conversation, Message, async_session_factory
|
|
from app.tools import resolve_tools
|
|
from app.tools.search import set_search_model
|
|
|
|
|
|
async def _ensure_conversation(conversation_id: str, first_message: str) -> None:
|
|
"""Create conversation and persist the user message."""
|
|
async with async_session_factory() as session:
|
|
existing = await session.get(Conversation, conversation_id)
|
|
if existing is None:
|
|
# Use first ~50 chars of message as title
|
|
title = first_message[:50].strip() or "New conversation"
|
|
conv = Conversation(id=conversation_id, title=title)
|
|
session.add(conv)
|
|
# Persist user message
|
|
msg = Message(
|
|
conversation_id=conversation_id,
|
|
role="human",
|
|
content=first_message,
|
|
)
|
|
session.add(msg)
|
|
await session.commit()
|
|
|
|
|
|
async def _persist_ai_message(conversation_id: str, content: str) -> None:
|
|
"""Persist the AI response message."""
|
|
async with async_session_factory() as session:
|
|
msg = Message(
|
|
conversation_id=conversation_id,
|
|
role="ai",
|
|
content=content,
|
|
)
|
|
session.add(msg)
|
|
await session.commit()
|
|
|
|
|
|
async def _stream_response(request: ChatRequest) -> AsyncIterator[bytes]:
|
|
"""Stream LLM response tokens via SSE."""
|
|
# Ensure conversation exists and persist user message
|
|
await _ensure_conversation(request.conversation_id, request.message)
|
|
|
|
# Resolve tools from frontend tool keys
|
|
active_tools = resolve_tools(request.tools)
|
|
|
|
# Set search strategy model so the web_search tool knows depth
|
|
set_search_model(request.model)
|
|
|
|
graph = await get_chat_graph(model=request.model, tools=active_tools)
|
|
|
|
config = {
|
|
"configurable": {"thread_id": request.conversation_id},
|
|
}
|
|
|
|
# When using ReAct agent (with tools), input is just messages.
|
|
# When using plain graph (no tools), input includes model key.
|
|
if active_tools:
|
|
input_data = {"messages": [HumanMessage(content=request.message)]}
|
|
else:
|
|
input_data = {
|
|
"messages": [HumanMessage(content=request.message)],
|
|
"model": request.model,
|
|
}
|
|
|
|
full_content: list[str] = []
|
|
|
|
async for event in graph.astream_events(
|
|
input_data,
|
|
config=config,
|
|
version="v2",
|
|
):
|
|
kind = event.get("event", "")
|
|
|
|
if kind == "on_chat_model_stream":
|
|
chunk = event.get("data", {}).get("chunk")
|
|
if chunk and hasattr(chunk, "content") and chunk.content:
|
|
# Only stream text content, skip tool call chunks
|
|
if isinstance(chunk.content, str):
|
|
full_content.append(chunk.content)
|
|
sse_data = json.dumps(
|
|
{"type": "token", "content": chunk.content},
|
|
ensure_ascii=False,
|
|
)
|
|
yield f"data: {sse_data}\n\n".encode("utf-8")
|
|
|
|
elif kind == "on_tool_start":
|
|
# Notify frontend that a tool is being called
|
|
tool_name = event.get("name", "unknown")
|
|
sse_data = json.dumps(
|
|
{"type": "tool_start", "tool": tool_name},
|
|
ensure_ascii=False,
|
|
)
|
|
yield f"data: {sse_data}\n\n".encode("utf-8")
|
|
|
|
elif kind == "on_tool_end":
|
|
tool_name = event.get("name", "unknown")
|
|
sse_data = json.dumps(
|
|
{"type": "tool_end", "tool": tool_name},
|
|
ensure_ascii=False,
|
|
)
|
|
yield f"data: {sse_data}\n\n".encode("utf-8")
|
|
|
|
# Persist AI response
|
|
ai_content = "".join(full_content)
|
|
if ai_content:
|
|
await _persist_ai_message(request.conversation_id, ai_content)
|
|
|
|
# Send done signal
|
|
done_data = json.dumps({"type": "done"})
|
|
yield f"data: {done_data}\n\n".encode("utf-8")
|
|
|
|
|
|
@post("/api/chat/stream")
|
|
async def stream_chat(data: ChatRequest) -> Stream:
|
|
"""POST /api/chat/stream - SSE streaming chat endpoint."""
|
|
if not data.conversation_id:
|
|
data.conversation_id = str(uuid.uuid4())
|
|
|
|
return Stream(
|
|
_stream_response(data),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
)
|