forked from HAL9000/cleveragents-core
430 lines
15 KiB
Python
430 lines
15 KiB
Python
"""Step definitions for vector store service coverage."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
from collections.abc import Iterable
|
|
from pathlib import Path
|
|
from typing import Any, ClassVar
|
|
from unittest.mock import patch
|
|
|
|
from behave import given, then, when
|
|
from behave.runner import Context
|
|
|
|
from cleveragents.application.services.vector_store_service import VectorStoreService
|
|
from cleveragents.config.settings import Settings
|
|
from cleveragents.core.exceptions import ConfigurationError
|
|
from cleveragents.domain.models.core.context import Context as ContextModel
|
|
from features.steps.service_steps import add_cleanup
|
|
|
|
|
|
class _StubContextRepository:
|
|
def __init__(self, backing_store: dict[int, list[ContextModel]]) -> None:
|
|
self._backing_store = backing_store
|
|
|
|
def get_for_plan(self, plan_id: int) -> Iterable[ContextModel]:
|
|
return list(self._backing_store.get(plan_id, []))
|
|
|
|
|
|
class _StubTransaction:
|
|
def __init__(self, unit: _StubUnitOfWork) -> None:
|
|
self._unit = unit
|
|
|
|
def __enter__(self) -> _StubUnitOfWork:
|
|
return self._unit
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> bool:
|
|
return False
|
|
|
|
|
|
class _StubUnitOfWork:
|
|
def __init__(self) -> None:
|
|
self._contexts_map: dict[int, list[ContextModel]] = {}
|
|
self.contexts = _StubContextRepository(self._contexts_map)
|
|
|
|
def transaction(self) -> _StubTransaction:
|
|
return _StubTransaction(self)
|
|
|
|
def set_contexts(self, plan_id: int, contexts: list[ContextModel]) -> None:
|
|
self._contexts_map[plan_id] = contexts
|
|
|
|
|
|
class _StubDocument:
|
|
def __init__(self, path: str, content: str) -> None:
|
|
self.metadata = {"path": path}
|
|
self.page_content = content
|
|
|
|
|
|
class RecordingFAISS:
|
|
"""Test double that records FAISS usage."""
|
|
|
|
last_from_texts_args: ClassVar[dict[str, Any] | None] = None
|
|
last_built_instance: ClassVar[RecordingFAISS | None] = None
|
|
load_local_should_raise: ClassVar[bool] = False
|
|
load_local_calls: ClassVar[list[str]] = []
|
|
default_similarity_payload: ClassVar[list[tuple[_StubDocument, float]]] = []
|
|
|
|
def __init__(self) -> None:
|
|
self.documents: list[str] = []
|
|
self.metadatas: list[dict[str, Any]] = []
|
|
self.saved_directory: str | None = None
|
|
self.similarity_payload: list[tuple[_StubDocument, float]] = []
|
|
self.last_limit: int | None = None
|
|
self.last_query: str | None = None
|
|
|
|
@classmethod
|
|
def reset(cls) -> None:
|
|
cls.last_from_texts_args = None
|
|
cls.last_built_instance = None
|
|
cls.load_local_should_raise = False
|
|
cls.load_local_calls = []
|
|
cls.default_similarity_payload = []
|
|
|
|
@classmethod
|
|
def from_texts(
|
|
cls,
|
|
documents: list[str],
|
|
*,
|
|
embedding: Any,
|
|
metadatas: list[dict[str, Any]],
|
|
) -> RecordingFAISS:
|
|
instance = cls()
|
|
instance.documents = list(documents)
|
|
instance.metadatas = list(metadatas)
|
|
instance.similarity_payload = list(cls.default_similarity_payload)
|
|
cls.last_from_texts_args = {
|
|
"documents": list(documents),
|
|
"metadatas": list(metadatas),
|
|
"embedding": embedding,
|
|
}
|
|
cls.last_built_instance = instance
|
|
return instance
|
|
|
|
def save_local(self, directory: str) -> None:
|
|
self.saved_directory = directory
|
|
|
|
def similarity_search_with_score(
|
|
self, query: str, k: int
|
|
) -> list[tuple[_StubDocument, float]]:
|
|
self.last_query = query
|
|
self.last_limit = k
|
|
return self.similarity_payload[:k]
|
|
|
|
@classmethod
|
|
def load_local(
|
|
cls,
|
|
directory: str,
|
|
embeddings: Any,
|
|
allow_dangerous_deserialization: bool = True,
|
|
) -> RecordingFAISS:
|
|
cls.load_local_calls.append(directory)
|
|
if cls.load_local_should_raise:
|
|
raise ValueError("Failed to load index")
|
|
instance = cls()
|
|
instance.saved_directory = directory
|
|
cls.last_built_instance = instance
|
|
return instance
|
|
|
|
|
|
class StubOpenAIEmbeddings:
|
|
"""Minimal stub to capture requested OpenAI embedding model."""
|
|
|
|
last_model: str | None = None
|
|
|
|
def __init__(self, *, model: str | None = None, **_) -> None:
|
|
StubOpenAIEmbeddings.last_model = model
|
|
|
|
|
|
def _create_vector_service(context: Context, *, enabled: bool) -> None:
|
|
temp_dir = tempfile.mkdtemp(prefix="vector-store-service-")
|
|
add_cleanup(context, lambda: shutil.rmtree(temp_dir, ignore_errors=True))
|
|
settings = Settings()
|
|
settings.vector_store_enabled = enabled
|
|
settings.vector_store_path = Path(temp_dir) / "vector_store"
|
|
settings.vector_embeddings_provider = "fake"
|
|
settings.vector_embeddings_dimension = 8
|
|
unit = _StubUnitOfWork()
|
|
context.vector_store_service = VectorStoreService(settings, unit)
|
|
context.stub_unit_of_work = unit
|
|
context.vector_store_root = temp_dir
|
|
context.error = None
|
|
context.refresh_result = None
|
|
context.search_results = None
|
|
context.stub_openai_class = None
|
|
|
|
|
|
def _set_plan_contexts(
|
|
context: Context,
|
|
plan_id: int,
|
|
entries: list[tuple[str, str | None]],
|
|
) -> None:
|
|
base = Path(context.vector_store_root)
|
|
contexts = [
|
|
ContextModel(plan_id=plan_id, path=str(base / rel_path), content=content)
|
|
for rel_path, content in entries
|
|
]
|
|
context.stub_unit_of_work.set_contexts(plan_id, contexts)
|
|
|
|
|
|
@given("a vector store service with search enabled")
|
|
def step_vector_service_enabled(context: Context) -> None:
|
|
_create_vector_service(context, enabled=True)
|
|
|
|
|
|
@given("a vector store service with search disabled")
|
|
def step_vector_service_disabled(context: Context) -> None:
|
|
_create_vector_service(context, enabled=False)
|
|
|
|
|
|
@given("plan {plan_id:d} has no context documents")
|
|
def step_plan_no_contexts(context: Context, plan_id: int) -> None:
|
|
context.stub_unit_of_work.set_contexts(plan_id, [])
|
|
|
|
|
|
@given("plan {plan_id:d} has contexts with stored content and blanks")
|
|
def step_plan_mixed_contexts(context: Context, plan_id: int) -> None:
|
|
_set_plan_contexts(
|
|
context,
|
|
plan_id,
|
|
[
|
|
(f"plan_{plan_id}_doc.md", "Context body"),
|
|
(f"plan_{plan_id}_empty.md", None),
|
|
],
|
|
)
|
|
|
|
|
|
@given("plan {plan_id:d} has contexts with stored content")
|
|
def step_plan_filled_contexts(context: Context, plan_id: int) -> None:
|
|
_set_plan_contexts(
|
|
context,
|
|
plan_id,
|
|
[
|
|
(f"plan_{plan_id}_alpha.md", "Alpha document"),
|
|
(f"plan_{plan_id}_beta.md", "Beta document"),
|
|
],
|
|
)
|
|
|
|
|
|
@given("plan {plan_id:d} already has persisted FAISS files")
|
|
def step_plan_has_files(context: Context, plan_id: int) -> None:
|
|
plan_dir = context.vector_store_service._plan_store_dir(plan_id)
|
|
for filename in ("index.faiss", "index.pkl"):
|
|
(plan_dir / filename).write_bytes(b"stub")
|
|
|
|
|
|
@given("FAISS interactions are recorded")
|
|
def step_patch_faiss(context: Context) -> None:
|
|
RecordingFAISS.reset()
|
|
patcher = patch(
|
|
"cleveragents.application.services.vector_store_service.FAISS",
|
|
RecordingFAISS,
|
|
)
|
|
patcher.start()
|
|
add_cleanup(context, patcher.stop)
|
|
context.faiss_class = RecordingFAISS
|
|
|
|
|
|
@given("future FAISS builds will return similarity hits")
|
|
def step_future_faiss_hits(context: Context) -> None:
|
|
assert hasattr(context, "faiss_class"), "FAISS interactions were not recorded"
|
|
payload = [(_StubDocument("doc.md", "C" * 600), 0.25)]
|
|
context.faiss_class.default_similarity_payload = payload
|
|
|
|
def cleanup() -> None:
|
|
context.faiss_class.default_similarity_payload = []
|
|
|
|
add_cleanup(context, cleanup)
|
|
|
|
|
|
@given("plan {plan_id:d} cache contains similarity results")
|
|
def step_plan_cache_with_results(context: Context, plan_id: int) -> None:
|
|
store = context.faiss_class()
|
|
store.similarity_payload = [
|
|
(_StubDocument("doc.md", "A" * 600), 0.25),
|
|
(_StubDocument("doc-two.md", "B" * 50), 0.9),
|
|
]
|
|
context.vector_store_service._cache[plan_id] = store
|
|
context.cached_store = store
|
|
|
|
|
|
@given("loading the plan {plan_id:d} index will fail")
|
|
def step_fail_load(context: Context, plan_id: int) -> None:
|
|
context.faiss_class.load_local_should_raise = True
|
|
|
|
|
|
@given('the embeddings provider flag is "{provider}"')
|
|
def step_set_embeddings_provider(context: Context, provider: str) -> None:
|
|
settings = context.vector_store_service.settings.model_copy(
|
|
update={"vector_embeddings_provider": provider}
|
|
)
|
|
context.vector_store_service.settings = settings
|
|
|
|
|
|
@given('the embeddings provider is "openai" using model "{model}"')
|
|
def step_set_openai_provider(context: Context, model: str) -> None:
|
|
settings = context.vector_store_service.settings.model_copy(
|
|
update={
|
|
"vector_embeddings_provider": "openai",
|
|
"vector_embeddings_model": model,
|
|
}
|
|
)
|
|
context.vector_store_service.settings = settings
|
|
|
|
|
|
@given("a stub OpenAI embeddings backend is available")
|
|
def step_stub_openai_backend(context: Context) -> None:
|
|
StubOpenAIEmbeddings.last_model = None
|
|
module = types.SimpleNamespace(OpenAIEmbeddings=StubOpenAIEmbeddings)
|
|
previous = sys.modules.get("langchain_openai")
|
|
sys.modules["langchain_openai"] = module
|
|
|
|
def cleanup() -> None:
|
|
if previous is None:
|
|
sys.modules.pop("langchain_openai", None)
|
|
else:
|
|
sys.modules["langchain_openai"] = previous
|
|
|
|
add_cleanup(context, cleanup)
|
|
context.stub_openai_class = StubOpenAIEmbeddings
|
|
|
|
|
|
@when("I attempt to refresh the vector store for plan {plan_id:d}")
|
|
@when("I refresh the vector store for plan {plan_id:d}")
|
|
def step_refresh_plan(context: Context, plan_id: int) -> None:
|
|
try:
|
|
context.refresh_result = context.vector_store_service.refresh_for_plan(plan_id)
|
|
context.error = None
|
|
except Exception as exc:
|
|
context.error = exc
|
|
context.refresh_result = None
|
|
|
|
|
|
@when('I search plan {plan_id:d} with the query "{query}" and limit {limit:d}')
|
|
def step_search_with_limit(
|
|
context: Context, plan_id: int, query: str, limit: int
|
|
) -> None:
|
|
try:
|
|
context.search_results = context.vector_store_service.search(
|
|
plan_id,
|
|
query,
|
|
top_k=limit,
|
|
refresh_if_missing=False,
|
|
)
|
|
context.error = None
|
|
except Exception as exc:
|
|
context.error = exc
|
|
context.search_results = None
|
|
|
|
|
|
@when('I search plan {plan_id:d} with the query "{query}"')
|
|
def step_search_auto_refresh(context: Context, plan_id: int, query: str) -> None:
|
|
context.search_results = context.vector_store_service.search(plan_id, query)
|
|
|
|
|
|
@when('I search plan {plan_id:d} with the query ""')
|
|
def step_search_empty_query(context: Context, plan_id: int) -> None:
|
|
step_search_auto_refresh(context, plan_id, "")
|
|
|
|
|
|
@when('I search plan {plan_id:d} with the query "{query}" and refresh disabled')
|
|
def step_search_no_refresh(context: Context, plan_id: int, query: str) -> None:
|
|
context.search_results = context.vector_store_service.search(
|
|
plan_id,
|
|
query,
|
|
refresh_if_missing=False,
|
|
)
|
|
|
|
|
|
@when("I invalidate the vector store without specifying a plan")
|
|
def step_invalidate_without_plan(context: Context) -> None:
|
|
context.vector_store_service.invalidate(None)
|
|
|
|
|
|
@then("the refresh result should be {expected:d} documents")
|
|
@then("the refresh result should be {expected:d} document")
|
|
def step_verify_refresh_result(context: Context, expected: int) -> None:
|
|
assert context.refresh_result == expected, context.refresh_result
|
|
|
|
|
|
@then("the plan {plan_id:d} cache should be empty")
|
|
def step_verify_cache_empty(context: Context, plan_id: int) -> None:
|
|
assert plan_id not in context.vector_store_service._cache
|
|
|
|
|
|
@then("the plan {plan_id:d} cache should still contain the cached store")
|
|
def step_verify_cache_still_present(context: Context, plan_id: int) -> None:
|
|
cached = context.vector_store_service._cache.get(plan_id)
|
|
assert cached is context.cached_store
|
|
|
|
|
|
@then("the plan {plan_id:d} persisted files should be removed")
|
|
def step_verify_files_removed(context: Context, plan_id: int) -> None:
|
|
plan_dir = context.vector_store_service._plan_store_dir(plan_id)
|
|
for filename in ("index.faiss", "index.pkl"):
|
|
assert not (plan_dir / filename).exists()
|
|
|
|
|
|
@then("FAISS should be built with {expected:d} cleaned document")
|
|
def step_verify_faiss_documents(context: Context, expected: int) -> None:
|
|
args = context.faiss_class.last_from_texts_args
|
|
assert args is not None
|
|
assert len(args["documents"]) == expected
|
|
assert len(args["metadatas"]) == expected
|
|
|
|
|
|
@then("the plan {plan_id:d} cache should hold the FAISS instance")
|
|
def step_verify_cache_holds_instance(context: Context, plan_id: int) -> None:
|
|
cached = context.vector_store_service._cache.get(plan_id)
|
|
assert cached is context.faiss_class.last_built_instance
|
|
|
|
|
|
@then("FAISS should be loaded for plan {plan_id:d}")
|
|
def step_verify_faiss_loaded(context: Context, plan_id: int) -> None:
|
|
assert hasattr(context, "faiss_class"), "FAISS interactions were not recorded"
|
|
expected_dir = str(context.vector_store_service._plan_store_dir(plan_id))
|
|
assert context.faiss_class.load_local_calls, "No FAISS load_local calls recorded"
|
|
assert expected_dir in context.faiss_class.load_local_calls
|
|
|
|
|
|
@then(
|
|
'the search results should include one formatted hit with path "{path}" and score {score:f}'
|
|
)
|
|
def step_verify_formatted_result(context: Context, path: str, score: float) -> None:
|
|
assert isinstance(context.search_results, list)
|
|
assert len(context.search_results) == 1
|
|
hit = context.search_results[0]
|
|
assert hit["path"] == path
|
|
assert abs(hit["score"] - score) < 1e-9
|
|
assert len(hit["snippet"]) == 500
|
|
|
|
|
|
@then("the FAISS similarity search limit should be {expected:d}")
|
|
def step_verify_limit(context: Context, expected: int) -> None:
|
|
assert context.cached_store.last_limit == expected
|
|
|
|
|
|
@then("the search results should be empty")
|
|
def step_verify_empty_results(context: Context) -> None:
|
|
assert context.search_results == []
|
|
|
|
|
|
@then("a configuration error should mention disabled vector store support")
|
|
def step_verify_disabled_error(context: Context) -> None:
|
|
assert isinstance(context.error, ConfigurationError)
|
|
assert "disabled" in str(context.error).lower()
|
|
|
|
|
|
@then("a configuration error should mention unsupported embeddings provider")
|
|
def step_verify_unsupported_provider(context: Context) -> None:
|
|
assert isinstance(context.error, ConfigurationError)
|
|
assert "unsupported" in str(context.error).lower()
|
|
|
|
|
|
@then('the last OpenAI embeddings model should be "{model}"')
|
|
def step_verify_openai_model(context: Context, model: str) -> None:
|
|
assert context.stub_openai_class is not None
|
|
assert context.stub_openai_class.last_model == model
|