Ingestion¶

def chunk_document(
    document: ExtractedDocument,
    *,
    target_tokens: int = DEFAULT_TARGET_TOKENS,
    overlap_tokens: int = DEFAULT_OVERLAP_TOKENS,
) -> list[Chunk]:
    """Chunk an ``ExtractedDocument`` — page by page when paginated (PDF),
    so each chunk carries its ``page_number``; whole-Markdown otherwise.

    ``pending_vision`` documents yield no chunks (indexed by the OCR/VLM
    tiers of a later release).
    """
    if document.status != "active":
        return []

    if document.pages:
        chunks: list[Chunk] = []
        for page in document.pages:
            chunks.extend(
                chunk_markdown(
                    page.text,
                    target_tokens=target_tokens,
                    overlap_tokens=overlap_tokens,
                    page_number=page.page_number,
                    start_index=len(chunks),
                )
            )
    else:
        chunks = chunk_markdown(
            document.markdown,
            target_tokens=target_tokens,
            overlap_tokens=overlap_tokens,
        )

    if chunks:
        avg_tokens = sum(chunk.token_count for chunk in chunks) / len(chunks)
        logger.info(
            f"Document chunked: {len(chunks)} chunks, "
            f"avg {avg_tokens:.0f} tokens/chunk "
            f"(extractor={document.metadata.get('extractor', '?')})"
        )
    else:
        logger.debug(
            f"Document chunked: 0 chunks " f"(extractor={document.metadata.get('extractor', '?')})"
        )
    return chunks

def chunk_markdown(
    text: str,
    *,
    target_tokens: int = DEFAULT_TARGET_TOKENS,
    overlap_tokens: int = DEFAULT_OVERLAP_TOKENS,
    page_number: int | None = None,
    start_index: int = 0,
) -> list[Chunk]:
    """Run the 3-pass chunker over one Markdown text. See module docstring."""
    if not text or not text.strip():
        return []

    # PostgreSQL text columns reject NUL bytes; some parsers emit them.
    text = text.replace("\x00", "")

    from langchain_text_splitters import MarkdownHeaderTextSplitter

    md_splitter = MarkdownHeaderTextSplitter(
        headers_to_split_on=HEADERS_TO_SPLIT_ON,
        strip_headers=False,
    )
    sub_splitter = _make_sub_splitter(target_tokens, overlap_tokens)

    chunks: list[Chunk] = []
    for section in md_splitter.split_text(text):
        breadcrumb = _heading_breadcrumb(section.metadata)

        sub_texts = sub_splitter.split_text(section.page_content)
        sub_texts = _reattach_table_headers(sub_texts)

        for sub_text in sub_texts:
            chunk_text = f"{breadcrumb}\n\n{sub_text}" if breadcrumb else sub_text

            chunks.append(
                Chunk(
                    chunk_index=start_index + len(chunks),
                    text=chunk_text,
                    token_count=count_tokens(chunk_text),
                    page_number=page_number,
                )
            )

    return chunks

def count_tokens(text: str) -> int:
    """Real e5 subword count (no special tokens) — never a chars/N estimate."""
    return len(_get_tokenizer().encode(text, add_special_tokens=False))

def clean_extracted_text(text: str) -> str:
    """Clean extracted text without destroying meaning. See module docstring."""
    if not text:
        return ""
    text = unicodedata.normalize("NFC", text)
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = _HYPHENATION_RE.sub(r"\1", text)
    text = "".join(ch for ch in text if _keep(ch))
    text = _HORIZONTAL_WS_RE.sub(" ", text)
    text = _NEWLINE_TRIM_RE.sub("\n", text)
    text = _PARAGRAPH_RE.sub("\n\n", text)
    return text.strip()

def download_state() -> dict:
    """The gate's view: on-disk presence + transient background-task flags."""
    return {
        "available": embedding_model_available(),
        "downloading": _state["downloading"],
        "error": _state["error"],
    }

def embedding_model_available() -> bool:
    """True iff EVERY required file is fully present in ``CACHE_DIR`` (no network)."""
    for filename in REQUIRED_FILES:
        cached = try_to_load_from_cache(
            EMBEDDING_MODEL_ID, filename, cache_dir=str(config.CACHE_DIR)
        )
        if not isinstance(cached, str):
            return False
    return True

def load_sentence_transformer(model_name: str, *, cache_folder: str, local_files_only: bool):
    """Load a SentenceTransformer on the best device the machine can actually use.

    Single entry point for every embedding-model load (the #146 download gate
    and the resident ``E5Embeddings`` singleton), so the MPS fallback is decided
    once and applies to both.
    """
    global _MPS_UNUSABLE
    from sentence_transformers import SentenceTransformer

    device = resolve_embedding_device()
    try:
        return SentenceTransformer(
            model_name,
            cache_folder=cache_folder,
            local_files_only=local_files_only,
            device=device,
        )
    except RuntimeError as exc:
        if device != "mps" or not _is_mps_allocation_failure(exc):
            raise
        _MPS_UNUSABLE = True
        logger.warning(
            f"MPS is available but the embedding model could not be allocated on it; "
            f"falling back to CPU: {exc}"
        )

    return SentenceTransformer(
        model_name,
        cache_folder=cache_folder,
        local_files_only=local_files_only,
        device="cpu",
    )

def resolve_embedding_device() -> Optional[str]:
    """Device for the embedding model. ``None`` = let sentence-transformers pick.

    Apple Silicon is the primary desktop platform, so MPS stays preferred
    whenever Metal genuinely works. ``is_available()`` is only a support flag
    though: a virtualised macOS host (the ``macos-14`` CI runner) answers True
    and then refuses the model's ~366 MiB allocation, and a real Mac under
    memory pressure fails the same way (#335). ``load_sentence_transformer``
    turns that failure into this cached CPU verdict.

    Off Apple Silicon the choice is left untouched (``None``), so a CUDA host
    keeps auto-selecting its GPU.
    """
    if _MPS_UNUSABLE:
        return "cpu"
    if _mps_is_available():
        return "mps"
    return None

def start_download() -> dict:
    """Idempotently kick off a background download. Returns the current state.

    No-op (returns current state) if the model is already present or a download
    is already running — guards double-clicks and concurrent callers.
    """
    with _lock:
        if not _state["downloading"] and not embedding_model_available():
            _state["downloading"] = True
            _state["error"] = None
            _spawn(_run_download)
    return download_state()

def build_embedding_text(*, file_name: str, chunk_text: str) -> str:
    """Prefix the chunk with its source document name (R&D POC format)."""
    return f"[document_name:{file_name}]\n{chunk_text}"

def rows_to_markdown_table(rows: list[list[str]]) -> str:
    """Render rows as a GitHub-style Markdown table (first row = header).

    Cells are flattened (newlines → spaces) and pipes escaped so the table
    survives the Markdown-header chunking pass intact.
    """
    if not rows:
        return ""
    width = max(len(row) for row in rows)

    def fmt(row: list[str]) -> str:
        cells = [str(cell).replace("|", "\\|").replace("\n", " ").strip() for cell in row]
        cells += [""] * (width - len(cells))
        return "| " + " | ".join(cells) + " |"

    header, *body = rows
    lines = [fmt(header), "| " + " | ".join(["---"] * width) + " |"]
    lines += [fmt(row) for row in body]
    return "\n".join(lines)

def add_kb_chunks(
    *,
    kb_id: int,
    document_id: int,
    source_file: str,
    chunks: list["Chunk"],
) -> list[str]:
    """Embed and store one document's chunks (sync — delegates to the
    PGEngine's internal loop; the P0bis-validated ingestion bridge).

    The ``[document_name:…]`` prefix is EMBEDDING-time text only (it boosts
    retrieval, POC pattern): vectors are computed over the prefixed text,
    but the STORED content is the clean chunk — it goes verbatim into the
    LLM prompt at generation time, and small models loop on the bracketed
    prefix (live-E2E finding)."""
    if not chunks:
        return []
    start_s = time.perf_counter()
    embedding_texts = [
        build_embedding_text(file_name=source_file, chunk_text=chunk.text) for chunk in chunks
    ]
    vectors = E5Embeddings().embed_documents(embedding_texts)
    texts = [chunk.text for chunk in chunks]
    metadatas = [
        {
            "kb_id": kb_id,
            "document_id": document_id,
            "source_file": source_file,
            "page": chunk.page_number,
            "chunk_index": chunk.chunk_index,
        }
        for chunk in chunks
    ]
    ids = get_kb_store().add_embeddings(texts=texts, embeddings=vectors, metadatas=metadatas)
    duration_ms = (time.perf_counter() - start_s) * 1000
    logger.info(
        f"KB chunks added: kb_id={kb_id}, document_id={document_id}, "
        f"source_file={source_file}, n_chunks={len(chunks)}, "
        f"duration_ms={duration_ms:.0f}"
    )
    return ids

def close_kb_store() -> None:
    """Release the store's PGEngine (lifespan shutdown, before cluster stop)."""
    global _pg_engine, _kb_store, _psycopg_url
    if _pg_engine is not None:
        _close_engine(_pg_engine)
    _pg_engine = None
    _kb_store = None
    _psycopg_url = None

def get_kb_store() -> PGVectorStore:
    """The resident store (set by the lifespan via ``init_kb_store``)."""
    if _kb_store is None:
        raise RuntimeError(
            "KB vector store not initialized — call init_kb_store() first "
            "(FastAPI lifespan does this at startup)."
        )
    return _kb_store

def init_kb_store(handle: "PostgresHandle") -> PGVectorStore:
    """Create/join ``rag.kb_chunks`` and build the resident hybrid store.

    Idempotent: safe on every boot (schema/table/FKs/HNSW all guarded).
    Must run AFTER the business tables exist (cross-schema FKs).
    """
    global _pg_engine, _kb_store, _psycopg_url

    with psycopg.connect(handle.psycopg_url, autocommit=True) as conn:
        conn.execute(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA_NAME}")
        table_exists = conn.execute(
            "SELECT 1 FROM information_schema.tables"
            " WHERE table_schema = %s AND table_name = %s",
            (SCHEMA_NAME, TABLE_NAME),
        ).fetchone()

    engine = PGEngine.from_connection_string(handle.sqlalchemy_url)

    if not table_exists:
        engine.init_vectorstore_table(
            table_name=TABLE_NAME,
            schema_name=SCHEMA_NAME,
            vector_size=EMBEDDING_DIMENSIONS,
            metadata_columns=[
                Column("kb_id", "INTEGER", nullable=False),
                Column("document_id", "INTEGER", nullable=False),
                Column("source_file", "TEXT", nullable=True),
                Column("page", "INTEGER", nullable=True),
                Column("chunk_index", "INTEGER", nullable=True),
            ],
            hybrid_search_config=_hybrid_config(),
        )

    with psycopg.connect(handle.psycopg_url, autocommit=True) as conn:
        _ensure_fk(
            conn,
            "fk_kb_chunks_kb_id",
            f"ALTER TABLE {SCHEMA_NAME}.{TABLE_NAME}"
            " ADD CONSTRAINT fk_kb_chunks_kb_id FOREIGN KEY (kb_id)"
            " REFERENCES public.knowledge_base(id) ON DELETE CASCADE",
        )
        _ensure_fk(
            conn,
            "fk_kb_chunks_document_id",
            f"ALTER TABLE {SCHEMA_NAME}.{TABLE_NAME}"
            " ADD CONSTRAINT fk_kb_chunks_document_id FOREIGN KEY (document_id)"
            " REFERENCES public.knowledge_documents(id) ON DELETE CASCADE",
        )

    store = PGVectorStore.create_sync(
        engine=engine,
        table_name=TABLE_NAME,
        schema_name=SCHEMA_NAME,
        embedding_service=E5Embeddings(),
        metadata_columns=METADATA_COLUMNS,
        hybrid_search_config=_hybrid_config(),
    )

    if not store.is_valid_index(HNSW_INDEX_NAME):
        store.apply_vector_index(
            HNSWIndex(
                name=HNSW_INDEX_NAME,
                distance_strategy=DistanceStrategy.COSINE_DISTANCE,
            )
        )

    # Swap-in last: never leave a half-initialized store visible.
    if _pg_engine is not None:
        _close_engine(_pg_engine)
    _pg_engine, _kb_store, _psycopg_url = engine, store, handle.psycopg_url
    logger.info(f"KB vector store ready ({SCHEMA_NAME}.{TABLE_NAME}, hybrid RRF)")
    return store

def search_kb_chunks_scored(
    query: str, *, kb_id: int, pool_k: int = POOL_K
) -> list[tuple["Document", float]]:
    """Hybrid candidate pool over one KB: documents in RRF order, each with
    its calibrated dense cosine similarity.

    The fusion overwrites per-row scores with RRF harmonics (rank algebra,
    no semantic scale), so the dense similarity is re-read from the stored
    vectors — that is what the adaptive cut upstream operates on. Fresh
    config per call (lib bug — see module docstring); the query is embedded
    once and shared by the search and the similarity read.
    """
    start_s = time.perf_counter()
    query_vector = E5Embeddings().embed_query(query)
    documents = get_kb_store().similarity_search_by_vector(
        query_vector,
        k=pool_k,
        filter={"kb_id": kb_id},
        hybrid_search_config=_hybrid_config(fts_query=query, pool_k=pool_k),
    )
    if not documents:
        duration_ms = (time.perf_counter() - start_s) * 1000
        logger.info(
            f"KB search: kb_id={kb_id}, k={pool_k}, hits=0, top_score=n/a, "
            f"duration_ms={duration_ms:.0f}, "
            f"query={truncate_for_log(query, 2000)}"
        )
        return []
    similarities = _dense_similarities([doc.id for doc in documents], query_vector)
    duration_ms = (time.perf_counter() - start_s) * 1000
    top_score = max(similarities.values()) if similarities else float("nan")
    logger.info(
        f"KB search: kb_id={kb_id}, k={pool_k}, hits={len(documents)}, "
        f"top_score={top_score:.4f}, duration_ms={duration_ms:.0f}, "
        f"query={truncate_for_log(query, 2000)}"
    )
    return [(doc, similarities[doc.id]) for doc in documents]