diff --git a/backend/src/kbqa/api/router.py b/backend/src/kbqa/api/router.py index 493fb7b..035f77c 100644 --- a/backend/src/kbqa/api/router.py +++ b/backend/src/kbqa/api/router.py @@ -1,6 +1,8 @@ from fastapi import APIRouter +from kbqa.chat.routes import router as sessions_router from kbqa.documents.routes import router as documents_router api_router = APIRouter(prefix="/api/v1") api_router.include_router(documents_router) +api_router.include_router(sessions_router) diff --git a/backend/src/kbqa/chat/repository.py b/backend/src/kbqa/chat/repository.py index b22b41d..1d7e503 100644 --- a/backend/src/kbqa/chat/repository.py +++ b/backend/src/kbqa/chat/repository.py @@ -6,6 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from kbqa.chat.models import ChatSession, Message, MessageSource from kbqa.database import utc_now +from kbqa.documents.models import Chunk, Document @dataclass(frozen=True, slots=True) @@ -16,6 +17,19 @@ class MessageSourceInput: score: float +@dataclass(frozen=True, slots=True) +class MessageSourceView: + message_id: str + document_id: str + chunk_id: str + citation_label: str + score: float + filename: str + order_index: int + page_number: int | None + content: str + + class ChatRepository: def __init__(self, session: AsyncSession) -> None: self.session = session @@ -67,17 +81,32 @@ class ChatRepository: ) return list(rows), int(total or 0) - async def sources_for_messages(self, message_ids: list[str]) -> dict[str, list[MessageSource]]: + async def sources_for_messages( + self, message_ids: list[str] + ) -> dict[str, list[MessageSourceView]]: if not message_ids: return {} - sources = await self.session.scalars( - select(MessageSource) + rows = await self.session.execute( + select(MessageSource, Chunk, Document) + .join(Chunk, Chunk.id == MessageSource.chunk_id) + .join(Document, Document.id == MessageSource.document_id) .where(MessageSource.message_id.in_(message_ids)) .order_by(MessageSource.citation_label) ) - result: dict[str, list[MessageSource]] = {} - for source in sources: - result.setdefault(source.message_id, []).append(source) + result: dict[str, list[MessageSourceView]] = {} + for source, chunk, document in rows: + view = MessageSourceView( + message_id=source.message_id, + document_id=source.document_id, + chunk_id=source.chunk_id, + citation_label=source.citation_label, + score=source.score, + filename=document.filename, + order_index=chunk.order_index, + page_number=chunk.page_number, + content=chunk.content, + ) + result.setdefault(source.message_id, []).append(view) return result async def begin_turn(self, session_id: str, query: str) -> tuple[Message, Message]: diff --git a/backend/src/kbqa/chat/routes.py b/backend/src/kbqa/chat/routes.py new file mode 100644 index 0000000..7fc528f --- /dev/null +++ b/backend/src/kbqa/chat/routes.py @@ -0,0 +1,93 @@ +from typing import Annotated + +from fastapi import APIRouter, Depends, Query, Response, status +from sqlalchemy.ext.asyncio import AsyncSession + +from kbqa.api.errors import AppError +from kbqa.chat.repository import ChatRepository +from kbqa.chat.schemas import ( + MessagePage, + MessageResponse, + MessageSourceResponse, + SessionCreate, + SessionPage, + SessionResponse, +) +from kbqa.database import get_session + +router = APIRouter(prefix="/sessions", tags=["sessions"]) +SessionDependency = Annotated[AsyncSession, Depends(get_session)] + + +@router.post("", response_model=SessionResponse, status_code=status.HTTP_201_CREATED) +async def create_session(payload: SessionCreate, session: SessionDependency) -> SessionResponse: + created = await ChatRepository(session).create_session(payload.title) + return SessionResponse.model_validate(created) + + +@router.get("", response_model=SessionPage) +async def list_sessions( + session: SessionDependency, + offset: Annotated[int, Query(ge=0)] = 0, + limit: Annotated[int, Query(ge=1, le=100)] = 20, +) -> SessionPage: + items, total = await ChatRepository(session).page_sessions(offset, limit) + return SessionPage( + items=[SessionResponse.model_validate(item) for item in items], + total=total, + offset=offset, + limit=limit, + ) + + +@router.get("/{session_id}/messages", response_model=MessagePage) +async def list_messages( + session_id: str, + session: SessionDependency, + offset: Annotated[int, Query(ge=0)] = 0, + limit: Annotated[int, Query(ge=1, le=100)] = 50, +) -> MessagePage: + repository = ChatRepository(session) + if await repository.get_session(session_id) is None: + raise AppError("SESSION_NOT_FOUND", "会话不存在", 404) + messages, total = await repository.page_messages(session_id, offset, limit) + sources_by_message = await repository.sources_for_messages([message.id for message in messages]) + responses: list[MessageResponse] = [] + for message in messages: + sources = sources_by_message.get(message.id, []) + source_labels = {source.citation_label for source in sources} + deleted = sorted(repository.citation_labels(message.content) - source_labels) + responses.append( + MessageResponse( + id=message.id, + session_id=message.session_id, + role=message.role, + content=message.content, + status=message.status, + error_code=message.error_code, + order_index=message.order_index, + created_at=message.created_at, + sources=[ + MessageSourceResponse( + label=source.citation_label, + document_id=source.document_id, + chunk_id=source.chunk_id, + filename=source.filename, + order_index=source.order_index, + page_number=source.page_number, + content=source.content, + score=source.score, + ) + for source in sources + ], + deleted_source_labels=deleted, + ) + ) + return MessagePage(items=responses, total=total, offset=offset, limit=limit) + + +@router.delete("/{session_id}", status_code=status.HTTP_204_NO_CONTENT) +async def delete_session(session_id: str, session: SessionDependency) -> Response: + if not await ChatRepository(session).delete_session(session_id): + raise AppError("SESSION_NOT_FOUND", "会话不存在", 404) + return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/backend/src/kbqa/chat/schemas.py b/backend/src/kbqa/chat/schemas.py new file mode 100644 index 0000000..7fc7698 --- /dev/null +++ b/backend/src/kbqa/chat/schemas.py @@ -0,0 +1,46 @@ +from datetime import datetime + +from pydantic import BaseModel, ConfigDict, Field + +from kbqa.documents.schemas import Page + + +class SessionCreate(BaseModel): + title: str | None = Field(default=None, max_length=200) + + +class SessionResponse(BaseModel): + model_config = ConfigDict(from_attributes=True) + + id: str + title: str | None + created_at: datetime + updated_at: datetime + + +class MessageSourceResponse(BaseModel): + label: str + document_id: str + chunk_id: str + filename: str + order_index: int + page_number: int | None + content: str + score: float + + +class MessageResponse(BaseModel): + id: str + session_id: str + role: str + content: str + status: str + error_code: str | None + order_index: int + created_at: datetime + sources: list[MessageSourceResponse] = Field(default_factory=list) + deleted_source_labels: list[str] = Field(default_factory=list) + + +SessionPage = Page[SessionResponse] +MessagePage = Page[MessageResponse]