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())