Optimalizácia RAG spracovania
This commit is contained in:
parent
bfe8cbba1b
commit
d8961f9db6
36
README.md
36
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.
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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
|
||||
|
||||
115
test/test_embedding_budget.py
Normal file
115
test/test_embedding_budget.py
Normal 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
|
||||
67
test/test_llm_context_budget.py
Normal file
67
test/test_llm_context_budget.py
Normal 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
|
||||
@ -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(
|
||||
|
||||
@ -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,
|
||||
|
||||
42
tools/rag/embedding_cache.py
Normal file
42
tools/rag/embedding_cache.py
Normal 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()))
|
||||
@ -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: "
|
||||
|
||||
@ -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=(',', ':'),
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
114
tools/rag/llm_context.py
Normal file
114
tools/rag/llm_context.py
Normal 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
|
||||
@ -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(
|
||||
|
||||
@ -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,
|
||||
}
|
||||
|
||||
@ -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),))
|
||||
|
||||
30
tools/text_budget.py
Normal file
30
tools/text_budget.py
Normal 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
|
||||
Loading…
Reference in New Issue
Block a user