convo-cli 0.1.0__py3-none-any.whl

convo/db/chunks.py

@@ -0,0 +1,132 @@
+import math
+import re
+from uuid import UUID
+
+from sqlalchemy import delete, func, select
+from sqlalchemy.orm import Session
+
+from convo.core.config import EMBEDDING_DIMENSIONS
+from convo.core.contracts import ChunkData, Hit
+from convo.db.models import ChunkRow, DocumentRow
+
+
+def keyword_query(question: str) -> str:
+    """Natural questions use optional content terms rather than requiring every word."""
+    ignored = set(
+        "a an the what when where why who which how is are was were do does did can could "
+        "would should will you me i we it this that these those file document please tell "
+        "explain describe summarize summarise about in on of for to and or".split()
+    )
+    terms = dict.fromkeys(
+        term.casefold()
+        for term in re.findall(r"\w+(?:[-.]\w+)*", question)
+        if term.casefold() not in ignored
+    )
+    return " OR ".join(list(terms)[:32]) or question
+
+
+def replace_chunks(
+    session: Session,
+    document_id: UUID,
+    chunks: list[ChunkData],
+    vectors: list[list[float]],
+) -> None:
+    if not chunks or len(chunks) != len(vectors):
+        raise ValueError("Chunks and vectors do not match")
+    if any(len(vector) != EMBEDDING_DIMENSIONS for vector in vectors):
+        raise ValueError("Incorrect embedding dimensions")
+    if any(not chunk.text.strip() or not chunk.locator for chunk in chunks):
+        raise ValueError("Chunks need text and a source locator")
+    if any(not math.isfinite(value) for vector in vectors for value in vector):
+        raise ValueError("Embeddings must contain finite values")
+    session.execute(delete(ChunkRow).where(ChunkRow.document_id == document_id))
+    session.add_all(
+        [
+            ChunkRow(
+                document_id=document_id,
+                ordinal=i,
+                text=chunk.text,
+                locator=chunk.locator,
+                embedding=vector,
+            )
+            for i, (chunk, vector) in enumerate(zip(chunks, vectors, strict=True))
+        ]
+    )
+
+
+def reciprocal_rank_fusion(rankings: list[list[Hit]], limit: int = 20) -> list[Hit]:
+    """Fuse lexical and semantic ranks without comparing incompatible raw scores."""
+    scores: dict[object, float] = {}
+    hits: dict[object, Hit] = {}
+    for ranking in rankings:
+        seen = set()
+        for rank, hit in enumerate(ranking, 1):
+            key = hit.chunk_id if hit.chunk_id is not None else (hit.text, hit.locator)
+            if key in seen:
+                continue
+            seen.add(key)
+            hits[key] = hit
+            scores[key] = scores.get(key, 0.0) + 1.0 / (60 + rank)
+    # Stable insertion order breaks equal scores deterministically.
+    return [hits[key] for key in sorted(scores, key=lambda key: -scores[key])[:limit]]
+
+
+def retrieve(
+    session: Session,
+    document_id: UUID,
+    signature: str,
+    query: list[float],
+    question: str,
+    limit: int = 20,
+    *,
+    overview: bool = False,
+) -> list[Hit]:
+    document = session.get(DocumentRow, document_id)
+    if document is None or document.index_signature != signature:
+        raise ValueError("Missing or outdated index; select the file again")
+    if len(query) != EMBEDDING_DIMENSIONS or any(not math.isfinite(value) for value in query):
+        raise ValueError("Invalid query embedding")
+    if not 1 <= limit <= 50 or not question.strip():
+        raise ValueError("Use a non-empty query and a candidate limit between 1 and 50")
+    distance = ChunkRow.embedding.cosine_distance(query).label("distance")
+    columns = (ChunkRow.id, ChunkRow.text, ChunkRow.locator, distance)
+    if overview:
+        # Broad summaries need coverage, not a passage ranker's QA relevance cutoff.
+        rows = list(
+            session.execute(
+                select(*columns)
+                .where(ChunkRow.document_id == document_id)
+                .order_by(ChunkRow.ordinal)
+            )
+        )
+        size = min(limit, len(rows))
+        indices = [round(index * (len(rows) - 1) / max(1, size - 1)) for index in range(size)]
+        return [
+            Hit(rows[index].text, rows[index].locator, float(rows[index].distance), rows[index].id)
+            for index in indices
+        ]
+    semantic = (
+        select(*columns)
+        .where(ChunkRow.document_id == document_id)
+        .order_by(distance, ChunkRow.ordinal)
+        .limit(limit)
+    )
+    lexical_query = func.websearch_to_tsquery("english", keyword_query(question))
+    lexical_rank = func.ts_rank_cd(ChunkRow.search_vector, lexical_query)
+    lexical = (
+        select(*columns)
+        .where(
+            ChunkRow.document_id == document_id,
+            ChunkRow.search_vector.op("@@")(lexical_query),
+        )
+        .order_by(lexical_rank.desc(), ChunkRow.ordinal)
+        .limit(limit)
+    )
+    rankings = [
+        [
+            Hit(text, locator, float(distance), identity)
+            for identity, text, locator, distance in session.execute(statement)
+        ]
+        for statement in (semantic, lexical)
+    ]
+    return reciprocal_rank_fusion(rankings, limit)

convo/db/conversations.py

@@ -0,0 +1,167 @@
+"""SQLAlchemy storage for sessions and atomic, completed conversation turns."""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from dataclasses import asdict
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal, cast
+from uuid import UUID, uuid4
+
+from sqlalchemy import func, select
+from sqlalchemy.exc import SQLAlchemyError
+
+from convo.core.config import Settings
+from convo.core.contracts import (
+    Answer,
+    Document,
+    Message,
+    SessionSnapshot,
+    SessionSummary,
+    Source,
+)
+from convo.core.exceptions import BackendNotConnected
+from convo.db.models import ChatSessionRow, DocumentRow, MessageRow
+from convo.db.session import make_session_factory
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Engine
+    from sqlalchemy.orm import Session
+
+    from convo.db.session import SessionFactory
+
+
+class Conversations:
+    """Created cheaply at startup; connect only when a session is requested."""
+
+    def __init__(self, settings: Settings) -> None:
+        self.settings = settings
+        self.engine: Engine | None = None
+        self.sessions: SessionFactory | None = None
+
+    @contextmanager
+    def transaction(self) -> Generator[Session]:
+        try:
+            if self.sessions is None:
+                self.engine, self.sessions = make_session_factory(self.settings)
+            with self.sessions.begin() as session:
+                yield session
+        except SQLAlchemyError as error:
+            raise BackendNotConnected(
+                "Could not access saved conversations. Check PostgreSQL/Docker/Podman "
+                "and run `convo --setup`."
+            ) from error
+
+    def create(self, document: Document | None = None) -> SessionSnapshot:
+        session_id = uuid4()
+        title = document.path.name[:200] if document else "New conversation"
+        with self.transaction() as session:
+            if document and session.get(DocumentRow, UUID(document.id)) is None:
+                raise BackendNotConnected("The document snapshot is missing. Select it with /file.")
+            session.add(
+                ChatSessionRow(
+                    id=session_id,
+                    document_id=UUID(document.id) if document else None,
+                    title=title,
+                )
+            )
+        return SessionSnapshot(str(session_id), title, document)
+
+    def list(self) -> tuple[SessionSummary, ...]:
+        with self.transaction() as session:
+            statement = (
+                select(ChatSessionRow, DocumentRow.source_path, func.count(MessageRow.id))
+                .outerjoin(DocumentRow, DocumentRow.id == ChatSessionRow.document_id)
+                .outerjoin(MessageRow, MessageRow.session_id == ChatSessionRow.id)
+                .group_by(ChatSessionRow.id, DocumentRow.source_path)
+                .order_by(ChatSessionRow.updated_at.desc(), ChatSessionRow.id)
+            )
+            return tuple(
+                SessionSummary(
+                    str(row.id), row.title, Path(path).name if path else None, row.updated_at, count
+                )
+                for row, path, count in session.execute(statement)
+            )
+
+    def load(self, session_id: str) -> SessionSnapshot:
+        try:
+            identity = UUID(session_id)
+        except ValueError as error:
+            raise BackendNotConnected("Choose a valid session with /sessions.") from error
+        with self.transaction() as session:
+            row = session.get(ChatSessionRow, identity)
+            if row is None:
+                raise BackendNotConnected("This conversation no longer exists. Use /sessions.")
+            document = None
+            if row.document_id:
+                stored = session.get(DocumentRow, row.document_id)
+                if stored is None:
+                    raise BackendNotConnected("The saved document snapshot is missing.")
+                document = Document(str(stored.id), Path(stored.source_path))
+            messages = session.scalars(
+                select(MessageRow)
+                .where(MessageRow.session_id == identity)
+                .order_by(MessageRow.ordinal)
+            ).all()
+            history = tuple(
+                Message(
+                    cast(Literal["user", "assistant"], message.role),
+                    message.content,
+                    tuple(Source(**source) for source in message.sources),
+                    message.structured,
+                )
+                for message in messages
+            )
+            return SessionSnapshot(str(row.id), row.title, document, history)
+
+    def append_turn(
+        self, session_id: str, question: str, answer: Answer, expected_messages: int
+    ) -> None:
+        """Commit both messages together; never save an incomplete or stale answer."""
+        identity = UUID(session_id)
+        with self.transaction() as session:
+            row = session.get(ChatSessionRow, identity, with_for_update=True)
+            if row is None:
+                raise BackendNotConnected("The conversation no longer exists. Use /sessions.")
+            count = (
+                session.scalar(
+                    select(func.count())
+                    .select_from(MessageRow)
+                    .where(MessageRow.session_id == identity)
+                )
+                or 0
+            )
+            if count != expected_messages:
+                raise BackendNotConnected(
+                    "This conversation changed in another process. Reopen it with /sessions and retry."
+                )
+            session.add_all(
+                [
+                    MessageRow(
+                        session_id=identity,
+                        ordinal=count,
+                        role="user",
+                        content=question,
+                        sources=[],
+                        structured=None,
+                    ),
+                    MessageRow(
+                        session_id=identity,
+                        ordinal=count + 1,
+                        role="assistant",
+                        content=answer.text,
+                        sources=[asdict(source) for source in answer.sources],
+                        structured=answer.structured,
+                    ),
+                ]
+            )
+            if count == 0:
+                row.title = question.strip()[:200]
+            row.updated_at = datetime.now(UTC)
+
+    def close(self) -> None:
+        if self.engine is not None:
+            self.engine.dispose()
+            self.engine, self.sessions = None, None

convo/db/documents.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+from uuid import UUID
+
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.orm import Session
+
+from convo.db.models import DocumentRow
+
+
+def upsert_document(
+    session: Session,
+    document_id: UUID,
+    path: Path,
+    content_hash: str,
+    signature: str,
+) -> None:
+    values = {
+        "id": document_id,
+        "source_path": str(path),
+        "content_hash": content_hash,
+        "index_signature": signature,
+    }
+    statement = insert(DocumentRow).values(**values)
+    session.execute(
+        statement.on_conflict_do_update(
+            index_elements=[DocumentRow.id],
+            set_={
+                "source_path": str(path),
+                "content_hash": content_hash,
+                "index_signature": signature,
+            },
+        )
+    )
+    # The upsert locks the document row until commit, serializing replacement.

convo/db/models.py

@@ -0,0 +1,81 @@
+from datetime import datetime
+from uuid import UUID, uuid4
+
+from pgvector.sqlalchemy import VECTOR
+from sqlalchemy import (
+    Computed,
+    DateTime,
+    ForeignKey,
+    Index,
+    Integer,
+    String,
+    Text,
+    UniqueConstraint,
+    func,
+)
+from sqlalchemy.dialects.postgresql import JSONB, TSVECTOR
+from sqlalchemy.dialects.postgresql import UUID as PG_UUID
+from sqlalchemy.orm import Mapped, mapped_column
+
+from convo.core.config import EMBEDDING_DIMENSIONS
+from convo.db.base import Base
+
+
+class DocumentRow(Base):
+    __tablename__ = "documents"
+
+    id: Mapped[UUID] = mapped_column(PG_UUID(as_uuid=True), primary_key=True)
+    source_path: Mapped[str] = mapped_column(Text, nullable=False)
+    content_hash: Mapped[str] = mapped_column(String(64), nullable=False)
+    index_signature: Mapped[str] = mapped_column(Text, nullable=False)
+
+
+class ChunkRow(Base):
+    __tablename__ = "chunks"
+    __table_args__ = (
+        UniqueConstraint("document_id", "ordinal"),
+        Index("ix_chunks_search_vector", "search_vector", postgresql_using="gin"),
+    )
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    document_id: Mapped[UUID] = mapped_column(
+        PG_UUID(as_uuid=True),
+        ForeignKey("documents.id", ondelete="CASCADE"),
+        nullable=False,
+    )
+    ordinal: Mapped[int] = mapped_column(Integer, nullable=False)
+    text: Mapped[str] = mapped_column(Text, nullable=False)
+    locator: Mapped[str] = mapped_column(Text, nullable=False)
+    embedding: Mapped[list[float]] = mapped_column(VECTOR(EMBEDDING_DIMENSIONS), nullable=False)
+    search_vector: Mapped[str] = mapped_column(
+        TSVECTOR, Computed("to_tsvector('english'::regconfig, text)", persisted=True)
+    )
+
+
+class ChatSessionRow(Base):
+    __tablename__ = "chat_sessions"
+    __table_args__ = (Index("ix_chat_sessions_updated_at", "updated_at"),)
+
+    id: Mapped[UUID] = mapped_column(PG_UUID(as_uuid=True), primary_key=True, default=uuid4)
+    document_id: Mapped[UUID | None] = mapped_column(
+        PG_UUID(as_uuid=True), ForeignKey("documents.id", ondelete="RESTRICT"), nullable=True
+    )
+    title: Mapped[str] = mapped_column(String(200), nullable=False)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
+
+
+class MessageRow(Base):
+    __tablename__ = "messages"
+    __table_args__ = (UniqueConstraint("session_id", "ordinal"),)
+
+    id: Mapped[UUID] = mapped_column(PG_UUID(as_uuid=True), primary_key=True, default=uuid4)
+    session_id: Mapped[UUID] = mapped_column(
+        PG_UUID(as_uuid=True), ForeignKey("chat_sessions.id", ondelete="CASCADE"), nullable=False
+    )
+    ordinal: Mapped[int] = mapped_column(Integer, nullable=False)
+    role: Mapped[str] = mapped_column(String(16), nullable=False)
+    content: Mapped[str] = mapped_column(Text, nullable=False)
+    sources: Mapped[list[dict[str, str]]] = mapped_column(JSONB, nullable=False, default=list)
+    structured: Mapped[dict[str, object] | None] = mapped_column(JSONB, nullable=True)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

convo/db/session.py

@@ -0,0 +1,27 @@
+"""Create the PostgreSQL engine and SQLAlchemy session factory.
+
+Each database operation gets its own short-lived session.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session, sessionmaker
+
+from convo.core.config import Settings
+
+if TYPE_CHECKING:
+    SessionFactory = sessionmaker[Session]
+
+
+def make_session_factory(settings: Settings) -> tuple[Engine, SessionFactory]:
+    engine = create_engine(
+        settings.database_url,
+        pool_pre_ping=True,
+        echo=False,
+        connect_args={"connect_timeout": 10},
+    )
+    return engine, sessionmaker(engine, expire_on_commit=False)