src/retrieval/base.py

"""Base classes for retrieval strategies."""

import logging
from abc import ABC, abstractmethod
from typing import Any

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever

logger = logging.getLogger(__name__)


class BaseRetrieverStrategy(ABC):
    """Abstract base class for all retrieval strategies.

    This provides a common interface for different retrieval approaches:
    - Vector-based (semantic similarity)
    - BM25 (lexical/keyword matching)
    - Hybrid combinations
    - Future strategies (PageIndex, GraphRAG, etc.)

    Subclasses must implement all abstract methods to be usable
    with the RetrieverFactory.
    """

    def __init__(self, config: dict[str, Any]):
        """Initialize the retrieval strategy.

        Args:
            config: Configuration dictionary for the strategy
        """
        self.config = config
        self._is_initialized = False

    @property
    @abstractmethod
    def name(self) -> str:
        """Return the unique identifier for this strategy.

        Returns:
            Strategy name (e.g., 'vector', 'bm25', 'bm25_vector')
        """
        pass

    @abstractmethod
    def build_index(self, documents: list[Document]) -> None:
        """Build or rebuild the index from documents.

        This method should process documents and create the necessary
        index structures for retrieval.

        Args:
            documents: List of Document objects to index
        """
        pass

    @abstractmethod
    def load_index(self) -> bool:
        """Load an existing index from persistent storage.

        Returns:
            True if index was loaded successfully, False otherwise
        """
        pass

    @abstractmethod
    def retrieve(self, query: str, k: int = 4) -> list[Document]:
        """Retrieve relevant documents for a query.

        Args:
            query: The search query string
            k: Number of documents to retrieve

        Returns:
            List of relevant Document objects
        """
        pass

    @abstractmethod
    def as_retriever(self, **kwargs: Any) -> BaseRetriever:
        """Return a LangChain-compatible retriever.

        Args:
            **kwargs: Additional arguments for the retriever

        Returns:
            A BaseRetriever instance compatible with LangChain chains
        """
        pass

    def is_index_ready(self) -> bool:
        """Check if the index is built and ready for queries.

        Returns:
            True if index is ready, False otherwise
        """
        return self._is_initialized

    def get_index_stats(self) -> dict[str, Any]:
        """Get statistics about the current index.

        Returns:
            Dictionary with index statistics (implementation-specific)
        """
        return {
            "strategy": self.name,
            "initialized": self._is_initialized,
        }


class StrategyRetriever(BaseRetriever):
    """LangChain-compatible retriever wrapper for BaseRetrieverStrategy.

    This adapter allows any BaseRetrieverStrategy to be used
    in LangChain LCEL chains and pipelines.
    """

    strategy: BaseRetrieverStrategy
    search_kwargs: dict[str, Any] = {}

    model_config = {"arbitrary_types_allowed": True}

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun | None = None,  # noqa: ARG002
    ) -> list[Document]:
        """Get documents relevant to a query.

        Args:
            query: String to find relevant documents for
            run_manager: Callback manager for the retriever run

        Returns:
            List of relevant documents
        """
        k = self.search_kwargs.get("k", 4)
        return self.strategy.retrieve(query, k=k)