Initial SQL-only 1C adapter baseline
This commit is contained in:
@@ -0,0 +1,207 @@
|
||||
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())
|
||||
Reference in New Issue
Block a user