23 lines
743 B
Python
23 lines
743 B
Python
import tempfile
|
|
|
|
import pytest
|
|
|
|
from companion.rag.vector_store import VectorStore
|
|
|
|
|
|
def test_vector_store_upsert_and_search():
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
store = VectorStore(uri=tmp, dimensions=4)
|
|
store.upsert(
|
|
ids=["a", "b"],
|
|
texts=["hello world", "goodbye world"],
|
|
embeddings=[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]],
|
|
metadatas=[
|
|
{"source_file": "a.md", "source_directory": "docs"},
|
|
{"source_file": "b.md", "source_directory": "docs"},
|
|
],
|
|
)
|
|
results = store.search(query_vector=[1.0, 0.0, 0.0, 0.0], top_k=1)
|
|
assert len(results) == 1
|
|
assert results[0]["source_file"] == "a.md"
|