llm/scripts/build_1c_rag_vector_index.py

from __future__ import annotations

import argparse
import json
import sqlite3
from datetime import datetime, timezone
from pathlib import Path

from common import corpus_content_hash, pack_float_vector, read_jsonl
from rag_embedding_providers import LOCAL_HASHING_MODEL, LOCAL_HASHING_PROVIDER, embed_texts, provider_metadata


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CORPUS = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_corpus.jsonl"
DEFAULT_OUTPUT = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_vector_index.sqlite"
SCHEMA_VERSION = 1
DEFAULT_EMBEDDING_MODEL = LOCAL_HASHING_MODEL


def connect_index(path: Path) -> sqlite3.Connection:
    path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(path)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")
    return conn


def reset_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        DROP TABLE IF EXISTS vector_documents;
        DROP TABLE IF EXISTS vector_meta;

        CREATE TABLE vector_meta (
            key TEXT PRIMARY KEY,
            value TEXT NOT NULL
        );

        CREATE TABLE vector_documents (
            id TEXT PRIMARY KEY,
            document_id TEXT,
            source_path TEXT,
            source_type TEXT,
            title TEXT,
            chunk_index INTEGER,
            content TEXT NOT NULL,
            metadata_json TEXT NOT NULL,
            vector BLOB NOT NULL
        );

        CREATE INDEX idx_vector_documents_source_type ON vector_documents(source_type);
        CREATE INDEX idx_vector_documents_source_path ON vector_documents(source_path);
        """
    )


def write_meta(conn: sqlite3.Connection, metadata: dict[str, object]) -> None:
    conn.executemany(
        "INSERT INTO vector_meta(key, value) VALUES(?, ?)",
        [(key, json.dumps(value, ensure_ascii=False, sort_keys=True)) for key, value in metadata.items()],
    )


def document_embedding_text(record: dict) -> str:
    title = str(record.get("title") or "").strip()
    content = str(record.get("content") or "").strip()
    metadata = record.get("metadata") if isinstance(record.get("metadata"), dict) else {}
    headings = metadata.get("headings") if isinstance(metadata.get("headings"), list) else []
    heading_text = "\n".join(str(item) for item in headings if str(item).strip())
    return "\n\n".join(part for part in (title, heading_text, content) if part)


def batched(items: list[dict], size: int) -> list[list[dict]]:
    return [items[index : index + size] for index in range(0, len(items), size)]


def build_vector_index(
    corpus_path: Path,
    output_path: Path,
    *,
    dimensions: int,
    embedding_model: str,
    embedding_provider: str = LOCAL_HASHING_PROVIDER,
    embedding_base_url: str = "",
    embedding_api_key_env: str = "OPENAI_API_KEY",
    batch_size: int = 16,
) -> dict:
    records = read_jsonl(corpus_path)
    corpus_hash = corpus_content_hash(records)
    if not records:
        inferred_dimensions = dimensions
    else:
        sample_vector = embed_texts(
            [document_embedding_text(records[0])],
            provider=embedding_provider,
            model=embedding_model,
            dimensions=dimensions,
            base_url=embedding_base_url,
            api_key_env=embedding_api_key_env,
        )[0]
        inferred_dimensions = len(sample_vector)
    conn = connect_index(output_path)
    try:
        with conn:
            reset_schema(conn)
            embedding_meta = provider_metadata(
                provider=embedding_provider,
                model=embedding_model,
                dimensions=inferred_dimensions,
                base_url=embedding_base_url,
            )
            write_meta(
                conn,
                {
                    "schema": "onec_rag_vector_index.v1",
                    "schema_version": SCHEMA_VERSION,
                    "type": "sqlite-vector-scan",
                    **embedding_meta,
                    "corpus_path": str(corpus_path),
                    "corpus_hash": corpus_hash,
                    "doc_count": len(records),
                    "built_at": datetime.now(timezone.utc).isoformat(),
                },
            )
            rows = []
            for batch in batched(records, max(int(batch_size or 1), 1)):
                texts = [document_embedding_text(record) for record in batch]
                vectors = embed_texts(
                    texts,
                    provider=embedding_provider,
                    model=embedding_model,
                    dimensions=inferred_dimensions,
                    base_url=embedding_base_url,
                    api_key_env=embedding_api_key_env,
                )
                for record, vector in zip(batch, vectors):
                    if len(vector) != inferred_dimensions:
                        raise ValueError(f"Embedding dimensions changed within the build: {len(vector)} != {inferred_dimensions}")
                    metadata = record.get("metadata") if isinstance(record.get("metadata"), dict) else {}
                    rows.append(
                        (
                            str(record.get("id") or ""),
                            str(record.get("document_id") or ""),
                            str(record.get("source_path") or ""),
                            str(metadata.get("source_type") or ""),
                            str(record.get("title") or ""),
                            int(record.get("chunk_index") or 0),
                            str(record.get("content") or ""),
                            json.dumps(metadata, ensure_ascii=False, sort_keys=True),
                            pack_float_vector(vector),
                        )
                    )
            conn.executemany(
                """
                INSERT INTO vector_documents(
                    id, document_id, source_path, source_type, title, chunk_index,
                    content, metadata_json, vector
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                rows,
            )
    finally:
        conn.close()
    return {
        "status": "ok",
        "index": str(output_path),
        "doc_count": len(records),
        "corpus_hash": corpus_hash,
        "embedding_provider": provider_metadata(provider=embedding_provider, model=embedding_model, dimensions=inferred_dimensions, base_url=embedding_base_url)["embedding_provider"],
        "embedding_model": embedding_model,
        "embedding_dimensions": inferred_dimensions,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Build a local SQLite vector index for the 1C RAG corpus.")
    parser.add_argument("--corpus", type=Path, default=DEFAULT_CORPUS)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--dimensions", type=int, default=384)
    parser.add_argument("--embedding-provider", default=LOCAL_HASHING_PROVIDER, choices=[LOCAL_HASHING_PROVIDER, "openai-compatible"])
    parser.add_argument("--embedding-model", default=DEFAULT_EMBEDDING_MODEL)
    parser.add_argument("--embedding-base-url", default="")
    parser.add_argument("--embedding-api-key-env", default="OPENAI_API_KEY")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    result = build_vector_index(
        args.corpus,
        args.output,
        dimensions=args.dimensions,
        embedding_model=args.embedding_model,
        embedding_provider=args.embedding_provider,
        embedding_base_url=args.embedding_base_url,
        embedding_api_key_env=args.embedding_api_key_env,
        batch_size=args.batch_size,
    )
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print(f"Wrote vector index with {result['doc_count']} document chunk(s) to {args.output}")
    return 0


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