feat: add search engine interface with embedding and filtering
This commit is contained in:
34
tests/test_search.py
Normal file
34
tests/test_search.py
Normal file
@@ -0,0 +1,34 @@
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from companion.rag.search import SearchEngine
|
||||
from companion.rag.vector_store import VectorStore
|
||||
|
||||
|
||||
@patch("companion.rag.search.OllamaEmbedder")
|
||||
def test_search_returns_results(mock_embedder_cls):
|
||||
mock_embedder = MagicMock()
|
||||
mock_embedder.embed.return_value = [[1.0, 0.0, 0.0, 0.0]]
|
||||
mock_embedder_cls.return_value = mock_embedder
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
store = VectorStore(uri=tmp, dimensions=4)
|
||||
store.upsert(
|
||||
ids=["a"],
|
||||
texts=["hello world"],
|
||||
embeddings=[[1.0, 0.0, 0.0, 0.0]],
|
||||
metadatas=[{"source_file": "a.md", "source_directory": "docs"}],
|
||||
)
|
||||
engine = SearchEngine(
|
||||
vector_store=store,
|
||||
embedder_base_url="http://localhost:11434",
|
||||
embedder_model="dummy",
|
||||
embedder_batch_size=32,
|
||||
default_top_k=5,
|
||||
similarity_threshold=0.0,
|
||||
hybrid_search_enabled=False,
|
||||
)
|
||||
results = engine.search("hello")
|
||||
assert len(results) == 1
|
||||
assert results[0]["source_file"] == "a.md"
|
||||
Reference in New Issue
Block a user