-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Expand file tree
/
Copy pathtest_vector_store.py
More file actions
62 lines (50 loc) · 2.04 KB
/
test_vector_store.py
File metadata and controls
62 lines (50 loc) · 2.04 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import numpy as np
from tests.example_repos.example_feature_repo_1 import document_embeddings
from tests.utils.rag_test_utils import MockVectorStore, example_feature_store
# showing intended use of example_feature_store fixture
_ = example_feature_store
def test_vector_store_initialization(example_feature_store):
"""Test vector store initialization."""
print("Testing vector store initialization...")
# Apply the feature view first
example_feature_store.apply([document_embeddings])
doc_view = example_feature_store.get_feature_view("document_embeddings")
store = MockVectorStore(
repo_path=str(example_feature_store.repo_path),
rag_view=doc_view,
features=[
"document_embeddings:content",
"document_embeddings:Embeddings",
"document_embeddings:item_id",
],
)
assert store.rag_view == doc_view
assert store.features == [
"document_embeddings:content",
"document_embeddings:Embeddings",
"document_embeddings:item_id",
]
def test_vector_store_query(example_feature_store):
"""Test vector store query method."""
print("Testing vector store query...")
# Apply the feature view first
example_feature_store.apply([document_embeddings])
doc_view = example_feature_store.get_feature_view("document_embeddings")
store = MockVectorStore(
repo_path=str(example_feature_store.repo_path),
rag_view=doc_view,
features=[
"document_embeddings:content",
"document_embeddings:Embeddings",
"document_embeddings:item_id",
],
)
# Test query with vector
query_vector = np.array([0.1] * 8) # 8-dimensional query vector to match schema
print("Querying with vector...")
response = store.query(query_vector=query_vector, query_string=None, top_k=5)
assert response is not None
# Test query with text
print("Querying with text...")
response = store.query(query_vector=None, query_string="test query", top_k=5)
assert response is not None