258 lines
8.5 KiB
Python
258 lines
8.5 KiB
Python
# Copyright 2026 zhaoxi826
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
|
||
import json
|
||
import asyncio
|
||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||
from fastapi.responses import StreamingResponse
|
||
from pydantic import BaseModel
|
||
from kilostar.utils.actor import get_actor
|
||
from kilostar.utils.access import Accessor, TokenData
|
||
from kilostar.core.individual.regulatory_node.template import (
|
||
MessageRequest,
|
||
MessageResponse,
|
||
)
|
||
|
||
chat_router = APIRouter(prefix="/api/v1/chat", tags=["chat"])
|
||
|
||
# 单次注入历史的最大轮数（user+assistant 算一轮），防止 token 爆炸。
|
||
_HISTORY_MAX_TURNS = 20
|
||
|
||
|
||
def _build_message_history(rows) -> list:
|
||
"""把 DB 中的 ChatHistoryMessage 列表转成 pydantic-ai message_history 格式。
|
||
|
||
历史按时间升序，截取末尾最多 _HISTORY_MAX_TURNS*2 条；user 消息映射为
|
||
``ModelRequest(parts=[UserPromptPart])``，assistant（``regulatory_node``）映射为
|
||
``ModelResponse(parts=[TextPart])``。其它 owner 跳过。
|
||
"""
|
||
from pydantic_ai.messages import (
|
||
ModelRequest, ModelResponse, UserPromptPart, TextPart,
|
||
)
|
||
|
||
trimmed = rows[-(_HISTORY_MAX_TURNS * 2):]
|
||
history: list = []
|
||
for row in trimmed:
|
||
owner = row.message_owner
|
||
text = row.message
|
||
if not text:
|
||
continue
|
||
if owner == "user":
|
||
history.append(ModelRequest(parts=[UserPromptPart(content=text)]))
|
||
elif owner == "regulatory_node":
|
||
history.append(ModelResponse(parts=[TextPart(content=text)]))
|
||
return history
|
||
|
||
|
||
async def _load_message_history(chat_id: str) -> list:
|
||
postgres_database = get_actor("postgres_database")
|
||
rows = await postgres_database.list_chat_messages(chat_id=chat_id)
|
||
return _build_message_history(rows or [])
|
||
|
||
|
||
def _extract_reply(resp: MessageResponse | None) -> str | None:
|
||
"""从 RegulatoryNode.working 的输出里取出对用户的回复文本。
|
||
|
||
RegulatoryNode 现在的 output_type 只剩 ``MessageResponse``（聊天/简单任务/汇报），
|
||
没有则视为节点降级为静默——上层不写回 chat history。
|
||
"""
|
||
if resp is None:
|
||
return None
|
||
return resp.reply_message
|
||
|
||
|
||
async def _ask_regulatory(
|
||
*, user_id: str, chat_id: str, message: str, message_history: list | None = None
|
||
) -> str | None:
|
||
"""统一封装 chat 入口对 RegulatoryNode 的调用。"""
|
||
regulatory_node = get_actor("regulatory_node")
|
||
payload = MessageRequest(
|
||
platform="client",
|
||
user_name=user_id,
|
||
platform_id=chat_id,
|
||
message=message,
|
||
)
|
||
resp: MessageResponse | None = await regulatory_node.working(
|
||
payload, message_history
|
||
)
|
||
return _extract_reply(resp)
|
||
|
||
|
||
class CreateChatRequest(BaseModel):
|
||
title: str = "新对话"
|
||
initial_message: str
|
||
|
||
|
||
class SendMessageRequest(BaseModel):
|
||
message: str
|
||
|
||
|
||
@chat_router.post("")
|
||
async def create_chat_session(
|
||
request: CreateChatRequest,
|
||
token_data: TokenData = Depends(Accessor.get_current_user),
|
||
):
|
||
postgres_database = get_actor("postgres_database")
|
||
chat = await postgres_database.create_chat_session(
|
||
user_id=token_data.user_id, title=request.title
|
||
)
|
||
|
||
# 存入用户消息
|
||
await postgres_database.add_chat_message(
|
||
chat_id=chat.chat_id, message=request.initial_message, message_owner="user"
|
||
)
|
||
|
||
# 调用监管节点处理简单任务/交流
|
||
response_msg = await _ask_regulatory(
|
||
user_id=token_data.user_id,
|
||
chat_id=chat.chat_id,
|
||
message=request.initial_message,
|
||
)
|
||
|
||
# 存入回复消息
|
||
if response_msg:
|
||
await postgres_database.add_chat_message(
|
||
chat_id=chat.chat_id, message=response_msg, message_owner="regulatory_node"
|
||
)
|
||
|
||
return {"chat_id": chat.chat_id, "reply": response_msg}
|
||
|
||
|
||
@chat_router.get("")
|
||
async def list_chat_sessions(
|
||
token_data: TokenData = Depends(Accessor.get_current_user),
|
||
):
|
||
postgres_database = get_actor("postgres_database")
|
||
sessions = await postgres_database.list_chat_sessions(
|
||
user_id=token_data.user_id
|
||
)
|
||
return {"sessions": sessions}
|
||
|
||
|
||
@chat_router.get("/{chat_id}")
|
||
async def get_chat_history(
|
||
chat_id: str, token_data: TokenData = Depends(Accessor.get_current_user)
|
||
):
|
||
postgres_database = get_actor("postgres_database")
|
||
messages = await postgres_database.list_chat_messages(chat_id=chat_id)
|
||
return {"messages": messages}
|
||
|
||
|
||
@chat_router.post("/{chat_id}/reply")
|
||
async def send_chat_message(
|
||
chat_id: str,
|
||
request: SendMessageRequest,
|
||
token_data: TokenData = Depends(Accessor.get_current_user),
|
||
):
|
||
postgres_database = get_actor("postgres_database")
|
||
# 先取历史（不含当前输入），再写入用户消息，避免历史里出现重复
|
||
message_history = await _load_message_history(chat_id)
|
||
await postgres_database.add_chat_message(
|
||
chat_id=chat_id, message=request.message, message_owner="user"
|
||
)
|
||
|
||
# 调用监管节点
|
||
response_msg = await _ask_regulatory(
|
||
user_id=token_data.user_id,
|
||
chat_id=chat_id,
|
||
message=request.message,
|
||
message_history=message_history,
|
||
)
|
||
|
||
# 存回复
|
||
if response_msg:
|
||
await postgres_database.add_chat_message(
|
||
chat_id=chat_id, message=response_msg, message_owner="regulatory_node"
|
||
)
|
||
|
||
return {"reply": response_msg}
|
||
|
||
|
||
@chat_router.delete("/{chat_id}")
|
||
async def delete_chat_session(
|
||
chat_id: str,
|
||
token_data: TokenData = Depends(Accessor.get_current_user),
|
||
):
|
||
postgres_database = get_actor("postgres_database")
|
||
session = await postgres_database.get_chat_session(chat_id=chat_id)
|
||
if not session:
|
||
raise HTTPException(status_code=404, detail="Chat session not found")
|
||
if session.user_id != token_data.user_id:
|
||
raise HTTPException(status_code=403, detail="Forbidden")
|
||
await postgres_database.delete_chat_session(chat_id=chat_id)
|
||
return {"message": "success"}
|
||
|
||
|
||
@chat_router.post("/{chat_id}/stream")
|
||
async def stream_chat_message(
|
||
chat_id: str,
|
||
request_body: SendMessageRequest,
|
||
request: Request,
|
||
token_data: TokenData = Depends(Accessor.get_current_user),
|
||
):
|
||
"""SSE 流式聊天端点：逐 token 流式输出。"""
|
||
postgres_database = get_actor("postgres_database")
|
||
|
||
message_history = await _load_message_history(chat_id)
|
||
|
||
await postgres_database.add_chat_message(
|
||
chat_id=chat_id, message=request_body.message, message_owner="user"
|
||
)
|
||
|
||
payload = MessageRequest(
|
||
platform="client",
|
||
user_name=token_data.user_id,
|
||
platform_id=chat_id,
|
||
message=request_body.message,
|
||
)
|
||
|
||
regulatory_node = get_actor("regulatory_node")
|
||
|
||
token_queue = asyncio.Queue()
|
||
stream_task = asyncio.create_task(
|
||
regulatory_node.stream_working(payload, token_queue, message_history)
|
||
)
|
||
|
||
async def event_generator():
|
||
full_response = ""
|
||
try:
|
||
while True:
|
||
if await request.is_disconnected():
|
||
stream_task.cancel()
|
||
break
|
||
try:
|
||
token = await asyncio.wait_for(token_queue.get(), timeout=0.5)
|
||
except asyncio.TimeoutError:
|
||
continue
|
||
if token is None:
|
||
break
|
||
full_response += token
|
||
yield f"data: {json.dumps({'token': token})}\n\n"
|
||
except Exception as e:
|
||
from kilostar.utils.logger import get_logger
|
||
get_logger("chat_stream").exception(f"Stream error: {e}")
|
||
if not full_response:
|
||
full_response = "抱歉，生成回复时出错。"
|
||
yield f"data: {json.dumps({'token': full_response})}\n\n"
|
||
|
||
if full_response:
|
||
await postgres_database.add_chat_message(
|
||
chat_id=chat_id,
|
||
message=full_response,
|
||
message_owner="regulatory_node",
|
||
)
|
||
yield f"data: {json.dumps({'done': True, 'full_message': full_response})}\n\n"
|
||
|
||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|