agent_search_gateway.providers.registry

 1"""Registration-order-preserving web provider registry."""
 2
 3from collections.abc import Callable
 4from dataclasses import dataclass
 5from typing import Literal
 6
 7from .contracts import ProviderCapabilities
 8
 9WebStage = Literal["search", "fetch"]
10WebProviderFactory = Callable[..., object]
11
12
13@dataclass(frozen=True, slots=True)
14class WebProviderRegistration:
15    name: str
16    capabilities: ProviderCapabilities
17    factory: WebProviderFactory
18    allowed_config_keys: frozenset[str]
19    requires_api_key: bool = True
20
21
22class ProviderRegistry:
23    def __init__(self) -> None:
24        self._registrations: dict[str, WebProviderRegistration] = {}
25
26    def register(self, registration: WebProviderRegistration) -> None:
27        if registration.name in self._registrations:
28            raise ValueError(f"provider already registered: {registration.name}")
29        self._registrations[registration.name] = registration
30
31    def get(self, name: str) -> WebProviderRegistration | None:
32        return self._registrations.get(name)
33
34    def require(self, name: str) -> WebProviderRegistration:
35        registration = self.get(name)
36        if registration is None:
37            raise KeyError(name)
38        return registration
39
40    def capabilities(self, name: str) -> ProviderCapabilities:
41        return self.require(name).capabilities
42
43    def list_in_registration_order(self) -> tuple[WebProviderRegistration, ...]:
44        return tuple(self._registrations.values())
45
46    def for_stage(self, stage: WebStage) -> tuple[WebProviderRegistration, ...]:
47        if stage == "search":
48            return tuple(item for item in self._registrations.values() if item.capabilities.search)
49        return tuple(item for item in self._registrations.values() if item.capabilities.fetch)