dp-zp-agent/scripts/build_sqlite_index.py
2026-07-28 23:56:03 +02:00

462 lines
11 KiB
Python

from __future__ import annotations
import json
import os
import sqlite3
import sys
from pathlib import Path
from typing import Any
from rich import print
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from scripts.common import CHUNKS_FILE, DB_FILE, DOCUMENTS_FILE, read_json
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:
"""Overí, či aktuálna SQLite knižnica podporuje FTS5."""
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 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 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 build_fts_index(conn: sqlite3.Connection) -> None:
"""Vytvorí FTS5 index nad chunkmi a ich metadátami."""
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(
f"SQLite integrity check 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]
if chunk_count != fts_count:
raise RuntimeError(
"Počet záznamov v chunks a chunks_fts sa nezhoduje: "
f"{chunk_count} != {fts_count}"
)
def get_counts(
conn: sqlite3.Connection,
) -> dict[str, int]:
return {
"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],
"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() -> dict[str, int]:
documents = read_json(DOCUMENTS_FILE)
chunks = read_json(CHUNKS_FILE)
DB_FILE.parent.mkdir(
parents=True,
exist_ok=True,
)
temporary_file = temporary_database_path(DB_FILE)
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)
validate_database(conn)
counts = get_counts(conn)
# Nová databáza nahradí starú až po úspešnom vytvorení.
os.replace(
temporary_file,
DB_FILE,
)
except Exception:
if temporary_file.exists():
temporary_file.unlink()
raise
print(
f"[green]SQLite index vytvorený:[/green] "
f"{DB_FILE}"
)
print(
f"Dokumentov: {counts['documents']}"
)
print(
f"Chunkov: {counts['chunks']}"
)
print(
f"FTS5 chunkov: {counts['fts_chunks']}"
)
print(
f"Tag záznamov: {counts['tags']}"
)
print(
f"Kategória záznamov: "
f"{counts['categories']}"
)
return counts
def main() -> None:
build_database()
if __name__ == "__main__":
main()