mirror of
https://github.com/semantica-agi/semantica.git
synced 2026-08-30 04:40:16 +00:00
86 lines
2.1 KiB
Python
86 lines
2.1 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from semantica.vector_store.faiss_store import FAISSStore
|
|
from semantica.vector_store.vector_store import VectorStore
|
|
|
|
# Fixtures
|
|
|
|
|
|
@pytest.fixture
|
|
def vector_dim():
|
|
return 768
|
|
|
|
|
|
@pytest.fixture
|
|
def random_vectors(vector_dim):
|
|
"""Generates a batch of 10,000 rando vectors."""
|
|
count = 10000
|
|
vectors = np.random.rand(count, vector_dim).astype(np.float32)
|
|
return vectors
|
|
|
|
|
|
@pytest.fixture
|
|
def populated_store(random_vectors, vector_dim):
|
|
"""
|
|
Returns a FAISS store bred with data.
|
|
"""
|
|
|
|
store = FAISSStore(dimension=vector_dim)
|
|
if hasattr(store, "progress_tracker"):
|
|
store.progress_tracker = MagicMock()
|
|
store.create_index(index_type="flat")
|
|
store.add_vectors(random_vectors)
|
|
return store
|
|
|
|
|
|
# Benchmarks
|
|
|
|
|
|
def test_faiss_insert_throughput(benchmark, random_vectors, vector_dim):
|
|
"""
|
|
Benchmarks raw Write speed to FAISS
|
|
"""
|
|
store = FAISSStore(dimension=vector_dim)
|
|
if hasattr(store, "progress_tracker"):
|
|
store.progress_tracker = MagicMock()
|
|
store.create_index(index_type="flat")
|
|
|
|
def insert_op():
|
|
store.add_vectors(random_vectors)
|
|
|
|
benchmark(insert_op)
|
|
|
|
assert len(store.index.vector_ids) >= 10000
|
|
|
|
|
|
def test_faiss_search_latency(benchmark, populated_store, vector_dim):
|
|
"""
|
|
Benchmarks Read/Search speed
|
|
"""
|
|
|
|
query = np.random.rand(1, vector_dim).astype(np.float32)
|
|
results = benchmark(populated_store.search_similar, query_vector=query, k=10)
|
|
assert len(results) == 10
|
|
|
|
|
|
def test_vector_storage_manager_overhead(benchmark, random_vectors, vector_dim):
|
|
"""
|
|
Benchmarks the overhead of the VectorStore class
|
|
"""
|
|
with patch(
|
|
"semantica.vector_store.vector_store.EmbeddingGenerator"
|
|
) as MockEmbedder:
|
|
manager = VectorStore(backend="faiss", dimension=vector_dim)
|
|
if hasattr(manager, "progress_tracker"):
|
|
manager.progress_tracker = MagicMock()
|
|
|
|
def store_op():
|
|
manager.store_vectors(random_vectors)
|
|
|
|
benchmark(store_op)
|
|
|
|
assert len(manager.vectors) >= 10000
|