feat: add persistent chat sessions
This commit is contained in:
@@ -1,6 +1,8 @@
|
|||||||
from fastapi import APIRouter
|
from fastapi import APIRouter
|
||||||
|
|
||||||
|
from kbqa.chat.routes import router as sessions_router
|
||||||
from kbqa.documents.routes import router as documents_router
|
from kbqa.documents.routes import router as documents_router
|
||||||
|
|
||||||
api_router = APIRouter(prefix="/api/v1")
|
api_router = APIRouter(prefix="/api/v1")
|
||||||
api_router.include_router(documents_router)
|
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.chat.models import ChatSession, Message, MessageSource
|
||||||
from kbqa.database import utc_now
|
from kbqa.database import utc_now
|
||||||
|
from kbqa.documents.models import Chunk, Document
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
@@ -16,6 +17,19 @@ class MessageSourceInput:
|
|||||||
score: float
|
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:
|
class ChatRepository:
|
||||||
def __init__(self, session: AsyncSession) -> None:
|
def __init__(self, session: AsyncSession) -> None:
|
||||||
self.session = session
|
self.session = session
|
||||||
@@ -67,17 +81,32 @@ class ChatRepository:
|
|||||||
)
|
)
|
||||||
return list(rows), int(total or 0)
|
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:
|
if not message_ids:
|
||||||
return {}
|
return {}
|
||||||
sources = await self.session.scalars(
|
rows = await self.session.execute(
|
||||||
select(MessageSource)
|
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))
|
.where(MessageSource.message_id.in_(message_ids))
|
||||||
.order_by(MessageSource.citation_label)
|
.order_by(MessageSource.citation_label)
|
||||||
)
|
)
|
||||||
result: dict[str, list[MessageSource]] = {}
|
result: dict[str, list[MessageSourceView]] = {}
|
||||||
for source in sources:
|
for source, chunk, document in rows:
|
||||||
result.setdefault(source.message_id, []).append(source)
|
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
|
return result
|
||||||
|
|
||||||
async def begin_turn(self, session_id: str, query: str) -> tuple[Message, Message]:
|
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