356 lines
13 KiB
Python
356 lines
13 KiB
Python
|
import json
|
|||
|
|
import logging
|
|||
|
|
from fastapi import APIRouter, Depends, HTTPException, Query
|
|||
|
|
from fastapi.responses import StreamingResponse
|
|||
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
|
from sqlalchemy import select, func
|
|||
|
|
from pydantic import BaseModel
|
|||
|
|
|
|||
|
|
from app.core.database import get_db
|
|||
|
|
from app.models.book import Book, Chapter, Segment
|
|||
|
|
from app.models.chat import Conversation, Message
|
|||
|
|
from app.llm.hub import resolve_provider
|
|||
|
|
from app.rag.engine import (
|
|||
|
|
search_segments, build_chat_prompt, parse_citations,
|
|||
|
|
estimate_tokens, match_chapters, _trim_messages,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
router = APIRouter(prefix="/chat", tags=["chat"])
|
|||
|
|
|
|||
|
|
|
|||
|
|
class CreateConversationRequest(BaseModel):
|
|||
|
|
book_id: str
|
|||
|
|
title: str | None = None
|
|||
|
|
model: str = "claude"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SendMessageRequest(BaseModel):
|
|||
|
|
content: str
|
|||
|
|
model: str | None = None
|
|||
|
|
# 被重试的用户消息 id：提供时不新建 user 消息，新 AI 回答作为该问题的新版本
|
|||
|
|
# （parent_message_id 指向它），历史上下文只取该问题之前的对话
|
|||
|
|
reply_to: str | None = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ConversationOut(BaseModel):
|
|||
|
|
id: str
|
|||
|
|
book_id: str
|
|||
|
|
title: str | None
|
|||
|
|
model: str
|
|||
|
|
created_at: str
|
|||
|
|
|
|||
|
|
class Config:
|
|||
|
|
from_attributes = True
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MessageOut(BaseModel):
|
|||
|
|
id: str
|
|||
|
|
conversation_id: str
|
|||
|
|
parent_message_id: str | None
|
|||
|
|
role: str
|
|||
|
|
content: str
|
|||
|
|
citations: list | None
|
|||
|
|
model: str | None
|
|||
|
|
created_at: str
|
|||
|
|
|
|||
|
|
class Config:
|
|||
|
|
from_attributes = True
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.get("/conversations", response_model=list[ConversationOut])
|
|||
|
|
async def list_conversations(book_id: str | None = None, db: AsyncSession = Depends(get_db)):
|
|||
|
|
q = select(Conversation).order_by(Conversation.updated_at.desc())
|
|||
|
|
if book_id:
|
|||
|
|
q = q.where(Conversation.book_id == book_id)
|
|||
|
|
result = await db.execute(q)
|
|||
|
|
convs = result.scalars().all()
|
|||
|
|
return [ConversationOut(
|
|||
|
|
id=c.id, book_id=c.book_id, title=c.title, model=c.model,
|
|||
|
|
created_at=c.created_at.isoformat() if c.created_at else "",
|
|||
|
|
) for c in convs]
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.post("/conversations", response_model=ConversationOut)
|
|||
|
|
async def create_conversation(req: CreateConversationRequest, db: AsyncSession = Depends(get_db)):
|
|||
|
|
book = await db.get(Book, req.book_id)
|
|||
|
|
if not book:
|
|||
|
|
raise HTTPException(404, "书不存在")
|
|||
|
|
conv = Conversation(book_id=req.book_id, title=req.title or f"与《{book.title}》的对话", model=req.model)
|
|||
|
|
db.add(conv)
|
|||
|
|
await db.commit()
|
|||
|
|
await db.refresh(conv)
|
|||
|
|
return ConversationOut(
|
|||
|
|
id=conv.id, book_id=conv.book_id, title=conv.title, model=conv.model,
|
|||
|
|
created_at=conv.created_at.isoformat() if conv.created_at else "",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.get("/conversations/{conv_id}/messages", response_model=list[MessageOut])
|
|||
|
|
async def get_messages(conv_id: str, db: AsyncSession = Depends(get_db)):
|
|||
|
|
conv = await db.get(Conversation, conv_id)
|
|||
|
|
if not conv:
|
|||
|
|
raise HTTPException(404, "会话不存在")
|
|||
|
|
result = await db.execute(
|
|||
|
|
select(Message).where(Message.conversation_id == conv_id).order_by(Message.created_at)
|
|||
|
|
)
|
|||
|
|
msgs = result.scalars().all()
|
|||
|
|
return [_msg_to_out(m) for m in msgs]
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.post("/conversations/{conv_id}/messages")
|
|||
|
|
async def send_message(conv_id: str, req: SendMessageRequest, db: AsyncSession = Depends(get_db)):
|
|||
|
|
conv = await db.get(Conversation, conv_id)
|
|||
|
|
if not conv:
|
|||
|
|
raise HTTPException(404, "会话不存在")
|
|||
|
|
|
|||
|
|
book = await db.get(Book, conv.book_id)
|
|||
|
|
if not book:
|
|||
|
|
raise HTTPException(404, "书不存在")
|
|||
|
|
|
|||
|
|
model_name = req.model or conv.model
|
|||
|
|
|
|||
|
|
# 重试：复用被重试的用户消息（不新建），新回答作为它的新版本
|
|||
|
|
reply_user: Message | None = None
|
|||
|
|
if req.reply_to:
|
|||
|
|
reply_user = await db.get(Message, req.reply_to)
|
|||
|
|
if not reply_user or reply_user.role != "user" or reply_user.conversation_id != conv_id:
|
|||
|
|
raise HTTPException(400, "reply_to 必须是本会话内的一条用户消息")
|
|||
|
|
user_msg = reply_user
|
|||
|
|
else:
|
|||
|
|
user_msg = Message(conversation_id=conv_id, role="user", content=req.content, model=model_name)
|
|||
|
|
db.add(user_msg)
|
|||
|
|
await db.commit()
|
|||
|
|
|
|||
|
|
chapters = (await db.execute(
|
|||
|
|
select(Chapter).where(Chapter.book_id == book.id).order_by(Chapter.order_idx)
|
|||
|
|
)).scalars().all()
|
|||
|
|
|
|||
|
|
# 章节标题匹配：命中章节时把整章原文注入 prompt（解决「第X章讲了什么」这类问题）
|
|||
|
|
matched = match_chapters(list(chapters), req.content)
|
|||
|
|
chapter_segments = []
|
|||
|
|
if matched:
|
|||
|
|
result = await db.execute(
|
|||
|
|
select(Segment)
|
|||
|
|
.where(Segment.chapter_id.in_([c.id for c in matched]))
|
|||
|
|
.order_by(Segment.order_idx)
|
|||
|
|
)
|
|||
|
|
chapter_segments = list(result.scalars().all())
|
|||
|
|
|
|||
|
|
segments = await search_segments(
|
|||
|
|
db, book.id, req.content, top_k=15, chapters=list(chapters)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
history_rows = (await db.execute(
|
|||
|
|
select(Message).where(Message.conversation_id == conv_id).order_by(Message.created_at)
|
|||
|
|
)).scalars().all()
|
|||
|
|
|
|||
|
|
# 历史上下文：重试时只取被重试问题之前的对话（旧版本回答不进 prompt）；
|
|||
|
|
# 正常发送时排除当前这轮 user 消息即可。
|
|||
|
|
if reply_user is not None:
|
|||
|
|
row_ids = {m.id for m in history_rows}
|
|||
|
|
if reply_user.id in row_ids:
|
|||
|
|
idx = next(i for i, m in enumerate(history_rows) if m.id == reply_user.id)
|
|||
|
|
history = [{"role": m.role, "content": m.content} for m in history_rows[:idx]]
|
|||
|
|
else:
|
|||
|
|
# 数据异常（如 SRU），退回为全量历史
|
|||
|
|
history = [{"role": m.role, "content": m.content} for m in history_rows]
|
|||
|
|
else:
|
|||
|
|
history = [{"role": m.role, "content": m.content} for m in history_rows[:-1]]
|
|||
|
|
|
|||
|
|
messages = build_chat_prompt(
|
|||
|
|
book, list(chapters), segments, history, req.content,
|
|||
|
|
chapter_segments=chapter_segments,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
total_tokens = estimate_tokens("".join(m["content"] for m in messages))
|
|||
|
|
provider = await resolve_provider(db, model_name)
|
|||
|
|
if provider is None:
|
|||
|
|
raise HTTPException(400, f"模型 {model_name} 未配置，请先在页面「管理模型」中添加或选择可用模型")
|
|||
|
|
if total_tokens > provider.max_context * 0.8:
|
|||
|
|
# 内容感知裁剪：保留 system（含检索原文 + 整章注入），从最旧历史开始丢
|
|||
|
|
messages = _trim_messages(messages, int(provider.max_context * 0.8))
|
|||
|
|
|
|||
|
|
async def stream_response():
|
|||
|
|
full_content = ""
|
|||
|
|
try:
|
|||
|
|
async for token in provider.chat_stream(messages, temperature=0.7):
|
|||
|
|
full_content += token
|
|||
|
|
yield f"data: {json.dumps({'type': 'token', 'content': token}, ensure_ascii=False)}\n\n"
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.exception("流式生成失败（conversation=%s）", conv_id)
|
|||
|
|
yield f"data: {json.dumps({'type': 'error', 'content': str(e)}, ensure_ascii=False)}\n\n"
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
parsed = parse_citations(full_content, segments)
|
|||
|
|
|
|||
|
|
assistant_msg = Message(
|
|||
|
|
conversation_id=conv_id,
|
|||
|
|
parent_message_id=user_msg.id,
|
|||
|
|
role="assistant",
|
|||
|
|
content=parsed["text"],
|
|||
|
|
citations=json.dumps(parsed["citations"], ensure_ascii=False),
|
|||
|
|
model=model_name,
|
|||
|
|
)
|
|||
|
|
db.add(assistant_msg)
|
|||
|
|
await db.commit()
|
|||
|
|
await db.refresh(assistant_msg)
|
|||
|
|
|
|||
|
|
yield f"data: {json.dumps({'type': 'done', 'message_id': assistant_msg.id, 'citations': parsed['citations']}, ensure_ascii=False)}\n\n"
|
|||
|
|
|
|||
|
|
return StreamingResponse(stream_response(), media_type="text/event-stream")
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.post("/messages/{msg_id}/retry")
|
|||
|
|
async def retry_message(msg_id: str, db: AsyncSession = Depends(get_db)):
|
|||
|
|
msg = await db.get(Message, msg_id)
|
|||
|
|
if not msg or msg.role != "user":
|
|||
|
|
raise HTTPException(400, "只能对用户消息重试")
|
|||
|
|
# 走 reply_to 语义：不新建 user 消息，而是给该问题追加一个新版本回答
|
|||
|
|
return await send_message(
|
|||
|
|
msg.conversation_id,
|
|||
|
|
SendMessageRequest(content=msg.content, model=msg.model, reply_to=msg.id),
|
|||
|
|
db,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.get("/segments/{segment_id}/context")
|
|||
|
|
async def segment_context(
|
|||
|
|
segment_id: str,
|
|||
|
|
before: int = Query(2, ge=0, le=6),
|
|||
|
|
after: int = Query(2, ge=0, le=6),
|
|||
|
|
db: AsyncSession = Depends(get_db),
|
|||
|
|
):
|
|||
|
|
"""引用段落的阅读上下文：章节标题 + 前后段落，供前端「阅读原文」弹层。"""
|
|||
|
|
seg = await db.get(Segment, segment_id)
|
|||
|
|
if not seg:
|
|||
|
|
raise HTTPException(404, "段落不存在")
|
|||
|
|
ch = await db.get(Chapter, seg.chapter_id)
|
|||
|
|
|
|||
|
|
prev_rows = await db.execute(
|
|||
|
|
select(Segment.id, Segment.text, Segment.order_idx, Segment.cfi)
|
|||
|
|
.where(Segment.chapter_id == seg.chapter_id, Segment.order_idx < seg.order_idx)
|
|||
|
|
.order_by(Segment.order_idx.desc())
|
|||
|
|
.limit(before)
|
|||
|
|
)
|
|||
|
|
next_rows = await db.execute(
|
|||
|
|
select(Segment.id, Segment.text, Segment.order_idx, Segment.cfi)
|
|||
|
|
.where(Segment.chapter_id == seg.chapter_id, Segment.order_idx > seg.order_idx)
|
|||
|
|
.order_by(Segment.order_idx)
|
|||
|
|
.limit(after)
|
|||
|
|
)
|
|||
|
|
prev_list = [{"id": r[0], "text": r[1], "order_idx": r[2], "cfi": r[3]} for r in prev_rows.all()]
|
|||
|
|
next_list = [{"id": r[0], "text": r[1], "order_idx": r[2], "cfi": r[3]} for r in next_rows.all()]
|
|||
|
|
prev_list.reverse() # 保持阅读顺序
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"segment": {"id": seg.id, "text": seg.text, "order_idx": seg.order_idx, "cfi": seg.cfi},
|
|||
|
|
"chapter": {"id": ch.id, "title": ch.title, "level": ch.level} if ch else None,
|
|||
|
|
"before": prev_list,
|
|||
|
|
"after": next_list,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.get("/search")
|
|||
|
|
async def search_book(book_id: str, q: str = Query(min_length=1), db: AsyncSession = Depends(get_db)):
|
|||
|
|
segments = await search_segments(db, book_id, q, top_k=10)
|
|||
|
|
return [
|
|||
|
|
{"id": s.id, "text": s.text[:300], "cfi": s.cfi, "char_count": s.char_count}
|
|||
|
|
for s in segments
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.delete("/messages/{msg_id}")
|
|||
|
|
async def delete_message(msg_id: str, db: AsyncSession = Depends(get_db)):
|
|||
|
|
msg = await db.get(Message, msg_id)
|
|||
|
|
if not msg:
|
|||
|
|
raise HTTPException(404, "消息不存在")
|
|||
|
|
await db.delete(msg)
|
|||
|
|
await db.commit()
|
|||
|
|
return {"deleted": msg_id}
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.delete("/conversations/{conv_id}")
|
|||
|
|
async def delete_conversation(conv_id: str, db: AsyncSession = Depends(get_db)):
|
|||
|
|
conv = await db.get(Conversation, conv_id)
|
|||
|
|
if not conv:
|
|||
|
|
raise HTTPException(404, "会话不存在")
|
|||
|
|
# ORM cascade="all, delete-orphan" 会一并删除该会话下的所有消息
|
|||
|
|
await db.delete(conv)
|
|||
|
|
await db.commit()
|
|||
|
|
return {"deleted": conv_id}
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.post("/conversations/{conv_id}/branch")
|
|||
|
|
async def branch_conversation(conv_id: str, from_message_id: str = Query(...), db: AsyncSession = Depends(get_db)):
|
|||
|
|
conv = await db.get(Conversation, conv_id)
|
|||
|
|
if not conv:
|
|||
|
|
raise HTTPException(404, "会话不存在")
|
|||
|
|
branch = Conversation(book_id=conv.book_id, title=f"{conv.title or '对话'} (分支)", model=conv.model)
|
|||
|
|
db.add(branch)
|
|||
|
|
await db.flush()
|
|||
|
|
|
|||
|
|
result = await db.execute(
|
|||
|
|
select(Message).where(Message.conversation_id == conv_id).order_by(Message.created_at)
|
|||
|
|
)
|
|||
|
|
msgs = result.scalars().all()
|
|||
|
|
for m in msgs:
|
|||
|
|
new_msg = Message(
|
|||
|
|
conversation_id=branch.id,
|
|||
|
|
role=m.role,
|
|||
|
|
content=m.content,
|
|||
|
|
citations=m.citations,
|
|||
|
|
model=m.model,
|
|||
|
|
)
|
|||
|
|
db.add(new_msg)
|
|||
|
|
if m.id == from_message_id:
|
|||
|
|
break
|
|||
|
|
await db.commit()
|
|||
|
|
await db.refresh(branch)
|
|||
|
|
return {
|
|||
|
|
"id": branch.id,
|
|||
|
|
"book_id": branch.book_id,
|
|||
|
|
"title": branch.title,
|
|||
|
|
"model": branch.model,
|
|||
|
|
"created_at": branch.created_at.isoformat() if branch.created_at else "",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
@router.get("/conversations/{conv_id}/export")
|
|||
|
|
async def export_conversation(conv_id: str, fmt: str = "markdown", db: AsyncSession = Depends(get_db)):
|
|||
|
|
conv = await db.get(Conversation, conv_id)
|
|||
|
|
if not conv:
|
|||
|
|
raise HTTPException(404, "会话不存在")
|
|||
|
|
result = await db.execute(
|
|||
|
|
select(Message).where(Message.conversation_id == conv_id).order_by(Message.created_at)
|
|||
|
|
)
|
|||
|
|
msgs = result.scalars().all()
|
|||
|
|
book = await db.get(Book, conv.book_id)
|
|||
|
|
|
|||
|
|
lines = [f"# {conv.title or '对话'}", f"书籍: {book.title if book else '未知'}\n"]
|
|||
|
|
for m in msgs:
|
|||
|
|
role_label = "🧑 用户" if m.role == "user" else "🤖 AI"
|
|||
|
|
lines.append(f"## {role_label}\n\n{m.content}\n")
|
|||
|
|
if m.citations:
|
|||
|
|
try:
|
|||
|
|
citations = json.loads(m.citations)
|
|||
|
|
for i, c in enumerate(citations):
|
|||
|
|
lines.append(f"> 引用 {i+1}: {c.get('excerpt', '')}\n")
|
|||
|
|
except (json.JSONDecodeError, TypeError):
|
|||
|
|
pass
|
|||
|
|
return {"content": "\n".join(lines), "format": fmt, "title": conv.title}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _msg_to_out(m: Message) -> MessageOut:
|
|||
|
|
citations = None
|
|||
|
|
if m.citations:
|
|||
|
|
try:
|
|||
|
|
citations = json.loads(m.citations)
|
|||
|
|
except (json.JSONDecodeError, TypeError):
|
|||
|
|
citations = None
|
|||
|
|
return MessageOut(
|
|||
|
|
id=m.id, conversation_id=m.conversation_id, parent_message_id=m.parent_message_id,
|
|||
|
|
role=m.role, content=m.content, citations=citations, model=m.model,
|
|||
|
|
created_at=m.created_at.isoformat() if m.created_at else "",
|
|||
|
|
)
|