175 lines
6.8 KiB
Python
175 lines
6.8 KiB
Python
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())
|