diff --git a/README.md b/README.md index bff97f4..979ff0c 100644 --- a/README.md +++ b/README.md @@ -983,6 +983,42 @@ tools/rag/evaluation/json_files/rag_answer_overrides.json 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 `POST /sync` uloží požiadavku do persistentnej fronty a ihneď vráti HTTP 202. diff --git a/app/routes.py b/app/routes.py index 5bb5ea1..34b2cbd 100644 --- a/app/routes.py +++ b/app/routes.py @@ -41,6 +41,7 @@ from tools.rag.common import ( ZPWIKI_ROOT, ) 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.wiki_git import expected_branch from tools.publication import DEFAULT_PUBLISHED_ONLY @@ -111,6 +112,7 @@ class SearchRequest( class RagRequest( StrictRequestModel ): + include_diagnostics: bool = False query: str = Field( ..., min_length=1, @@ -230,6 +232,8 @@ class RagResponse(RetrievalResponse): RagSource ] index_version: dict[str, Any] | None = None + context_budget: dict[str, Any] | None = None + diagnostics: dict[str, Any] | None = None class SearchResponse(RetrievalResponse): @@ -577,7 +581,7 @@ def rag( ), ) http_request.app.state.embedding_ready = True - return result + return compact_rag_payload(result, include_diagnostics=request.include_diagnostics) except Exception as error: raise_retrieval_error( diff --git a/docker-compose.yml b/docker-compose.yml index b489da2..f72d75b 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -14,6 +14,8 @@ services: CHUNK_OVERLAP_TOKENS: "70" CHUNK_MIN_TOKENS: "80" 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 GIT_SSH_COMMAND: >- ssh -F /dev/null -i /run/secrets/zpwiki_ssh_key diff --git a/test/test_embedding_budget.py b/test/test_embedding_budget.py new file mode 100644 index 0000000..68c795e --- /dev/null +++ b/test/test_embedding_budget.py @@ -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 diff --git a/test/test_llm_context_budget.py b/test/test_llm_context_budget.py new file mode 100644 index 0000000..41d8000 --- /dev/null +++ b/test/test_llm_context_budget.py @@ -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 diff --git a/tools/rag/build_chunks.py b/tools/rag/build_chunks.py index a277ad6..7dff528 100644 --- a/tools/rag/build_chunks.py +++ b/tools/rag/build_chunks.py @@ -23,6 +23,7 @@ from tools.rag.common import ( load_zpwiki_page, write_json, ) +from tools.text_budget import split_to_fit MAX_TOKENS = int(os.getenv("CHUNK_MAX_TOKENS", "450")) @@ -51,9 +52,12 @@ class ChunkDraft: context_lines: list[str] units: list[str] heading_paths: list[list[str]] + literal_text: str | None = None @property def text(self) -> str: + if self.literal_text is not None: + return self.literal_text context = "\n".join(self.context_lines).strip() body = join_units(self.units) @@ -284,7 +288,8 @@ def split_by_words(text: str, max_tokens: int) -> list[str]: if 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]: @@ -587,6 +592,10 @@ def chunk_section( context_tokens = TOKEN_COUNTER.count( "\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( 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( diff --git a/tools/rag/build_sqlite_index.py b/tools/rag/build_sqlite_index.py index aeb26d7..3e31dee 100644 --- a/tools/rag/build_sqlite_index.py +++ b/tools/rag/build_sqlite_index.py @@ -3,6 +3,7 @@ from __future__ import annotations import json import os import sqlite3 +import numpy as np import sys from pathlib import Path from typing import Any @@ -27,7 +28,10 @@ from tools.rag.embedding_utils import ( 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" @@ -133,16 +137,20 @@ def create_tables( ); 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, 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); @@ -348,78 +356,43 @@ def insert_chunks( ) -def insert_embeddings( - conn: sqlite3.Connection, - chunks: list[dict], -) -> None: +def insert_embeddings(conn: sqlite3.Connection, chunks: list[dict], *, cache_file: Path) -> dict: if not chunks: - return - - chunk_ids: list[str] = [] - texts: list[str] = [] - + 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: - chunk_id = str( - chunk.get("chunk_id") or "" - ).strip() - - if not chunk_id: - raise ValueError( - "Chunk bez chunk_id nemôže " - "mať embedding" - ) - - chunk_ids.append( - chunk_id - ) - - texts.append( - build_embedding_text( - chunk - ) - ) - - vectors = embed_passages( - texts - ) - - if len(vectors) != len(chunk_ids): - raise RuntimeError( - "Počet embeddingov sa " - "nezhoduje s počtom chunkov" - ) - - 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, - ) + 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( @@ -536,7 +509,7 @@ def validate_database( embedding_count = conn.execute( """ - SELECT COUNT(*) + SELECT COUNT(DISTINCT chunk_id) FROM chunk_embeddings """ ).fetchone()[0] @@ -579,6 +552,7 @@ 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(*) @@ -599,7 +573,7 @@ def get_counts( ).fetchone()[0], "embedding_chunks": conn.execute( """ - SELECT COUNT(*) + SELECT COUNT(DISTINCT chunk_id) FROM chunk_embeddings """ ).fetchone()[0], @@ -627,7 +601,7 @@ def temporary_database_path( 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 documents = read_json( documents_file or DOCUMENTS_FILE @@ -685,18 +659,13 @@ def build_database(*, documents_file: Path | None = None, chunks_file: Path | No conn ) - insert_embeddings( - conn, - chunks, - ) + embedding_stats = insert_embeddings(conn, chunks, cache_file=cache_file or target.parent / "embedding_cache.sqlite") validate_database( conn ) - counts = get_counts( - conn - ) + counts = {**get_counts(conn), **embedding_stats} os.replace( temporary_file, diff --git a/tools/rag/embedding_cache.py b/tools/rag/embedding_cache.py new file mode 100644 index 0000000..d6f2bfb --- /dev/null +++ b/tools/rag/embedding_cache.py @@ -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())) diff --git a/tools/rag/embedding_utils.py b/tools/rag/embedding_utils.py index 3a9a922..b84a7ee 100644 --- a/tools/rag/embedding_utils.py +++ b/tools/rag/embedding_utils.py @@ -1,10 +1,13 @@ from __future__ import annotations import os +import hashlib +import json from functools import lru_cache from typing import Any import numpy as np +from tools.text_budget import split_to_fit DEFAULT_EMBEDDING_MODEL = ( @@ -44,19 +47,66 @@ def embedding_batch_size() -> int: ) -@lru_cache(maxsize=2) def get_embedding_model( 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 ( 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( chunk: dict[str, Any], ) -> str: @@ -156,9 +206,11 @@ def embed_passages( prepared = [ "passage: " - + text.strip() + + text 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( prepared, @@ -197,6 +249,11 @@ def embed_query( 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( [ "query: " diff --git a/tools/rag/evaluation/rag_runner.py b/tools/rag/evaluation/rag_runner.py index a7204d0..e0c41a5 100644 --- a/tools/rag/evaluation/rag_runner.py +++ b/tools/rag/evaluation/rag_runner.py @@ -1,4 +1,5 @@ from __future__ import annotations +from tools.rag.llm_context import compact_rag_payload import json import os @@ -861,6 +862,7 @@ def run_question( tool_latency ) + rag_result = compact_rag_payload(rag_result) current_urls = [source['source_url'] for source in rag_result.get('sources', []) if source.get('source_url')] retrieval_coverage.append({key: rag_result.get(key) for key in @@ -910,6 +912,7 @@ def run_question( "content": json.dumps( rag_result, ensure_ascii=False, + separators=(',', ':'), ), } ) diff --git a/tools/rag/llm_context.py b/tools/rag/llm_context.py new file mode 100644 index 0000000..a9d483a --- /dev/null +++ b/tools/rag/llm_context.py @@ -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 diff --git a/tools/rag/search_core.py b/tools/rag/search_core.py index 7d82d02..6bc478c 100644 --- a/tools/rag/search_core.py +++ b/tools/rag/search_core.py @@ -11,6 +11,7 @@ from tools.rag.embedding_utils import ( blob_to_vector, cosine_similarity, embed_query, + model_profile, ) @@ -305,6 +306,10 @@ def embedding_index_info( raise RuntimeError( "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 ( str(row["model"]), @@ -545,9 +550,12 @@ def run_vector_query( ) ) - return candidates[ - :candidate_limit - ] + # A chunk can have several lossless embedding segments. Rank by its best + # 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( diff --git a/tools/rag/tests/test_database.py b/tools/rag/tests/test_database.py index ae6991e..a5436a8 100644 --- a/tools/rag/tests/test_database.py +++ b/tools/rag/tests/test_database.py @@ -144,6 +144,9 @@ def test_database_contains_documents_chunks_metadata_and_fts( "chunks": 1, "fts_chunks": 1, "embedding_chunks": 1, + "embedding_segments": 1, + "embedding_reused": 0, + "embedding_computed": 1, "tags": 2, "categories": 1, } diff --git a/tools/sync_indexes.py b/tools/sync_indexes.py index 1abe68a..41b22ed 100644 --- a/tools/sync_indexes.py +++ b/tools/sync_indexes.py @@ -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') chunks = build_chunks(wiki_root=frozen_wiki, output=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, 'created_at_utc': datetime.now(timezone.utc).isoformat(), 'source_sha256': hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest(), 'documents_sha256': digest(stage / 'documents.json'), 'chunks_sha256': digest(stage / 'chunks.json'), '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: 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),)) diff --git a/tools/text_budget.py b/tools/text_budget.py new file mode 100644 index 0000000..812f1d4 --- /dev/null +++ b/tools/text_budget.py @@ -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