forked from cleveragents/cleveragents-core
aditya
eff446f5e8
Add FAISS-backed ACMS read and write adapters on top of the shared VectorStoreService, wire them through the DI container, and cover indexing, scoped search, removal, and benchmark behavior with Behave and ASV. Address review feedback by keeping the commit scoped to the FAISS backend work and replacing the loose vector-store cache typing with explicit FAISS store protocols instead of Any. ISSUES CLOSED: #871
103 lines
3.4 KiB
Python
103 lines
3.4 KiB
Python
"""ASV benchmarks for the FAISS ACMS vector backend adapters."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import ClassVar
|
|
|
|
_SRC = str(Path(__file__).resolve().parents[1] / "src")
|
|
if _SRC not in sys.path:
|
|
sys.path.insert(0, _SRC)
|
|
|
|
import cleveragents # noqa: E402
|
|
|
|
importlib.reload(cleveragents)
|
|
|
|
from cleveragents.application.services.config_service import ConfigService # noqa: E402
|
|
from cleveragents.application.services.faiss_vector_backend import ( # noqa: E402
|
|
FAISSVectorBackend,
|
|
FAISSVectorIndexBackend,
|
|
)
|
|
from cleveragents.application.services.vector_store_service import ( # noqa: E402
|
|
FAISS,
|
|
VectorStoreService,
|
|
)
|
|
from cleveragents.config.settings import Settings # noqa: E402
|
|
from cleveragents.infrastructure.database.unit_of_work import UnitOfWork # noqa: E402
|
|
|
|
|
|
class VectorSearchLatencySuite:
|
|
"""Benchmark FAISS-backed ACMS vector indexing and search."""
|
|
|
|
params: ClassVar[list[int]] = [250, 1000]
|
|
param_names: ClassVar[list[str]] = ["documents"]
|
|
timeout = 120
|
|
|
|
def setup(self, documents: int) -> None:
|
|
if FAISS is None:
|
|
raise RuntimeError("FAISS is required for this benchmark")
|
|
|
|
self._saved_env = {
|
|
key: os.environ.get(key)
|
|
for key in (
|
|
"CLEVERAGENTS_INDEX_VECTOR_BACKEND",
|
|
"CLEVERAGENTS_INDEX_VECTOR_DIR",
|
|
"CLEVERAGENTS_EMBEDDING_PROVIDER",
|
|
"CLEVERAGENTS_EMBEDDING_DIMENSIONS",
|
|
)
|
|
}
|
|
self._tempdir = tempfile.mkdtemp(prefix="asv-faiss-vector-")
|
|
os.environ["CLEVERAGENTS_INDEX_VECTOR_BACKEND"] = "faiss"
|
|
os.environ["CLEVERAGENTS_INDEX_VECTOR_DIR"] = self._tempdir
|
|
os.environ["CLEVERAGENTS_EMBEDDING_PROVIDER"] = "fake"
|
|
os.environ["CLEVERAGENTS_EMBEDDING_DIMENSIONS"] = "32"
|
|
|
|
service = VectorStoreService(
|
|
Settings(),
|
|
UnitOfWork("sqlite:///:memory:", require_confirmation=False),
|
|
ConfigService(),
|
|
)
|
|
self._index_backend = FAISSVectorIndexBackend(service)
|
|
self._vector_backend = FAISSVectorBackend(service)
|
|
self._query = [1.0, 2.0, 3.0]
|
|
self._scope = frozenset({f"RES{i:04d}" for i in range(min(documents, 32))})
|
|
|
|
for i in range(documents):
|
|
vector = [float(i % 11), float((i * 3) % 17), float((i * 5) % 19)]
|
|
self._index_backend.index_embedding(
|
|
"local/bench",
|
|
f"uko:{i:04d}",
|
|
vector,
|
|
{
|
|
"resource_id": f"RES{i:04d}",
|
|
"location": f"src/file_{i:04d}.py",
|
|
"resource_type": "python",
|
|
},
|
|
)
|
|
|
|
def teardown(self, documents: int) -> None:
|
|
_ = documents
|
|
shutil.rmtree(self._tempdir, ignore_errors=True)
|
|
for key, value in self._saved_env.items():
|
|
if value is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = value
|
|
|
|
def time_project_search(self, documents: int) -> None:
|
|
_ = documents
|
|
self._index_backend.search_similar("local/bench", self._query, limit=20)
|
|
|
|
def time_scoped_similarity_search(self, documents: int) -> None:
|
|
_ = documents
|
|
self._vector_backend.similarity_search(
|
|
self._query,
|
|
scope=self._scope,
|
|
top_k=20,
|
|
)
|