convo-cli 0.1.0__py3-none-any.whl

convo/migrations/versions/a42e9d8f0b31_hybrid_search_and_chat_sessions.py

@@ -0,0 +1,66 @@
+"""Add full-text search and durable conversations without replacing existing documents.
+
+Revision ID: a42e9d8f0b31
+Revises: 5c3b75f9c5c1
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+revision = "a42e9d8f0b31"
+down_revision = "5c3b75f9c5c1"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "chunks",
+        sa.Column(
+            "search_vector",
+            postgresql.TSVECTOR(),
+            sa.Computed("to_tsvector('english'::regconfig, text)", persisted=True),
+            nullable=False,
+        ),
+    )
+    op.create_index("ix_chunks_search_vector", "chunks", ["search_vector"], postgresql_using="gin")
+    op.create_table(
+        "chat_sessions",
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("document_id", sa.UUID(), nullable=True),
+        sa.Column("title", sa.String(200), nullable=False),
+        sa.Column(
+            "created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
+        ),
+        sa.Column(
+            "updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
+        ),
+        sa.ForeignKeyConstraint(["document_id"], ["documents.id"], ondelete="RESTRICT"),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index("ix_chat_sessions_updated_at", "chat_sessions", ["updated_at"])
+    op.create_table(
+        "messages",
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("session_id", sa.UUID(), nullable=False),
+        sa.Column("ordinal", sa.Integer(), nullable=False),
+        sa.Column("role", sa.String(16), nullable=False),
+        sa.Column("content", sa.Text(), nullable=False),
+        sa.Column("sources", postgresql.JSONB(), nullable=False),
+        sa.Column("structured", postgresql.JSONB(), nullable=True),
+        sa.Column(
+            "created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
+        ),
+        sa.ForeignKeyConstraint(["session_id"], ["chat_sessions.id"], ondelete="CASCADE"),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("session_id", "ordinal"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("messages")
+    op.drop_index("ix_chat_sessions_updated_at", table_name="chat_sessions")
+    op.drop_table("chat_sessions")
+    op.drop_index("ix_chunks_search_vector", table_name="chunks")
+    op.drop_column("chunks", "search_vector")

convo/pipeline/chat.py

@@ -0,0 +1,129 @@
+"""One chat flow: route, retrieve/rerank, generate a validated answer, persist the turn."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from convo.core.config import Settings
+from convo.core.contracts import (
+    Answer,
+    ChatModel,
+    Document,
+    GenerationSelection,
+    Message,
+    ReasoningEffort,
+    SessionSnapshot,
+    SessionSummary,
+)
+from convo.llm.auth import ChatGPTAuth
+from convo.llm.client import ChatGPTProvider
+from convo.llm.credentials import KeyringCredentialStore
+from convo.pipeline.routing import conversation_reply, search_query
+from convo.pipeline.worker import IndexWorker
+
+if TYPE_CHECKING:
+    from convo.db.conversations import Conversations
+
+
+class RagBackend:
+    def __init__(self, settings: Settings) -> None:
+        self.settings = settings
+        self.documents = IndexWorker(settings)
+        self.auth = ChatGPTAuth(KeyringCredentialStore())
+        self.chat = ChatGPTProvider(self.auth)
+        self._conversations: Conversations | None = None
+        self._session_lock = asyncio.Lock()
+
+    def _store(self) -> Conversations:
+        if self._conversations is None:
+            from convo.db.conversations import Conversations
+
+            self._conversations = Conversations(self.settings)
+        return self._conversations
+
+    async def prepare(self, path: Path) -> Document:
+        return await self.documents.prepare(path)
+
+    async def prepare_with_progress(self, path: Path, progress: Callable[[str], None]) -> Document:
+        return await self.documents.prepare(path, progress)
+
+    async def _answer(
+        self,
+        document: Document | None,
+        question: str,
+        history: tuple[Message, ...],
+        progress: Callable[[str], None] | None = None,
+    ) -> Answer:
+        if not question.strip() or len(question) > 4_000:
+            from convo.core.exceptions import BackendNotConnected
+
+            raise BackendNotConnected("Use a question between 1 and 4,000 characters.")
+        conversational = conversation_reply(question)
+        if conversational:
+            return Answer(conversational)
+        if document is None:
+            return Answer("Choose a file with /file, then ask me a question about it.")
+        query = search_query(question, history)
+        hits = await self.documents.search(document, query, progress)
+        if progress:
+            progress("writing a grounded answer…")
+        return await self.chat.answer(document.path.name, question, history, hits)
+
+    async def answer(
+        self, document: Document, question: str, history: tuple[Message, ...]
+    ) -> Answer:
+        return await self._answer(document, question, history)
+
+    async def new_session(self, document: Document | None = None) -> SessionSnapshot:
+        async with self._session_lock:
+            return await asyncio.to_thread(lambda: self._store().create(document))
+
+    async def list_sessions(self) -> tuple[SessionSummary, ...]:
+        async with self._session_lock:
+            return await asyncio.to_thread(lambda: self._store().list())
+
+    async def load_session(self, session_id: str) -> SessionSnapshot:
+        async with self._session_lock:
+            return await asyncio.to_thread(lambda: self._store().load(session_id))
+
+    async def answer_session(
+        self, session_id: str, question: str, progress: Callable[[str], None] | None = None
+    ) -> Answer:
+        # No database transaction stays open during embedding or the LLM request.
+        snapshot = await self.load_session(session_id)
+        answer = await self._answer(snapshot.document, question, snapshot.history, progress)
+        async with self._session_lock:
+            await asyncio.to_thread(
+                lambda: self._store().append_turn(
+                    session_id, question, answer, expected_messages=len(snapshot.history)
+                )
+            )
+        return answer
+
+    async def login_openai(self) -> str:
+        account = await self.auth.login()
+        self.chat.reset_catalog()
+        return account
+
+    def generation_selection(self) -> GenerationSelection:
+        return self.chat.generation_selection()
+
+    async def list_models(self) -> tuple[ChatModel, ...]:
+        return await self.chat.list_models()
+
+    async def select_model(self, model_id: str) -> GenerationSelection:
+        return await self.chat.select_model(model_id)
+
+    async def cycle_reasoning(self) -> GenerationSelection:
+        return await self.chat.cycle_reasoning()
+
+    async def select_reasoning(self, effort: ReasoningEffort | None) -> GenerationSelection:
+        return await self.chat.select_reasoning(effort)
+
+    def close(self) -> None:
+        self.documents.close()
+        if self._conversations is not None:
+            self._conversations.close()

convo/pipeline/chunking.py

@@ -0,0 +1,42 @@
+"""Convert reader documents into the application's small chunk contract."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from llama_index.core.schema import Document
+
+from convo.core.contracts import ChunkData
+
+MAX_CHUNKS = 2_000
+
+
+def split_documents(
+    documents: list[Document], tokenizer: Callable[[str], list[int]]
+) -> list[ChunkData]:
+    """Split reader output using the embedding model's tokenizer and chunk size."""
+    from llama_index.core.node_parser import SentenceSplitter
+    from llama_index.core.schema import MetadataMode
+
+    splitter = SentenceSplitter(
+        chunk_size=512,
+        chunk_overlap=64,
+        tokenizer=tokenizer,
+        include_metadata=False,
+        include_prev_next_rel=False,
+    )
+    chunks: list[ChunkData] = []
+    for section, document in enumerate(documents, 1):
+        page = document.metadata.get("page_label")
+        locator = f"page {page}" if page is not None else f"reader section {section}"
+        for node in splitter.get_nodes_from_documents([document]):
+            text = node.get_content(metadata_mode=MetadataMode.NONE).strip()
+            if text:
+                chunks.append(ChunkData(text=text, locator=locator))
+                if len(chunks) > MAX_CHUNKS:
+                    raise ValueError("File produced too many chunks")
+    if not chunks:
+        raise ValueError("File did not produce any chunks")
+    return chunks

convo/pipeline/embeddings.py

@@ -0,0 +1,83 @@
+"""Generate Qwen3 vectors locally with sentence-transformers."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
+from convo.core.config import EMBEDDING_DIMENSIONS
+
+
+class Embedder:
+    def __init__(self, model_name: str) -> None:
+        self.name = model_name
+        self._model: SentenceTransformer | None = None
+
+    @property
+    def model(self) -> SentenceTransformer:
+        if self._model is None:
+            from sentence_transformers import SentenceTransformer
+
+            # A cached model should not make HTTP checks on every new CLI session.
+            model = None
+            try:
+                model = SentenceTransformer(self.name, local_files_only=True)
+            except (OSError, TypeError):
+                # A partial ST cache can return an empty pooling config and raise
+                # TypeError instead of OSError. Let the Hub fill missing assets.
+                pass
+            # Leave the exception scope first so a partial model and its traceback
+            # can be released before allocating the replacement weights.
+            if model is None:
+                model = SentenceTransformer(self.name)
+            self._model = model
+            self._model.max_seq_length = 2048
+        return self._model
+
+    @property
+    def signature(self) -> str:
+        # Bump v1 after changing reader behavior, tokenizer, prompts or model revision.
+        return f"{self.name}:qwen-tokenizer:512tokens:64overlap:normalized:2048max:v1"
+
+    def tokenize(self, text: str) -> list[int]:
+        return list(self.model.tokenizer.encode(text, add_special_tokens=False))
+
+    def documents(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            raise ValueError("Cannot embed an empty chunk list")
+        vectors = np.asarray(
+            self.model.encode(
+                texts,
+                batch_size=8,
+                normalize_embeddings=True,
+                convert_to_numpy=True,
+                show_progress_bar=False,
+            ),
+            dtype=np.float32,
+        )
+        if vectors.shape != (len(texts), EMBEDDING_DIMENSIONS):
+            raise ValueError("Unexpected embedding dimensions")
+        if not np.isfinite(vectors).all():
+            raise ValueError("Invalid embeddings")
+        return vectors.tolist()
+
+    def query(self, question: str) -> list[float]:
+        if not question.strip():
+            raise ValueError("Cannot embed an empty question")
+        vector = np.asarray(
+            self.model.encode(
+                [question],
+                prompt_name="query",
+                normalize_embeddings=True,
+                convert_to_numpy=True,
+                show_progress_bar=False,
+            )[0],
+            dtype=np.float32,
+        )
+        if vector.shape != (EMBEDDING_DIMENSIONS,) or not np.isfinite(vector).all():
+            raise ValueError("Unexpected query embedding")
+        return vector.tolist()

convo/pipeline/index.py

@@ -0,0 +1,160 @@
+"""Own file preparation and document-scoped pgvector search."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import tempfile
+from collections.abc import Callable
+from pathlib import Path
+from typing import TypeVar
+from uuid import NAMESPACE_URL, UUID, uuid5
+
+from sqlalchemy import text
+from sqlalchemy.exc import OperationalError, SQLAlchemyError
+
+from convo.core.config import Settings
+from convo.core.contracts import Document, Hit
+from convo.core.exceptions import BackendNotConnected
+from convo.db.chunks import replace_chunks, retrieve
+from convo.db.documents import upsert_document
+from convo.db.session import make_session_factory
+from convo.pipeline.chunking import split_documents
+from convo.pipeline.embeddings import Embedder
+from convo.pipeline.ingestion import MAX_BYTES, read_file
+from convo.pipeline.reranking import Reranker
+from convo.pipeline.routing import is_overview
+
+_Result = TypeVar("_Result")
+
+
+class DocumentIndex:
+    """Coordinate parsing, embedding, PostgreSQL writes, and scoped retrieval."""
+
+    def __init__(self, settings: Settings) -> None:
+        self.engine, self.sessions = make_session_factory(settings)
+        self.embedder = Embedder(settings.embedding_model)
+        self.reranker = Reranker(settings.reranker_model, settings.rerank_min_score)
+        self.candidate_limit = settings.retrieval_candidates
+        self.lock = asyncio.Lock()
+        self._closed = False
+        self.progress: Callable[[str], None] = lambda _stage: None
+
+    def _check_database(self) -> None:
+        """Fail before parsing or downloading weights when storage isn't ready."""
+        try:
+            with self.engine.connect() as connection:
+                connection.execute(text("SELECT id, embedding, search_vector FROM chunks LIMIT 0"))
+        except OperationalError as error:
+            raise BackendNotConnected(
+                "Cannot connect to PostgreSQL. Check DATABASE_URL and start your Podman "
+                "machine/container (`podman start convo-postgres`), then select the file again."
+            ) from error
+        except SQLAlchemyError as error:
+            raise BackendNotConnected(
+                "The RAG database is not ready. Run `convo --setup`, then select the file again."
+            ) from error
+
+    def _prepare(self, path: Path) -> Document:
+        self.progress("checking PostgreSQL…")
+        self._check_database()
+        path = path.resolve(strict=True)
+        before = path.stat()
+        if before.st_size > MAX_BYTES:
+            raise ValueError("File exceeds 10 MB")
+        with path.open("rb") as stream:
+            raw = stream.read(MAX_BYTES + 1)
+        if len(raw) > MAX_BYTES:
+            raise ValueError("File grew past the size limit")
+        content_hash = hashlib.sha256(raw).hexdigest()
+
+        # The reader parses the exact bytes that are hashed, even if the source changes later.
+        self.progress("reading file…")
+        with tempfile.TemporaryDirectory(prefix="convo-reader-") as temporary_directory:
+            snapshot = Path(temporary_directory) / path.name
+            snapshot.write_bytes(raw)
+            documents = read_file(snapshot)
+        self.progress("loading Qwen3 (first use downloads model weights)…")
+        # Loading here gives the UI a distinct stage before tokenization starts.
+        self.embedder.model
+        self.progress("splitting text…")
+        chunks = split_documents(documents, tokenizer=self.embedder.tokenize)
+        self.progress(f"embedding {len(chunks)} chunks…")
+        vectors = self.embedder.documents([chunk.text for chunk in chunks])
+
+        after = path.stat()
+        with path.open("rb") as stream:
+            current_raw = stream.read(MAX_BYTES + 1)
+        current_hash = hashlib.sha256(current_raw).hexdigest()
+        if (
+            len(current_raw) > MAX_BYTES
+            or current_hash != content_hash
+            or (before.st_mtime_ns, before.st_size) != (after.st_mtime_ns, after.st_size)
+        ):
+            raise ValueError("File changed during indexing; select it again")
+
+        signature = self.embedder.signature
+        # Snapshot identity includes path, exact content, and reader/embedding recipe.
+        document_id = uuid5(NAMESPACE_URL, f"{path}\n{content_hash}\n{signature}")
+        # Complete parsing/model work before opening the transaction; only row writes are atomic.
+        self.progress("saving index to PostgreSQL…")
+        with self.sessions.begin() as session:
+            upsert_document(session, document_id, path, content_hash, signature)
+            replace_chunks(session, document_id, chunks, vectors)
+        return Document(id=str(document_id), path=path)
+
+    async def prepare(self, path: Path) -> Document:
+        self._ensure_open()
+        async with self.lock:
+            return await self._run_worker(self._prepare, path)
+
+    def _retrieve(self, document: Document, question: str) -> list[Hit]:
+        self.progress("embedding the question…")
+        vector = self.embedder.query(question)
+        overview = is_overview(question)
+        self.progress(
+            "sampling sections for an overview…" if overview else "searching vectors and keywords…"
+        )
+        with self.sessions() as session:
+            candidates = retrieve(
+                session,
+                UUID(document.id),
+                self.embedder.signature,
+                vector,
+                question,
+                limit=5 if overview else self.candidate_limit,
+                overview=overview,
+            )
+        # The database session is closed before model loading/reranking.
+        self.progress("reranking candidates (first use downloads a small model)…")
+        return self.reranker.rerank(question, candidates, filter_weak=not overview)
+
+    async def search(self, document: Document, question: str) -> list[Hit]:
+        self._ensure_open()
+        if not question.strip() or len(question) > 4_000:
+            raise ValueError("Use a question between 1 and 4,000 characters")
+        async with self.lock:
+            return await self._run_worker(self._retrieve, document, question)
+
+    async def _run_worker(self, function: Callable[..., _Result], *args: object) -> _Result:
+        """Keep a worker under the lock until its thread has really stopped."""
+        task = asyncio.create_task(asyncio.to_thread(function, *args))
+        try:
+            return await asyncio.shield(task)
+        except asyncio.CancelledError:
+            # Cancelling asyncio.to_thread cannot stop its thread. Wait here so a later
+            # operation cannot use the embedder or engine concurrently with that thread.
+            try:
+                await asyncio.shield(task)
+            except Exception:
+                pass
+            raise
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("Document index is closed")
+
+    def close(self) -> None:
+        if not self._closed:
+            self.engine.dispose()
+            self._closed = True

convo/pipeline/ingestion.py

@@ -0,0 +1,46 @@
+"""Read one selected file with LlamaIndex's file readers only."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+from zipfile import ZipFile
+
+if TYPE_CHECKING:
+    from llama_index.core.schema import Document
+
+SUPPORTED = {".txt", ".md", ".pdf", ".docx", ".csv", ".json", ".py", ".rst", ".log"}
+MAX_BYTES = 10 * 1024 * 1024
+MAX_EXTRACTED_CHARACTERS = 500_000
+MAX_DOCX_EXPANDED_BYTES = 50 * 1024 * 1024
+
+
+def read_file(path: Path) -> list[Document]:
+    """Load text and reader metadata for one file; never scan its directory."""
+    if not path.is_file() or path.suffix.lower() not in SUPPORTED:
+        raise ValueError("Choose a supported file")
+    if path.stat().st_size > MAX_BYTES:
+        raise ValueError("File exceeds 10 MB")
+    if path.suffix.lower() == ".docx":
+        with ZipFile(path) as archive:
+            if sum(item.file_size for item in archive.infolist()) > MAX_DOCX_EXPANDED_BYTES:
+                raise ValueError("Expanded DOCX is too large")
+
+    from llama_index.core import SimpleDirectoryReader
+    from llama_index.readers.file import PDFReader
+
+    documents = SimpleDirectoryReader(
+        input_files=[str(path)],
+        encoding="utf-8",
+        errors="strict",
+        raise_on_error=True,
+        # Keep PDF pages separate, preserving page_label metadata for sources.
+        file_extractor={".pdf": PDFReader(return_full_document=False)},
+    ).load_data()
+    documents = [document for document in documents if document.text.strip()]
+    if (
+        not documents
+        or sum(len(document.text) for document in documents) > MAX_EXTRACTED_CHARACTERS
+    ):
+        raise ValueError("Empty file or too much extracted text")
+    return documents

convo/pipeline/reranking.py

@@ -0,0 +1,63 @@
+"""Rerank the hybrid candidate pool with a small, local cross-encoder."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from convo.core.contracts import Hit
+
+if TYPE_CHECKING:
+    from sentence_transformers import CrossEncoder
+
+
+class Reranker:
+    def __init__(self, name: str, min_score: float = -3.0) -> None:
+        self.name, self.min_score = name, min_score
+        self._model: CrossEncoder | None = None
+
+    @property
+    def model(self) -> CrossEncoder:
+        if self._model is None:
+            from sentence_transformers import CrossEncoder
+            from torch.nn import Identity
+
+            model = None
+            try:
+                model = CrossEncoder(
+                    self.name, local_files_only=True, max_length=512, activation_fn=Identity()
+                )
+            except OSError:
+                pass
+            if model is None:
+                model = CrossEncoder(self.name, max_length=512, activation_fn=Identity())
+            self._model = model
+        return self._model
+
+    def rerank(
+        self, question: str, hits: list[Hit], limit: int = 5, *, filter_weak: bool = True
+    ) -> list[Hit]:
+        if not hits:
+            return []
+        scores = np.asarray(
+            self.model.predict(
+                [(question, hit.text) for hit in hits],
+                batch_size=8,
+                show_progress_bar=False,
+                convert_to_numpy=True,
+            ),
+            dtype=np.float32,
+        ).reshape(-1)
+        if len(scores) != len(hits) or not np.isfinite(scores).all():
+            raise ValueError("Reranker returned invalid scores")
+        ranked = sorted(
+            zip(hits, scores, strict=True), key=lambda item: float(item[1]), reverse=True
+        )
+        # Raw model logits are ranking scores, not calibrated probabilities.
+        return [
+            replace(hit, score=float(score))
+            for hit, score in ranked
+            if not filter_weak or float(score) >= self.min_score
+        ][:limit]

convo/pipeline/routing.py

@@ -0,0 +1,64 @@
+"""Conservative conversational routing and history-aware follow-up search queries."""
+
+import re
+
+from convo.core.contracts import Message
+
+_GREETING = re.compile(
+    r"(?:hi|hello|hey|hey there|hello there|good morning|good evening|good afternoon)[!.\s]*", re.I
+)
+_THANKS = re.compile(r"(?:thanks|thank you|thanks a lot|thank you so much)[!.\s]*", re.I)
+
+
+def conversation_reply(question: str) -> str | None:
+    if _GREETING.fullmatch(question.strip()):
+        return "Hi! Ask me about your selected file, or use /file to choose one."
+    if _THANKS.fullmatch(question.strip()):
+        return "You're welcome! Ask another question whenever you're ready."
+    return None
+
+
+def is_overview(question: str) -> bool:
+    normalized = re.sub(r"[^\w\s]", "", question.casefold()).strip()
+    normalized = re.sub(r"^(?:please|can you|could you)\s+", "", normalized)
+    return normalized in {
+        "summarize",
+        "summarise",
+        "summary",
+        "summarize this",
+        "summarise this",
+        "summarize this file",
+        "summarise this file",
+        "summarize the file",
+        "summarize this document",
+        "summarise this document",
+        "summarize the document",
+        "what is this file about",
+        "what is this document about",
+        "explain this file",
+        "explain this document",
+        "give me an overview of this file",
+        "give me a summary of this file",
+        "give me a summary of this document",
+    }
+
+
+def search_query(question: str, history: tuple[Message, ...]) -> str:
+    # Mixed greetings/questions remain document queries. Short references use the
+    # last user question to avoid embedding an ambiguous 'what about that?' alone.
+    if is_overview(question):
+        return question
+    if len(question.split()) <= 16 and re.search(
+        r"\b(it|that|those|they|them|more|further|why|explain|elaborate)\b", question, re.I
+    ):
+        previous = next(
+            (
+                message.content
+                for message in reversed(history)
+                if message.role == "user" and conversation_reply(message.content) is None
+            ),
+            None,
+        )
+        if previous:
+            return f"{previous[:600]}\n{question}"[:4_000]
+    return question