921 lines
32 KiB
Python
921 lines
32 KiB
Python
"""
|
||
知识库对话相关 API
|
||
"""
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import re
|
||
from typing import Optional
|
||
|
||
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy import select, delete
|
||
|
||
from app.core.database import get_db
|
||
from app.core.deps import get_current_user
|
||
from app.models.llm_model_config import LLMModelConfig
|
||
from app.models.project import Project
|
||
from app.models.user import User
|
||
from app.models.chat_session import ChatSession, ChatMessage
|
||
from app.schemas.response import success_response
|
||
from app.services.project_service import (
|
||
get_project_or_404,
|
||
require_project_read_access,
|
||
require_project_write_access,
|
||
)
|
||
from app.services.rag_service import rag_service
|
||
from app.services.zvec_service import zvec_service
|
||
from app.services.project_vectorization_task_service import project_vectorization_task_service
|
||
from pydantic import BaseModel
|
||
|
||
|
||
router = APIRouter()
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
CHAT_HISTORY_MESSAGE_LIMIT = 20
|
||
# 生成被客户端中断时写入助手消息的统一标记，前端据此识别“已停止生成”
|
||
INTERRUPTED_RESPONSE_MARKER = "interrupt"
|
||
|
||
|
||
def _is_interrupted_content(content: str) -> bool:
|
||
"""判断助手消息内容是否带中断标记。
|
||
|
||
中断标记要么单独成文（内容为空时），要么以 \\n\\n 追加在内容末尾。
|
||
"""
|
||
if not content:
|
||
return False
|
||
return content == INTERRUPTED_RESPONSE_MARKER or content.endswith(
|
||
f"\n\n{INTERRUPTED_RESPONSE_MARKER}"
|
||
)
|
||
|
||
|
||
class VectorizeRequest(BaseModel):
|
||
"""批量向量化请求"""
|
||
force: bool = False # False=增量（跳过未变更文件）, True=全量重建
|
||
|
||
|
||
@router.get("/projects/{project_id}/vectorize/progress", response_model=dict)
|
||
async def get_vectorize_progress(
|
||
project_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""查询项目向量化进度"""
|
||
project = await get_project_or_404(db, project_id)
|
||
await require_project_read_access(db, project_id, current_user)
|
||
|
||
progress = await zvec_service.get_progress(db, project_id, project.storage_key)
|
||
return success_response(data=progress)
|
||
|
||
|
||
@router.post("/projects/{project_id}/vectorize", response_model=dict)
|
||
async def vectorize_project(
|
||
project_id: int,
|
||
req: VectorizeRequest,
|
||
background_tasks: BackgroundTasks,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""创建项目文件夹向量化后台任务（增量或全量）"""
|
||
await get_project_or_404(db, project_id)
|
||
await require_project_write_access(db, project_id, current_user)
|
||
|
||
if await zvec_service.get_embedding_config(db) is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="尚未配置可用的向量模型，请先在「模型配置 - 向量模型」中添加并启用。",
|
||
)
|
||
|
||
running_task = await project_vectorization_task_service.get_running_task(db, project_id)
|
||
if running_task:
|
||
task = running_task
|
||
else:
|
||
task = await project_vectorization_task_service.create_task(
|
||
db,
|
||
project_id,
|
||
current_user.id,
|
||
force=req.force,
|
||
)
|
||
background_tasks.add_task(project_vectorization_task_service.run_task, task.task_id)
|
||
|
||
return success_response(
|
||
data=project_vectorization_task_service.serialize_task(task),
|
||
message="向量化任务已提交",
|
||
)
|
||
|
||
|
||
@router.get("/projects/{project_id}/vectorize/tasks/latest", response_model=dict)
|
||
async def get_latest_vectorize_task(
|
||
project_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""查询项目最近一次向量化任务"""
|
||
await require_project_read_access(db, project_id, current_user)
|
||
task = await project_vectorization_task_service.get_latest_task(db, project_id)
|
||
return success_response(
|
||
data=project_vectorization_task_service.serialize_task(task) if task else None
|
||
)
|
||
|
||
|
||
@router.get("/projects/{project_id}/vectorize/tasks/{task_id}", response_model=dict)
|
||
async def get_vectorize_task(
|
||
project_id: int,
|
||
task_id: str,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""查询项目向量化任务状态"""
|
||
await require_project_read_access(db, project_id, current_user)
|
||
task = await project_vectorization_task_service.get_task(db, project_id, task_id)
|
||
if not task:
|
||
raise HTTPException(status_code=404, detail="向量化任务不存在")
|
||
return success_response(data=project_vectorization_task_service.serialize_task(task))
|
||
|
||
|
||
class ChatCreateRequest(BaseModel):
|
||
"""创建对话会话请求"""
|
||
project_id: int
|
||
llm_config_id: int
|
||
title: str = "新对话"
|
||
|
||
|
||
class ChatMessageRequest(BaseModel):
|
||
"""发送聊天消息请求"""
|
||
session_id: int
|
||
message: str
|
||
# 是否同时写入一条新的用户消息。重新生成旧回复时应传 false，避免用户提问重复入库。
|
||
insert_user_message: bool = True
|
||
|
||
|
||
class ChatSessionUpdateRequest(BaseModel):
|
||
"""更新对话会话请求"""
|
||
title: str
|
||
|
||
|
||
class ChatResponse(BaseModel):
|
||
"""对话响应"""
|
||
session_id: int
|
||
user_message: str
|
||
assistant_message: str
|
||
|
||
|
||
def _parse_refs(value):
|
||
if not value:
|
||
return []
|
||
try:
|
||
return json.loads(value)
|
||
except (ValueError, TypeError):
|
||
return []
|
||
|
||
|
||
def _normalize_refs(refs):
|
||
"""将存储的引用归一化为 [{citation_id, file_path, anchor_text, excerpt}] 形式。
|
||
|
||
兼容旧格式（纯文件路径数组）与新格式（对象数组）。
|
||
"""
|
||
normalized = []
|
||
for idx, ref in enumerate(refs or [], 1):
|
||
if isinstance(ref, dict):
|
||
file_path = ref.get("file_path")
|
||
citation_id = ref.get("citation_id", idx)
|
||
anchor_text = ref.get("anchor_text") or ""
|
||
excerpt = ref.get("excerpt") or ""
|
||
else:
|
||
file_path = ref
|
||
citation_id = idx
|
||
anchor_text = ""
|
||
excerpt = ""
|
||
if file_path:
|
||
normalized.append({
|
||
"citation_id": citation_id,
|
||
"file_path": file_path,
|
||
"anchor_text": anchor_text,
|
||
"excerpt": excerpt,
|
||
})
|
||
return normalized
|
||
|
||
|
||
def _build_reference_items(refs, project_id=None):
|
||
return [
|
||
{
|
||
"citation_id": ref["citation_id"],
|
||
"file_path": ref["file_path"],
|
||
"file_name": ref["file_path"].rsplit("/", 1)[-1],
|
||
"anchor_text": ref.get("anchor_text") or "",
|
||
"excerpt": ref.get("excerpt") or "",
|
||
"project_id": project_id,
|
||
}
|
||
for ref in _normalize_refs(refs)
|
||
]
|
||
|
||
|
||
def _canonicalize_message_citations(content: str, refs):
|
||
"""按文件合并引用并将正文引用重新编号，兼容历史重复分块数据。"""
|
||
normalized = _normalize_refs(refs)
|
||
if not normalized:
|
||
return content, []
|
||
|
||
canonical_refs = []
|
||
file_to_id = {}
|
||
old_to_new = {}
|
||
for ref in normalized:
|
||
file_path = ref["file_path"]
|
||
old_id = int(ref["citation_id"])
|
||
new_id = file_to_id.get(file_path)
|
||
if new_id is None:
|
||
new_id = len(canonical_refs) + 1
|
||
file_to_id[file_path] = new_id
|
||
canonical_refs.append({**ref, "citation_id": new_id})
|
||
else:
|
||
canonical_ref = canonical_refs[new_id - 1]
|
||
current_excerpt = canonical_ref.get("excerpt", "").strip()
|
||
incoming_excerpt = ref.get("excerpt", "").strip()
|
||
if incoming_excerpt and incoming_excerpt not in current_excerpt:
|
||
canonical_ref["excerpt"] = (
|
||
f"{current_excerpt}\n\n{incoming_excerpt}"
|
||
if current_excerpt else incoming_excerpt
|
||
)
|
||
old_to_new[old_id] = new_id
|
||
|
||
def replace_citation(match):
|
||
old_id = int(match.group(1))
|
||
return f"[{old_to_new.get(old_id, old_id)}]"
|
||
|
||
normalized_content = _CITATION_PATTERN.sub(replace_citation, content or "")
|
||
normalized_content = re.sub(r"\[(\d+)\](?:\s*\[\1\])+", r"[\1]", normalized_content)
|
||
return normalized_content, canonical_refs
|
||
|
||
|
||
def _stream_event(event: str, data) -> str:
|
||
return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||
|
||
|
||
_CITATION_PATTERN = re.compile(r"\[(\d+)\]")
|
||
|
||
|
||
def _compact_cited_refs(answer: str, retrieved_docs):
|
||
"""仅保留实际使用的文档，并按正文首次引用顺序连续编号。"""
|
||
docs_by_id = {
|
||
int(doc.get("citation_id", index)): doc
|
||
for index, doc in enumerate(retrieved_docs, 1)
|
||
}
|
||
old_to_new = {}
|
||
file_to_new = {}
|
||
refs = []
|
||
|
||
for match in _CITATION_PATTERN.finditer(answer or ""):
|
||
old_id = int(match.group(1))
|
||
doc = docs_by_id.get(old_id)
|
||
if doc is None or old_id in old_to_new:
|
||
continue
|
||
file_path = doc.get("file_path")
|
||
if not file_path:
|
||
continue
|
||
|
||
new_id = file_to_new.get(file_path)
|
||
if new_id is None:
|
||
new_id = len(refs) + 1
|
||
file_to_new[file_path] = new_id
|
||
refs.append({
|
||
"citation_id": new_id,
|
||
"file_path": file_path,
|
||
"anchor_text": doc.get("anchor_text") or "",
|
||
# 优先使用命中文档的上下文片段（命中点前后扩展），避免分块边界
|
||
# 的短文本（如文档末尾残余块）导致引用预览内容过少
|
||
"excerpt": doc.get("content") or doc.get("excerpt") or doc.get("anchor_text") or "",
|
||
})
|
||
old_to_new[old_id] = new_id
|
||
|
||
def replace_citation(match):
|
||
old_id = int(match.group(1))
|
||
return f"[{old_to_new.get(old_id, old_id)}]"
|
||
|
||
normalized_answer = _CITATION_PATTERN.sub(replace_citation, answer or "")
|
||
normalized_answer = re.sub(
|
||
r"\[(\d+)\](?:\s*\[\1\])+", r"[\1]", normalized_answer
|
||
)
|
||
return normalized_answer, refs
|
||
|
||
|
||
MAX_SESSION_TITLE_LENGTH = 60
|
||
|
||
|
||
def _normalize_session_title(title: Optional[str]) -> str:
|
||
"""会话标题统一处理：取首行、去掉多余引号与空白，超长截断，空值回退。"""
|
||
title = (title or "").strip().strip('"').strip("「」").strip()
|
||
title = title.splitlines()[0].strip() if title else ""
|
||
if not title:
|
||
return "新对话"
|
||
return title[:MAX_SESSION_TITLE_LENGTH]
|
||
|
||
|
||
@router.post("/sessions", response_model=dict)
|
||
async def create_chat_session(
|
||
req: ChatCreateRequest,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""创建新的对话会话"""
|
||
await get_project_or_404(db, req.project_id)
|
||
await require_project_read_access(db, req.project_id, current_user)
|
||
model_result = await db.execute(
|
||
select(LLMModelConfig).where(
|
||
LLMModelConfig.config_id == req.llm_config_id,
|
||
LLMModelConfig.model_type == "chat",
|
||
LLMModelConfig.is_active == True,
|
||
)
|
||
)
|
||
if model_result.scalar_one_or_none() is None:
|
||
raise HTTPException(status_code=400, detail="请选择有效的对话模型")
|
||
|
||
session = ChatSession(
|
||
user_id=current_user.id,
|
||
project_id=req.project_id,
|
||
llm_config_id=req.llm_config_id,
|
||
title=_normalize_session_title(req.title),
|
||
)
|
||
db.add(session)
|
||
await db.commit()
|
||
await db.refresh(session)
|
||
|
||
return success_response(data={
|
||
"session_id": session.id,
|
||
"project_id": session.project_id,
|
||
"llm_config_id": session.llm_config_id,
|
||
"title": session.title,
|
||
"created_at": session.created_at.isoformat(),
|
||
"updated_at": session.updated_at.isoformat(),
|
||
})
|
||
|
||
|
||
@router.get("/sessions", response_model=dict)
|
||
async def list_chat_sessions(
|
||
project_id: Optional[int] = None,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""列出当前用户的对话会话，可按知识库筛选"""
|
||
stmt = select(ChatSession).where(ChatSession.user_id == current_user.id)
|
||
if project_id is not None:
|
||
await require_project_read_access(db, project_id, current_user)
|
||
stmt = stmt.where(ChatSession.project_id == project_id)
|
||
stmt = stmt.order_by(ChatSession.updated_at.desc(), ChatSession.created_at.desc())
|
||
result = await db.execute(stmt)
|
||
sessions = result.scalars().all()
|
||
|
||
return success_response(data=[{
|
||
"session_id": s.id,
|
||
"project_id": s.project_id,
|
||
"llm_config_id": s.llm_config_id,
|
||
"title": s.title,
|
||
"created_at": s.created_at.isoformat(),
|
||
"updated_at": s.updated_at.isoformat(),
|
||
} for s in sessions])
|
||
|
||
|
||
@router.get("/sessions/{session_id}/messages", response_model=dict)
|
||
async def get_session_messages(
|
||
session_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""获取对话会话的所有消息"""
|
||
stmt = select(ChatSession).where(ChatSession.id == session_id)
|
||
result = await db.execute(stmt)
|
||
session = result.scalar_one_or_none()
|
||
|
||
if not session:
|
||
raise HTTPException(status_code=404, detail="对话会话不存在")
|
||
|
||
if session.user_id != current_user.id:
|
||
raise HTTPException(status_code=403, detail="无权访问该对话")
|
||
|
||
await require_project_read_access(db, session.project_id, current_user)
|
||
|
||
msg_stmt = select(ChatMessage).where(
|
||
ChatMessage.session_id == session_id
|
||
).order_by(ChatMessage.created_at.asc())
|
||
msg_result = await db.execute(msg_stmt)
|
||
messages = msg_result.scalars().all()
|
||
|
||
data = []
|
||
for message in messages:
|
||
stored_refs = _parse_refs(message.referenced_files)
|
||
content = message.content
|
||
refs = stored_refs
|
||
if message.role == "assistant":
|
||
content, refs = _canonicalize_message_citations(content, stored_refs)
|
||
data.append({
|
||
"id": message.id,
|
||
"role": message.role,
|
||
"content": content,
|
||
"referenced_files": refs,
|
||
"references": _build_reference_items(refs, session.project_id) if message.role == "assistant" else [],
|
||
"created_at": message.created_at.isoformat(),
|
||
})
|
||
|
||
return success_response(data=data)
|
||
|
||
|
||
@router.get("/search", response_model=dict)
|
||
async def search_chat_messages(
|
||
keyword: str,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""搜索聊天消息内容"""
|
||
keyword = (keyword or "").strip()
|
||
if not keyword:
|
||
return success_response(data=[])
|
||
|
||
escaped_keyword = (
|
||
keyword.replace("\\", "\\\\")
|
||
.replace("%", "\\%")
|
||
.replace("_", "\\_")
|
||
)
|
||
stmt = (
|
||
select(
|
||
ChatMessage.id,
|
||
ChatMessage.session_id,
|
||
ChatMessage.role,
|
||
ChatMessage.content,
|
||
ChatMessage.created_at,
|
||
ChatSession.title,
|
||
ChatSession.project_id,
|
||
ChatSession.llm_config_id,
|
||
Project.name.label("project_name"),
|
||
LLMModelConfig.model_name.label("model_name"),
|
||
)
|
||
.join(ChatSession, ChatMessage.session_id == ChatSession.id)
|
||
.join(Project, Project.id == ChatSession.project_id)
|
||
.join(LLMModelConfig, LLMModelConfig.config_id == ChatSession.llm_config_id)
|
||
.where(
|
||
ChatSession.user_id == current_user.id,
|
||
ChatMessage.content.ilike(f"%{escaped_keyword}%", escape="\\"),
|
||
)
|
||
.order_by(ChatMessage.created_at.desc())
|
||
.limit(50)
|
||
)
|
||
result = await db.execute(stmt)
|
||
rows = result.all()
|
||
|
||
def _build_snippet(content: str, term: str) -> str:
|
||
text = content or ""
|
||
if not text:
|
||
return ""
|
||
lowered = text.lower()
|
||
needle = term.lower()
|
||
idx = lowered.find(needle)
|
||
if idx < 0:
|
||
return text[:120]
|
||
start = max(0, idx - 24)
|
||
end = min(len(text), idx + len(term) + 48)
|
||
prefix = "..." if start > 0 else ""
|
||
suffix = "..." if end < len(text) else ""
|
||
return f"{prefix}{text[start:end].strip()}{suffix}"
|
||
|
||
return success_response(data=[
|
||
{
|
||
"result_id": row.id,
|
||
"session_id": row.session_id,
|
||
"message_id": row.id,
|
||
"role": row.role,
|
||
"content": row.content,
|
||
"snippet": _build_snippet(row.content, keyword),
|
||
"session_title": row.title,
|
||
"project_id": row.project_id,
|
||
"project_name": row.project_name,
|
||
"llm_config_id": row.llm_config_id,
|
||
"model_name": row.model_name,
|
||
"created_at": row.created_at.isoformat() if row.created_at else None,
|
||
}
|
||
for row in rows
|
||
])
|
||
|
||
|
||
@router.post("/send", response_model=dict)
|
||
async def send_chat_message(
|
||
req: ChatMessageRequest,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""发送对话消息并获取回复"""
|
||
question = (req.message or "").strip()
|
||
if not question:
|
||
raise HTTPException(status_code=400, detail="消息内容不能为空")
|
||
|
||
stmt = select(ChatSession).where(ChatSession.id == req.session_id)
|
||
result = await db.execute(stmt)
|
||
session = result.scalar_one_or_none()
|
||
|
||
if not session:
|
||
raise HTTPException(status_code=404, detail="对话会话不存在")
|
||
|
||
if session.user_id != current_user.id:
|
||
raise HTTPException(status_code=403, detail="无权访问该对话")
|
||
|
||
await require_project_read_access(db, session.project_id, current_user)
|
||
|
||
query_message = ChatMessage(
|
||
session_id=req.session_id,
|
||
role="user",
|
||
content=question,
|
||
)
|
||
db.add(query_message)
|
||
await db.flush()
|
||
|
||
try:
|
||
msg_stmt = (
|
||
select(ChatMessage)
|
||
.where(
|
||
ChatMessage.session_id == req.session_id,
|
||
ChatMessage.id != query_message.id,
|
||
)
|
||
.order_by(ChatMessage.created_at.desc())
|
||
.limit(CHAT_HISTORY_MESSAGE_LIMIT)
|
||
)
|
||
msg_result = await db.execute(msg_stmt)
|
||
prev_messages = list(reversed(msg_result.scalars().all()))
|
||
|
||
conversation_history = [
|
||
{"role": m.role, "content": m.content}
|
||
for m in prev_messages
|
||
# 被中断的助手消息（含统一中断标记）不参与上下文，避免污染后续模型输入
|
||
if not (
|
||
m.role == "assistant"
|
||
and _is_interrupted_content(m.content or "")
|
||
)
|
||
]
|
||
|
||
retrieved_docs = await rag_service.retrieve_documents(
|
||
db,
|
||
session.project_id,
|
||
question,
|
||
top_k=5,
|
||
)
|
||
|
||
assistant_response = await rag_service.generate_response(
|
||
db,
|
||
question,
|
||
session.project_id,
|
||
session.llm_config_id,
|
||
retrieved_docs,
|
||
conversation_history,
|
||
)
|
||
|
||
# 仅保留实际引用的文档，并把引用编号压缩为从 1 开始的连续序列。
|
||
assistant_response, cited_refs = _compact_cited_refs(
|
||
assistant_response, retrieved_docs
|
||
)
|
||
|
||
assistant_message = ChatMessage(
|
||
session_id=req.session_id,
|
||
role="assistant",
|
||
content=assistant_response,
|
||
referenced_files=json.dumps(cited_refs, ensure_ascii=False) if cited_refs else None,
|
||
)
|
||
db.add(assistant_message)
|
||
|
||
# 累加会话消息计数（用户 + 助手 共 2 条）
|
||
session.message_count = (session.message_count or 0) + 2
|
||
|
||
await db.commit()
|
||
|
||
return success_response(data={
|
||
"session_id": req.session_id,
|
||
"user_message": question,
|
||
"assistant_message": assistant_response,
|
||
"referenced_files": [ref["file_path"] for ref in cited_refs],
|
||
"references": _build_reference_items(cited_refs, session.project_id),
|
||
})
|
||
|
||
except Exception as e:
|
||
await db.rollback()
|
||
raise HTTPException(
|
||
status_code=500,
|
||
detail=f"生成对话回复失败: {str(e)}"
|
||
)
|
||
|
||
|
||
@router.post("/send/stream", response_model=dict)
|
||
async def send_chat_message_stream(
|
||
req: ChatMessageRequest,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""发送对话消息并流式返回回复"""
|
||
question = (req.message or "").strip()
|
||
if not question:
|
||
raise HTTPException(status_code=400, detail="消息内容不能为空")
|
||
|
||
stmt = select(ChatSession).where(ChatSession.id == req.session_id)
|
||
result = await db.execute(stmt)
|
||
session = result.scalar_one_or_none()
|
||
|
||
if not session:
|
||
raise HTTPException(status_code=404, detail="对话会话不存在")
|
||
|
||
if session.user_id != current_user.id:
|
||
raise HTTPException(status_code=403, detail="无权访问该对话")
|
||
|
||
await require_project_read_access(db, session.project_id, current_user)
|
||
|
||
# 先读取历史消息（不含本轮提问），用于多轮上下文与判断是否首轮
|
||
msg_stmt = (
|
||
select(ChatMessage)
|
||
.where(ChatMessage.session_id == req.session_id)
|
||
.order_by(ChatMessage.created_at.desc())
|
||
.limit(CHAT_HISTORY_MESSAGE_LIMIT)
|
||
)
|
||
msg_result = await db.execute(msg_stmt)
|
||
prev_messages = list(reversed(msg_result.scalars().all()))
|
||
|
||
conversation_history = [
|
||
{"role": m.role, "content": m.content}
|
||
for m in prev_messages
|
||
# 被中断的助手消息（含统一中断标记）不参与上下文，避免污染后续模型输入
|
||
if not (
|
||
m.role == "assistant"
|
||
and _is_interrupted_content(m.content or "")
|
||
)
|
||
]
|
||
|
||
# 先持久化用户提问和空的助手消息占位行，再执行向量检索。
|
||
# 占位行先入库后，即使客户端在检索阶段断开（停止生成），
|
||
# 前端也能拿到真实消息 id 并正确标记“已停止”，而不会留下悬空占位。
|
||
assistant_message = ChatMessage(
|
||
session_id=req.session_id,
|
||
role="assistant",
|
||
content="",
|
||
)
|
||
if req.insert_user_message:
|
||
query_message = ChatMessage(
|
||
session_id=req.session_id,
|
||
role="user",
|
||
content=question,
|
||
)
|
||
db.add(query_message)
|
||
db.add(assistant_message)
|
||
session.message_count = (session.message_count or 0) + (2 if req.insert_user_message else 1)
|
||
await db.commit()
|
||
|
||
user_message_id = query_message.id if req.insert_user_message else None
|
||
assistant_message_id = assistant_message.id
|
||
|
||
async def _mark_interrupted_assistant():
|
||
"""把中断标记写入仍为空的助手占位行（断开连接时任务可能正被取消）。"""
|
||
try:
|
||
msg = await db.get(ChatMessage, assistant_message_id)
|
||
if msg is not None and not (msg.content or "").strip():
|
||
msg.content = INTERRUPTED_RESPONSE_MARKER
|
||
await db.commit()
|
||
except Exception as save_exc: # noqa: BLE001
|
||
logger.warning("保存中断状态失败: %s", save_exc)
|
||
|
||
try:
|
||
retrieved_docs = await rag_service.retrieve_documents(
|
||
db,
|
||
session.project_id,
|
||
question,
|
||
top_k=5,
|
||
)
|
||
except asyncio.CancelledError:
|
||
# 客户端在向量检索阶段断开（停止生成）：任务被取消，占位行需写入中断标记。
|
||
await _mark_interrupted_assistant()
|
||
raise
|
||
except Exception:
|
||
# 占位行已入库：把失败原因写进助手消息，避免留下永远显示“正在思考”的空行。
|
||
try:
|
||
msg = await db.get(ChatMessage, assistant_message_id)
|
||
if msg is not None:
|
||
msg.content = "回答生成失败，请稍后重试。"
|
||
await db.commit()
|
||
except Exception as save_exc: # noqa: BLE001
|
||
logger.warning("保存检索失败状态失败: %s", save_exc)
|
||
raise
|
||
|
||
llm_config_id = session.llm_config_id
|
||
project_id = session.project_id
|
||
# 每累计这么多字符就回写一次占位行，平衡「刷新可见性」与「写库频率」
|
||
FLUSH_EVERY_CHARS = 120
|
||
|
||
async def _persist_assistant(parts, *, completed):
|
||
"""把已生成内容回写到占位的助手消息行。
|
||
|
||
completed=True 时附带引用解析；False 表示流式中途的增量回写。
|
||
"""
|
||
assistant_response = "".join(parts)
|
||
cited_refs = []
|
||
if completed:
|
||
assistant_response, cited_refs = _compact_cited_refs(
|
||
assistant_response, retrieved_docs
|
||
)
|
||
references = _build_reference_items(cited_refs, project_id)
|
||
|
||
msg = await db.get(ChatMessage, assistant_message_id)
|
||
if msg is not None:
|
||
msg.content = assistant_response
|
||
if completed:
|
||
msg.referenced_files = (
|
||
json.dumps(cited_refs, ensure_ascii=False) if cited_refs else None
|
||
)
|
||
|
||
await db.commit()
|
||
return assistant_response, references
|
||
|
||
async def event_generator():
|
||
assistant_response_parts = []
|
||
chars_since_flush = 0
|
||
finished = False
|
||
try:
|
||
# 先回传两条消息的真实 id，使前端无需刷新即可获得删除入口等能力
|
||
yield _stream_event("ids", {
|
||
"session_id": req.session_id,
|
||
"user_message_id": user_message_id,
|
||
"assistant_message_id": assistant_message_id,
|
||
})
|
||
|
||
async for chunk in rag_service.generate_response_stream(
|
||
db,
|
||
question,
|
||
project_id,
|
||
llm_config_id,
|
||
retrieved_docs,
|
||
conversation_history,
|
||
):
|
||
assistant_response_parts.append(chunk)
|
||
chars_since_flush += len(chunk)
|
||
yield _stream_event("chunk", {"content": chunk})
|
||
|
||
# 边生成边落库：达到阈值就把当前进度回写到占位行
|
||
if chars_since_flush >= FLUSH_EVERY_CHARS:
|
||
chars_since_flush = 0
|
||
await _persist_assistant(assistant_response_parts, completed=False)
|
||
|
||
assistant_response, references = await _persist_assistant(
|
||
assistant_response_parts, completed=True
|
||
)
|
||
finished = True
|
||
|
||
yield _stream_event("references", references or [])
|
||
yield _stream_event("done", {
|
||
"session_id": req.session_id,
|
||
"user_message_id": user_message_id,
|
||
"assistant_message_id": assistant_message_id,
|
||
"content": assistant_response,
|
||
"references": references or [],
|
||
})
|
||
except (asyncio.CancelledError, GeneratorExit):
|
||
# 客户端中途断开（如刷新页面）：尽力把已生成的部分回写到占位行。
|
||
# 由于提问与占位行已入库，刷新后至少能看到提问与已生成内容。
|
||
if not finished:
|
||
if not assistant_response_parts:
|
||
assistant_response_parts.append(INTERRUPTED_RESPONSE_MARKER)
|
||
else:
|
||
assistant_response_parts.append(f"\n\n{INTERRUPTED_RESPONSE_MARKER}")
|
||
try:
|
||
await db.rollback()
|
||
await _persist_assistant(assistant_response_parts, completed=False)
|
||
except Exception as save_exc: # noqa: BLE001
|
||
logger.warning("保存中断的助手回复失败: %s", save_exc)
|
||
raise
|
||
except Exception as exc:
|
||
await db.rollback()
|
||
if not assistant_response_parts:
|
||
assistant_response_parts.append("回答生成失败，请稍后重试。")
|
||
try:
|
||
await _persist_assistant(assistant_response_parts, completed=False)
|
||
except Exception as save_exc: # noqa: BLE001
|
||
logger.warning("保存失败的助手回复状态失败: %s", save_exc)
|
||
yield _stream_event("error", {"detail": f"生成对话回复失败: {str(exc)}"})
|
||
|
||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
||
|
||
|
||
@router.delete("/sessions/{session_id}", response_model=dict)
|
||
async def delete_chat_session(
|
||
session_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""删除对话会话"""
|
||
stmt = select(ChatSession).where(ChatSession.id == session_id)
|
||
result = await db.execute(stmt)
|
||
session = result.scalar_one_or_none()
|
||
|
||
if not session:
|
||
raise HTTPException(status_code=404, detail="对话会话不存在")
|
||
|
||
if session.user_id != current_user.id:
|
||
raise HTTPException(status_code=403, detail="无权删除该对话")
|
||
|
||
await db.delete(session)
|
||
msg_stmt = delete(ChatMessage).where(ChatMessage.session_id == session_id)
|
||
await db.execute(msg_stmt)
|
||
await db.commit()
|
||
|
||
return success_response(message="对话会话已删除")
|
||
|
||
|
||
@router.delete("/messages/{message_id}", response_model=dict)
|
||
async def delete_chat_message(
|
||
message_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""删除单条对话消息"""
|
||
stmt = (
|
||
select(ChatMessage, ChatSession)
|
||
.join(ChatSession, ChatMessage.session_id == ChatSession.id)
|
||
.where(ChatMessage.id == message_id)
|
||
)
|
||
result = await db.execute(stmt)
|
||
row = result.first()
|
||
|
||
if not row:
|
||
raise HTTPException(status_code=404, detail="消息不存在")
|
||
|
||
msg, session = row
|
||
if session.user_id != current_user.id:
|
||
raise HTTPException(status_code=403, detail="无权删除该消息")
|
||
|
||
await db.delete(msg)
|
||
if session.message_count and session.message_count > 0:
|
||
session.message_count = max(0, session.message_count - 1)
|
||
await db.commit()
|
||
|
||
return success_response(message="消息已删除")
|
||
|
||
|
||
@router.post("/messages/{message_id}/interrupt", response_model=dict)
|
||
async def mark_message_interrupted(
|
||
message_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""把生成被中断的助手消息标记为已停止。
|
||
|
||
前端在点击“停止生成”后立即调用本接口，把统一中断标记写入数据库，
|
||
避免依赖连接断开后生成器的兜底写入（存在时间差），保证切换会话或刷新后
|
||
仍能识别该消息为“已停止生成”。
|
||
"""
|
||
stmt = (
|
||
select(ChatMessage, ChatSession)
|
||
.join(ChatSession, ChatMessage.session_id == ChatSession.id)
|
||
.where(ChatMessage.id == message_id)
|
||
)
|
||
result = await db.execute(stmt)
|
||
row = result.first()
|
||
|
||
if not row:
|
||
raise HTTPException(status_code=404, detail="消息不存在")
|
||
|
||
msg, session = row
|
||
if session.user_id != current_user.id:
|
||
raise HTTPException(status_code=403, detail="无权操作该消息")
|
||
if msg.role != "assistant":
|
||
raise HTTPException(status_code=400, detail="仅助手消息可标记为已停止")
|
||
|
||
content = msg.content or ""
|
||
if not content.strip():
|
||
msg.content = INTERRUPTED_RESPONSE_MARKER
|
||
elif not _is_interrupted_content(content):
|
||
msg.content = f"{content.rstrip()}\n\n{INTERRUPTED_RESPONSE_MARKER}"
|
||
await db.commit()
|
||
|
||
return success_response(data={
|
||
"message_id": msg.id,
|
||
"content": msg.content,
|
||
})
|
||
|
||
|
||
@router.put("/sessions/{session_id}", response_model=dict)
|
||
async def update_chat_session(
|
||
session_id: int,
|
||
req: ChatSessionUpdateRequest,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
"""更新对话会话标题"""
|
||
stmt = select(ChatSession).where(ChatSession.id == session_id)
|
||
result = await db.execute(stmt)
|
||
session = result.scalar_one_or_none()
|
||
|
||
if not session:
|
||
raise HTTPException(status_code=404, detail="对话会话不存在")
|
||
|
||
if session.user_id != current_user.id:
|
||
raise HTTPException(status_code=403, detail="无权修改该对话")
|
||
|
||
title = (req.title or "").strip()
|
||
if not title:
|
||
raise HTTPException(status_code=400, detail="标题不能为空")
|
||
|
||
session.title = title
|
||
await db.commit()
|
||
|
||
return success_response(message="标题已更新", data={"session_id": session_id, "title": title})
|