feat: add LanceDB vector store with upsert, delete, and search
This commit is contained in:
22
tests/test_vector_store.py
Normal file
22
tests/test_vector_store.py
Normal file
@@ -0,0 +1,22 @@
|
||||
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"
|
||||
Reference in New Issue
Block a user