temp/features/steps/faiss_vector_backend_steps.py

"""Step definitions for the FAISS ACMS vector backend feature."""

from __future__ import annotations

import os
import shutil
import tempfile
from pathlib import Path
from typing import Any, ClassVar, cast
from unittest.mock import patch

from behave import given, then, when
from langchain_community.embeddings import FakeEmbeddings

from cleveragents.application.container import get_container, reset_container
from cleveragents.application.services.config_service import ConfigService
from cleveragents.application.services.faiss_vector_backend import (
    FAISSVectorBackend,
    FAISSVectorIndexBackend,
)
from cleveragents.application.services.vector_store_service import VectorStoreService
from cleveragents.config.settings import Settings
from cleveragents.domain.models.acms.index_stubs import InMemoryVectorIndexBackend
from cleveragents.domain.models.acms.stubs import InMemoryVectorBackend
from cleveragents.infrastructure.database.unit_of_work import UnitOfWork
from features.steps.service_steps import add_cleanup


class _StubDocument:
    def __init__(self, page_content: str, metadata: dict[str, str]) -> None:
        self.page_content = page_content
        self.metadata = dict(metadata)


class _StubDocStore:
    def __init__(self) -> None:
        self._docs: dict[str, _StubDocument] = {}

    def search(self, store_id: str) -> _StubDocument:
        return self._docs[store_id]


class RecordingFAISS:
    """Deterministic FAISS stand-in for ACMS backend tests."""

    saved_snapshots: ClassVar[
        dict[str, list[tuple[str, str, list[float], dict[str, str]]]]
    ] = {}
    last_embedding_backend: ClassVar[Any | None] = None

    def __init__(self) -> None:
        self.docstore = _StubDocStore()
        self.index_to_docstore_id: dict[int, str] = {}
        self._vectors: dict[str, list[float]] = {}

    @classmethod
    def reset(cls) -> None:
        cls.saved_snapshots = {}
        cls.last_embedding_backend = None

    @classmethod
    def from_embeddings(
        cls,
        text_embeddings: list[tuple[str, list[float]]],
        embedding: Any,
        metadatas: list[dict[str, str]] | None = None,
        ids: list[str] | None = None,
        **_: Any,
    ) -> RecordingFAISS:
        instance = cls()
        cls.last_embedding_backend = embedding
        instance.add_embeddings(text_embeddings, metadatas=metadatas, ids=ids)
        return instance

    def add_embeddings(
        self,
        text_embeddings: list[tuple[str, list[float]]],
        metadatas: list[dict[str, str]] | None = None,
        ids: list[str] | None = None,
        **_: Any,
    ) -> list[str]:
        stored_ids: list[str] = []
        metadata_list = metadatas or [{} for _ in text_embeddings]
        for index, (text, vector) in enumerate(text_embeddings):
            store_id = ids[index] if ids is not None else f"doc-{len(self._vectors)}"
            self.docstore._docs[store_id] = _StubDocument(text, metadata_list[index])
            self._vectors[store_id] = list(vector)
            stored_ids.append(store_id)
        self._rebuild_index_map()
        return stored_ids

    def similarity_search_with_score_by_vector(
        self,
        embedding: list[float],
        k: int = 4,
        filter: Any | None = None,
        fetch_k: int = 20,
        **_: Any,
    ) -> list[tuple[_StubDocument, float]]:
        matches: list[tuple[_StubDocument, float]] = []
        for store_id, vector in self._vectors.items():
            document = self.docstore.search(store_id)
            if filter is not None and not filter(document.metadata):
                continue
            distance = sum(
                abs(left - right)
                for left, right in zip(embedding, vector, strict=False)
            )
            matches.append((document, float(distance)))
        matches.sort(key=lambda item: item[1])
        return matches[: min(fetch_k, k)]

    def save_local(self, directory: str, index_name: str = "index") -> None:
        _ = index_name
        RecordingFAISS.saved_snapshots[directory] = [
            (
                store_id,
                self.docstore.search(store_id).page_content,
                list(self._vectors[store_id]),
                dict(self.docstore.search(store_id).metadata),
            )
            for store_id in self.index_to_docstore_id.values()
        ]
        target = Path(directory)
        target.mkdir(parents=True, exist_ok=True)
        (target / "index.faiss").write_bytes(b"stub")
        (target / "index.pkl").write_bytes(b"stub")

    @classmethod
    def load_local(
        cls,
        directory: str,
        embeddings: Any,
        index_name: str = "index",
        *,
        allow_dangerous_deserialization: bool = False,
        **_: Any,
    ) -> RecordingFAISS:
        _unused = (index_name, allow_dangerous_deserialization)
        cls.last_embedding_backend = embeddings
        if directory not in cls.saved_snapshots:
            raise FileNotFoundError(directory)
        instance = cls()
        for store_id, text, vector, metadata in cls.saved_snapshots[directory]:
            instance.docstore._docs[store_id] = _StubDocument(text, metadata)
            instance._vectors[store_id] = list(vector)
        instance._rebuild_index_map()
        return instance

    def delete(self, ids: list[str] | None = None, **_: Any) -> bool:
        for store_id in ids or []:
            self.docstore._docs.pop(store_id, None)
            self._vectors.pop(store_id, None)
        self._rebuild_index_map()
        return True

    def _rebuild_index_map(self) -> None:
        ordered_ids = sorted(self._vectors)
        self.index_to_docstore_id = {
            position: store_id for position, store_id in enumerate(ordered_ids)
        }


def _set_env(context: Any, key: str, value: str | None) -> None:
    original = os.environ.get(key)
    if value is None:
        os.environ.pop(key, None)
    else:
        os.environ[key] = value

    def cleanup() -> None:
        if original is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = original

    add_cleanup(context, cleanup)


def _parse_vector(raw: str) -> list[float]:
    return [float(part.strip()) for part in raw.split(",") if part.strip()]


@given("the ACMS FAISS backend is configured to use fake embeddings")
def step_configure_acms_faiss(context: Any) -> None:
    index_dir = tempfile.mkdtemp(prefix="acms-faiss-index-")
    add_cleanup(context, lambda: shutil.rmtree(index_dir, ignore_errors=True))
    _set_env(context, "CLEVERAGENTS_INDEX_VECTOR_BACKEND", "faiss")
    _set_env(context, "CLEVERAGENTS_INDEX_VECTOR_DIR", index_dir)
    _set_env(context, "CLEVERAGENTS_EMBEDDING_PROVIDER", "fake")
    _set_env(context, "CLEVERAGENTS_EMBEDDING_DIMENSIONS", "8")
    reset_container()


@given("FAISS interactions are recorded for ACMS backends")
def step_patch_acms_faiss(context: Any) -> None:
    RecordingFAISS.reset()
    patchers = [
        patch(
            "cleveragents.application.services.vector_store_service.FAISS",
            RecordingFAISS,
        ),
        patch(
            "cleveragents.application.services.faiss_vector_backend.FAISS",
            RecordingFAISS,
        ),
    ]
    for patcher in patchers:
        patcher.start()
        add_cleanup(context, patcher.stop)


@given("FAISS is unavailable for ACMS backends")
def step_disable_acms_faiss(context: Any) -> None:
    patchers = [
        patch("cleveragents.application.services.vector_store_service.FAISS", None),
        patch("cleveragents.application.services.faiss_vector_backend.FAISS", None),
    ]
    for patcher in patchers:
        patcher.start()
        add_cleanup(context, patcher.stop)
    reset_container()


@given("an ACMS FAISS backend pair")
def step_acms_backend_pair(context: Any) -> None:
    settings = Settings()
    unit_of_work = UnitOfWork("sqlite:///:memory:", require_confirmation=False)
    service = VectorStoreService(settings, unit_of_work, ConfigService())
    context.acms_service = service
    context.acms_vector_backend = FAISSVectorBackend(service)
    context.acms_index_backend = FAISSVectorIndexBackend(service)


@given(
    'I index ACMS embedding "{doc_id}" in project "{project}" for resource "{resource_id}" with vector "{vector}"'
)
@when(
    'I index ACMS embedding "{doc_id}" in project "{project}" for resource "{resource_id}" with vector "{vector}"'
)
def step_index_acms_embedding(
    context: Any,
    doc_id: str,
    project: str,
    resource_id: str,
    vector: str,
) -> None:
    context.acms_index_backend.index_embedding(
        project,
        doc_id,
        _parse_vector(vector),
        {
            "resource_id": resource_id,
            "location": f"src/{resource_id.lower()}.py",
            "resource_type": "python",
        },
    )


@when('I search ACMS vectors in project "{project}" for vector "{vector}"')
def step_search_project_vectors(context: Any, project: str, vector: str) -> None:
    context.index_results = context.acms_index_backend.search_similar(
        project,
        _parse_vector(vector),
        limit=5,
    )


@when('I search ACMS vectors scoped to resource "{resource_id}" for vector "{vector}"')
def step_search_scoped_vectors(context: Any, resource_id: str, vector: str) -> None:
    context.vector_results = context.acms_vector_backend.similarity_search(
        _parse_vector(vector),
        scope=frozenset({resource_id}),
        top_k=5,
    )


@when('I remove ACMS embedding "{doc_id}" from project "{project}"')
def step_remove_acms_embedding(context: Any, doc_id: str, project: str) -> None:
    context.acms_index_backend.remove_embedding(project, doc_id)


@when("I resolve ACMS vector backends from the DI container")
def step_resolve_container_backends(context: Any) -> None:
    reset_container()
    container = get_container()
    context.container_vector_backend = container.vector_backend()
    context.container_index_vector_backend = container.index_vector_backend()


@then('the project vector search should include doc_id "{doc_id}"')
def step_assert_project_result(context: Any, doc_id: str) -> None:
    assert any(result.doc_id == doc_id for result in context.index_results)


@then('the scoped vector search should include uko_uri "{uko_uri}"')
def step_assert_scoped_result(context: Any, uko_uri: str) -> None:
    assert any(result.uko_uri == uko_uri for result in context.vector_results)


@then("the scoped vector search should be empty")
def step_assert_scoped_empty(context: Any) -> None:
    assert context.vector_results == []


@then('project "{project}" should have {count:d} ACMS embeddings')
def step_assert_embedding_count(context: Any, project: str, count: int) -> None:
    assert context.acms_service.acms_count(project=project) == count


@then("the ACMS embeddings provider should be FakeEmbeddings sized {size:d}")
def step_assert_fake_embeddings(context: Any, size: int) -> None:
    backend = RecordingFAISS.last_embedding_backend
    assert isinstance(backend, FakeEmbeddings)
    typed_backend = cast(Any, backend)
    assert typed_backend.size == size


@then("the container vector backend should be a FAISSVectorBackend")
def step_assert_container_vector_backend(context: Any) -> None:
    assert isinstance(context.container_vector_backend, FAISSVectorBackend)


@then("the container index vector backend should be a FAISSVectorIndexBackend")
def step_assert_container_index_backend(context: Any) -> None:
    assert isinstance(context.container_index_vector_backend, FAISSVectorIndexBackend)


@then("the container vector backend should be an InMemoryVectorBackend")
def step_assert_container_vector_fallback(context: Any) -> None:
    assert isinstance(context.container_vector_backend, InMemoryVectorBackend)


@then("the container index vector backend should be an InMemoryVectorIndexBackend")
def step_assert_container_index_fallback(context: Any) -> None:
    assert isinstance(
        context.container_index_vector_backend, InMemoryVectorIndexBackend
    )