pv-agent/agent/bootstrap.py

"""Container-Bootstrap: vollständigen Produktionsindex vor API-Start sicherstellen."""
from __future__ import annotations

import json
import sqlite3
from pathlib import Path

from .config import Config
from .ingest import build_index
from .ollama_client import OllamaClient


def index_is_ready(cfg: Config) -> bool:
    path = Path(cfg.db_path)
    if not path.is_file() or path.stat().st_size == 0:
        return False
    try:
        con = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
        try:
            n_chunks = con.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
            if n_chunks <= 0:
                return False
            if cfg.embed_off:
                return True
            n_dense = con.execute(
                "SELECT COUNT(*) FROM chunks c JOIN vectors v "
                "ON v.content_hash = c.content_hash AND v.model = ?",
                (cfg.embed_model,),
            ).fetchone()[0]
            return n_dense == n_chunks
        finally:
            con.close()
    except (OSError, sqlite3.Error):
        return False


def ensure_index(cfg: Config) -> dict:
    if index_is_ready(cfg):
        result = {"event": "index_ready", "db_path": cfg.db_path}
        print(json.dumps(result, ensure_ascii=False), flush=True)
        return result

    print(
        json.dumps(
            {
                "event": "index_bootstrap_started",
                "db_path": cfg.db_path,
                "embed_model": None if cfg.embed_off else cfg.embed_model,
            },
            ensure_ascii=False,
        ),
        flush=True,
    )
    client = None
    if not cfg.embed_off:
        client = OllamaClient(
            cfg.ollama_url,
            embed_timeout_s=cfg.embed_timeout_s,
            chat_timeout_s=cfg.chat_timeout_s,
        )
    try:
        stats = build_index(cfg, client=client)
    finally:
        if client is not None:
            client.close()
    if stats.embed_error:
        raise RuntimeError(
            "Index-Bootstrap ohne vollständige Embeddings abgebrochen: "
            + stats.embed_error
        )
    if not index_is_ready(cfg):
        raise RuntimeError("Index-Bootstrap abgeschlossen, Index ist aber unvollständig")
    result = {"event": "index_bootstrap_completed", **stats.as_dict()}
    print(json.dumps(result, ensure_ascii=False), flush=True)
    return result


def main() -> int:
    ensure_index(Config.from_env())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())