from __future__ import annotations import argparse import json import sqlite3 from pathlib import Path from typing import Any from common import cosine_similarity, read_jsonl, corpus_content_hash, unpack_float_vector from rag_embedding_providers import LOCAL_HASHING_PROVIDER, embed_texts from rag_profiles import resolve_rag_profile ROOT = Path(__file__).resolve().parents[1] DEFAULT_CORPUS = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_corpus.jsonl" DEFAULT_INDEX = ROOT / "plugins" / "1c" / "datasets" / "prepared" / "rag_vector_index.sqlite" def load_meta(conn: sqlite3.Connection) -> dict[str, Any]: meta: dict[str, Any] = {} for key, value in conn.execute("SELECT key, value FROM vector_meta"): try: meta[str(key)] = json.loads(value) except json.JSONDecodeError: meta[str(key)] = value return meta def freshness(index_meta: dict[str, Any], corpus_path: Path | None) -> dict[str, Any]: if not corpus_path: return {"status": "unknown", "reason": "corpus path was not provided"} if not corpus_path.exists(): return {"status": "unknown", "reason": "corpus file is missing", "corpus": str(corpus_path)} current_hash = corpus_content_hash(read_jsonl(corpus_path)) index_hash = str(index_meta.get("corpus_hash") or "") return { "status": "fresh" if current_hash == index_hash else "stale", "corpus": str(corpus_path), "corpus_hash": current_hash, "index_corpus_hash": index_hash, } def row_document(row: sqlite3.Row) -> dict[str, Any]: try: metadata = json.loads(row["metadata_json"] or "{}") except json.JSONDecodeError: metadata = {} return { "id": row["id"], "document_id": row["document_id"], "source_path": row["source_path"], "title": row["title"], "chunk_index": row["chunk_index"], "content": row["content"], "metadata": metadata, } def search_vector_index( index_path: Path, query: str, *, limit: int, candidate_limit: int | None = None, min_score: float = 0.0, source_types: list[str] | None = None, corpus_path: Path | None = DEFAULT_CORPUS, embedding_base_url: str = "", embedding_api_key_env: str = "OPENAI_API_KEY", ) -> dict[str, Any]: conn = sqlite3.connect(index_path) conn.row_factory = sqlite3.Row try: meta = load_meta(conn) dimensions = int(meta.get("embedding_dimensions") or 0) if dimensions <= 0: raise ValueError("Vector index metadata is missing embedding_dimensions") provider = str(meta.get("embedding_provider") or LOCAL_HASHING_PROVIDER) model = str(meta.get("embedding_model") or "") query_vector = embed_texts( [query], provider=provider, model=model, dimensions=dimensions, base_url=embedding_base_url or str(meta.get("embedding_base_url") or ""), api_key_env=embedding_api_key_env, )[0] if len(query_vector) != dimensions: raise ValueError(f"Query embedding dimensions {len(query_vector)} do not match index dimensions {dimensions}") allowed_source_types = {value.lower() for value in source_types or []} sql = "SELECT * FROM vector_documents" params: list[Any] = [] if allowed_source_types: placeholders = ",".join("?" for _ in allowed_source_types) sql += f" WHERE lower(source_type) IN ({placeholders})" params.extend(sorted(allowed_source_types)) results = [] for row in conn.execute(sql, params): score = cosine_similarity(query_vector, unpack_float_vector(row["vector"], dimensions)) if score > min_score: results.append({"score": score, "document": row_document(row)}) results.sort(key=lambda item: item["score"], reverse=True) if candidate_limit: results = results[:candidate_limit] return { "schema": "onec_rag_vector_search.v1", "status": "ok", "query": query, "index": str(index_path), "meta": meta, "freshness": freshness(meta, corpus_path), "results": results[:limit], "counts": {"results": min(len(results), limit), "candidates": len(results)}, } finally: conn.close() def print_result(result: dict, index: int) -> None: document = result["document"] content = (document.get("content") or "").strip().replace("\n", " ") if len(content) > 500: content = content[:497].rstrip() + "..." print(f"{index}. score={result['score']:.4f}") print(f" source={document.get('source_path')} chunk={document.get('chunk_index')}") print(f" title={document.get('title')}") if (document.get("metadata") or {}).get("heading"): print(f" heading={(document.get('metadata') or {}).get('heading')}") print(f" content={content}") def main() -> int: parser = argparse.ArgumentParser(description="Search the local 1C RAG SQLite vector index.") parser.add_argument("query") parser.add_argument("--index", type=Path, default=DEFAULT_INDEX) parser.add_argument("--corpus", type=Path, default=DEFAULT_CORPUS) parser.add_argument("--profile", default="auto") parser.add_argument("--limit", type=int) parser.add_argument("--candidate-limit", type=int) parser.add_argument("--min-score", type=float, default=0.0) parser.add_argument("--source-type", action="append", dest="source_types") parser.add_argument("--embedding-base-url", default="") parser.add_argument("--embedding-api-key-env", default="OPENAI_API_KEY") parser.add_argument("--json", action="store_true") args = parser.parse_args() profile = resolve_rag_profile(args.profile, args.query) source_types = args.source_types if args.source_types is not None else profile["source_types"] result = search_vector_index( args.index, args.query, limit=args.limit or int(profile["limit"]), candidate_limit=args.candidate_limit or int(profile["candidate_limit"]), min_score=float(args.min_score), source_types=source_types, corpus_path=args.corpus, embedding_base_url=args.embedding_base_url, embedding_api_key_env=args.embedding_api_key_env, ) if args.json: print(json.dumps(result, ensure_ascii=False, indent=2)) elif result["results"]: freshness_status = (result.get("freshness") or {}).get("status") print(f"freshness={freshness_status} embedding_model={result['meta'].get('embedding_model')}") for position, item in enumerate(result["results"], start=1): print_result(item, position) else: print("No matches.") return 0 if (result.get("freshness") or {}).get("status") != "stale" else 1 if __name__ == "__main__": raise SystemExit(main())