devplacepy/tests/unit/services/deepsearch/store.py

# retoor <retoor@molodetz.nl>

from devplacepy.services.deepsearch.embeddings import local_embed
from devplacepy.services.deepsearch.store import Chunk, VectorStore
from tests.conftest import run_async


def _chunks():
    texts = [
        "The transistor was invented at Bell Labs in nineteen forty seven.",
        "Silicon is the primary semiconductor material used in chips.",
        "Quantum tunnelling limits how small a transistor can become.",
    ]
    chunks = [
        Chunk(uid=f"c{i}", text=text, url=f"https://s{i}.example", title=f"Source {i}")
        for i, text in enumerate(texts)
    ]
    return chunks


def test_add_and_count():
    store = VectorStore("ds_store_test_count")
    try:
        chunks = _chunks()
        vectors = local_embed([c.text for c in chunks]).vectors
        run_async(store.add(chunks, vectors))
        assert run_async(store.count()) == 3
    finally:
        store.drop()


def test_hybrid_search_returns_relevant_chunk():
    store = VectorStore("ds_store_test_hybrid")
    try:
        chunks = _chunks()
        vectors = local_embed([c.text for c in chunks]).vectors
        run_async(store.add(chunks, vectors))
        query = "where was the transistor invented"
        query_vector = local_embed([query]).vectors[0]
        results = run_async(store.hybrid_search(query, query_vector, top_k=2))
        assert results
        assert any("transistor" in r.text.lower() for r in results)
    finally:
        store.drop()


def test_keyword_scores_rank_match_higher():
    store = VectorStore("ds_store_test_keyword")
    chunks = _chunks()
    scores = store.keyword_scores("silicon semiconductor", chunks)
    assert scores
    best = max(scores, key=scores.get)
    assert "silicon" in next(c.text for c in chunks if c.uid == best).lower()


def test_dims_reflects_embedding_width():
    store = VectorStore("ds_store_test_dims")
    try:
        chunks = _chunks()
        vectors = local_embed([c.text for c in chunks]).vectors
        run_async(store.add(chunks, vectors))
        assert store.dims == len(vectors[0])
    finally:
        store.drop()


def test_coverage_analytics_empty_collection():
    store = VectorStore("ds_store_test_cov_empty")
    try:
        analytics = run_async(store.coverage_analytics())
        assert analytics == {
            "chunks": 0,
            "domains": 0,
            "sources": 0,
            "avg_chunk_chars": 0,
        }
    finally:
        store.drop()


def test_coverage_analytics_counts_chunks_and_sources():
    store = VectorStore("ds_store_test_cov")
    try:
        chunks = _chunks()
        for index, chunk in enumerate(chunks):
            chunk.source = "httpx" if index == 0 else "playwright"
        vectors = local_embed([c.text for c in chunks]).vectors
        run_async(store.add(chunks, vectors))
        analytics = run_async(store.coverage_analytics())
        assert analytics["chunks"] == 3
        assert analytics["sources"] == 2
        assert analytics["avg_chunk_chars"] > 0
    finally:
        store.drop()