Spaces:
Running on Zero
Running on Zero
File size: 3,400 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 | """Factory for creating retrieval strategies."""
import logging
from typing import Any
from .base import BaseRetrieverStrategy
logger = logging.getLogger(__name__)
class RetrieverFactory:
"""Factory for creating and managing retrieval strategies.
This factory uses a registry pattern to allow dynamic registration
of new retrieval strategies. Strategies are registered using the
@RetrieverFactory.register decorator.
Example:
@RetrieverFactory.register("my_strategy")
class MyStrategy(BaseRetrieverStrategy):
...
# Later:
strategy = RetrieverFactory.create("my_strategy", config)
"""
_strategies: dict[str, type[BaseRetrieverStrategy]] = {}
@classmethod
def register(cls, name: str):
"""Decorator to register a retrieval strategy.
Args:
name: Unique identifier for the strategy
Returns:
Decorator function that registers the strategy class
Example:
@RetrieverFactory.register("vector")
class VectorStrategy(BaseRetrieverStrategy):
...
"""
def decorator(strategy_class: type[BaseRetrieverStrategy]):
if name in cls._strategies:
logger.warning(f"Overwriting existing strategy: {name}")
cls._strategies[name] = strategy_class
logger.debug(f"Registered retrieval strategy: {name}")
return strategy_class
return decorator
@classmethod
def create(cls, strategy_name: str, config: dict[str, Any]) -> BaseRetrieverStrategy:
"""Create a retrieval strategy instance by name.
Args:
strategy_name: Name of the strategy to create
config: Configuration dictionary for the strategy
Returns:
An instance of the requested strategy
Raises:
ValueError: If strategy_name is not registered
"""
if strategy_name not in cls._strategies:
available = ", ".join(cls._strategies.keys()) or "none"
msg = f"Unknown retrieval strategy: '{strategy_name}'. Available: {available}"
raise ValueError(msg)
strategy_class = cls._strategies[strategy_name]
logger.info(f"Creating retrieval strategy: {strategy_name}")
return strategy_class(config)
@classmethod
def available_strategies(cls) -> list[str]:
"""List all registered strategy names.
Returns:
List of registered strategy names
"""
return list(cls._strategies.keys())
@classmethod
def is_registered(cls, name: str) -> bool:
"""Check if a strategy is registered.
Args:
name: Strategy name to check
Returns:
True if strategy is registered, False otherwise
"""
return name in cls._strategies
@classmethod
def get_strategy_class(cls, name: str) -> type[BaseRetrieverStrategy] | None:
"""Get the strategy class by name without instantiating.
Args:
name: Strategy name
Returns:
Strategy class or None if not found
"""
return cls._strategies.get(name)
def get_retriever_factory() -> type[RetrieverFactory]:
"""Get the RetrieverFactory class.
Returns:
The RetrieverFactory class
"""
return RetrieverFactory
|