158 lines
4.5 KiB
Python
158 lines
4.5 KiB
Python
"""Testes do banco vetorial e dos embeddings."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from app.bible_parser import VerseChunk # noqa: E402
|
|
from app.embedding import HashEmbedder, normalize # noqa: E402
|
|
from app.vector_store import InMemoryVectorStore, VectorStore # noqa: E402
|
|
|
|
|
|
def make_chunk(verse: int, book: str = "Gênesis", chapter: int = 1, file: str = "biblia.pdf"):
|
|
return VerseChunk(
|
|
book=book,
|
|
chapter=chapter,
|
|
verse=verse,
|
|
page=6,
|
|
file=file,
|
|
text=f"Versículo {verse} de {book} {chapter}",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def embedder():
|
|
return HashEmbedder(dimensions=32)
|
|
|
|
|
|
@pytest.fixture
|
|
def store(embedder):
|
|
return InMemoryVectorStore(embedder=embedder)
|
|
|
|
|
|
def test_normalize_produces_unit_vector():
|
|
vector = normalize([3.0, 4.0])
|
|
assert vector == [0.6, 0.8]
|
|
|
|
|
|
def test_normalize_of_zero_vector_is_safe():
|
|
assert normalize([0.0, 0.0]) == [0.0, 0.0]
|
|
|
|
|
|
def test_hash_embedder_is_deterministic():
|
|
embedder = HashEmbedder(dimensions=16)
|
|
assert embedder.encode(["olá mundo"]) == embedder.encode(["olá mundo"])
|
|
|
|
|
|
def test_hash_embedder_produces_requested_dimensions():
|
|
embedder = HashEmbedder(dimensions=7)
|
|
vector = embedder.encode(["texto"])[0]
|
|
assert len(vector) == 7
|
|
|
|
|
|
def test_upsert_then_query_finds_the_matching_chunk(store, embedder):
|
|
chunks = [make_chunk(1), make_chunk(2), make_chunk(3)]
|
|
vectors = embedder.encode([chunk.text for chunk in chunks])
|
|
|
|
store.upsert(chunks, vectors)
|
|
|
|
query_vector = embedder.encode(["Versículo 2 de Gênesis 1"])[0]
|
|
results = store.query(query_vector, n_results=3, min_similarity=0.0)
|
|
|
|
assert results
|
|
assert results[0].verse == 2
|
|
assert results[0].book == "Gênesis"
|
|
|
|
|
|
def test_query_respects_min_similarity(store, embedder):
|
|
chunks = [make_chunk(1)]
|
|
vectors = embedder.encode([chunk.text for chunk in chunks])
|
|
store.upsert(chunks, vectors)
|
|
|
|
query_vector = embedder.encode(["palavra completamente diferente xyz"])[0]
|
|
results = store.query(query_vector, n_results=5, min_similarity=0.99)
|
|
|
|
assert results == [], "nada deveria passar com limiar 0.99"
|
|
|
|
|
|
def test_query_respects_n_results(store, embedder):
|
|
chunks = [make_chunk(i) for i in range(1, 11)]
|
|
vectors = embedder.encode([chunk.text for chunk in chunks])
|
|
store.upsert(chunks, vectors)
|
|
|
|
query_vector = embedder.encode(["Versículo"])[0]
|
|
results = store.query(query_vector, n_results=3, min_similarity=0.0)
|
|
|
|
assert len(results) == 3
|
|
|
|
|
|
def test_delete_file_removes_only_that_file(store, embedder):
|
|
chunks = [make_chunk(1, file="a.pdf"), make_chunk(2, file="b.pdf")]
|
|
vectors = embedder.encode([chunk.text for chunk in chunks])
|
|
store.upsert(chunks, vectors)
|
|
|
|
store.delete_file("a.pdf")
|
|
|
|
assert store.count() == 1
|
|
query_vector = embedder.encode(["Versículo"])[0]
|
|
results = store.query(query_vector, n_results=10, min_similarity=0.0)
|
|
assert {result.file for result in results} == {"b.pdf"}
|
|
|
|
|
|
def test_get_by_reference_returns_exact_verse(store, embedder):
|
|
chunks = [make_chunk(1), make_chunk(2, chapter=2)]
|
|
vectors = embedder.encode([chunk.text for chunk in chunks])
|
|
store.upsert(chunks, vectors)
|
|
|
|
results = store.get_by_reference("Gênesis", 2, 2)
|
|
|
|
assert len(results) == 1
|
|
assert results[0].chapter == 2
|
|
assert results[0].verse == 2
|
|
|
|
|
|
def test_get_by_reference_returns_empty_when_missing(store, embedder):
|
|
chunks = [make_chunk(1)]
|
|
vectors = embedder.encode([chunk.text for chunk in chunks])
|
|
store.upsert(chunks, vectors)
|
|
|
|
assert store.get_by_reference("João", 3, 16) == []
|
|
|
|
|
|
def test_count_reflects_upserts(store, embedder):
|
|
chunks = [make_chunk(1), make_chunk(2)]
|
|
vectors = embedder.encode([chunk.text for chunk in chunks])
|
|
|
|
store.upsert(chunks, vectors)
|
|
|
|
assert store.count() == 2
|
|
|
|
|
|
def test_upsert_is_idempotent(store, embedder):
|
|
chunk = make_chunk(1)
|
|
vector = embedder.encode([chunk.text])[0]
|
|
|
|
store.upsert([chunk], [vector])
|
|
store.upsert([chunk], [vector])
|
|
|
|
assert store.count() == 1
|
|
|
|
|
|
def test_chroma_store_roundtrip(tmp_path, embedder):
|
|
"""Garante que a implementação Chroma real tem a mesma interface."""
|
|
store = VectorStore(persist_dir=tmp_path, embedder=embedder)
|
|
chunk = make_chunk(1)
|
|
vector = embedder.encode([chunk.text])[0]
|
|
|
|
store.upsert([chunk], [vector])
|
|
|
|
assert store.count() == 1
|
|
results = store.get_by_reference("Gênesis", 1, 1)
|
|
assert len(results) == 1
|
|
assert results[0].text == chunk.text
|