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 "",
|
||
) |