forked from HAL9000/cleveragents-core
233 lines
7.6 KiB
Python
233 lines
7.6 KiB
Python
"""Step definitions for context vector search integration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from behave import given, then, when
|
|
|
|
from cleveragents.application.services.context_service import ContextService
|
|
from cleveragents.config.settings import Settings
|
|
from cleveragents.core.exceptions import ConfigurationError
|
|
from cleveragents.domain.models.core import Context, Plan, Project
|
|
from features.steps.service_steps import add_cleanup
|
|
|
|
|
|
class RecordingContextRepository:
|
|
"""In-memory context repository used for vector integration tests."""
|
|
|
|
def __init__(self) -> None:
|
|
self._store: dict[int, list[Context]] = {}
|
|
self._next_id = 1
|
|
|
|
def get_for_plan(self, plan_id: int) -> list[Context]:
|
|
return [*self._store.get(plan_id, [])]
|
|
|
|
def add(self, context: Context) -> None:
|
|
context.id = context.id or self._next_id
|
|
self._next_id += 1
|
|
self._store.setdefault(context.plan_id, []).append(context)
|
|
|
|
def remove(self, context_id: int) -> None:
|
|
for plan_id, contexts in self._store.items():
|
|
self._store[plan_id] = [c for c in contexts if c.id != context_id]
|
|
|
|
def clear_for_plan(self, plan_id: int) -> None:
|
|
self._store.pop(plan_id, None)
|
|
|
|
|
|
class RecordingPlansRepository:
|
|
def __init__(self, plan: Plan) -> None:
|
|
self._plan = plan
|
|
|
|
def get_current_for_project(self, project_id: int | None) -> Plan | None:
|
|
return self._plan if project_id else None
|
|
|
|
|
|
class RecordingTransaction:
|
|
def __init__(
|
|
self,
|
|
plan_repo: RecordingPlansRepository,
|
|
context_repo: RecordingContextRepository,
|
|
) -> None:
|
|
self.plans = plan_repo
|
|
self.contexts = context_repo
|
|
|
|
def __enter__(self) -> RecordingTransaction:
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> bool:
|
|
return False
|
|
|
|
|
|
class RecordingUnitOfWork:
|
|
def __init__(self, plan: Plan) -> None:
|
|
self._plans = RecordingPlansRepository(plan)
|
|
self._contexts = RecordingContextRepository()
|
|
|
|
def transaction(self) -> RecordingTransaction:
|
|
return RecordingTransaction(self._plans, self._contexts)
|
|
|
|
|
|
class RecordingVectorStoreStub:
|
|
"""Test double that records vector store usage."""
|
|
|
|
def __init__(self, *, enabled: bool) -> None:
|
|
self.enabled = enabled
|
|
self.refresh_calls: list[int] = []
|
|
self.search_calls: list[tuple[int, str, int, bool]] = []
|
|
self.results: list[dict[str, Any]] = []
|
|
self.refresh_error: Exception | None = None
|
|
self.search_error: Exception | None = None
|
|
|
|
def is_enabled(self) -> bool:
|
|
return self.enabled
|
|
|
|
def refresh_for_plan(self, plan_id: int) -> int:
|
|
self.refresh_calls.append(plan_id)
|
|
if self.refresh_error:
|
|
raise self.refresh_error
|
|
return len(self.results)
|
|
|
|
def search(
|
|
self,
|
|
plan_id: int,
|
|
query: str,
|
|
*,
|
|
top_k: int = 5,
|
|
refresh_if_missing: bool = True,
|
|
) -> list[dict[str, Any]]:
|
|
self.search_calls.append((plan_id, query, top_k, refresh_if_missing))
|
|
if self.search_error:
|
|
raise self.search_error
|
|
return [*self.results]
|
|
|
|
|
|
def _setup_vector_context_service(context, *, enabled: bool) -> None:
|
|
temp_dir = Path(tempfile.mkdtemp(prefix="context-vector-"))
|
|
add_cleanup(context, lambda: shutil.rmtree(temp_dir, ignore_errors=True))
|
|
|
|
plan = Plan(
|
|
id=101,
|
|
project_id=1,
|
|
name="vector-plan",
|
|
prompt="Vector integration",
|
|
current=True,
|
|
)
|
|
project = Project(id=1, name="vector-project", path=temp_dir)
|
|
|
|
context.temp_dir = temp_dir
|
|
context.plan = plan
|
|
context.project = project
|
|
context.unit_of_work = RecordingUnitOfWork(plan)
|
|
context.vector_stub = RecordingVectorStoreStub(enabled=enabled)
|
|
context.settings = Settings()
|
|
context.context_service = ContextService(
|
|
settings=context.settings,
|
|
unit_of_work=context.unit_of_work,
|
|
vector_store_service=context.vector_stub,
|
|
)
|
|
context.created_files: dict[str, Path] = {}
|
|
|
|
|
|
@given("a context service with vector search enabled")
|
|
def step_context_service_vector_enabled(context) -> None:
|
|
_setup_vector_context_service(context, enabled=True)
|
|
|
|
|
|
@given("a context service with vector search disabled")
|
|
def step_context_service_vector_disabled(context) -> None:
|
|
_setup_vector_context_service(context, enabled=False)
|
|
|
|
|
|
@given('the vector store stub will return "{path}" with score {score:f}')
|
|
def step_stub_vector_results(context, path: str, score: float) -> None:
|
|
context.vector_stub.results = [
|
|
{"path": path, "score": score, "snippet": f"match:{path}"}
|
|
]
|
|
|
|
|
|
@given("the vector test project loses its identifier")
|
|
def step_vector_project_loses_identifier(context) -> None:
|
|
assert hasattr(context, "project"), "Vector test project is not initialized"
|
|
context.project.id = None
|
|
|
|
|
|
@given("the vector store refresh fails with configuration error")
|
|
def step_vector_refresh_config_error(context) -> None:
|
|
context.vector_stub.refresh_error = ConfigurationError(
|
|
"Vector refresh configuration is invalid"
|
|
)
|
|
|
|
|
|
@given("the vector store search fails with configuration error")
|
|
def step_vector_search_config_error(context) -> None:
|
|
context.vector_stub.search_error = ConfigurationError(
|
|
"Vector search configuration is invalid"
|
|
)
|
|
|
|
|
|
@given('I have a vector context file "{filename}" with content:')
|
|
def step_create_vector_context_file(context, filename: str) -> None:
|
|
file_path = context.temp_dir / filename
|
|
file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
file_path.write_text(context.text.strip("\n"), encoding="utf-8")
|
|
context.created_files[filename] = file_path
|
|
|
|
|
|
@when('I add the file "{filename}" to the vector-aware context')
|
|
def step_add_file_vector_context(context, filename: str) -> None:
|
|
file_path = context.created_files[filename]
|
|
context.context_service.add_to_context(context.project, file_path)
|
|
|
|
|
|
@when('I remove the file "{filename}" from the vector-aware context')
|
|
def step_remove_file_vector_context(context, filename: str) -> None:
|
|
file_path = context.created_files[filename]
|
|
context.context_service.remove_from_context(context.project, file_path)
|
|
|
|
|
|
@when("I clear the vector-aware context")
|
|
def step_clear_vector_context(context) -> None:
|
|
context.context_service.clear_context(context.project)
|
|
|
|
|
|
@when('I search the context for "{query}" with limit {limit:d}')
|
|
def step_search_context(context, query: str, limit: int) -> None:
|
|
context.search_results = context.context_service.search_context(
|
|
context.project,
|
|
query,
|
|
limit=limit,
|
|
)
|
|
|
|
|
|
@then('the semantic search results should include "{path}" with score {score:f}')
|
|
def step_assert_semantic_result(context, path: str, score: float) -> None:
|
|
assert context.search_results, "No semantic search results recorded"
|
|
first = context.search_results[0]
|
|
assert first["path"] == path, first
|
|
assert abs(first["score"] - score) < 1e-9, first
|
|
|
|
|
|
@then("the vector store query count should be {expected:d}")
|
|
def step_assert_query_count(context, expected: int) -> None:
|
|
assert len(context.vector_stub.search_calls) == expected, (
|
|
context.vector_stub.search_calls
|
|
)
|
|
|
|
|
|
@then("the semantic search results should be empty")
|
|
def step_semantic_results_empty(context) -> None:
|
|
assert context.search_results == []
|
|
|
|
|
|
@then("the vector index should have refreshed {expected:d} time")
|
|
@then("the vector index should have refreshed {expected:d} times")
|
|
def step_assert_refresh_count(context, expected: int) -> None:
|
|
assert len(context.vector_stub.refresh_calls) == expected, (
|
|
context.vector_stub.refresh_calls
|
|
)
|