Edit on GitHub

agent_search_gateway.providers.openai_chat

OpenAI-compatible chat-completions LLM adapter.

  1"""OpenAI-compatible chat-completions LLM adapter."""
  2
  3import asyncio
  4import json
  5import logging
  6from collections.abc import Awaitable, Callable, Mapping, Sequence
  7
  8from ..concurrency import CapacityGate
  9from ..errors import ErrorCode, ExecutionFailure, ProtocolFailure
 10from ..models import LLMInvocation, RetryPolicy
 11from ..observability import SecretValue, log_event
 12from ..retry import retry_async
 13from .contracts import ChatMessage
 14from .http import HttpJsonExecutor
 15
 16_RESERVED_EXTRA_BODY_KEYS = frozenset({"model", "messages"})
 17
 18
 19class OpenAIChatCompletionsClient:
 20    def __init__(
 21        self,
 22        *,
 23        name: str,
 24        api_url: str,
 25        secret: SecretValue,
 26        executor: HttpJsonExecutor,
 27        quota: CapacityGate,
 28        retry_policy: RetryPolicy,
 29        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
 30        logger: logging.Logger | None = None,
 31    ) -> None:
 32        self.name = name
 33        self._endpoint = f"{api_url.rstrip('/')}/v1/chat/completions"
 34        self._secret = secret
 35        self._executor = executor
 36        self._quota = quota
 37        self._retry_policy = retry_policy
 38        self._sleep = sleep
 39        self._logger = logger or logging.getLogger(__name__)
 40
 41    async def complete_text(
 42        self,
 43        invocation: LLMInvocation,
 44        messages: Sequence[ChatMessage],
 45    ) -> str:
 46        self._validate_invocation(invocation)
 47
 48        async def operation() -> str:
 49            payload = await self._request_once(invocation, messages)
 50            return self._extract_content(payload)
 51
 52        async with self._quota.lease():
 53            result = await retry_async(
 54                self._retry_policy,
 55                operation,
 56                is_retryable=lambda exc: isinstance(exc, ProtocolFailure),
 57                sleep=self._sleep,
 58                before_attempt=lambda attempt: self._log_protocol_attempt(
 59                    invocation,
 60                    messages,
 61                    attempt,
 62                ),
 63                on_retry=lambda attempt, exc, delay: self._log_protocol_retry(
 64                    invocation,
 65                    messages,
 66                    attempt,
 67                    exc,
 68                    delay,
 69                ),
 70            )
 71        self._log_protocol_completed(invocation, messages, output_chars=len(result))
 72        return result
 73
 74    async def complete_json(
 75        self,
 76        invocation: LLMInvocation,
 77        messages: Sequence[ChatMessage],
 78    ) -> Mapping[str, object]:
 79        self._validate_invocation(invocation)
 80
 81        async def operation() -> Mapping[str, object]:
 82            payload = await self._request_once(invocation, messages)
 83            content = self._extract_content(payload)
 84            try:
 85                decoded = json.loads(content)
 86            except json.JSONDecodeError as exc:
 87                raise self._protocol_failure("content was not valid JSON") from exc
 88            if not isinstance(decoded, dict):
 89                raise self._protocol_failure("JSON content must be an object")
 90            return decoded
 91
 92        async with self._quota.lease():
 93            result = await retry_async(
 94                self._retry_policy,
 95                operation,
 96                is_retryable=lambda exc: isinstance(exc, ProtocolFailure),
 97                sleep=self._sleep,
 98                before_attempt=lambda attempt: self._log_protocol_attempt(
 99                    invocation,
100                    messages,
101                    attempt,
102                ),
103                on_retry=lambda attempt, exc, delay: self._log_protocol_retry(
104                    invocation,
105                    messages,
106                    attempt,
107                    exc,
108                    delay,
109                ),
110            )
111        output_chars = len(json.dumps(result, ensure_ascii=False, separators=(",", ":")))
112        self._log_protocol_completed(invocation, messages, output_chars=output_chars)
113        return result
114
115    def _log_protocol_attempt(
116        self,
117        invocation: LLMInvocation,
118        messages: Sequence[ChatMessage],
119        attempt: int,
120    ) -> None:
121        message_count, input_chars, extra_body_keys = self._protocol_metadata(invocation, messages)
122        log_event(
123            self._logger,
124            logging.DEBUG,
125            "provider_started",
126            provider=self.name,
127            stage="llm",
128            model=invocation.model,
129            message_count=message_count,
130            input_chars=input_chars,
131            extra_body_keys=extra_body_keys,
132            attempt=attempt,
133        )
134
135    def _log_protocol_retry(
136        self,
137        invocation: LLMInvocation,
138        messages: Sequence[ChatMessage],
139        attempt: int,
140        exc: BaseException,
141        delay: float,
142    ) -> None:
143        message_count, input_chars, extra_body_keys = self._protocol_metadata(invocation, messages)
144        log_event(
145            self._logger,
146            logging.DEBUG,
147            "provider_failed",
148            provider=self.name,
149            stage="llm",
150            model=invocation.model,
151            message_count=message_count,
152            input_chars=input_chars,
153            extra_body_keys=extra_body_keys,
154            attempt=attempt,
155            error_type=type(exc).__name__,
156            delay_ms=max(0, int(delay * 1000)),
157            reason="protocol_retry",
158        )
159
160    def _log_protocol_completed(
161        self,
162        invocation: LLMInvocation,
163        messages: Sequence[ChatMessage],
164        *,
165        output_chars: int,
166    ) -> None:
167        message_count, input_chars, extra_body_keys = self._protocol_metadata(invocation, messages)
168        log_event(
169            self._logger,
170            logging.DEBUG,
171            "provider_completed",
172            provider=self.name,
173            stage="llm",
174            model=invocation.model,
175            message_count=message_count,
176            input_chars=input_chars,
177            extra_body_keys=extra_body_keys,
178            output_chars=output_chars,
179        )
180
181    @staticmethod
182    def _protocol_metadata(
183        invocation: LLMInvocation,
184        messages: Sequence[ChatMessage],
185    ) -> tuple[int, int, str]:
186        return (
187            len(messages),
188            sum(len(value) for message in messages for value in message.values()),
189            ",".join(sorted(invocation.extra_body)) or "-",
190        )
191
192    async def aclose(self) -> None:
193        await self._executor.aclose()
194
195    def _validate_invocation(self, invocation: LLMInvocation) -> None:
196        if invocation.provider != self.name:
197            raise ExecutionFailure(
198                ErrorCode.LLM_STAGE_FAILED,
199                f"LLM invocation provider {invocation.provider} does not match {self.name}",
200            )
201        collisions = _RESERVED_EXTRA_BODY_KEYS.intersection(invocation.extra_body)
202        if collisions:
203            raise ExecutionFailure(
204                ErrorCode.LLM_STAGE_FAILED,
205                f"LLM extra_body cannot override: {', '.join(sorted(collisions))}",
206            )
207
208    async def _request_once(
209        self,
210        invocation: LLMInvocation,
211        messages: Sequence[ChatMessage],
212    ) -> object:
213        body: dict[str, object] = {
214            "model": invocation.model,
215            "messages": [dict(message) for message in messages],
216        }
217        body.update(invocation.extra_body)
218        return await self._executor.request_json(
219            "POST",
220            self._endpoint,
221            stage="llm",
222            headers={
223                "Authorization": f"Bearer {self._secret.reveal()}",
224                "Content-Type": "application/json",
225            },
226            json_body=body,
227        )
228
229    def _extract_content(self, payload: object) -> str:
230        if not isinstance(payload, dict):
231            raise self._protocol_failure("response must be an object")
232        choices = payload.get("choices")
233        if not isinstance(choices, list) or not choices:
234            raise self._protocol_failure("response choices are missing")
235        first = choices[0]
236        if not isinstance(first, dict):
237            raise self._protocol_failure("response choice is invalid")
238        message = first.get("message")
239        if not isinstance(message, dict):
240            raise self._protocol_failure("response message is missing")
241        content = message.get("content")
242        if not isinstance(content, str) or not content.strip():
243            raise self._protocol_failure("response content is empty")
244        return content
245
246    def _protocol_failure(self, reason: str) -> ProtocolFailure:
247        return ProtocolFailure(ErrorCode.PROTOCOL_ERROR, f"{self.name}/llm: {reason}")
class OpenAIChatCompletionsClient: 
 20class OpenAIChatCompletionsClient:
 21    def __init__(
 22        self,
 23        *,
 24        name: str,
 25        api_url: str,
 26        secret: SecretValue,
 27        executor: HttpJsonExecutor,
 28        quota: CapacityGate,
 29        retry_policy: RetryPolicy,
 30        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
 31        logger: logging.Logger | None = None,
 32    ) -> None:
 33        self.name = name
 34        self._endpoint = f"{api_url.rstrip('/')}/v1/chat/completions"
 35        self._secret = secret
 36        self._executor = executor
 37        self._quota = quota
 38        self._retry_policy = retry_policy
 39        self._sleep = sleep
 40        self._logger = logger or logging.getLogger(__name__)
 41
 42    async def complete_text(
 43        self,
 44        invocation: LLMInvocation,
 45        messages: Sequence[ChatMessage],
 46    ) -> str:
 47        self._validate_invocation(invocation)
 48
 49        async def operation() -> str:
 50            payload = await self._request_once(invocation, messages)
 51            return self._extract_content(payload)
 52
 53        async with self._quota.lease():
 54            result = await retry_async(
 55                self._retry_policy,
 56                operation,
 57                is_retryable=lambda exc: isinstance(exc, ProtocolFailure),
 58                sleep=self._sleep,
 59                before_attempt=lambda attempt: self._log_protocol_attempt(
 60                    invocation,
 61                    messages,
 62                    attempt,
 63                ),
 64                on_retry=lambda attempt, exc, delay: self._log_protocol_retry(
 65                    invocation,
 66                    messages,
 67                    attempt,
 68                    exc,
 69                    delay,
 70                ),
 71            )
 72        self._log_protocol_completed(invocation, messages, output_chars=len(result))
 73        return result
 74
 75    async def complete_json(
 76        self,
 77        invocation: LLMInvocation,
 78        messages: Sequence[ChatMessage],
 79    ) -> Mapping[str, object]:
 80        self._validate_invocation(invocation)
 81
 82        async def operation() -> Mapping[str, object]:
 83            payload = await self._request_once(invocation, messages)
 84            content = self._extract_content(payload)
 85            try:
 86                decoded = json.loads(content)
 87            except json.JSONDecodeError as exc:
 88                raise self._protocol_failure("content was not valid JSON") from exc
 89            if not isinstance(decoded, dict):
 90                raise self._protocol_failure("JSON content must be an object")
 91            return decoded
 92
 93        async with self._quota.lease():
 94            result = await retry_async(
 95                self._retry_policy,
 96                operation,
 97                is_retryable=lambda exc: isinstance(exc, ProtocolFailure),
 98                sleep=self._sleep,
 99                before_attempt=lambda attempt: self._log_protocol_attempt(
100                    invocation,
101                    messages,
102                    attempt,
103                ),
104                on_retry=lambda attempt, exc, delay: self._log_protocol_retry(
105                    invocation,
106                    messages,
107                    attempt,
108                    exc,
109                    delay,
110                ),
111            )
112        output_chars = len(json.dumps(result, ensure_ascii=False, separators=(",", ":")))
113        self._log_protocol_completed(invocation, messages, output_chars=output_chars)
114        return result
115
116    def _log_protocol_attempt(
117        self,
118        invocation: LLMInvocation,
119        messages: Sequence[ChatMessage],
120        attempt: int,
121    ) -> None:
122        message_count, input_chars, extra_body_keys = self._protocol_metadata(invocation, messages)
123        log_event(
124            self._logger,
125            logging.DEBUG,
126            "provider_started",
127            provider=self.name,
128            stage="llm",
129            model=invocation.model,
130            message_count=message_count,
131            input_chars=input_chars,
132            extra_body_keys=extra_body_keys,
133            attempt=attempt,
134        )
135
136    def _log_protocol_retry(
137        self,
138        invocation: LLMInvocation,
139        messages: Sequence[ChatMessage],
140        attempt: int,
141        exc: BaseException,
142        delay: float,
143    ) -> None:
144        message_count, input_chars, extra_body_keys = self._protocol_metadata(invocation, messages)
145        log_event(
146            self._logger,
147            logging.DEBUG,
148            "provider_failed",
149            provider=self.name,
150            stage="llm",
151            model=invocation.model,
152            message_count=message_count,
153            input_chars=input_chars,
154            extra_body_keys=extra_body_keys,
155            attempt=attempt,
156            error_type=type(exc).__name__,
157            delay_ms=max(0, int(delay * 1000)),
158            reason="protocol_retry",
159        )
160
161    def _log_protocol_completed(
162        self,
163        invocation: LLMInvocation,
164        messages: Sequence[ChatMessage],
165        *,
166        output_chars: int,
167    ) -> None:
168        message_count, input_chars, extra_body_keys = self._protocol_metadata(invocation, messages)
169        log_event(
170            self._logger,
171            logging.DEBUG,
172            "provider_completed",
173            provider=self.name,
174            stage="llm",
175            model=invocation.model,
176            message_count=message_count,
177            input_chars=input_chars,
178            extra_body_keys=extra_body_keys,
179            output_chars=output_chars,
180        )
181
182    @staticmethod
183    def _protocol_metadata(
184        invocation: LLMInvocation,
185        messages: Sequence[ChatMessage],
186    ) -> tuple[int, int, str]:
187        return (
188            len(messages),
189            sum(len(value) for message in messages for value in message.values()),
190            ",".join(sorted(invocation.extra_body)) or "-",
191        )
192
193    async def aclose(self) -> None:
194        await self._executor.aclose()
195
196    def _validate_invocation(self, invocation: LLMInvocation) -> None:
197        if invocation.provider != self.name:
198            raise ExecutionFailure(
199                ErrorCode.LLM_STAGE_FAILED,
200                f"LLM invocation provider {invocation.provider} does not match {self.name}",
201            )
202        collisions = _RESERVED_EXTRA_BODY_KEYS.intersection(invocation.extra_body)
203        if collisions:
204            raise ExecutionFailure(
205                ErrorCode.LLM_STAGE_FAILED,
206                f"LLM extra_body cannot override: {', '.join(sorted(collisions))}",
207            )
208
209    async def _request_once(
210        self,
211        invocation: LLMInvocation,
212        messages: Sequence[ChatMessage],
213    ) -> object:
214        body: dict[str, object] = {
215            "model": invocation.model,
216            "messages": [dict(message) for message in messages],
217        }
218        body.update(invocation.extra_body)
219        return await self._executor.request_json(
220            "POST",
221            self._endpoint,
222            stage="llm",
223            headers={
224                "Authorization": f"Bearer {self._secret.reveal()}",
225                "Content-Type": "application/json",
226            },
227            json_body=body,
228        )
229
230    def _extract_content(self, payload: object) -> str:
231        if not isinstance(payload, dict):
232            raise self._protocol_failure("response must be an object")
233        choices = payload.get("choices")
234        if not isinstance(choices, list) or not choices:
235            raise self._protocol_failure("response choices are missing")
236        first = choices[0]
237        if not isinstance(first, dict):
238            raise self._protocol_failure("response choice is invalid")
239        message = first.get("message")
240        if not isinstance(message, dict):
241            raise self._protocol_failure("response message is missing")
242        content = message.get("content")
243        if not isinstance(content, str) or not content.strip():
244            raise self._protocol_failure("response content is empty")
245        return content
246
247    def _protocol_failure(self, reason: str) -> ProtocolFailure:
248        return ProtocolFailure(ErrorCode.PROTOCOL_ERROR, f"{self.name}/llm: {reason}")
OpenAIChatCompletionsClient( *, name: str, api_url: str, secret: agent_search_gateway.observability.SecretValue, executor: agent_search_gateway.providers.http.HttpJsonExecutor, quota: agent_search_gateway.concurrency.CapacityGate, retry_policy: agent_search_gateway.models.RetryPolicy, sleep: Callable[[float], Awaitable[None]] = <function sleep>, logger: logging.Logger | None = None) 
21    def __init__(
22        self,
23        *,
24        name: str,
25        api_url: str,
26        secret: SecretValue,
27        executor: HttpJsonExecutor,
28        quota: CapacityGate,
29        retry_policy: RetryPolicy,
30        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
31        logger: logging.Logger | None = None,
32    ) -> None:
33        self.name = name
34        self._endpoint = f"{api_url.rstrip('/')}/v1/chat/completions"
35        self._secret = secret
36        self._executor = executor
37        self._quota = quota
38        self._retry_policy = retry_policy
39        self._sleep = sleep
40        self._logger = logger or logging.getLogger(__name__)
name
async def complete_text( self, invocation: agent_search_gateway.models.LLMInvocation, messages: Sequence[Mapping[str, str]]) -> str: 
42    async def complete_text(
43        self,
44        invocation: LLMInvocation,
45        messages: Sequence[ChatMessage],
46    ) -> str:
47        self._validate_invocation(invocation)
48
49        async def operation() -> str:
50            payload = await self._request_once(invocation, messages)
51            return self._extract_content(payload)
52
53        async with self._quota.lease():
54            result = await retry_async(
55                self._retry_policy,
56                operation,
57                is_retryable=lambda exc: isinstance(exc, ProtocolFailure),
58                sleep=self._sleep,
59                before_attempt=lambda attempt: self._log_protocol_attempt(
60                    invocation,
61                    messages,
62                    attempt,
63                ),
64                on_retry=lambda attempt, exc, delay: self._log_protocol_retry(
65                    invocation,
66                    messages,
67                    attempt,
68                    exc,
69                    delay,
70                ),
71            )
72        self._log_protocol_completed(invocation, messages, output_chars=len(result))
73        return result
async def complete_json( self, invocation: agent_search_gateway.models.LLMInvocation, messages: Sequence[Mapping[str, str]]) -> Mapping[str, object]: 
 75    async def complete_json(
 76        self,
 77        invocation: LLMInvocation,
 78        messages: Sequence[ChatMessage],
 79    ) -> Mapping[str, object]:
 80        self._validate_invocation(invocation)
 81
 82        async def operation() -> Mapping[str, object]:
 83            payload = await self._request_once(invocation, messages)
 84            content = self._extract_content(payload)
 85            try:
 86                decoded = json.loads(content)
 87            except json.JSONDecodeError as exc:
 88                raise self._protocol_failure("content was not valid JSON") from exc
 89            if not isinstance(decoded, dict):
 90                raise self._protocol_failure("JSON content must be an object")
 91            return decoded
 92
 93        async with self._quota.lease():
 94            result = await retry_async(
 95                self._retry_policy,
 96                operation,
 97                is_retryable=lambda exc: isinstance(exc, ProtocolFailure),
 98                sleep=self._sleep,
 99                before_attempt=lambda attempt: self._log_protocol_attempt(
100                    invocation,
101                    messages,
102                    attempt,
103                ),
104                on_retry=lambda attempt, exc, delay: self._log_protocol_retry(
105                    invocation,
106                    messages,
107                    attempt,
108                    exc,
109                    delay,
110                ),
111            )
112        output_chars = len(json.dumps(result, ensure_ascii=False, separators=(",", ":")))
113        self._log_protocol_completed(invocation, messages, output_chars=output_chars)
114        return result
async def aclose(self) -> None: 
193    async def aclose(self) -> None:
194        await self._executor.aclose()