Spaces:
Running on Zero
Running on Zero
File size: 3,879 Bytes
0828c2c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """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)
|