diff --git a/backend/alembic.ini b/backend/alembic.ini new file mode 100644 index 0000000..3f89fe0 --- /dev/null +++ b/backend/alembic.ini @@ -0,0 +1,39 @@ +[alembic] +script_location = %(here)s/alembic +prepend_sys_path = %(here)s/src +path_separator = os +sqlalchemy.url = sqlite+aiosqlite:///./data/kbqa.sqlite3 + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/backend/alembic/env.py b/backend/alembic/env.py new file mode 100644 index 0000000..2e1812e --- /dev/null +++ b/backend/alembic/env.py @@ -0,0 +1,58 @@ +from logging.config import fileConfig + +from alembic import context +from sqlalchemy import pool +from sqlalchemy.ext.asyncio import async_engine_from_config + +from kbqa.config import get_settings +from kbqa.database import Base +import kbqa.models # noqa: F401 + +config = context.config +config.set_main_option("sqlalchemy.url", get_settings().database_url) + +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +target_metadata = Base.metadata + + +def run_migrations_offline() -> None: + context.configure( + url=config.get_main_option("sqlalchemy.url"), + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + render_as_batch=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +def do_run_migrations(connection) -> None: + context.configure(connection=connection, target_metadata=target_metadata, render_as_batch=True) + with context.begin_transaction(): + context.run_migrations() + + +async def run_async_migrations() -> None: + connectable = async_engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + async with connectable.connect() as connection: + await connection.run_sync(do_run_migrations) + await connectable.dispose() + + +def run_migrations_online() -> None: + import asyncio + + asyncio.run(run_async_migrations()) + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/backend/alembic/versions/0001_initial.py b/backend/alembic/versions/0001_initial.py new file mode 100644 index 0000000..8c0d444 --- /dev/null +++ b/backend/alembic/versions/0001_initial.py @@ -0,0 +1,154 @@ +"""Create initial SQLite schema. + +Revision ID: 0001_initial +Revises: +Create Date: 2026-07-13 +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +revision: str = "0001_initial" +down_revision: str | None = None +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.create_table( + "documents", + sa.Column("id", sa.String(36), primary_key=True), + sa.Column("filename", sa.Text(), nullable=False), + sa.Column("stored_path", sa.Text(), nullable=False), + sa.Column("file_type", sa.String(16), nullable=False), + sa.Column("mime_type", sa.String(128), nullable=False), + sa.Column("file_size", sa.Integer(), nullable=False), + sa.Column("sha256", sa.String(64), nullable=False), + sa.Column("status", sa.String(16), nullable=False), + sa.Column("chunk_count", sa.Integer(), nullable=False), + sa.Column("error_code", sa.String(64)), + sa.Column("error_message", sa.Text()), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.CheckConstraint( + "status IN ('pending','indexing','indexed','failed','deleting')", + name="ck_documents_status", + ), + ) + op.create_index("ix_documents_status_updated_at", "documents", ["status", "updated_at"]) + + op.create_table( + "chat_sessions", + sa.Column("id", sa.String(36), primary_key=True), + sa.Column("title", sa.Text()), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + ) + op.create_index("ix_chat_sessions_updated_at", "chat_sessions", ["updated_at"]) + + op.create_table( + "index_jobs", + sa.Column("id", sa.String(36), primary_key=True), + sa.Column( + "document_id", + sa.String(36), + sa.ForeignKey("documents.id", ondelete="CASCADE"), + nullable=False, + unique=True, + ), + sa.Column("status", sa.String(16), nullable=False), + sa.Column("attempt_count", sa.Integer(), nullable=False), + sa.Column("last_error", sa.Text()), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("started_at", sa.DateTime(timezone=True)), + sa.Column("finished_at", sa.DateTime(timezone=True)), + sa.CheckConstraint( + "status IN ('queued','running','succeeded','failed')", + name="ck_index_jobs_status", + ), + ) + + op.create_table( + "chunks", + sa.Column("id", sa.String(36), primary_key=True), + sa.Column( + "document_id", + sa.String(36), + sa.ForeignKey("documents.id", ondelete="CASCADE"), + nullable=False, + ), + sa.Column("order_index", sa.Integer(), nullable=False), + sa.Column("content", sa.Text(), nullable=False), + sa.Column("page_number", sa.Integer()), + sa.Column("char_count", sa.Integer(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.UniqueConstraint("document_id", "order_index", name="uq_chunks_document_order"), + ) + op.create_index("ix_chunks_document_id", "chunks", ["document_id"]) + + op.create_table( + "messages", + sa.Column("id", sa.String(36), primary_key=True), + sa.Column( + "session_id", + sa.String(36), + sa.ForeignKey("chat_sessions.id", ondelete="CASCADE"), + nullable=False, + ), + sa.Column("role", sa.String(16), nullable=False), + sa.Column("content", sa.Text(), nullable=False), + sa.Column("status", sa.String(16), nullable=False), + sa.Column("error_code", sa.String(64)), + sa.Column("order_index", sa.Integer(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.CheckConstraint("role IN ('user','assistant')", name="ck_messages_role"), + sa.CheckConstraint( + "status IN ('generating','completed','failed')", name="ck_messages_status" + ), + sa.UniqueConstraint("session_id", "order_index", name="uq_messages_session_order"), + ) + op.create_index("ix_messages_session_created_at", "messages", ["session_id", "created_at"]) + + op.create_table( + "message_sources", + sa.Column("id", sa.String(36), primary_key=True), + sa.Column( + "message_id", + sa.String(36), + sa.ForeignKey("messages.id", ondelete="CASCADE"), + nullable=False, + ), + sa.Column( + "document_id", + sa.String(36), + sa.ForeignKey("documents.id", ondelete="CASCADE"), + nullable=False, + ), + sa.Column( + "chunk_id", + sa.String(36), + sa.ForeignKey("chunks.id", ondelete="CASCADE"), + nullable=False, + ), + sa.Column("citation_label", sa.String(8), nullable=False), + sa.Column("score", sa.Float(), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.UniqueConstraint("message_id", "chunk_id", name="uq_message_sources_message_chunk"), + ) + op.create_index("ix_message_sources_message_id", "message_sources", ["message_id"]) + + +def downgrade() -> None: + op.drop_index("ix_message_sources_message_id", table_name="message_sources") + op.drop_table("message_sources") + op.drop_index("ix_messages_session_created_at", table_name="messages") + op.drop_table("messages") + op.drop_index("ix_chunks_document_id", table_name="chunks") + op.drop_table("chunks") + op.drop_table("index_jobs") + op.drop_index("ix_chat_sessions_updated_at", table_name="chat_sessions") + op.drop_table("chat_sessions") + op.drop_index("ix_documents_status_updated_at", table_name="documents") + op.drop_table("documents") diff --git a/backend/src/kbqa/chat/__init__.py b/backend/src/kbqa/chat/__init__.py new file mode 100644 index 0000000..63c179c --- /dev/null +++ b/backend/src/kbqa/chat/__init__.py @@ -0,0 +1 @@ +"""Chat persistence and streaming feature.""" diff --git a/backend/src/kbqa/chat/models.py b/backend/src/kbqa/chat/models.py new file mode 100644 index 0000000..b32b236 --- /dev/null +++ b/backend/src/kbqa/chat/models.py @@ -0,0 +1,78 @@ +from datetime import datetime + +from sqlalchemy import ( + CheckConstraint, + DateTime, + Float, + ForeignKey, + Index, + Integer, + String, + Text, + UniqueConstraint, +) +from sqlalchemy.orm import Mapped, mapped_column + +from kbqa.database import Base, utc_now +from kbqa.documents.models import new_uuid + + +class ChatSession(Base): + __tablename__ = "chat_sessions" + __table_args__ = (Index("ix_chat_sessions_updated_at", "updated_at"),) + + id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid) + title: Mapped[str | None] = mapped_column(Text) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=utc_now + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=utc_now, onupdate=utc_now + ) + + +class Message(Base): + __tablename__ = "messages" + __table_args__ = ( + CheckConstraint("role IN ('user','assistant')", name="ck_messages_role"), + CheckConstraint("status IN ('generating','completed','failed')", name="ck_messages_status"), + UniqueConstraint("session_id", "order_index", name="uq_messages_session_order"), + Index("ix_messages_session_created_at", "session_id", "created_at"), + ) + + id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid) + session_id: Mapped[str] = mapped_column( + String(36), ForeignKey("chat_sessions.id", ondelete="CASCADE"), nullable=False + ) + role: Mapped[str] = mapped_column(String(16), nullable=False) + content: Mapped[str] = mapped_column(Text, nullable=False, default="") + status: Mapped[str] = mapped_column(String(16), nullable=False) + error_code: Mapped[str | None] = mapped_column(String(64)) + order_index: Mapped[int] = mapped_column(Integer, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=utc_now + ) + + +class MessageSource(Base): + __tablename__ = "message_sources" + __table_args__ = ( + UniqueConstraint("message_id", "chunk_id", name="uq_message_sources_message_chunk"), + Index("ix_message_sources_message_id", "message_id"), + ) + + id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid) + message_id: Mapped[str] = mapped_column( + String(36), ForeignKey("messages.id", ondelete="CASCADE"), nullable=False + ) + document_id: Mapped[str] = mapped_column( + String(36), ForeignKey("documents.id", ondelete="CASCADE"), nullable=False + ) + chunk_id: Mapped[str] = mapped_column( + String(36), ForeignKey("chunks.id", ondelete="CASCADE"), nullable=False + ) + citation_label: Mapped[str] = mapped_column(String(8), nullable=False) + score: Mapped[float] = mapped_column(Float, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=utc_now + ) diff --git a/backend/src/kbqa/chat/repository.py b/backend/src/kbqa/chat/repository.py new file mode 100644 index 0000000..b22b41d --- /dev/null +++ b/backend/src/kbqa/chat/repository.py @@ -0,0 +1,164 @@ +import re +from dataclasses import dataclass + +from sqlalchemy import delete, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from kbqa.chat.models import ChatSession, Message, MessageSource +from kbqa.database import utc_now + + +@dataclass(frozen=True, slots=True) +class MessageSourceInput: + document_id: str + chunk_id: str + citation_label: str + score: float + + +class ChatRepository: + def __init__(self, session: AsyncSession) -> None: + self.session = session + + async def create_session(self, title: str | None) -> ChatSession: + session = ChatSession(title=title.strip() if title and title.strip() else None) + self.session.add(session) + await self.session.commit() + await self.session.refresh(session) + return session + + async def get_session(self, session_id: str) -> ChatSession | None: + return await self.session.get(ChatSession, session_id) + + async def page_sessions(self, offset: int, limit: int) -> tuple[list[ChatSession], int]: + total = await self.session.scalar(select(func.count(ChatSession.id))) + rows = await self.session.scalars( + select(ChatSession).order_by(ChatSession.updated_at.desc()).offset(offset).limit(limit) + ) + return list(rows), int(total or 0) + + async def load_recent_messages(self, session_id: str, limit: int) -> list[Message]: + rows = list( + await self.session.scalars( + select(Message) + .where(Message.session_id == session_id, Message.status == "completed") + .order_by(Message.order_index.desc()) + .limit(limit) + ) + ) + rows.reverse() + return rows + + async def page_messages( + self, + session_id: str, + offset: int, + limit: int, + ) -> tuple[list[Message], int]: + total = await self.session.scalar( + select(func.count(Message.id)).where(Message.session_id == session_id) + ) + rows = await self.session.scalars( + select(Message) + .where(Message.session_id == session_id) + .order_by(Message.order_index) + .offset(offset) + .limit(limit) + ) + return list(rows), int(total or 0) + + async def sources_for_messages(self, message_ids: list[str]) -> dict[str, list[MessageSource]]: + if not message_ids: + return {} + sources = await self.session.scalars( + select(MessageSource) + .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) + return result + + async def begin_turn(self, session_id: str, query: str) -> tuple[Message, Message]: + chat_session = await self.session.get(ChatSession, session_id) + if chat_session is None: + raise LookupError("session not found") + maximum = await self.session.scalar( + select(func.max(Message.order_index)).where(Message.session_id == session_id) + ) + next_order = int(maximum if maximum is not None else -1) + 1 + user = Message( + session_id=session_id, + role="user", + content=query, + status="completed", + order_index=next_order, + ) + assistant = Message( + session_id=session_id, + role="assistant", + content="", + status="generating", + order_index=next_order + 1, + ) + if chat_session.title is None: + chat_session.title = query.strip()[:30] + chat_session.updated_at = utc_now() + self.session.add_all([user, assistant]) + await self.session.commit() + await self.session.refresh(user) + await self.session.refresh(assistant) + return user, assistant + + async def complete_answer( + self, + message_id: str, + content: str, + sources: list[MessageSourceInput], + ) -> None: + message = await self.session.get(Message, message_id) + if message is None: + raise LookupError("message not found") + await self.session.execute( + delete(MessageSource).where(MessageSource.message_id == message_id) + ) + message.content = content + message.status = "completed" + message.error_code = None + self.session.add_all( + [ + MessageSource( + message_id=message_id, + document_id=source.document_id, + chunk_id=source.chunk_id, + citation_label=source.citation_label, + score=source.score, + ) + for source in sources + ] + ) + chat_session = await self.session.get(ChatSession, message.session_id) + if chat_session is not None: + chat_session.updated_at = utc_now() + await self.session.commit() + + async def fail_answer(self, message_id: str, error_code: str) -> None: + message = await self.session.get(Message, message_id) + if message is None: + return + message.status = "failed" + message.error_code = error_code + await self.session.commit() + + async def delete_session(self, session_id: str) -> bool: + chat_session = await self.session.get(ChatSession, session_id) + if chat_session is None: + return False + await self.session.delete(chat_session) + await self.session.commit() + return True + + @staticmethod + def citation_labels(content: str) -> set[str]: + return set(re.findall(r"\[(S\d+)]", content)) diff --git a/backend/src/kbqa/database.py b/backend/src/kbqa/database.py index 3bd3c9e..8d6f727 100644 --- a/backend/src/kbqa/database.py +++ b/backend/src/kbqa/database.py @@ -1,4 +1,5 @@ from collections.abc import AsyncIterator +from datetime import UTC, datetime from fastapi import Request from sqlalchemy import event @@ -16,6 +17,10 @@ class Base(DeclarativeBase): pass +def utc_now() -> datetime: + return datetime.now(UTC) + + class Database: def __init__(self, url: str) -> None: self.engine: AsyncEngine = create_async_engine(url, pool_pre_ping=True) diff --git a/backend/src/kbqa/documents/__init__.py b/backend/src/kbqa/documents/__init__.py new file mode 100644 index 0000000..b642b3b --- /dev/null +++ b/backend/src/kbqa/documents/__init__.py @@ -0,0 +1 @@ +"""Document management feature.""" diff --git a/backend/src/kbqa/documents/models.py b/backend/src/kbqa/documents/models.py new file mode 100644 index 0000000..d00f8fd --- /dev/null +++ b/backend/src/kbqa/documents/models.py @@ -0,0 +1,92 @@ +from datetime import datetime +from uuid import uuid4 + +from sqlalchemy import ( + CheckConstraint, + DateTime, + ForeignKey, + Index, + Integer, + String, + Text, + UniqueConstraint, +) +from sqlalchemy.orm import Mapped, mapped_column + +from kbqa.database import Base, utc_now + + +def new_uuid() -> str: + return str(uuid4()) + + +class Document(Base): + __tablename__ = "documents" + __table_args__ = ( + CheckConstraint( + "status IN ('pending','indexing','indexed','failed','deleting')", + name="ck_documents_status", + ), + Index("ix_documents_status_updated_at", "status", "updated_at"), + ) + + id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid) + filename: Mapped[str] = mapped_column(Text, nullable=False) + stored_path: Mapped[str] = mapped_column(Text, nullable=False) + file_type: Mapped[str] = mapped_column(String(16), nullable=False) + mime_type: Mapped[str] = mapped_column(String(128), nullable=False) + file_size: Mapped[int] = mapped_column(Integer, nullable=False) + sha256: Mapped[str] = mapped_column(String(64), nullable=False) + status: Mapped[str] = mapped_column(String(16), nullable=False, default="pending") + chunk_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + error_code: Mapped[str | None] = mapped_column(String(64)) + error_message: Mapped[str | None] = mapped_column(Text) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=utc_now + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=utc_now, onupdate=utc_now + ) + + +class IndexJob(Base): + __tablename__ = "index_jobs" + __table_args__ = ( + CheckConstraint( + "status IN ('queued','running','succeeded','failed')", + name="ck_index_jobs_status", + ), + ) + + id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid) + document_id: Mapped[str] = mapped_column( + String(36), ForeignKey("documents.id", ondelete="CASCADE"), nullable=False, unique=True + ) + status: Mapped[str] = mapped_column(String(16), nullable=False, default="queued") + attempt_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + last_error: Mapped[str | None] = mapped_column(Text) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=utc_now + ) + started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + + +class Chunk(Base): + __tablename__ = "chunks" + __table_args__ = ( + UniqueConstraint("document_id", "order_index", name="uq_chunks_document_order"), + Index("ix_chunks_document_id", "document_id"), + ) + + id: Mapped[str] = mapped_column(String(36), primary_key=True) + document_id: Mapped[str] = mapped_column( + String(36), ForeignKey("documents.id", ondelete="CASCADE"), nullable=False + ) + order_index: Mapped[int] = mapped_column(Integer, nullable=False) + content: Mapped[str] = mapped_column(Text, nullable=False) + page_number: Mapped[int | None] = mapped_column(Integer) + char_count: Mapped[int] = mapped_column(Integer, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, default=utc_now + ) diff --git a/backend/src/kbqa/documents/repository.py b/backend/src/kbqa/documents/repository.py new file mode 100644 index 0000000..c1add13 --- /dev/null +++ b/backend/src/kbqa/documents/repository.py @@ -0,0 +1,195 @@ +from dataclasses import dataclass + +from sqlalchemy import delete, func, select +from sqlalchemy.ext.asyncio import AsyncSession + +from kbqa.database import utc_now +from kbqa.documents.models import Chunk, Document, IndexJob, new_uuid + + +@dataclass(frozen=True, slots=True) +class StoredUpload: + document_id: str + filename: str + stored_path: str + file_type: str + mime_type: str + file_size: int + sha256: str + + +class DocumentRepository: + def __init__(self, session: AsyncSession) -> None: + self.session = session + + async def create_pending(self, upload: StoredUpload) -> tuple[Document, IndexJob]: + document = Document( + id=upload.document_id, + filename=upload.filename, + stored_path=upload.stored_path, + file_type=upload.file_type, + mime_type=upload.mime_type, + file_size=upload.file_size, + sha256=upload.sha256, + status="pending", + ) + job = IndexJob(id=new_uuid(), document_id=document.id, status="queued") + self.session.add_all([document, job]) + await self.session.commit() + await self.session.refresh(document) + await self.session.refresh(job) + return document, job + + async def page( + self, + status: str | None, + offset: int, + limit: int, + ) -> tuple[list[Document], int]: + filters = [Document.status == status] if status else [] + total = await self.session.scalar(select(func.count(Document.id)).where(*filters)) + result = await self.session.scalars( + select(Document) + .where(*filters) + .order_by(Document.created_at.desc()) + .offset(offset) + .limit(limit) + ) + return list(result), int(total or 0) + + async def get(self, document_id: str) -> Document | None: + return await self.session.get(Document, document_id) + + async def get_job_for_document(self, document_id: str) -> IndexJob | None: + return await self.session.scalar( + select(IndexJob).where(IndexJob.document_id == document_id) + ) + + async def page_chunks( + self, + document_id: str, + offset: int, + limit: int, + ) -> tuple[list[Chunk], int]: + total = await self.session.scalar( + select(func.count(Chunk.id)).where(Chunk.document_id == document_id) + ) + result = await self.session.scalars( + select(Chunk) + .where(Chunk.document_id == document_id) + .order_by(Chunk.order_index) + .offset(offset) + .limit(limit) + ) + return list(result), int(total or 0) + + async def claim_job(self, job_id: str) -> IndexJob | None: + job = await self.session.get(IndexJob, job_id) + if job is None or job.status != "queued": + return None + document = await self.session.get(Document, job.document_id) + if document is None or document.status not in {"pending", "failed"}: + return None + now = utc_now() + job.status = "running" + job.attempt_count += 1 + job.started_at = now + job.finished_at = None + job.last_error = None + document.status = "indexing" + document.error_code = None + document.error_message = None + document.updated_at = now + await self.session.commit() + return job + + async def replace_chunks(self, document_id: str, chunks: list[Chunk]) -> None: + await self.session.execute(delete(Chunk).where(Chunk.document_id == document_id)) + self.session.add_all(chunks) + await self.session.commit() + + async def mark_indexed(self, document_id: str, chunk_count: int) -> None: + document = await self.session.get(Document, document_id) + job = await self.get_job_for_document(document_id) + if document is None or job is None: + return + now = utc_now() + document.status = "indexed" + document.chunk_count = chunk_count + document.updated_at = now + document.error_code = None + document.error_message = None + job.status = "succeeded" + job.finished_at = now + job.last_error = None + await self.session.commit() + + async def mark_failed(self, document_id: str, code: str, message: str) -> None: + document = await self.session.get(Document, document_id) + job = await self.get_job_for_document(document_id) + if document is None or job is None: + return + now = utc_now() + document.status = "failed" + document.error_code = code + document.error_message = message + document.updated_at = now + job.status = "failed" + job.finished_at = now + job.last_error = message + await self.session.commit() + + async def requeue_unfinished(self) -> list[str]: + jobs = list( + await self.session.scalars( + select(IndexJob).where(IndexJob.status.in_(("queued", "running"))) + ) + ) + now = utc_now() + for job in jobs: + job.status = "queued" + job.started_at = None + job.finished_at = None + document = await self.session.get(Document, job.document_id) + if document is not None: + document.status = "pending" + document.updated_at = now + await self.session.commit() + return [job.id for job in jobs] + + async def requeue_failed(self, document_id: str) -> IndexJob | None: + document = await self.session.get(Document, document_id) + job = await self.get_job_for_document(document_id) + if document is None or job is None or document.status != "failed": + return None + document.status = "pending" + document.error_code = None + document.error_message = None + document.updated_at = utc_now() + job.status = "queued" + job.last_error = None + job.finished_at = None + await self.session.commit() + return job + + async def mark_deleting(self, document_id: str) -> Document | None: + document = await self.session.get(Document, document_id) + if document is None: + return None + document.status = "deleting" + document.updated_at = utc_now() + await self.session.commit() + return document + + async def deleting_documents(self) -> list[Document]: + return list( + await self.session.scalars(select(Document).where(Document.status == "deleting")) + ) + + async def delete_document(self, document_id: str) -> bool: + document = await self.session.get(Document, document_id) + if document is None: + return False + await self.session.delete(document) + await self.session.commit() + return True diff --git a/backend/src/kbqa/models.py b/backend/src/kbqa/models.py new file mode 100644 index 0000000..911d8b2 --- /dev/null +++ b/backend/src/kbqa/models.py @@ -0,0 +1,6 @@ +"""Import all ORM models so Alembic can discover metadata.""" + +from kbqa.chat.models import ChatSession, Message, MessageSource +from kbqa.documents.models import Chunk, Document, IndexJob + +__all__ = ["ChatSession", "Chunk", "Document", "IndexJob", "Message", "MessageSource"]