agent_search_gateway.runtime

  1"""Resolved runtime assembly for the foreground daemon."""
  2
  3import logging
  4from collections.abc import Callable, Mapping
  5from typing import cast
  6
  7import httpx
  8
  9from .academic.aggregator import PaperAggregator
 10from .concurrency import ProviderQuotaManager
 11from .config import (
 12    ResolvedAcademicProviderConfig,
 13    ResolvedConfig,
 14    ResolvedOAResolverConfig,
 15    ResolvedWebProviderConfig,
 16)
 17from .errors import ConfigFailure, ErrorCode
 18from .llm.stages import LLMStages
 19from .observability import log_event
 20from .orchestrators.fetch import FetchOrchestrator
 21from .orchestrators.paper import PaperSearchOrchestrator
 22from .orchestrators.search import SearchOrchestrator
 23from .paths import RuntimePaths
 24from .providers.academic.defaults import (
 25    build_default_academic_registry,
 26    build_default_oa_resolver_registry,
 27)
 28from .providers.academic.registry import AcademicProviderRegistry, OAResolverRegistry
 29from .providers.contracts import (
 30    AcademicSearchProvider,
 31    KeywordSearchProvider,
 32    LLMClient,
 33    OAResolver,
 34    URLFetchProvider,
 35)
 36from .providers.defaults import build_default_registry
 37from .providers.http import HttpJsonExecutor
 38from .providers.openai_chat import OpenAIChatCompletionsClient
 39from .providers.registry import ProviderRegistry
 40from .result_writer import ResultWriter
 41from .scheduler.fetch import FetchScheduler
 42from .url_store import URLStore
 43
 44HttpClientFactory = Callable[[], httpx.AsyncClient]
 45_RESERVED_WEB_ADAPTER_KWARGS = frozenset({"name", "http_executor", "secret"})
 46_RESERVED_ACADEMIC_ADAPTER_KWARGS = frozenset({"name", "executor", "api_key", "contact_email"})
 47
 48
 49class Runtime:
 50    """Owns one in-memory gateway runtime and its transport clients."""
 51
 52    def __init__(
 53        self,
 54        *,
 55        quotas: ProviderQuotaManager,
 56        web_search_providers: tuple[KeywordSearchProvider, ...],
 57        web_fetch_providers: tuple[URLFetchProvider, ...],
 58        academic_search_providers: tuple[AcademicSearchProvider, ...],
 59        oa_resolver: OAResolver | None,
 60        llm_clients: Mapping[str, LLMClient],
 61        store: URLStore,
 62        search_orchestrator: SearchOrchestrator,
 63        paper_search_orchestrator: PaperSearchOrchestrator,
 64        fetch_orchestrator: FetchOrchestrator,
 65        web_http_executors: tuple[HttpJsonExecutor, ...],
 66        academic_http_executors: tuple[HttpJsonExecutor, ...],
 67    ) -> None:
 68        self.quotas = quotas
 69        self.web_search_providers = web_search_providers
 70        self.web_fetch_providers = web_fetch_providers
 71        self.academic_search_providers = academic_search_providers
 72        self.oa_resolver = oa_resolver
 73        self.llm_clients = dict(llm_clients)
 74        self.store = store
 75        self.search_orchestrator = search_orchestrator
 76        self.paper_search_orchestrator = paper_search_orchestrator
 77        self.fetch_orchestrator = fetch_orchestrator
 78        self._web_http_executors = web_http_executors
 79        self._academic_http_executors = academic_http_executors
 80        self._closed = False
 81
 82    @classmethod
 83    def build(
 84        cls,
 85        config: ResolvedConfig,
 86        paths: RuntimePaths,
 87        *,
 88        registry: ProviderRegistry | None = None,
 89        academic_registry: AcademicProviderRegistry | None = None,
 90        oa_resolver_registry: OAResolverRegistry | None = None,
 91        http_client_factory: HttpClientFactory = httpx.AsyncClient,
 92    ) -> "Runtime":
 93        provider_registry = registry or build_default_registry()
 94        academic_provider_registry = academic_registry or build_default_academic_registry()
 95        resolver_registry = oa_resolver_registry or build_default_oa_resolver_registry()
 96        enabled_web = tuple(
 97            item for item in config.web.providers if item.enable_search or item.enable_fetch
 98        )
 99        enabled_academic = tuple(item for item in config.academic.providers if item.enabled)
100        quotas = ProviderQuotaManager(
101            web_limits={item.name: item.max_concurrency for item in enabled_web},
102            llm_limits={item.name: item.max_concurrency for item in config.llm.providers},
103            academic_limits={item.name: item.max_concurrency for item in enabled_academic},
104        )
105        web_search, web_fetch, web_executors = cls._build_web_providers(
106            enabled_web,
107            provider_registry,
108            config,
109            http_client_factory,
110        )
111        academic_search, academic_executors = cls._build_academic_providers(
112            enabled_academic,
113            academic_provider_registry,
114            config,
115            http_client_factory,
116        )
117        oa_resolver, resolver_executors = cls._build_oa_resolver(
118            config.oa_resolver,
119            resolver_registry,
120            config,
121            http_client_factory,
122        )
123        llm_clients = cls._build_llm_clients(config, quotas, http_client_factory)
124        stages = LLMStages(
125            llm_clients,
126            judge=config.llm.judge,
127            safety=config.llm.safety,
128            content_clean=config.llm.content_clean,
129            focus_summary=config.llm.focus_summary,
130        )
131        store = URLStore()
132        result_writer = ResultWriter(paths.results_dir)
133        paper_aggregator = PaperAggregator(
134            (
135                *(provider.name for provider in academic_search),
136                *(f"llm:{item.provider}" for item in config.llm.search_invocations),
137            )
138        )
139        search_orchestrator = SearchOrchestrator(
140            keyword_providers=web_search,
141            llm_invocations=config.llm.search_invocations,
142            quotas=quotas,
143            stages=stages,
144            store=store,
145            result_writer=result_writer,
146            paper_aggregator=paper_aggregator,
147            paper_resolver=oa_resolver,
148        )
149        paper_search_orchestrator = PaperSearchOrchestrator(
150            providers=academic_search,
151            quotas=quotas,
152            aggregator=paper_aggregator,
153            resolver=oa_resolver,
154            store=store,
155            result_writer=result_writer,
156        )
157        fetch_orchestrator = FetchOrchestrator(
158            store=store,
159            scheduler=FetchScheduler(web_fetch, quotas, stages),
160            stages=stages,
161        )
162        runtime = cls(
163            quotas=quotas,
164            web_search_providers=web_search,
165            web_fetch_providers=web_fetch,
166            academic_search_providers=academic_search,
167            oa_resolver=oa_resolver,
168            llm_clients=llm_clients,
169            store=store,
170            search_orchestrator=search_orchestrator,
171            paper_search_orchestrator=paper_search_orchestrator,
172            fetch_orchestrator=fetch_orchestrator,
173            web_http_executors=web_executors,
174            academic_http_executors=(*academic_executors, *resolver_executors),
175        )
176        log_event(
177            logging.getLogger(__name__),
178            logging.DEBUG,
179            "runtime_built",
180            web_providers=",".join(item.name for item in enabled_web) or "-",
181            web_provider_count=len(enabled_web),
182            llm_providers=",".join(item.name for item in config.llm.providers) or "-",
183            llm_provider_count=len(config.llm.providers),
184            academic_providers=",".join(item.name for item in enabled_academic) or "-",
185            academic_provider_count=len(enabled_academic),
186            oa_resolver=config.oa_resolver.name if config.oa_resolver is not None else "-",
187            web_limits=",".join(f"{item.name}:{item.max_concurrency}" for item in enabled_web)
188            or "-",
189            llm_limits=",".join(
190                f"{item.name}:{item.max_concurrency}" for item in config.llm.providers
191            )
192            or "-",
193            academic_limits=",".join(
194                f"{item.name}:{item.max_concurrency}" for item in enabled_academic
195            )
196            or "-",
197        )
198        return runtime
199
200    @classmethod
201    def _build_web_providers(
202        cls,
203        providers: tuple[ResolvedWebProviderConfig, ...],
204        registry: ProviderRegistry,
205        config: ResolvedConfig,
206        http_client_factory: HttpClientFactory,
207    ) -> tuple[
208        tuple[KeywordSearchProvider, ...],
209        tuple[URLFetchProvider, ...],
210        tuple[HttpJsonExecutor, ...],
211    ]:
212        search: list[KeywordSearchProvider] = []
213        fetch: list[URLFetchProvider] = []
214        executors: list[HttpJsonExecutor] = []
215        for provider_config in providers:
216            registration = registry.get(provider_config.name)
217            if registration is None:
218                raise ConfigFailure(
219                    ErrorCode.CONFIG_ERROR,
220                    f"Invalid enabled web provider: {provider_config.name}",
221                )
222            credential_kwargs: dict[str, object] = {}
223            if registration.requires_api_key:
224                private_value = provider_config.secret
225                if private_value is None:
226                    raise ConfigFailure(
227                        ErrorCode.CONFIG_ERROR,
228                        f"Invalid enabled web provider: {provider_config.name}",
229                    )
230                credential_kwargs["secret"] = private_value
231            reserved = set(provider_config.options) & _RESERVED_WEB_ADAPTER_KWARGS
232            if reserved:
233                names = ", ".join(sorted(reserved))
234                raise ConfigFailure(
235                    ErrorCode.CONFIG_ERROR,
236                    f"Reserved config key(s) for web provider {provider_config.name}: {names}",
237                )
238            executor = HttpJsonExecutor(
239                http_client_factory(),
240                config.retry,
241                provider_name=provider_config.name,
242            )
243            kwargs: dict[str, object] = {
244                "name": provider_config.name,
245                "http_executor": executor,
246            }
247            kwargs.update(credential_kwargs)
248            kwargs.update(provider_config.options)
249            try:
250                adapter = registration.factory(**kwargs)
251            except TypeError as exc:
252                raise ConfigFailure(
253                    ErrorCode.CONFIG_ERROR,
254                    f"Invalid configuration for web provider {provider_config.name}",
255                ) from exc
256            executors.append(executor)
257            if provider_config.enable_search:
258                search.append(cast(KeywordSearchProvider, adapter))
259            if provider_config.enable_fetch:
260                fetch.append(cast(URLFetchProvider, adapter))
261        return tuple(search), tuple(fetch), tuple(executors)
262
263    @classmethod
264    def _build_academic_providers(
265        cls,
266        providers: tuple[ResolvedAcademicProviderConfig, ...],
267        registry: AcademicProviderRegistry,
268        config: ResolvedConfig,
269        http_client_factory: HttpClientFactory,
270    ) -> tuple[tuple[AcademicSearchProvider, ...], tuple[HttpJsonExecutor, ...]]:
271        search: list[AcademicSearchProvider] = []
272        executors: list[HttpJsonExecutor] = []
273        for provider_config in providers:
274            registration = registry.get(provider_config.name)
275            if registration is None:
276                raise ConfigFailure(
277                    ErrorCode.CONFIG_ERROR,
278                    f"Invalid enabled academic provider: {provider_config.name}",
279                )
280            reserved = set(provider_config.options) & _RESERVED_ACADEMIC_ADAPTER_KWARGS
281            if reserved:
282                names = ", ".join(sorted(reserved))
283                raise ConfigFailure(
284                    ErrorCode.CONFIG_ERROR,
285                    f"Reserved config key(s) for academic provider {provider_config.name}: {names}",
286                )
287            executor = HttpJsonExecutor(
288                http_client_factory(),
289                config.retry,
290                provider_name=provider_config.name,
291            )
292            kwargs: dict[str, object] = {"executor": executor}
293            if provider_config.api_key is not None:
294                kwargs["api_key"] = provider_config.api_key
295            if provider_config.contact_email is not None:
296                kwargs["contact_email"] = provider_config.contact_email
297            kwargs.update(provider_config.options)
298            try:
299                adapter = registration.factory(**kwargs)
300            except TypeError as exc:
301                raise ConfigFailure(
302                    ErrorCode.CONFIG_ERROR,
303                    f"Invalid configuration for academic provider {provider_config.name}",
304                ) from exc
305            search.append(cast(AcademicSearchProvider, adapter))
306            executors.append(executor)
307        return tuple(search), tuple(executors)
308
309    @classmethod
310    def _build_oa_resolver(
311        cls,
312        resolver_config: ResolvedOAResolverConfig | None,
313        registry: OAResolverRegistry,
314        config: ResolvedConfig,
315        http_client_factory: HttpClientFactory,
316    ) -> tuple[OAResolver | None, tuple[HttpJsonExecutor, ...]]:
317        if resolver_config is None:
318            return None, ()
319        registration = registry.get(resolver_config.name)
320        if registration is None:
321            raise ConfigFailure(
322                ErrorCode.CONFIG_ERROR,
323                f"Invalid enabled OA resolver: {resolver_config.name}",
324            )
325        reserved = set(resolver_config.options) & _RESERVED_ACADEMIC_ADAPTER_KWARGS
326        if reserved:
327            names = ", ".join(sorted(reserved))
328            raise ConfigFailure(
329                ErrorCode.CONFIG_ERROR,
330                f"Reserved config key(s) for OA resolver {resolver_config.name}: {names}",
331            )
332        executor = HttpJsonExecutor(
333            http_client_factory(),
334            config.retry,
335            provider_name=resolver_config.name,
336        )
337        kwargs: dict[str, object] = {"executor": executor}
338        if resolver_config.api_key is not None:
339            kwargs["api_key"] = resolver_config.api_key
340        if resolver_config.contact_email is not None:
341            kwargs["contact_email"] = resolver_config.contact_email
342        kwargs.update(resolver_config.options)
343        try:
344            adapter = registration.factory(**kwargs)
345        except TypeError as exc:
346            raise ConfigFailure(
347                ErrorCode.CONFIG_ERROR,
348                f"Invalid configuration for OA resolver {resolver_config.name}",
349            ) from exc
350        return cast(OAResolver, adapter), (executor,)
351
352    @staticmethod
353    def _build_llm_clients(
354        config: ResolvedConfig,
355        quotas: ProviderQuotaManager,
356        http_client_factory: HttpClientFactory,
357    ) -> dict[str, LLMClient]:
358        clients: dict[str, LLMClient] = {}
359        for provider_config in config.llm.providers:
360            private_value = provider_config.secret
361            executor = HttpJsonExecutor(
362                http_client_factory(),
363                config.retry,
364                provider_name=provider_config.name,
365            )
366            clients[provider_config.name] = OpenAIChatCompletionsClient(
367                name=provider_config.name,
368                api_url=provider_config.api_url,
369                secret=private_value,
370                executor=executor,
371                quota=quotas.get_llm(provider_config.name),
372                retry_policy=config.retry,
373            )
374        return clients
375
376    async def aclose(self) -> None:
377        if self._closed:
378            return
379        self._closed = True
380        for executor in self._web_http_executors:
381            await executor.aclose()
382        for executor in self._academic_http_executors:
383            await executor.aclose()
384        for client in self.llm_clients.values():
385            await client.aclose()
386
387    def __repr__(self) -> str:
388        resolver_name = self.oa_resolver.name if self.oa_resolver is not None else "-"
389        return (
390            "Runtime("
391            f"web_search={len(self.web_search_providers)}, "
392            f"web_fetch={len(self.web_fetch_providers)}, "
393            f"academic_search={len(self.academic_search_providers)}, "
394            f"oa_resolver={resolver_name}, "
395            f"llm_clients={len(self.llm_clients)}"
396            ")"
397        )