feat: add persistent chat sessions
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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)
|
||||
@@ -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]
|
||||
Reference in New Issue
Block a user