Optimalizácia RAG spracovania

This commit is contained in:
Ján Pták 2026-10-26 13:38:15 +01:00
parent bfe8cbba1b
commit d8961f9db6
15 changed files with 560 additions and 91 deletions

View File

@ -983,6 +983,42 @@ tools/rag/evaluation/json_files/rag_answer_overrides.json
Pôvodný frozen `questions.json` tým ostáva nezmenený. Pôvodný frozen `questions.json` tým ostáva nezmenený.
## Tokenové limity, segmenty a embeddingová cache
Chunkovanie dodrží `CHUNK_MAX_TOKENS` aj pri dlhom slove, riadku kódu,
tabuľke alebo nadpise. Posledné delenie zachová Unicode znaky a obsah;
celá cesta nadpisov zostáva v metadátach.
Finálny embeddingový vstup obsahuje názov, autora, tagy, kategórie, obsah
aj prefix `passage: `. Jeho veľkosť sa meria skutočným tokenizérom modelu
vrátane špeciálnych tokenov. Nadlimitný vstup sa bez straty obsahu rozdelí
na segmenty. Pôvodný chunk, FTS a graf ostávajú zachované. Vektorové
vyhľadávanie vyberá najlepšiu zhodu segmentu a každý chunk vráti iba raz.
Encoder nadlimitný nesegmentovaný vstup odmietne namiesto tichého orezania.
`data/embedding_cache.sqlite` uchováva vektory podľa presného vstupu a
odtlačku váh, tokenizéra a konfigurácie modelu. Pri zmene relevantných
metadát alebo modelu sa vektor prepočíta. Kontrolný súčet zachytí poškodený
záznam cache. Prvé vytvorenie novej cache prepočíta všetko; ďalšia obnova
nezmenených vstupov použije cache. Výsledok synchronizácie obsahuje
`embedding_segments`, `embedding_computed` a `embedding_reused`; počet
`embedding_chunks` stále označuje počet pokrytých pôvodných chunkov.
Cache je obnoviteľná optimalizácia; jej odstránenie pri zastavenej
synchronizácii vyvolá nový výpočet, nezmení aktívne indexy.
`POST /rag` predvolene vracia dôkazy iba raz v `context` a stručné citačné
metadáta v `sources`. Diagnostické expanzie a ich texty sa vrátia len pri
explicitnom `include_diagnostics: true`, oddelene v `diagnostics`.
RAG evaluačný klient diagnostiku pred odoslaním modelu vždy odstráni.
`RAG_LLM_MAX_TOKENS` (predvolene 12000) obmedzuje celý JSON podkladu
vrátane kontextu, pokynov, zdrojov a verzie dát. Nie je to rozpočet celej
konverzácie ani výstupu modelu. `RAG_LLM_TOKEN_ENCODING` má predvolene
`cl100k_base`; ide o explicitný odhad, nie presný tokenizér všetkých
vzdialených modelov. Výstup `context_budget` uvádza použitý tokenizér,
limit a skutočný počet podľa tohto odhadu. Pri orezaní nástroj nastaví
`complete=false`, `truncated=true` a dôvod `llm_context_budget`.
## Manuálna synchronizácia a stav úloh ## Manuálna synchronizácia a stav úloh
`POST /sync` uloží požiadavku do persistentnej fronty a ihneď vráti HTTP 202. `POST /sync` uloží požiadavku do persistentnej fronty a ihneď vráti HTTP 202.

View File

@ -41,6 +41,7 @@ from tools.rag.common import (
ZPWIKI_ROOT, ZPWIKI_ROOT,
) )
from tools.rag.rag_utils import build_rag_context from tools.rag.rag_utils import build_rag_context
from tools.rag.llm_context import compact_rag_payload
from tools.rag.search_utils import search_database from tools.rag.search_utils import search_database
from tools.wiki_git import expected_branch from tools.wiki_git import expected_branch
from tools.publication import DEFAULT_PUBLISHED_ONLY from tools.publication import DEFAULT_PUBLISHED_ONLY
@ -111,6 +112,7 @@ class SearchRequest(
class RagRequest( class RagRequest(
StrictRequestModel StrictRequestModel
): ):
include_diagnostics: bool = False
query: str = Field( query: str = Field(
..., ...,
min_length=1, min_length=1,
@ -230,6 +232,8 @@ class RagResponse(RetrievalResponse):
RagSource RagSource
] ]
index_version: dict[str, Any] | None = None index_version: dict[str, Any] | None = None
context_budget: dict[str, Any] | None = None
diagnostics: dict[str, Any] | None = None
class SearchResponse(RetrievalResponse): class SearchResponse(RetrievalResponse):
@ -577,7 +581,7 @@ def rag(
), ),
) )
http_request.app.state.embedding_ready = True http_request.app.state.embedding_ready = True
return result return compact_rag_payload(result, include_diagnostics=request.include_diagnostics)
except Exception as error: except Exception as error:
raise_retrieval_error( raise_retrieval_error(

View File

@ -14,6 +14,8 @@ services:
CHUNK_OVERLAP_TOKENS: "70" CHUNK_OVERLAP_TOKENS: "70"
CHUNK_MIN_TOKENS: "80" CHUNK_MIN_TOKENS: "80"
CHUNK_TOKEN_ENCODING: cl100k_base CHUNK_TOKEN_ENCODING: cl100k_base
RAG_LLM_MAX_TOKENS: ${RAG_LLM_MAX_TOKENS:-12000}
RAG_LLM_TOKEN_ENCODING: ${RAG_LLM_TOKEN_ENCODING:-cl100k_base}
HF_HOME: /cache/huggingface HF_HOME: /cache/huggingface
GIT_SSH_COMMAND: >- GIT_SSH_COMMAND: >-
ssh -F /dev/null -i /run/secrets/zpwiki_ssh_key ssh -F /dev/null -i /run/secrets/zpwiki_ssh_key

View File

@ -0,0 +1,115 @@
import copy
import json
import sqlite3
from pathlib import Path
import numpy as np
import pytest
from tools.rag import build_chunks as chunking
from tools.rag import build_sqlite_index as indexer
from tools.rag import embedding_utils as embedding
from tools.rag import search_core
from tools.rag.embedding_cache import EmbeddingCache, cache_key
from tools.text_budget import split_to_fit
def test_unicode_split_is_lossless_and_respects_exact_counter():
text = 'Príliš dlhý text 🧑🏽‍💻 漢字\n' * 30
parts = split_to_fit(text, lambda value: chunking.TOKEN_COUNTER.count(value) <= 15)
assert ''.join(parts) == text
assert all(chunking.TOKEN_COUNTER.count(part) <= 15 for part in parts)
assert '\ufffd' not in ''.join(parts)
@pytest.mark.parametrize('body', ['a' * 12000, '```python\n' + 'x' * 12000 + '\n```',
'| ' + 'stĺpec' * 2500 + ' |\n| --- |\n| hodnota |'])
def test_pathological_chunks_have_hard_limit_without_lost_body(body):
drafts = chunking.chunk_markdown(body, document_title='Test')
assert drafts and max(map(chunking.token_count, drafts)) <= chunking.MAX_TOKENS
combined = ''.join(draft.literal_text if draft.literal_text is not None else
''.join(draft.units) for draft in drafts)
# Newlines/context can be repeated by natural block splitting; every long
# unbroken payload is retained in its original order by the fallback.
payload = 'a' * 12000 if body.startswith('a') else ('x' * 12000 if '```' in body else 'stĺpec' * 2500)
assert payload in combined
def test_long_heading_is_preserved_in_content_and_metadata():
heading = 'dlhy_nadpis_' * 700
drafts = chunking.chunk_markdown('# ' + heading + '\n\nKoniec obsahu.', document_title='Test')
assert max(map(chunking.token_count, drafts)) <= chunking.MAX_TOKENS
assert heading in ''.join(draft.text for draft in drafts)
assert 'Koniec obsahu.' in ''.join(draft.text for draft in drafts)
assert all(draft.heading_paths == [[heading]] for draft in drafts)
def test_real_embedding_tokenizer_counts_metadata_prefix_and_special_tokens():
chunk = {'title': 'Dlhý názov ' * 160, 'tags': ['tag'] * 60, 'text': 'Text práce 🐈 ' * 230}
text = embedding.build_embedding_text(chunk)
parts = embedding.split_passage(text)
assert len(parts) > 1 and ''.join(parts) == text
assert max(map(embedding.passage_token_count, parts)) <= embedding.get_embedding_model(embedding.embedding_model_name()).max_seq_length
with pytest.raises(ValueError, match='exceeds model limit'):
embedding.embed_passages([text])
def test_model_fingerprint_is_independent_of_last_tokenizer_call():
name = embedding.embedding_model_name()
embedding._model_profile.cache_clear()
before = embedding.model_profile(name)['fingerprint']
model = embedding.get_embedding_model(name)
model.tokenizer(['query: krátky text', 'query: dlhší text na overenie'],
truncation=True, padding=True, max_length=model.max_seq_length)
embedding._model_profile.cache_clear()
assert embedding.model_profile(name)['fingerprint'] == before
def test_oversized_query_is_explicitly_rejected_instead_of_silently_truncated():
with pytest.raises(ValueError, match='Query prekračuje tokenový limit'):
embedding.embed_query('🧑🏽‍💻' * 200)
def test_cache_key_binds_metadata_and_exact_model_and_checks_corruption(tmp_path):
chunk = {'text': 'obsah', 'title': 'názov', 'author': 'autor', 'tags': ['tag'], 'categories': ['category']}
profile = {'fingerprint': 'weights-and-tokenizer-a'}
key = cache_key(profile, embedding.build_embedding_text(chunk))
for field in chunk:
changed = copy.deepcopy(chunk)
changed[field] = ['other'] if isinstance(chunk[field], list) else 'other'
assert cache_key(profile, embedding.build_embedding_text(changed)) != key
assert cache_key({'fingerprint': 'new-model'}, embedding.build_embedding_text(chunk)) != key
cache = EmbeddingCache(tmp_path / 'cache.sqlite')
cache.save({key: np.array([1., 0.], dtype=np.float32)})
assert key in cache.load([key])
with sqlite3.connect(cache.path) as conn:
conn.execute('UPDATE embeddings SET vector=?', (b'corrupted',))
assert cache.load([key]) == {}
def test_multi_segment_index_reuses_cache_and_returns_each_chunk_once(tmp_path, monkeypatch):
calls = []
monkeypatch.setattr(indexer, 'model_profile', lambda _: {'fingerprint': 'test-model'})
monkeypatch.setattr(indexer, 'split_passage', lambda text, **kw: [text[:8], text[8:]])
def encode(texts):
calls.append(texts)
return np.array([[1., 0.] for _ in texts], dtype=np.float32)
monkeypatch.setattr(indexer, 'embed_passages', encode)
monkeypatch.setattr(search_core, 'model_profile', lambda _: {'fingerprint': 'test-model'})
monkeypatch.setattr(search_core, 'embed_query', lambda *a, **kw: np.array([1., 0.]))
docs = [{'path': 'pages/test/README.md', 'title': 'Test', 'published': True}]
chunks = [{'chunk_id': 'test::0', 'document_path': docs[0]['path'], 'chunk_index': 0,
'text': 'celý dlhý obsah', 'title': 'Test'}]
cache = tmp_path / 'cache.sqlite'
for run in range(2):
with sqlite3.connect(':memory:') as conn:
conn.row_factory = sqlite3.Row
indexer.create_tables(conn)
indexer.insert_documents(conn, docs)
indexer.insert_chunks(conn, chunks)
stats = indexer.insert_embeddings(conn, chunks, cache_file=cache)
assert conn.execute('SELECT COUNT(*) FROM chunk_embeddings').fetchone()[0] == 2
assert stats == {'embedding_reused': 2 if run else 0, 'embedding_computed': 0 if run else 2}
result = search_core.run_vector_query(conn, 'obsah', 5, True)
assert len(result) == 1 and result[0]['chunk_id'] == 'test::0'
assert len(calls) == 1

View File

@ -0,0 +1,67 @@
import copy
import json
from tools.rag.build_chunks import TOKEN_COUNTER
from tools.rag.llm_context import compact_rag_payload
def response(text='Unikátny dôkaz o práci osoby. '):
return {'status': 'ok', 'query': 'Práca?', 'engine': 'test', 'strategies': ['test'],
'complete': None, 'truncated': False, 'truncation_reasons': [], 'index_version': {'version': 'test'},
'instructions': ['Cituj iba dodané zdroje.'], 'answer_format': {},
'sources': [{'source_id': 'S1', 'document_path': 'pages/eva/README.md', 'title': 'Eva',
'source_url': 'https://example.test/eva', 'text': text,
'query_evidence': {'huge_diagnostic': text * 10},
'retrieval': {'rank': 1}}], 'context': text}
def test_compact_payload_contains_evidence_once_and_preserves_sources():
original = response()
before = copy.deepcopy(original)
compact = compact_rag_payload(original)
assert original == before
assert json.dumps(compact, ensure_ascii=False).count(original['sources'][0]['text']) == 1
assert 'text' not in compact['sources'][0] and 'query_evidence' not in compact['sources'][0]
assert compact['sources'][0]['source_url'] == 'https://example.test/eva'
assert compact['index_version'] == original['index_version']
assert compact_rag_payload(compact) == compact
def test_entire_tool_json_is_bounded_and_marks_omitted_evidence():
original = response('Veľmi dlhý dôkaz 🧑🏽‍💻. ' * 2000)
original['sources'].append({**original['sources'][0], 'source_id': 'S2', 'text': 'Ďalší dôkaz'})
compact = compact_rag_payload(original, max_tokens=700)
used = TOKEN_COUNTER.count(json.dumps(compact, ensure_ascii=False, separators=(',', ':')))
assert used <= 700 and compact['context_budget']['used_tokens'] == used
assert compact['complete'] is False and compact['truncated'] is True
assert 'llm_context_budget' in compact['truncation_reasons']
assert 'KONIEC ZDROJA S1' in compact['context']
assert compact['source_count'] == 1 and '[OBSAH ZDROJA OREZANÝ]' in compact['context']
assert json.loads(compact['context'].splitlines()[1])['truncated'] is True
def test_diagnostics_are_explicit_and_never_forwarded_by_projection():
verbose = compact_rag_payload(response(), include_diagnostics=True)
assert verbose['diagnostics']['sources'][0]['query_evidence']
projected = compact_rag_payload(verbose)
assert 'diagnostics' not in projected
assert 'huge_diagnostic' not in json.dumps(projected)
def test_http_serialization_stays_within_budget_and_diagnostics_are_opt_in(monkeypatch):
from fastapi.testclient import TestClient
import app.main as main
import app.routes as routes
original = response('Dlhý dôkaz 🐈 ' * 1800)
monkeypatch.setattr(main, 'embed_query', lambda _: None)
monkeypatch.setattr(routes, 'build_rag_context', lambda *a, **kw: original)
monkeypatch.setenv('RAG_LLM_MAX_TOKENS', '700')
with TestClient(main.app) as client:
headers = {'X-API-Key': 'a' * 64}
compact = client.post('/rag', json={'query': 'Práca?'}, headers=headers)
assert compact.status_code == 200
assert TOKEN_COUNTER.count(compact.text) <= 700
assert 'diagnostics' not in compact.json()
assert 'text' not in compact.json()['sources'][0]
verbose = client.post('/rag', json={'query': 'Práca?', 'include_diagnostics': True}, headers=headers)
assert verbose.status_code == 200 and verbose.json()['diagnostics'] == original

View File

@ -23,6 +23,7 @@ from tools.rag.common import (
load_zpwiki_page, load_zpwiki_page,
write_json, write_json,
) )
from tools.text_budget import split_to_fit
MAX_TOKENS = int(os.getenv("CHUNK_MAX_TOKENS", "450")) MAX_TOKENS = int(os.getenv("CHUNK_MAX_TOKENS", "450"))
@ -51,9 +52,12 @@ class ChunkDraft:
context_lines: list[str] context_lines: list[str]
units: list[str] units: list[str]
heading_paths: list[list[str]] heading_paths: list[list[str]]
literal_text: str | None = None
@property @property
def text(self) -> str: def text(self) -> str:
if self.literal_text is not None:
return self.literal_text
context = "\n".join(self.context_lines).strip() context = "\n".join(self.context_lines).strip()
body = join_units(self.units) body = join_units(self.units)
@ -284,7 +288,8 @@ def split_by_words(text: str, max_tokens: int) -> list[str]:
if current: if current:
pieces.append(" ".join(current)) pieces.append(" ".join(current))
return pieces return [part for piece in pieces for part in split_to_fit(
piece, lambda value: TOKEN_COUNTER.count(value) <= max_tokens)]
def split_code_block(text: str, max_tokens: int) -> list[str]: def split_code_block(text: str, max_tokens: int) -> list[str]:
@ -587,6 +592,10 @@ def chunk_section(
context_tokens = TOKEN_COUNTER.count( context_tokens = TOKEN_COUNTER.count(
"\n".join(context_lines) "\n".join(context_lines)
) )
if context_tokens > MAX_TOKENS - 16:
# Preserve even enormous titles/headings once in the content. The final
# token fallback splits this draft and retains the full heading metadata.
return [ChunkDraft(context_lines, [section.body], heading_path_list(section.heading_path))]
body_limit = max( body_limit = max(
1, 1,
@ -908,7 +917,14 @@ def chunk_markdown(
) )
) )
return balance_short_chunks(drafts) bounded = []
for draft in balance_short_chunks(drafts):
if token_count(draft) <= MAX_TOKENS:
bounded.append(draft)
else:
for piece in split_to_fit(draft.text, lambda value: TOKEN_COUNTER.count(value) <= MAX_TOKENS):
bounded.append(ChunkDraft([], [piece], draft.heading_paths, literal_text=piece))
return bounded
def content_hash( def content_hash(

View File

@ -3,6 +3,7 @@ from __future__ import annotations
import json import json
import os import os
import sqlite3 import sqlite3
import numpy as np
import sys import sys
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@ -27,7 +28,10 @@ from tools.rag.embedding_utils import (
embed_passages, embed_passages,
embedding_model_name, embedding_model_name,
vector_to_blob, vector_to_blob,
model_profile,
split_passage,
) )
from tools.rag.embedding_cache import EmbeddingCache, cache_key
FTS_TOKENIZER = "unicode61 remove_diacritics 2" FTS_TOKENIZER = "unicode61 remove_diacritics 2"
@ -133,16 +137,20 @@ def create_tables(
); );
CREATE TABLE chunk_embeddings ( CREATE TABLE chunk_embeddings (
chunk_id TEXT PRIMARY KEY, chunk_id TEXT NOT NULL,
segment_index INTEGER NOT NULL DEFAULT 0,
model TEXT NOT NULL, model TEXT NOT NULL,
dimensions INTEGER NOT NULL, dimensions INTEGER NOT NULL,
embedding BLOB NOT NULL, embedding BLOB NOT NULL,
PRIMARY KEY(chunk_id, segment_index),
FOREIGN KEY(chunk_id) FOREIGN KEY(chunk_id)
REFERENCES chunks(chunk_id) REFERENCES chunks(chunk_id)
ON UPDATE CASCADE ON UPDATE CASCADE
ON DELETE CASCADE ON DELETE CASCADE
); );
CREATE TABLE embedding_model_info (metadata_json TEXT NOT NULL);
CREATE INDEX idx_documents_path CREATE INDEX idx_documents_path
ON documents(path); ON documents(path);
@ -348,78 +356,43 @@ def insert_chunks(
) )
def insert_embeddings( def insert_embeddings(conn: sqlite3.Connection, chunks: list[dict], *, cache_file: Path) -> dict:
conn: sqlite3.Connection,
chunks: list[dict],
) -> None:
if not chunks: if not chunks:
return return {'embedding_reused': 0, 'embedding_computed': 0}
model_name = embedding_model_name()
chunk_ids: list[str] = [] profile = model_profile(model_name)
texts: list[str] = [] conn.execute('INSERT INTO embedding_model_info VALUES(?)', (json.dumps(profile, sort_keys=True),))
segments = []
for chunk in chunks: for chunk in chunks:
chunk_id = str( if not chunk.get('chunk_id'):
chunk.get("chunk_id") or "" raise ValueError('Chunk bez chunk_id nemôže mať embedding')
).strip() text = build_embedding_text(chunk)
for index, part in enumerate(split_passage(text, model_name=model_name)):
if not chunk_id: segments.append((chunk['chunk_id'], index, part, cache_key(profile, part)))
raise ValueError( cache = EmbeddingCache(cache_file)
"Chunk bez chunk_id nemôže " found = cache.load([row[3] for row in segments])
"mať embedding" missing = {row[3]: row[2] for row in segments if row[3] not in found}
) computed = {}
if missing:
chunk_ids.append( vectors = embed_passages(list(missing.values()))
chunk_id if len(vectors) != len(missing):
) raise RuntimeError('Počet embeddingov sa nezhoduje s počtom vstupov')
for key, vector in zip(missing, vectors):
texts.append( vector = np.asarray(vector, dtype=np.float32)
build_embedding_text( if vector.ndim != 1 or not len(vector) or not np.isfinite(vector).all() or np.linalg.norm(vector) == 0:
chunk raise RuntimeError('Embedding nie je platný nenulový vektor')
) computed[key] = vector
) cache.save(computed)
all_vectors = {**found, **computed}
vectors = embed_passages( dimensions = {len(vector) for vector in all_vectors.values()}
texts if len(dimensions) != 1:
) raise RuntimeError('Embeddingy majú rozdielne rozmery')
conn.executemany('''INSERT INTO chunk_embeddings
if len(vectors) != len(chunk_ids): (chunk_id,segment_index,model,dimensions,embedding) VALUES(?,?,?,?,?)''',
raise RuntimeError( [(chunk_id, index, model_name, len(all_vectors[key]), vector_to_blob(all_vectors[key]))
"Počet embeddingov sa " for chunk_id, index, _, key in segments])
"nezhoduje s počtom chunkov" return {'embedding_reused': sum(row[3] in found for row in segments),
) 'embedding_computed': len(computed)}
model_name = (
embedding_model_name()
)
rows: list[tuple] = []
for chunk_id, vector in zip(
chunk_ids,
vectors,
):
rows.append(
(
chunk_id,
model_name,
int(vector.shape[0]),
vector_to_blob(vector),
)
)
conn.executemany(
"""
INSERT INTO chunk_embeddings (
chunk_id,
model,
dimensions,
embedding
)
VALUES (?, ?, ?, ?)
""",
rows,
)
def build_fts_index( def build_fts_index(
@ -536,7 +509,7 @@ def validate_database(
embedding_count = conn.execute( embedding_count = conn.execute(
""" """
SELECT COUNT(*) SELECT COUNT(DISTINCT chunk_id)
FROM chunk_embeddings FROM chunk_embeddings
""" """
).fetchone()[0] ).fetchone()[0]
@ -579,6 +552,7 @@ def get_counts(
conn: sqlite3.Connection, conn: sqlite3.Connection,
) -> dict[str, int]: ) -> dict[str, int]:
return { return {
"embedding_segments": conn.execute("SELECT COUNT(*) FROM chunk_embeddings").fetchone()[0],
"documents": conn.execute( "documents": conn.execute(
""" """
SELECT COUNT(*) SELECT COUNT(*)
@ -599,7 +573,7 @@ def get_counts(
).fetchone()[0], ).fetchone()[0],
"embedding_chunks": conn.execute( "embedding_chunks": conn.execute(
""" """
SELECT COUNT(*) SELECT COUNT(DISTINCT chunk_id)
FROM chunk_embeddings FROM chunk_embeddings
""" """
).fetchone()[0], ).fetchone()[0],
@ -627,7 +601,7 @@ def temporary_database_path(
def build_database(*, documents_file: Path | None = None, chunks_file: Path | None = None, def build_database(*, documents_file: Path | None = None, chunks_file: Path | None = None,
db_file: Path | None = None) -> dict[str, int]: db_file: Path | None = None, cache_file: Path | None = None) -> dict[str, int]:
target = db_file or DB_FILE target = db_file or DB_FILE
documents = read_json( documents = read_json(
documents_file or DOCUMENTS_FILE documents_file or DOCUMENTS_FILE
@ -685,18 +659,13 @@ def build_database(*, documents_file: Path | None = None, chunks_file: Path | No
conn conn
) )
insert_embeddings( embedding_stats = insert_embeddings(conn, chunks, cache_file=cache_file or target.parent / "embedding_cache.sqlite")
conn,
chunks,
)
validate_database( validate_database(
conn conn
) )
counts = get_counts( counts = {**get_counts(conn), **embedding_stats}
conn
)
os.replace( os.replace(
temporary_file, temporary_file,

View File

@ -0,0 +1,42 @@
"""Reusable vectors keyed by the exact input and model fingerprint."""
import hashlib
import json
import sqlite3
from contextlib import closing
from pathlib import Path
import numpy as np
def cache_key(profile: dict, text: str) -> str:
return hashlib.sha256(json.dumps({'profile': profile['fingerprint'], 'text': text},
ensure_ascii=False, sort_keys=True).encode()).hexdigest()
class EmbeddingCache:
def __init__(self, path: Path):
self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True)
with closing(sqlite3.connect(self.path)) as conn, conn:
conn.execute('''CREATE TABLE IF NOT EXISTS embeddings (
key TEXT PRIMARY KEY, dimensions INTEGER NOT NULL, vector BLOB NOT NULL,
vector_sha256 TEXT NOT NULL)''')
def load(self, keys: list[str]) -> dict[str, np.ndarray]:
found = {}
with closing(sqlite3.connect(self.path)) as conn:
for key in set(keys):
row = conn.execute('SELECT dimensions,vector,vector_sha256 FROM embeddings WHERE key=?', (key,)).fetchone()
if not row or hashlib.sha256(row[1]).hexdigest() != row[2]:
continue
vector = np.frombuffer(row[1], dtype=np.float32)
if len(vector) == row[0] and np.isfinite(vector).all() and np.linalg.norm(vector) > 0:
found[key] = vector.copy()
return found
def save(self, vectors: dict[str, np.ndarray]):
with closing(sqlite3.connect(self.path)) as conn, conn:
for key, vector in vectors.items():
blob = np.asarray(vector, dtype=np.float32).tobytes()
conn.execute('INSERT OR REPLACE INTO embeddings VALUES(?,?,?,?)',
(key, len(vector), blob, hashlib.sha256(blob).hexdigest()))

View File

@ -1,10 +1,13 @@
from __future__ import annotations from __future__ import annotations
import os import os
import hashlib
import json
from functools import lru_cache from functools import lru_cache
from typing import Any from typing import Any
import numpy as np import numpy as np
from tools.text_budget import split_to_fit
DEFAULT_EMBEDDING_MODEL = ( DEFAULT_EMBEDDING_MODEL = (
@ -44,19 +47,66 @@ def embedding_batch_size() -> int:
) )
@lru_cache(maxsize=2)
def get_embedding_model( def get_embedding_model(
model_name: str, model_name: str,
): ):
return _load_embedding_model(model_name, os.getenv('EMBEDDING_REVISION') or None)
@lru_cache(maxsize=2)
def _load_embedding_model(model_name: str, revision: str | None):
from sentence_transformers import ( from sentence_transformers import (
SentenceTransformer, SentenceTransformer,
) )
return SentenceTransformer( return SentenceTransformer(
model_name model_name,
revision=revision,
) )
def model_profile(model_name: str) -> dict:
"""Fingerprint the loaded weights, tokenizer and pooling, not a mutable alias."""
return _model_profile(model_name, get_embedding_model(model_name))
@lru_cache(maxsize=2)
def _model_profile(model_name, model):
import torch
weights = hashlib.sha256()
for name, tensor in sorted(model.state_dict().items()):
weights.update(name.encode())
weights.update(str((tuple(tensor.shape), tensor.dtype)).encode())
weights.update(tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes())
tokenizer_config = json.loads(model.tokenizer.backend_tokenizer.to_str())
# Fast tokenizers store the last call's padding/truncation in the backend.
# These mutable execution settings must not change the vocabulary fingerprint.
tokenizer_config.pop('padding', None)
tokenizer_config.pop('truncation', None)
tokenizer = json.dumps(tokenizer_config, sort_keys=True, separators=(',', ':'), ensure_ascii=False)
modules = {name: module.get_config_dict() for name, module in model._modules.items()
if hasattr(module, 'get_config_dict')}
profile = {'model': model_name, 'weights_sha256': weights.hexdigest(),
'tokenizer_sha256': hashlib.sha256(tokenizer.encode()).hexdigest(),
'max_seq_length': int(model.max_seq_length), 'modules': modules,
'passage_prefix': 'passage: ', 'query_prefix': 'query: ',
'normalize_embeddings': True, 'dtype': 'float32', 'segment_policy': 'lossless-v1'}
profile['fingerprint'] = hashlib.sha256(json.dumps(profile, sort_keys=True).encode()).hexdigest()
return profile
def passage_token_count(text: str, *, model_name: str | None = None) -> int:
model = get_embedding_model(model_name or embedding_model_name())
return len(model.tokenizer('passage: ' + text, truncation=False, add_special_tokens=True,
verbose=False)['input_ids'])
def split_passage(text: str, *, model_name: str | None = None) -> list[str]:
selected = model_name or embedding_model_name()
limit = int(get_embedding_model(selected).max_seq_length)
return split_to_fit(text, lambda value: passage_token_count(value, model_name=selected) <= limit)
def build_embedding_text( def build_embedding_text(
chunk: dict[str, Any], chunk: dict[str, Any],
) -> str: ) -> str:
@ -156,9 +206,11 @@ def embed_passages(
prepared = [ prepared = [
"passage: " "passage: "
+ text.strip() + text
for text in texts for text in texts
] ]
if any(passage_token_count(text, model_name=selected_model) > model.max_seq_length for text in texts):
raise ValueError('Embedding input exceeds model limit; split it before encoding.')
vectors = model.encode( vectors = model.encode(
prepared, prepared,
@ -197,6 +249,11 @@ def embed_query(
selected_model selected_model
) )
length = len(model.tokenizer('query: ' + clean_query, truncation=False,
add_special_tokens=True, verbose=False)['input_ids'])
if length > model.max_seq_length:
raise ValueError('Query prekračuje tokenový limit embeddingového modelu')
vector = model.encode( vector = model.encode(
[ [
"query: " "query: "

View File

@ -1,4 +1,5 @@
from __future__ import annotations from __future__ import annotations
from tools.rag.llm_context import compact_rag_payload
import json import json
import os import os
@ -861,6 +862,7 @@ def run_question(
tool_latency tool_latency
) )
rag_result = compact_rag_payload(rag_result)
current_urls = [source['source_url'] for source in rag_result.get('sources', []) current_urls = [source['source_url'] for source in rag_result.get('sources', [])
if source.get('source_url')] if source.get('source_url')]
retrieval_coverage.append({key: rag_result.get(key) for key in retrieval_coverage.append({key: rag_result.get(key) for key in
@ -910,6 +912,7 @@ def run_question(
"content": json.dumps( "content": json.dumps(
rag_result, rag_result,
ensure_ascii=False, ensure_ascii=False,
separators=(',', ':'),
), ),
} }
) )

114
tools/rag/llm_context.py Normal file
View File

@ -0,0 +1,114 @@
"""A bounded tool message with one copy of evidence and no retrieval diagnostics."""
from __future__ import annotations
import copy
import json
import os
from tools.rag.build_chunks import TokenCounter
from tools.text_budget import split_to_fit
SOURCE_FIELDS = ('source_id', 'title', 'author', 'document_path', 'source_url', 'published',
'categories', 'tags', 'section')
FIELDS = ('status', 'reason', 'query', 'engine', 'strategies', 'instructions', 'answer_format',
'complete', 'truncated', 'truncation_reasons', 'index_version', 'source_count')
def compact_rag_payload(data: dict, *, max_tokens: int | None = None,
include_diagnostics: bool = False) -> dict:
from tools.rag.rag_utils import build_context_text
limit = max_tokens if max_tokens is not None else int(os.getenv('RAG_LLM_MAX_TOKENS', '12000'))
if limit < 256:
raise ValueError('RAG_LLM_MAX_TOKENS musí byť aspoň 256')
encoding = os.getenv('RAG_LLM_TOKEN_ENCODING', 'cl100k_base')
counter = TokenCounter(encoding)
if counter.encoding is None:
raise RuntimeError('Kontextový rozpočet vyžaduje platný tiktoken tokenizér')
base = {key: copy.deepcopy(data[key]) for key in FIELDS if key in data}
base.setdefault('status', 'ok')
base.setdefault('complete', None)
base.setdefault('truncated', False)
base.setdefault('truncation_reasons', [])
input_sources = [{**source, 'source_id': source.get('source_id') or f'S{number}'}
for number, source in enumerate(data.get('sources', []), 1)]
references = [{key: source[key] for key in SOURCE_FIELDS if key in source}
for source in input_sources]
def stamp(payload):
payload['context_budget'] = {'max_tokens': limit, 'used_tokens': 0,
'tokenizer': counter.name, 'tokens_are_estimate': True}
for _ in range(8):
used = counter.count(json.dumps(payload, ensure_ascii=False, separators=(',', ':')))
if payload['context_budget']['used_tokens'] == used:
break
payload['context_budget']['used_tokens'] = used
return payload
def render(sources, *, truncated=False, raw_context=None):
payload = copy.deepcopy(base)
if truncated:
payload['complete'] = False
payload['truncated'] = True
payload['truncation_reasons'] = list(dict.fromkeys(
[*payload['truncation_reasons'], 'llm_context_budget']))
payload['sources'] = [{key: source[key] for key in SOURCE_FIELDS if key in source} for source in sources]
payload['source_count'] = len(sources)
coverage = {key: payload.get(key) for key in ['complete', 'truncated', 'truncation_reasons', 'index_version']}
context = raw_context if raw_context is not None else build_context_text(sources)
payload['context'] = 'METADÁTA VYHĽADANIA\n' + json.dumps(coverage, ensure_ascii=False) + '\n\n' + context
return stamp(payload)
def fits(payload):
return payload['context_budget']['used_tokens'] <= limit
# An already projected server response needs no reconstruction by its client.
if data.get('context_budget'):
result = stamp({**base, 'sources': references, 'context': data.get('context', '')})
if not fits(result):
parts = split_to_fit(data.get('context', ''), lambda text:
fits(render(references, truncated=True, raw_context=text)))
result = render(references, truncated=True, raw_context=parts[0])
elif any('text' in source for source in input_sources):
sources = input_sources
result = render(sources)
if not fits(result):
selected = []
if not fits(render([], truncated=True)):
raise ValueError('Tokenový rozpočet je menší ako povinné pokyny a metadáta')
for source in sources:
candidate = render([*selected, source], truncated=True)
if fits(candidate):
selected.append(source)
continue
body = source.get('text') or ''
suffix = '\n[OBSAH ZDROJA OREZANÝ]'
def body_fits(text):
return fits(render([*selected, {**source, 'text': text + suffix}], truncated=True))
if body and body_fits(body[:1]):
# Only the first fitting prefix is delivered; the complete
# original remains available through the diagnostics option.
low, high = 1, len(body)
while low < high:
middle = (low + high + 1) // 2
if body_fits(body[:middle]): low = middle
else: high = middle - 1
selected.append({**source, 'text': body[:low] + suffix})
break
result = render(selected, truncated=True)
if not selected:
result['status'] = 'no_answer'
result['reason'] = 'llm_context_budget_exhausted'
result = stamp(result)
else:
# Compatibility with older tool servers that expose only context + refs.
context = data.get('context', '')
result = render(references, raw_context=context)
if not fits(result):
parts = split_to_fit(context, lambda text: fits(render(references, truncated=True, raw_context=text)))
result = render(references, truncated=True, raw_context=parts[0])
if not fits(result):
raise ValueError('Povinné metadáta prekračujú kontextový rozpočet')
if include_diagnostics:
result['diagnostics'] = copy.deepcopy(data)
return result

View File

@ -11,6 +11,7 @@ from tools.rag.embedding_utils import (
blob_to_vector, blob_to_vector,
cosine_similarity, cosine_similarity,
embed_query, embed_query,
model_profile,
) )
@ -305,6 +306,10 @@ def embedding_index_info(
raise RuntimeError( raise RuntimeError(
"Embedding index je prázdny" "Embedding index je prázdny"
) )
if conn.execute("SELECT 1 FROM sqlite_master WHERE name='embedding_model_info'").fetchone():
metadata = conn.execute('SELECT metadata_json FROM embedding_model_info').fetchone()
if metadata and json.loads(metadata[0])['fingerprint'] != model_profile(str(row['model']))['fingerprint']:
raise RuntimeError('Verzia embeddingového modelu nezodpovedá indexu; obnovte index.')
return ( return (
str(row["model"]), str(row["model"]),
@ -545,9 +550,12 @@ def run_vector_query(
) )
) )
return candidates[ # A chunk can have several lossless embedding segments. Rank by its best
:candidate_limit # matching segment, and never spend several result slots on the same chunk.
] unique = {}
for item in candidates:
unique.setdefault(item['chunk_id'], item)
return list(unique.values())[:candidate_limit]
def tokens_match_for_title( def tokens_match_for_title(

View File

@ -144,6 +144,9 @@ def test_database_contains_documents_chunks_metadata_and_fts(
"chunks": 1, "chunks": 1,
"fts_chunks": 1, "fts_chunks": 1,
"embedding_chunks": 1, "embedding_chunks": 1,
"embedding_segments": 1,
"embedding_reused": 0,
"embedding_computed": 1,
"tags": 2, "tags": 2,
"categories": 1, "categories": 1,
} }

View File

@ -64,12 +64,15 @@ def prepare_snapshot(wiki_root: Path, stage: Path) -> tuple[dict, dict, dict]:
documents = scan_pages(wiki_root=frozen_wiki, output=stage / 'documents.json') documents = scan_pages(wiki_root=frozen_wiki, output=stage / 'documents.json')
chunks = build_chunks(wiki_root=frozen_wiki, output=stage / 'chunks.json') chunks = build_chunks(wiki_root=frozen_wiki, output=stage / 'chunks.json')
counts = build_database(documents_file=stage / 'documents.json', chunks_file=stage / 'chunks.json', counts = build_database(documents_file=stage / 'documents.json', chunks_file=stage / 'chunks.json',
db_file=stage / 'zp_index.sqlite') db_file=stage / 'zp_index.sqlite', cache_file=stage.parent / 'embedding_cache.sqlite')
version = {'version': uuid.uuid4().hex, 'git_revision': revision, version = {'version': uuid.uuid4().hex, 'git_revision': revision,
'created_at_utc': datetime.now(timezone.utc).isoformat(), 'created_at_utc': datetime.now(timezone.utc).isoformat(),
'source_sha256': hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest(), 'source_sha256': hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest(),
'documents_sha256': digest(stage / 'documents.json'), 'chunks_sha256': digest(stage / 'chunks.json'), 'documents_sha256': digest(stage / 'documents.json'), 'chunks_sha256': digest(stage / 'chunks.json'),
'embedding_model': embedding_model_name(), 'publication_policy': PUBLICATION_POLICY_VERSION} 'embedding_model': embedding_model_name(), 'publication_policy': PUBLICATION_POLICY_VERSION}
if chunks:
from tools.rag.embedding_utils import model_profile
version['embedding_fingerprint'] = model_profile(embedding_model_name())['fingerprint']
with closing(sqlite3.connect(stage / 'zp_index.sqlite')) as conn, conn: with closing(sqlite3.connect(stage / 'zp_index.sqlite')) as conn, conn:
conn.execute('CREATE TABLE index_state (id INTEGER PRIMARY KEY CHECK(id=1), metadata_json TEXT NOT NULL)') conn.execute('CREATE TABLE index_state (id INTEGER PRIMARY KEY CHECK(id=1), metadata_json TEXT NOT NULL)')
conn.execute('INSERT INTO index_state VALUES (1, ?)', (json.dumps(version, sort_keys=True),)) conn.execute('INSERT INTO index_state VALUES (1, ?)', (json.dumps(version, sort_keys=True),))

30
tools/text_budget.py Normal file
View File

@ -0,0 +1,30 @@
"""Lossless character boundaries constrained by an actual token counter."""
from collections.abc import Callable
def split_to_fit(text: str, fits: Callable[[str], bool]) -> list[str]:
"""Keep every character; use natural boundaries when they fit near the end."""
pieces = []
while text:
if fits(text):
pieces.append(text)
break
low, high = 0, len(text)
while low < high:
middle = (low + high + 1) // 2
if fits(text[:middle]):
low = middle
else:
high = middle - 1
if low == 0:
raise ValueError('Token budget cannot contain one Unicode character.')
# Avoid breaking a word unless its length makes that unavoidable.
boundary = max(text.rfind('\n', 0, low), text.rfind(' ', 0, low)) + 1
cut = boundary if boundary >= low * 0.7 else low
while cut and not fits(text[:cut]):
cut -= 1
if not cut:
raise ValueError('No valid token boundary.')
pieces.append(text[:cut])
text = text[cut:]
return pieces