mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-31 03:12:16 +00:00
324 lines
10 KiB
Python
324 lines
10 KiB
Python
"""Tests for TwoStageRetriever — BM25 recall + optional semantic reranking."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import List
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from openjarvis.connectors.retriever import ColBERTReranker, Reranker, TwoStageRetriever
|
|
from openjarvis.connectors.store import KnowledgeStore
|
|
from openjarvis.tools.storage._stubs import RetrievalResult
|
|
|
|
|
|
def _has_torch() -> bool:
|
|
try:
|
|
import torch # noqa: F401
|
|
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture()
|
|
def store(tmp_path: Path) -> KnowledgeStore:
|
|
"""KnowledgeStore pre-populated with 4 items across different sources/authors."""
|
|
ks = KnowledgeStore(db_path=tmp_path / "test_retriever.db")
|
|
|
|
ks.store(
|
|
content="AI research on neural networks and deep learning advances.",
|
|
source="gmail",
|
|
doc_type="email",
|
|
author="alice@example.com",
|
|
)
|
|
ks.store(
|
|
content="AI research on distributed systems and cloud infrastructure.",
|
|
source="obsidian",
|
|
doc_type="note",
|
|
author="bob@example.com",
|
|
)
|
|
ks.store(
|
|
content="AI research on natural language processing and transformers.",
|
|
source="gmail",
|
|
doc_type="email",
|
|
author="alice@example.com",
|
|
)
|
|
ks.store(
|
|
content="AI research on reinforcement learning for robotics.",
|
|
source="slack",
|
|
doc_type="message",
|
|
author="carol@example.com",
|
|
)
|
|
|
|
return ks
|
|
|
|
|
|
@pytest.fixture()
|
|
def retriever(store: KnowledgeStore) -> TwoStageRetriever:
|
|
return TwoStageRetriever(store)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 1: retrieve_returns_results — basic search works
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_retrieve_returns_results(retriever: TwoStageRetriever) -> None:
|
|
"""A matching query returns at least one result."""
|
|
results = retriever.retrieve("AI research")
|
|
assert isinstance(results, list)
|
|
assert len(results) >= 1
|
|
assert all(isinstance(r, RetrievalResult) for r in results)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 2: retrieve_respects_top_k — limits results
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_retrieve_respects_top_k(retriever: TwoStageRetriever) -> None:
|
|
"""retrieve() returns at most top_k results."""
|
|
results = retriever.retrieve("AI research", top_k=2)
|
|
assert len(results) <= 2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 3: retrieve_with_source_filter — filters pass through to store
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_retrieve_with_source_filter(
|
|
store: KnowledgeStore,
|
|
) -> None:
|
|
"""source= filter is forwarded to the KnowledgeStore."""
|
|
ret = TwoStageRetriever(store)
|
|
results = ret.retrieve("AI research", source="gmail")
|
|
assert len(results) >= 1
|
|
for r in results:
|
|
assert r.metadata.get("source") == "gmail"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 4: retrieve_with_author_filter — filters pass through
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_retrieve_with_author_filter(
|
|
store: KnowledgeStore,
|
|
) -> None:
|
|
"""author= filter is forwarded to the KnowledgeStore."""
|
|
ret = TwoStageRetriever(store)
|
|
results = ret.retrieve("AI research", author="alice@example.com")
|
|
assert len(results) >= 1
|
|
for r in results:
|
|
assert r.metadata.get("author") == "alice@example.com"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 5: retrieve_bm25_only_when_no_colbert — works without reranker
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_retrieve_bm25_only_when_no_reranker(
|
|
store: KnowledgeStore,
|
|
) -> None:
|
|
"""Without a reranker the retriever returns BM25 results directly."""
|
|
ret = TwoStageRetriever(store, reranker=None)
|
|
results = ret.retrieve("AI research", top_k=3)
|
|
assert isinstance(results, list)
|
|
assert len(results) <= 3
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 6: retrieve_with_colbert_reranking — mock reranker is called
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_retrieve_with_colbert_reranking(
|
|
store: KnowledgeStore,
|
|
) -> None:
|
|
"""When a reranker is provided it is called to reorder candidates."""
|
|
mock_reranker = MagicMock(spec=Reranker)
|
|
# The mock reranker returns a pair of results to distinguish from BM25 output
|
|
reranked = [
|
|
RetrievalResult(
|
|
content="Reranked result A",
|
|
score=0.99,
|
|
source="gmail",
|
|
metadata={"reranked": True},
|
|
),
|
|
RetrievalResult(
|
|
content="Reranked result B",
|
|
score=0.95,
|
|
source="slack",
|
|
metadata={"reranked": True},
|
|
),
|
|
]
|
|
mock_reranker.rerank.return_value = reranked
|
|
|
|
# recall_k=4 fetches all 4 docs; top_k=2 so len(candidates)=4 > top_k=2
|
|
ret = TwoStageRetriever(store, reranker=mock_reranker, recall_k=4)
|
|
results = ret.retrieve("AI research", top_k=2)
|
|
|
|
# The reranker must have been called with top_k=2
|
|
mock_reranker.rerank.assert_called_once()
|
|
call_args = mock_reranker.rerank.call_args
|
|
assert call_args[1].get("top_k") == 2
|
|
|
|
# Results come from the reranker
|
|
assert results == reranked
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 7: retrieve_no_results — empty for nonexistent query
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_retrieve_no_results(retriever: TwoStageRetriever) -> None:
|
|
"""A query matching nothing returns an empty list."""
|
|
results = retriever.retrieve("xyzzy_nonexistent_zqjwkm")
|
|
assert results == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 8: retrieve_recall_k_larger_than_top_k — BM25 fetches more than final k
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_retrieve_recall_k_larger_than_top_k(
|
|
store: KnowledgeStore,
|
|
) -> None:
|
|
"""Stage-1 recall fetches max(recall_k, top_k*3) candidates."""
|
|
call_log: List[int] = []
|
|
|
|
class SpyStore(KnowledgeStore):
|
|
def retrieve(self, query, *, top_k=5, **kwargs): # type: ignore[override]
|
|
call_log.append(top_k)
|
|
return super().retrieve(query, top_k=top_k, **kwargs)
|
|
|
|
spy = SpyStore(db_path=":memory:")
|
|
spy.store(
|
|
content="Deep learning research in computer vision tasks.",
|
|
source="gmail",
|
|
doc_type="email",
|
|
author="alice@example.com",
|
|
)
|
|
|
|
ret = TwoStageRetriever(spy, recall_k=50)
|
|
ret.retrieve("deep learning research", top_k=2)
|
|
|
|
# Stage-1 must have requested more candidates than the final top_k
|
|
assert len(call_log) == 1
|
|
# recall_k=50 vs top_k*3=6 → should use 50
|
|
assert call_log[0] == 50
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 9: reranker_uses_cached_embeddings — EmbeddingStore integration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _has_torch(),
|
|
reason="torch required for embedding tests",
|
|
)
|
|
def test_reranker_uses_cached_embeddings() -> None:
|
|
"""ColBERTReranker checks EmbeddingStore.get() before calling docFromText().
|
|
|
|
When a cached embedding is found, docFromText() should NOT be called for
|
|
that candidate.
|
|
"""
|
|
import torch # type: ignore[import]
|
|
|
|
# Create a fake cached embedding (T=20 tokens, dim=128)
|
|
cached_emb = torch.randn(20, 128)
|
|
|
|
# Mock the embedding store
|
|
mock_store = MagicMock()
|
|
mock_store.get.return_value = cached_emb
|
|
|
|
# Create a reranker with the mocked embedding store
|
|
reranker = ColBERTReranker(embedding_store=mock_store)
|
|
|
|
# Mock the ColBERT model so _load_model succeeds
|
|
mock_model = MagicMock()
|
|
# queryFromText returns a (Q, dim) tensor
|
|
mock_model.queryFromText.return_value = [torch.randn(32, 128)]
|
|
reranker._model = mock_model # bypass _load_model()
|
|
|
|
candidates = [
|
|
RetrievalResult(
|
|
content="Some document about AI",
|
|
score=1.0,
|
|
source="gmail",
|
|
metadata={"chunk_id": "chunk-123"},
|
|
),
|
|
]
|
|
|
|
results = reranker.rerank("AI research", candidates, top_k=1)
|
|
|
|
# Embedding store should have been consulted
|
|
mock_store.get.assert_called_once_with("chunk-123")
|
|
|
|
# docFromText should NOT have been called (cache hit)
|
|
mock_model.docFromText.assert_not_called()
|
|
|
|
# We should still get a result back
|
|
assert len(results) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 10: reranker_caches_new_embeddings — stores after compute
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _has_torch(),
|
|
reason="torch required for embedding tests",
|
|
)
|
|
def test_reranker_caches_new_embeddings() -> None:
|
|
"""When EmbeddingStore.get() returns None, docFromText() is called and
|
|
the result is stored back via EmbeddingStore.store()."""
|
|
import torch # type: ignore[import]
|
|
|
|
mock_store = MagicMock()
|
|
mock_store.get.return_value = None # cache miss
|
|
|
|
reranker = ColBERTReranker(embedding_store=mock_store)
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.queryFromText.return_value = [torch.randn(32, 128)]
|
|
# docFromText returns (1, T, dim) which the reranker squeezes
|
|
doc_emb = torch.randn(1, 20, 128)
|
|
mock_model.docFromText.return_value = [doc_emb]
|
|
reranker._model = mock_model
|
|
|
|
candidates = [
|
|
RetrievalResult(
|
|
content="Document about neural nets",
|
|
score=1.0,
|
|
source="obsidian",
|
|
metadata={"chunk_id": "chunk-456"},
|
|
),
|
|
]
|
|
|
|
results = reranker.rerank("neural nets", candidates, top_k=1)
|
|
|
|
# docFromText should have been called (cache miss)
|
|
mock_model.docFromText.assert_called_once()
|
|
|
|
# The computed embedding should be stored back
|
|
mock_store.store.assert_called_once()
|
|
call_args = mock_store.store.call_args
|
|
assert call_args[0][0] == "chunk-456" # chunk_id
|
|
assert call_args[0][1].shape == (20, 128) # squeezed tensor
|
|
|
|
assert len(results) == 1
|