""" Main LLM client — presents the OpenAI Responses API interface and routes to provider adapters. All methods are async for non-blocking I/O. """ from __future__ import annotations import asyncio import logging from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any, TypeVar from omnigent.llms._responses_to_chat import ( chat_response_to_response, chat_stream_to_response_events, responses_input_to_chat_messages, ) from omnigent.llms._usage_observer import notify as _notify_usage from omnigent.llms.adapters import get_adapter from omnigent.llms.adapters.openai import OpenAIAdapter from omnigent.llms.errors import ( PermanentLLMError, RetryableLLMError, ) from omnigent.llms.routing import parse_model_string from omnigent.llms.types import ( Response, ResponseCompletedEvent, ResponseStreamEvent, ) from omnigent.reasoning_effort import OPENAI_EFFORTS, validate_effort_or_llm_error from omnigent.runtime.llm_retry import classify_llm_error from omnigent.spec.types import RetryPolicy _logger = logging.getLogger(__name__) _T = TypeVar("_T") def _emit_usage_from_response(response: Response) -> None: usage = response.usage if usage is None: return _notify_usage( model=response.model, input_tokens=int(usage.input_tokens or 0), output_tokens=int(usage.output_tokens or 0), total_tokens=int(usage.total_tokens or 0), ) async def _tee_stream_for_usage( stream: AsyncIterator[ResponseStreamEvent], ) -> AsyncIterator[ResponseStreamEvent]: async for event in stream: if isinstance(event, ResponseCompletedEvent): _emit_usage_from_response(event.response) yield event class _ResponsesNamespace: """ Namespace providing ``client.responses.create()`` to mirror the OpenAI SDK interface. :param client: The parent :class:`Client` instance. """ def __init__(self, client: Client) -> None: self._client = client async def create( self, *, input: list[dict[str, Any]], instructions: str | None = None, model: str, tools: list[dict[str, Any]] | None = None, reasoning: dict[str, str] | None = None, stream: bool = False, connection_params: dict[str, str] | None = None, timeout: int | None = None, retry: RetryPolicy | None = None, **kwargs: Any, ) -> Response | AsyncIterator[ResponseStreamEvent]: """ Create a response from the LLM, routing to the appropriate provider based on the model string. :param input: Responses API input items, e.g. ``[{"role": "user", "content": "Hello"}]``. :param instructions: System instructions string. :param model: Provider-prefixed model string, e.g. ``"anthropic/claude-sonnet-4-20250514"`` or ``"gpt-5.4"``. :param tools: OpenAI-format tool schemas, or ``None``. :param reasoning: Reasoning configuration dict, e.g. ``{"effort": "high", "summary": "concise"}``. :param stream: If ``True``, return an async iterator of streaming events. If ``False``, return a :class:`Response`. :param connection_params: Per-call connection overrides. Keys are provider-specific, e.g. ``{"api_key": "...", "base_url": "..."}`` for OpenAI-compatible providers, or ``{"aws_region": "us-west-2"}`` for Bedrock. ``None`` uses the adapter's default credentials. :param timeout: Request timeout in seconds. ``None`` uses the adapter's default (120s non-streaming, 300s streaming). :param retry: Retry policy for transient failures (timeouts, rate limits). ``None`` disables client-level retries. Useful for standalone calls outside the workflow engine. :param kwargs: Additional provider-specific kwargs (e.g. ``temperature``, ``max_tokens``). :returns: A :class:`Response` when ``stream=False``, or an async iterator of :data:`ResponseStreamEvent` when ``stream=True``. :raises PermanentLLMError: On non-retryable errors. :raises RetryableLLMError: When all retry attempts are exhausted. """ async def call_fn() -> Response | AsyncIterator[ResponseStreamEvent]: """ Dispatch to the adapter. :returns: Response or streaming event iterator. """ return await self._do_create( input=input, instructions=instructions, model=model, tools=tools, reasoning=reasoning, stream=stream, connection_params=connection_params, timeout=timeout, **kwargs, ) if retry is None: result = await call_fn() else: result = await _execute_with_retry(call_fn, retry) if isinstance(result, Response): _emit_usage_from_response(result) return result return _tee_stream_for_usage(result) async def _do_create( self, *, input: list[dict[str, Any]], instructions: str | None, model: str, tools: list[dict[str, Any]] | None, reasoning: dict[str, str] | None, stream: bool, connection_params: dict[str, str] | None, timeout: int | None, **kwargs: Any, ) -> Response | AsyncIterator[ResponseStreamEvent]: """ Route the LLM call to the appropriate provider adapter. :param input: Responses API input items. :param instructions: System instructions string. :param model: Provider-prefixed model string. :param tools: Tool schemas or ``None``. :param reasoning: Reasoning config or ``None``. :param stream: Enable streaming. :param connection_params: Connection overrides or ``None``. :param timeout: Timeout in seconds or ``None``. :param kwargs: Additional provider-specific kwargs. :returns: Response or async streaming event iterator. """ routed = parse_model_string(model) adapter = get_adapter(routed.provider) # OpenAI supports the Responses API natively — use it # directly so reasoning token events flow through # unmodified. if isinstance(adapter, OpenAIAdapter): if reasoning and reasoning.get("effort"): effort = validate_effort_or_llm_error( reasoning.get("effort"), "OpenAI Responses", OPENAI_EFFORTS ) if effort == "none": reasoning = {"effort": effort} return await adapter.responses_create( input=input, instructions=instructions, model=routed.model, tools=tools, reasoning=reasoning, stream=stream, connection_params=connection_params, timeout=timeout, **kwargs, ) messages = responses_input_to_chat_messages( input, instructions, ) extra: dict[str, Any] = dict(kwargs) # Translate Responses API ``text`` (structured output) to the # Chat Completions ``response_format`` parameter. The Responses # API shape is ``text={"format": {"type": "json_schema", ...}}``; # the Chat Completions equivalent is # ``response_format={"type": "json_schema", "json_schema": ...}``. # Without this, the ``text`` kwarg is sent as-is in the Chat # Completions body and rejected with 400 by providers that don't # recognise it (e.g. Databricks). text_param = extra.pop("text", None) if isinstance(text_param, dict): fmt = text_param.get("format") if isinstance(fmt, dict) and fmt.get("type") == "json_schema": extra["response_format"] = { "type": "json_schema", "json_schema": {k: v for k, v in fmt.items() if k != "type"}, } if reasoning: extra["reasoning_effort"] = reasoning.get("effort") if stream: chunks = await adapter.chat_completions( messages, routed.model, tools, True, extra, connection_params=connection_params, timeout=timeout, ) assert not isinstance(chunks, dict) return chat_stream_to_response_events( chunks, model=routed.model, ) result = await adapter.chat_completions( messages, routed.model, tools, False, extra, connection_params=connection_params, timeout=timeout, ) assert isinstance(result, dict) return chat_response_to_response(result) async def _execute_with_retry( call_fn: Callable[[], Awaitable[_T]], retry_config: RetryPolicy, ) -> _T: """ Execute ``call_fn`` with retry on transient failures. Standalone retry logic for the LLM client. Uses ``asyncio.sleep`` for backoff so the event loop stays free. :param call_fn: Zero-argument async callable that performs the LLM call. :param retry_config: Retry policy (max_attempts, backoff, etc.). :returns: The successful result from ``call_fn``. :raises PermanentLLMError: On non-retryable errors. :raises RetryableLLMError: When all retry attempts are exhausted. """ last_error: RetryableLLMError | None = None total_tries = retry_config.max_retries + 1 for attempt in range(total_tries): try: return await call_fn() except (PermanentLLMError, RetryableLLMError): raise except Exception as exc: classified = classify_llm_error( exc, retry_config.retryable_status_codes, ) if isinstance(classified, PermanentLLMError): raise classified from exc last_error = classified if attempt + 1 < total_tries: await _backoff_sleep(attempt, retry_config) assert last_error is not None raise last_error async def _backoff_sleep( attempt: int, config: RetryPolicy, ) -> None: """ Sleep with exponential backoff and jitter. Uses ``asyncio.sleep`` for non-blocking backoff. :param attempt: Zero-based attempt index (0 = first attempt). :param config: Retry policy with backoff parameters. """ delay = config.compute_backoff_delay(retry_index=attempt + 1) total_tries = config.max_retries + 1 _logger.info( "LLM retry %d/%d after %.1fs", attempt + 2, total_tries, delay, ) await _sleep(delay) async def _sleep(seconds: float) -> None: """ Indirection point for the LLM retry backoff sleep. Exists so tests can stub the retry delay without patching ``asyncio.sleep`` globally (patching ``omnigent.llms.client.asyncio.sleep`` walks the dotted path into the real ``asyncio`` module singleton and leaks the mock into every other test in the process). :param seconds: Delay in seconds. """ await asyncio.sleep(seconds) class Client: """ Multi-provider async LLM client. Provides ``await client.responses.create()`` matching the OpenAI SDK interface, routing to any supported provider based on the model string prefix. Usage:: client = Client() resp = await client.responses.create( input=[{"role": "user", "content": "Hello"}], instructions="You are helpful.", model="anthropic/claude-sonnet-4-20250514", ) """ def __init__(self) -> None: """Initialize the client with a responses namespace.""" self.responses = _ResponsesNamespace(self)