116 lines
5.8 KiB
Python
116 lines
5.8 KiB
Python
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
|