151 lines
6.1 KiB
Python
151 lines
6.1 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.retrieval.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, 0.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, 0.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, 0.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
|