726 lines
17 KiB
Python
726 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import numpy as np
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from rich import print
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
|
|
from tools.rag.common import (
|
|
CHUNKS_FILE,
|
|
DB_FILE,
|
|
DOCUMENTS_FILE,
|
|
read_json,
|
|
)
|
|
from tools.rag.embedding_utils import (
|
|
build_embedding_text,
|
|
embed_passages,
|
|
embedding_model_name,
|
|
vector_to_blob,
|
|
model_profile,
|
|
split_passage,
|
|
)
|
|
from tools.rag.embedding_cache import EmbeddingCache, cache_key
|
|
|
|
|
|
FTS_TOKENIZER = "unicode61 remove_diacritics 2"
|
|
FTS_PREFIXES = "3 4 5"
|
|
|
|
|
|
def published_to_db(value: Any) -> int | None:
|
|
if value is True:
|
|
return 1
|
|
|
|
if value is False:
|
|
return 0
|
|
|
|
return None
|
|
|
|
|
|
def verify_fts5(
|
|
conn: sqlite3.Connection,
|
|
) -> None:
|
|
try:
|
|
conn.execute(
|
|
"CREATE VIRTUAL TABLE temp.fts5_check "
|
|
"USING fts5(value)"
|
|
)
|
|
conn.execute(
|
|
"DROP TABLE temp.fts5_check"
|
|
)
|
|
|
|
except sqlite3.OperationalError as error:
|
|
raise RuntimeError(
|
|
"Táto inštalácia SQLite nemá dostupné FTS5. "
|
|
"Použi Python/SQLite zostavenie s podporou "
|
|
"SQLITE_ENABLE_FTS5."
|
|
) from error
|
|
|
|
|
|
def create_tables(
|
|
conn: sqlite3.Connection,
|
|
) -> None:
|
|
conn.executescript(
|
|
f"""
|
|
PRAGMA foreign_keys = ON;
|
|
|
|
CREATE TABLE documents (
|
|
id INTEGER PRIMARY KEY,
|
|
path TEXT UNIQUE NOT NULL,
|
|
title TEXT,
|
|
author TEXT,
|
|
published INTEGER
|
|
CHECK (
|
|
published IN (0, 1)
|
|
OR published IS NULL
|
|
),
|
|
content_length INTEGER
|
|
NOT NULL DEFAULT 0,
|
|
metadata_json TEXT
|
|
NOT NULL DEFAULT '{{}}'
|
|
);
|
|
|
|
CREATE TABLE chunks (
|
|
id INTEGER PRIMARY KEY,
|
|
chunk_id TEXT UNIQUE NOT NULL,
|
|
document_path TEXT NOT NULL,
|
|
title TEXT,
|
|
author TEXT,
|
|
published INTEGER
|
|
CHECK (
|
|
published IN (0, 1)
|
|
OR published IS NULL
|
|
),
|
|
chunk_index INTEGER NOT NULL,
|
|
heading_paths_json TEXT
|
|
NOT NULL DEFAULT '[]',
|
|
text TEXT NOT NULL,
|
|
text_length INTEGER
|
|
NOT NULL DEFAULT 0,
|
|
token_count INTEGER,
|
|
content_hash TEXT,
|
|
FOREIGN KEY(document_path)
|
|
REFERENCES documents(path)
|
|
ON UPDATE CASCADE
|
|
ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE chunk_tags (
|
|
chunk_id TEXT NOT NULL,
|
|
tag TEXT NOT NULL,
|
|
PRIMARY KEY(chunk_id, tag),
|
|
FOREIGN KEY(chunk_id)
|
|
REFERENCES chunks(chunk_id)
|
|
ON UPDATE CASCADE
|
|
ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE chunk_categories (
|
|
chunk_id TEXT NOT NULL,
|
|
category TEXT NOT NULL,
|
|
PRIMARY KEY(chunk_id, category),
|
|
FOREIGN KEY(chunk_id)
|
|
REFERENCES chunks(chunk_id)
|
|
ON UPDATE CASCADE
|
|
ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE chunk_embeddings (
|
|
chunk_id TEXT NOT NULL,
|
|
segment_index INTEGER NOT NULL DEFAULT 0,
|
|
model TEXT NOT NULL,
|
|
dimensions INTEGER NOT NULL,
|
|
embedding BLOB NOT NULL,
|
|
PRIMARY KEY(chunk_id, segment_index),
|
|
FOREIGN KEY(chunk_id)
|
|
REFERENCES chunks(chunk_id)
|
|
ON UPDATE CASCADE
|
|
ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE embedding_model_info (metadata_json TEXT NOT NULL);
|
|
|
|
CREATE INDEX idx_documents_path
|
|
ON documents(path);
|
|
|
|
CREATE INDEX idx_documents_published
|
|
ON documents(published);
|
|
|
|
CREATE INDEX idx_chunks_document_path
|
|
ON chunks(document_path);
|
|
|
|
CREATE INDEX idx_chunks_title
|
|
ON chunks(title);
|
|
|
|
CREATE INDEX idx_chunks_author
|
|
ON chunks(author);
|
|
|
|
CREATE INDEX idx_chunks_published
|
|
ON chunks(published);
|
|
|
|
CREATE INDEX idx_chunk_tags_tag
|
|
ON chunk_tags(tag);
|
|
|
|
CREATE INDEX idx_chunk_categories_category
|
|
ON chunk_categories(category);
|
|
|
|
CREATE INDEX idx_chunk_embeddings_model
|
|
ON chunk_embeddings(model);
|
|
|
|
CREATE VIRTUAL TABLE chunks_fts USING fts5(
|
|
chunk_id UNINDEXED,
|
|
title,
|
|
author,
|
|
document_path,
|
|
tags,
|
|
categories,
|
|
text,
|
|
tokenize='{FTS_TOKENIZER}',
|
|
prefix='{FTS_PREFIXES}'
|
|
);
|
|
"""
|
|
)
|
|
|
|
|
|
def insert_documents(
|
|
conn: sqlite3.Connection,
|
|
documents: list[dict],
|
|
) -> None:
|
|
rows = [
|
|
(
|
|
document.get("path"),
|
|
document.get("title"),
|
|
document.get("author"),
|
|
published_to_db(
|
|
document.get("published")
|
|
),
|
|
int(
|
|
document.get("content_length")
|
|
or 0
|
|
),
|
|
json.dumps(
|
|
document.get("metadata") or {},
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
),
|
|
)
|
|
for document in documents
|
|
]
|
|
|
|
conn.executemany(
|
|
"""
|
|
INSERT INTO documents (
|
|
path,
|
|
title,
|
|
author,
|
|
published,
|
|
content_length,
|
|
metadata_json
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
""",
|
|
rows,
|
|
)
|
|
|
|
|
|
def insert_chunks(
|
|
conn: sqlite3.Connection,
|
|
chunks: list[dict],
|
|
) -> None:
|
|
chunk_rows: list[tuple] = []
|
|
tag_rows: list[tuple[str, str]] = []
|
|
category_rows: list[
|
|
tuple[str, str]
|
|
] = []
|
|
|
|
for chunk in chunks:
|
|
chunk_id = str(
|
|
chunk.get("chunk_id") or ""
|
|
).strip()
|
|
|
|
if not chunk_id:
|
|
raise ValueError(
|
|
"Chunk bez chunk_id nie je "
|
|
"možné indexovať"
|
|
)
|
|
|
|
text = chunk.get("text") or ""
|
|
|
|
chunk_rows.append(
|
|
(
|
|
chunk_id,
|
|
chunk.get("document_path"),
|
|
chunk.get("title"),
|
|
chunk.get("author"),
|
|
published_to_db(
|
|
chunk.get("published")
|
|
),
|
|
int(
|
|
chunk.get("chunk_index")
|
|
or 0
|
|
),
|
|
json.dumps(
|
|
chunk.get(
|
|
"heading_paths"
|
|
)
|
|
or [],
|
|
ensure_ascii=False,
|
|
),
|
|
text,
|
|
int(
|
|
chunk.get("text_length")
|
|
or len(text)
|
|
),
|
|
chunk.get("token_count"),
|
|
chunk.get("content_hash"),
|
|
)
|
|
)
|
|
|
|
for tag in chunk.get("tags") or []:
|
|
value = str(tag).strip()
|
|
|
|
if value:
|
|
tag_rows.append(
|
|
(
|
|
chunk_id,
|
|
value,
|
|
)
|
|
)
|
|
|
|
for category in (
|
|
chunk.get("categories") or []
|
|
):
|
|
value = str(
|
|
category
|
|
).strip()
|
|
|
|
if value:
|
|
category_rows.append(
|
|
(
|
|
chunk_id,
|
|
value,
|
|
)
|
|
)
|
|
|
|
conn.executemany(
|
|
"""
|
|
INSERT INTO chunks (
|
|
chunk_id,
|
|
document_path,
|
|
title,
|
|
author,
|
|
published,
|
|
chunk_index,
|
|
heading_paths_json,
|
|
text,
|
|
text_length,
|
|
token_count,
|
|
content_hash
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
chunk_rows,
|
|
)
|
|
|
|
conn.executemany(
|
|
"""
|
|
INSERT OR IGNORE INTO chunk_tags (
|
|
chunk_id,
|
|
tag
|
|
)
|
|
VALUES (?, ?)
|
|
""",
|
|
tag_rows,
|
|
)
|
|
|
|
conn.executemany(
|
|
"""
|
|
INSERT OR IGNORE INTO chunk_categories (
|
|
chunk_id,
|
|
category
|
|
)
|
|
VALUES (?, ?)
|
|
""",
|
|
category_rows,
|
|
)
|
|
|
|
|
|
def insert_embeddings(conn: sqlite3.Connection, chunks: list[dict], *, cache_file: Path) -> dict:
|
|
if not chunks:
|
|
return {'embedding_reused': 0, 'embedding_computed': 0}
|
|
model_name = embedding_model_name()
|
|
profile = model_profile(model_name)
|
|
conn.execute('INSERT INTO embedding_model_info VALUES(?)', (json.dumps(profile, sort_keys=True),))
|
|
segments = []
|
|
for chunk in chunks:
|
|
if not chunk.get('chunk_id'):
|
|
raise ValueError('Chunk bez chunk_id nemôže mať embedding')
|
|
text = build_embedding_text(chunk)
|
|
for index, part in enumerate(split_passage(text, model_name=model_name)):
|
|
segments.append((chunk['chunk_id'], index, part, cache_key(profile, part)))
|
|
cache = EmbeddingCache(cache_file)
|
|
found = cache.load([row[3] for row in segments])
|
|
missing = {row[3]: row[2] for row in segments if row[3] not in found}
|
|
computed = {}
|
|
if missing:
|
|
vectors = embed_passages(list(missing.values()))
|
|
if len(vectors) != len(missing):
|
|
raise RuntimeError('Počet embeddingov sa nezhoduje s počtom vstupov')
|
|
for key, vector in zip(missing, vectors):
|
|
vector = np.asarray(vector, dtype=np.float32)
|
|
if vector.ndim != 1 or not len(vector) or not np.isfinite(vector).all() or np.linalg.norm(vector) == 0:
|
|
raise RuntimeError('Embedding nie je platný nenulový vektor')
|
|
computed[key] = vector
|
|
cache.save(computed)
|
|
all_vectors = {**found, **computed}
|
|
dimensions = {len(vector) for vector in all_vectors.values()}
|
|
if len(dimensions) != 1:
|
|
raise RuntimeError('Embeddingy majú rozdielne rozmery')
|
|
conn.executemany('''INSERT INTO chunk_embeddings
|
|
(chunk_id,segment_index,model,dimensions,embedding) VALUES(?,?,?,?,?)''',
|
|
[(chunk_id, index, model_name, len(all_vectors[key]), vector_to_blob(all_vectors[key]))
|
|
for chunk_id, index, _, key in segments])
|
|
return {'embedding_reused': sum(row[3] in found for row in segments),
|
|
'embedding_computed': len(computed)}
|
|
|
|
|
|
def build_fts_index(
|
|
conn: sqlite3.Connection,
|
|
) -> None:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO chunks_fts (
|
|
rowid,
|
|
chunk_id,
|
|
title,
|
|
author,
|
|
document_path,
|
|
tags,
|
|
categories,
|
|
text
|
|
)
|
|
SELECT
|
|
chunks.id,
|
|
chunks.chunk_id,
|
|
COALESCE(
|
|
chunks.title,
|
|
''
|
|
),
|
|
COALESCE(
|
|
chunks.author,
|
|
''
|
|
),
|
|
chunks.document_path,
|
|
COALESCE(
|
|
tags.values_text,
|
|
''
|
|
),
|
|
COALESCE(
|
|
categories.values_text,
|
|
''
|
|
),
|
|
chunks.text
|
|
FROM chunks
|
|
LEFT JOIN (
|
|
SELECT
|
|
chunk_id,
|
|
GROUP_CONCAT(
|
|
tag,
|
|
' '
|
|
) AS values_text
|
|
FROM chunk_tags
|
|
GROUP BY chunk_id
|
|
) AS tags
|
|
ON tags.chunk_id
|
|
= chunks.chunk_id
|
|
LEFT JOIN (
|
|
SELECT
|
|
chunk_id,
|
|
GROUP_CONCAT(
|
|
category,
|
|
' '
|
|
) AS values_text
|
|
FROM chunk_categories
|
|
GROUP BY chunk_id
|
|
) AS categories
|
|
ON categories.chunk_id
|
|
= chunks.chunk_id
|
|
ORDER BY chunks.id
|
|
"""
|
|
)
|
|
|
|
conn.execute(
|
|
"INSERT INTO chunks_fts"
|
|
"(chunks_fts) "
|
|
"VALUES('optimize')"
|
|
)
|
|
|
|
|
|
def validate_database(
|
|
conn: sqlite3.Connection,
|
|
) -> None:
|
|
integrity = conn.execute(
|
|
"PRAGMA integrity_check"
|
|
).fetchone()[0]
|
|
|
|
if integrity != "ok":
|
|
raise RuntimeError(
|
|
"SQLite integrity check "
|
|
f"zlyhal: {integrity}"
|
|
)
|
|
|
|
foreign_key_errors = (
|
|
conn.execute(
|
|
"PRAGMA foreign_key_check"
|
|
).fetchall()
|
|
)
|
|
|
|
if foreign_key_errors:
|
|
raise RuntimeError(
|
|
"Databáza obsahuje chyby "
|
|
"cudzích kľúčov: "
|
|
f"{foreign_key_errors[:5]}"
|
|
)
|
|
|
|
conn.execute(
|
|
"INSERT INTO chunks_fts"
|
|
"(chunks_fts) "
|
|
"VALUES('integrity-check')"
|
|
)
|
|
|
|
chunk_count = conn.execute(
|
|
"SELECT COUNT(*) FROM chunks"
|
|
).fetchone()[0]
|
|
|
|
fts_count = conn.execute(
|
|
"SELECT COUNT(*) FROM chunks_fts"
|
|
).fetchone()[0]
|
|
|
|
embedding_count = conn.execute(
|
|
"""
|
|
SELECT COUNT(DISTINCT chunk_id)
|
|
FROM chunk_embeddings
|
|
"""
|
|
).fetchone()[0]
|
|
|
|
if chunk_count != fts_count:
|
|
raise RuntimeError(
|
|
"Počet záznamov v chunks "
|
|
"a chunks_fts sa nezhoduje: "
|
|
f"{chunk_count} != {fts_count}"
|
|
)
|
|
|
|
if chunk_count != embedding_count:
|
|
raise RuntimeError(
|
|
"Počet chunkov a embeddingov "
|
|
"sa nezhoduje: "
|
|
f"{chunk_count} != "
|
|
f"{embedding_count}"
|
|
)
|
|
|
|
model_count = conn.execute(
|
|
"""
|
|
SELECT COUNT(
|
|
DISTINCT model
|
|
)
|
|
FROM chunk_embeddings
|
|
"""
|
|
).fetchone()[0]
|
|
|
|
if (
|
|
embedding_count > 0
|
|
and model_count != 1
|
|
):
|
|
raise RuntimeError(
|
|
"Embedding index obsahuje "
|
|
"viac modelov"
|
|
)
|
|
|
|
|
|
def get_counts(
|
|
conn: sqlite3.Connection,
|
|
) -> dict[str, int]:
|
|
return {
|
|
"embedding_segments": conn.execute("SELECT COUNT(*) FROM chunk_embeddings").fetchone()[0],
|
|
"documents": conn.execute(
|
|
"""
|
|
SELECT COUNT(*)
|
|
FROM documents
|
|
"""
|
|
).fetchone()[0],
|
|
"chunks": conn.execute(
|
|
"""
|
|
SELECT COUNT(*)
|
|
FROM chunks
|
|
"""
|
|
).fetchone()[0],
|
|
"fts_chunks": conn.execute(
|
|
"""
|
|
SELECT COUNT(*)
|
|
FROM chunks_fts
|
|
"""
|
|
).fetchone()[0],
|
|
"embedding_chunks": conn.execute(
|
|
"""
|
|
SELECT COUNT(DISTINCT chunk_id)
|
|
FROM chunk_embeddings
|
|
"""
|
|
).fetchone()[0],
|
|
"tags": conn.execute(
|
|
"""
|
|
SELECT COUNT(*)
|
|
FROM chunk_tags
|
|
"""
|
|
).fetchone()[0],
|
|
"categories": conn.execute(
|
|
"""
|
|
SELECT COUNT(*)
|
|
FROM chunk_categories
|
|
"""
|
|
).fetchone()[0],
|
|
}
|
|
|
|
|
|
def temporary_database_path(
|
|
db_file: Path,
|
|
) -> Path:
|
|
return db_file.with_name(
|
|
f".{db_file.name}.tmp"
|
|
)
|
|
|
|
|
|
def build_database(*, documents_file: Path | None = None, chunks_file: Path | None = None,
|
|
db_file: Path | None = None, cache_file: Path | None = None) -> dict[str, int]:
|
|
target = db_file or DB_FILE
|
|
documents = read_json(
|
|
documents_file or DOCUMENTS_FILE
|
|
)
|
|
|
|
chunks = read_json(
|
|
chunks_file or CHUNKS_FILE
|
|
)
|
|
|
|
target.parent.mkdir(
|
|
parents=True,
|
|
exist_ok=True,
|
|
)
|
|
|
|
temporary_file = (
|
|
temporary_database_path(
|
|
target
|
|
)
|
|
)
|
|
|
|
if temporary_file.exists():
|
|
temporary_file.unlink()
|
|
|
|
try:
|
|
with sqlite3.connect(
|
|
temporary_file
|
|
) as conn:
|
|
conn.execute(
|
|
"PRAGMA foreign_keys = ON"
|
|
)
|
|
conn.execute(
|
|
"PRAGMA temp_store = MEMORY"
|
|
)
|
|
|
|
verify_fts5(
|
|
conn
|
|
)
|
|
|
|
with conn:
|
|
create_tables(
|
|
conn
|
|
)
|
|
|
|
insert_documents(
|
|
conn,
|
|
documents,
|
|
)
|
|
|
|
insert_chunks(
|
|
conn,
|
|
chunks,
|
|
)
|
|
|
|
build_fts_index(
|
|
conn
|
|
)
|
|
|
|
embedding_stats = insert_embeddings(conn, chunks, cache_file=cache_file or target.parent / "embedding_cache.sqlite")
|
|
|
|
validate_database(
|
|
conn
|
|
)
|
|
|
|
counts = {**get_counts(conn), **embedding_stats}
|
|
|
|
os.replace(
|
|
temporary_file,
|
|
target,
|
|
)
|
|
|
|
except Exception:
|
|
if temporary_file.exists():
|
|
temporary_file.unlink()
|
|
|
|
raise
|
|
|
|
print(
|
|
"[green]SQLite index "
|
|
"vytvorený:[/green] "
|
|
f"{target}"
|
|
)
|
|
|
|
print(
|
|
f"Dokumentov: "
|
|
f"{counts['documents']}"
|
|
)
|
|
|
|
print(
|
|
f"Chunkov: "
|
|
f"{counts['chunks']}"
|
|
)
|
|
|
|
print(
|
|
f"FTS5 chunkov: "
|
|
f"{counts['fts_chunks']}"
|
|
)
|
|
|
|
print(
|
|
f"Embedding chunkov: "
|
|
f"{counts['embedding_chunks']}"
|
|
)
|
|
|
|
print(
|
|
f"Tag záznamov: "
|
|
f"{counts['tags']}"
|
|
)
|
|
|
|
print(
|
|
"Kategória záznamov: "
|
|
f"{counts['categories']}"
|
|
)
|
|
|
|
return counts
|
|
|
|
|
|
def main() -> None:
|
|
build_database()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|