diff --git a/docs/en_US/ai_tools.rst b/docs/en_US/ai_tools.rst index 68416bb46ba..5363dc547f2 100644 --- a/docs/en_US/ai_tools.rst +++ b/docs/en_US/ai_tools.rst @@ -56,7 +56,8 @@ Select your preferred LLM provider from the dropdown: **OpenAI** Use GPT models from OpenAI, or any OpenAI-compatible API provider (e.g., - LiteLLM, LM Studio, EXO, or other local inference servers). + LiteLLM, LM Studio, EXO, hosted services such as GLM (Z.ai), or local + inference servers). * **API URL**: Custom API endpoint URL (leave empty for default: https://api.openai.com/v1). Include the ``/v1`` path prefix if required by your provider. * **API Key File**: Path to a file containing your OpenAI API key (obtain from https://platform.openai.com/). This path refers to the filesystem where the pgAdmin server is running (e.g., inside the container if using Docker). The ``~`` prefix is expanded to the home directory of the user running the pgAdmin server process. Optional when using a custom URL with a provider that does not require authentication. @@ -77,7 +78,10 @@ Select your preferred LLM provider from the dropdown: .. note:: You can also use the *OpenAI* provider with a custom API URL for any OpenAI-compatible endpoint, including Docker Model Runner and other local - inference servers. + inference servers. pgAdmin sends ``max_tokens`` by default and learns to use + ``max_completion_tokens`` when a server explicitly rejects it. When the + ``ALLOWED_LLM_API_URLS`` configuration setting is non-empty, custom API URLs must match one of its entries; an + empty list permits any URL. After configuring your provider, click *Save* to apply the changes. diff --git a/docs/en_US/preferences.rst b/docs/en_US/preferences.rst index 57a68692b63..29d0de8c2a3 100644 --- a/docs/en_US/preferences.rst +++ b/docs/en_US/preferences.rst @@ -66,7 +66,8 @@ Use the fields on the *AI* panel to configure your LLM provider: * Use the *API URL* field to set a custom API endpoint URL. Leave empty to use the default OpenAI API (``https://api.openai.com/v1``). Set a custom URL to - use any OpenAI-compatible API provider (e.g., LiteLLM, LM Studio, EXO). + use any OpenAI-compatible API provider (e.g., LiteLLM, LM Studio, EXO, + GLM (Z.ai)). Include the ``/v1`` path prefix if required by your provider (e.g., ``http://localhost:1234/v1``). @@ -101,8 +102,14 @@ Use the fields on the *AI* panel to configure your LLM provider: .. note:: You can also use the *OpenAI* provider with a custom API URL for any OpenAI-compatible endpoint, including Docker Model Runner, LM Studio, EXO, - and other local inference servers. This can be useful when you want to use - a provider that isn't explicitly listed but supports the OpenAI API format. + hosted services such as GLM (Z.ai), and other local inference servers. This + can be useful when you want to use a provider that isn't explicitly listed + but supports the OpenAI API format. + +.. note:: With the default allowlist enabled, an administrator must add any + API host not already in ``ALLOWED_LLM_API_URLS`` before pgAdmin can connect + to it. By default the list allows the OpenAI and Anthropic APIs and loopback + addresses on any port. An empty allowlist disables this restriction. The Browser Node **************** diff --git a/web/pgadmin/llm/models.py b/web/pgadmin/llm/models.py index ba20245cb76..b9bceddbecd 100644 --- a/web/pgadmin/llm/models.py +++ b/web/pgadmin/llm/models.py @@ -194,6 +194,8 @@ class LLMError: code: Optional[str] = None provider: Optional[str] = None retryable: bool = False + http_status: Optional[int] = None + param: Optional[str] = None def __str__(self) -> str: if self.code: diff --git a/web/pgadmin/llm/providers/openai.py b/web/pgadmin/llm/providers/openai.py index 075652e3224..e272ee36ef3 100644 --- a/web/pgadmin/llm/providers/openai.py +++ b/web/pgadmin/llm/providers/openai.py @@ -10,10 +10,12 @@ """OpenAI GPT LLM client implementation.""" import json +import logging import socket import ssl import urllib.request import urllib.error +from urllib.parse import urlparse from collections.abc import Generator from typing import Optional, Union import uuid @@ -39,6 +41,11 @@ # Default model if none specified DEFAULT_MODEL = 'gpt-4o' +# Learned server/model capabilities survive per-request client instances. +_adaptive_cache = {} +_logger = logging.getLogger(__name__) +_MAX_SENDS = 3 + # Default API base URL DEFAULT_API_BASE_URL = 'https://api.openai.com/v1' @@ -76,7 +83,6 @@ def __init__(self, api_key: Optional[str] = None, base_url = base_url[:-len(suffix)].rstrip('/') break self._base_url = base_url - self._use_responses_api = False @property def _api_url(self) -> str: @@ -85,6 +91,14 @@ def _api_url(self) -> str: return f'{self._base_url}/responses' return f'{self._base_url}/chat/completions' + @property + def _is_default_host(self) -> bool: + """True when talking to the official OpenAI API host.""" + # Compare hostnames, not authorities: an explicit default port + # (https://api.openai.com:443/v1) is still the official endpoint. + return (urlparse(self._base_url).hostname == + urlparse(DEFAULT_API_BASE_URL).hostname) + @property def provider_name(self) -> str: return 'openai' @@ -97,9 +111,7 @@ def is_available(self) -> bool: """Check if the client is properly configured.""" # API key is required for the default OpenAI endpoint, but optional # for custom endpoints (e.g., local LLM servers). - if self._base_url.rstrip('/').startswith( - DEFAULT_API_BASE_URL.rstrip('/') - ): + if self._is_default_host: return bool(self._api_key) return True @@ -127,25 +139,67 @@ def chat( Raises: LLMClientError: If the request fails. """ - if self._use_responses_api: - return self._chat_responses( - messages, tools, system_prompt, max_tokens - ) + for attempt in range(_MAX_SENDS): + budget = max_tokens * (3 if self._capabilities.get('reasoning') + else 1) + try: + if self._use_responses_api: + return self._chat_responses( + messages, tools, system_prompt, budget) + return self._chat_completions( + messages, tools, system_prompt, budget) + except LLMClientError as error: + if not self._adjust_after_error(error, attempt): + raise - # Try Chat Completions API first - payload = self._build_chat_payload( - messages, tools, system_prompt, max_tokens - ) + @property + def _capabilities(self) -> dict: + """Return the learned capabilities for this server and model.""" + return _adaptive_cache.setdefault((self._base_url, self._model), {}) + @property + def _use_responses_api(self) -> bool: + """Keep routing decisions across per-request client instances.""" + return self._capabilities.get('responses', False) + + def _remember(self, capability: str, value=True): + """Log each new server/model adjustment once.""" + if self._capabilities.get(capability) != value: + self._capabilities[capability] = value + # Do not log API keys, URLs, prompts or provider error bodies. + _logger.debug('OpenAI provider learned %s adjustment', capability) + + def _adjust_after_error(self, error: LLMClientError, + attempt: int) -> bool: + """Apply each learned adjustment once, with at most three sends.""" + adjusted = False + if (self._is_unsupported_token_param_error(error) and + not self._capabilities.get('completion_tokens')): + self._remember('completion_tokens') + adjusted = True + elif (self._should_use_responses_api(error) and + not self._use_responses_api): + self._remember('responses') + adjusted = True + elif (error.error.code == 'reasoning_budget' and + not self._capabilities.get('reasoning')): + self._remember('reasoning') + adjusted = True + return adjusted and attempt + 1 < _MAX_SENDS + + def _chat_completions( + self, + messages: list[Message], + tools: Optional[list[Tool]], + system_prompt: Optional[str], + max_tokens: int + ) -> LLMResponse: + """Send one Chat Completions request; the caller bounds retries.""" + payload = self._build_chat_payload( + messages, tools, system_prompt, max_tokens) try: - response_data = self._make_request(payload) - return self._parse_response(response_data) - except LLMClientError as e: - if self._should_use_responses_api(e): - self._use_responses_api = True - return self._chat_responses( - messages, tools, system_prompt, max_tokens - ) + return self._parse_response(self._make_request(payload)) + except LLMClientError: raise except Exception as e: raise LLMClientError(LLMError( @@ -178,10 +232,20 @@ def _chat_responses( def _should_use_responses_api(self, error: LLMClientError) -> bool: """Check if the error indicates we should use the Responses API.""" - error_msg = str(error).lower() - return ('v1/responses' in error_msg or - 'not supported in the v1/chat/completions' in error_msg or - 'not a chat model' in error_msg) + detail = error.error + return (getattr(detail, 'http_status', None) in (400, 404) and + getattr(detail, 'param', None) in (None, '', 'model') and + 'only supported in' in detail.message.lower() and + 'v1/responses' in detail.message.lower()) + + def _is_unsupported_token_param_error( + self, error: LLMClientError + ) -> bool: + """Require the server's structured rejection of max_tokens.""" + detail = error.error + return (getattr(detail, 'http_status', None) == 400 and + getattr(detail, 'param', None) == 'max_tokens' and + detail.code == 'unsupported_parameter') def _build_chat_payload( self, @@ -199,10 +263,13 @@ def _build_chat_payload( 'content': system_prompt }) + token_param = ('max_completion_tokens' + if self._capabilities.get('completion_tokens') + else 'max_tokens') payload = { 'model': self._model, 'messages': converted_messages, - 'max_completion_tokens': max_tokens, + token_param: max_tokens, } if tools: @@ -357,6 +424,45 @@ def _convert_tools_responses(self, tools: list[Tool]) -> list[dict]: for tool in tools ] + def _http_error_to_llm_error( + self, e: urllib.error.HTTPError + ) -> LLMError: + """Convert an HTTP error response to an LLMError. + + Preserves the structured error details (code, parameter) so + callers can react to them - e.g. the unsupported-parameter + fallback - even when the message itself is generic. + """ + error_body = e.read().decode('utf-8') + error_msg = None + error_code = None + error_param = None + try: + error_data = json.loads(error_body) + except json.JSONDecodeError: + error_data = None + # Compatible servers do not all use OpenAI's error shape; only + # read the structured fields when they are actually mappings. + if isinstance(error_data, dict): + error_info = error_data.get('error') + if isinstance(error_info, dict): + error_msg = error_info.get('message') + error_param = error_info.get('param') + error_code = error_info.get('code') + elif isinstance(error_info, str): + error_msg = error_info + if not error_msg: + error_msg = error_body or str(e) + + return LLMError( + message=error_msg, + code=error_code or str(e.code), + provider=self.provider_name, + retryable=e.code in (429, 500, 502, 503, 504), + http_status=e.code, + param=error_param + ) + def _make_request(self, payload: dict) -> dict: """Make an HTTP request to the OpenAI API.""" headers = { @@ -379,19 +485,7 @@ def _make_request(self, payload: dict) -> dict: ) as response: return json.loads(response.read().decode('utf-8')) except urllib.error.HTTPError as e: - error_body = e.read().decode('utf-8') - try: - error_data = json.loads(error_body) - error_msg = error_data.get('error', {}).get('message', str(e)) - except json.JSONDecodeError: - error_msg = error_body or str(e) - - raise LLMClientError(LLMError( - message=error_msg, - code=str(e.code), - provider=self.provider_name, - retryable=e.code in (429, 500, 502, 503, 504) - )) + raise LLMClientError(self._http_error_to_llm_error(e)) except urllib.error.URLError as e: raise LLMClientError(LLMError( message=f"Connection error: {e.reason}", @@ -423,6 +517,16 @@ def _raise_max_tokens_error(self, input_tokens: int): retryable=False )) + def _raise_reasoning_budget_error(self): + """Explain empty output caused by a model's reasoning budget.""" + raise LLMClientError(LLMError( + message=f'Reasoning model {self._model} exhausted its output ' + f'budget on internal reasoning and returned no content.', + code='reasoning_budget', + provider=self.provider_name, + retryable=True + )) + def _parse_response(self, data: dict) -> LLMResponse: """Parse the Chat Completions API response into an LLMResponse.""" # Check for API-level errors in the response @@ -490,9 +594,20 @@ def _parse_response(self, data: dict) -> LLMResponse: total_tokens=usage_data.get('total_tokens', 0) ) + details = usage_data.get('completion_tokens_details') or {} + reasoning = (bool(message.get('reasoning_content')) or + (details.get('reasoning_tokens') or 0) > 0) # Check for problematic responses if not content and not tool_calls: if stop_reason == StopReason.MAX_TOKENS: + details = usage_data.get('completion_tokens_details') or {} + reasoning_tokens = details.get('reasoning_tokens') or 0 + if (message.get('reasoning_content') or + reasoning_tokens > 0): + # A reasoning model exhausted its output budget on + # internal reasoning before producing any visible + # content. + self._raise_reasoning_budget_error() self._raise_max_tokens_error(usage.input_tokens) elif finish_reason and finish_reason not in ('stop', 'tool_calls'): raise LLMClientError(LLMError( @@ -503,6 +618,9 @@ def _parse_response(self, data: dict) -> LLMResponse: retryable=False )) + if reasoning: + self._remember('reasoning') + return LLMResponse( content=content, tool_calls=tool_calls, @@ -577,9 +695,14 @@ def _parse_responses_response(self, data: dict) -> LLMResponse: total_tokens=usage_data.get('total_tokens', 0) ) + details = usage_data.get('output_tokens_details') or {} + reasoning = (details.get('reasoning_tokens') or 0) > 0 + # Check for problematic responses if not content and not tool_calls: if stop_reason == StopReason.MAX_TOKENS: + if reasoning: + self._raise_reasoning_budget_error() self._raise_max_tokens_error(usage.input_tokens) elif stop_reason == StopReason.STOP_SEQUENCE: raise LLMClientError(LLMError( @@ -589,6 +712,9 @@ def _parse_responses_response(self, data: dict) -> LLMResponse: retryable=False )) + if reasoning: + self._remember('reasoning') + return LLMResponse( content=content, tool_calls=tool_calls, @@ -608,46 +734,32 @@ def chat_stream( **kwargs ) -> Generator[Union[str, LLMResponse], None, None]: """Stream a chat response from OpenAI.""" - if self._use_responses_api: - payload = self._build_responses_payload( - messages, tools, system_prompt, max_tokens - ) + for attempt in range(_MAX_SENDS): + budget = max_tokens * (3 if self._capabilities.get('reasoning') + else 1) + if self._use_responses_api: + payload = self._build_responses_payload( + messages, tools, system_prompt, budget) + else: + payload = self._build_chat_payload( + messages, tools, system_prompt, budget) + payload['stream_options'] = {'include_usage': True} payload['stream'] = True + yielded = False try: - yield from self._process_stream(payload) - except LLMClientError: - raise + for item in self._process_stream(payload): + yielded = True + yield item + return + except LLMClientError as error: + # Never replay a stream once visible output has escaped. + if yielded or not self._adjust_after_error(error, attempt): + raise except Exception as e: raise LLMClientError(LLMError( message=f"Streaming request failed: {e!s}", provider=self.provider_name )) from e - return - - # Try Chat Completions API first - payload = self._build_chat_payload( - messages, tools, system_prompt, max_tokens - ) - payload['stream'] = True - payload['stream_options'] = {'include_usage': True} - - try: - yield from self._process_stream(payload) - except LLMClientError as e: - if self._should_use_responses_api(e): - self._use_responses_api = True - payload = self._build_responses_payload( - messages, tools, system_prompt, max_tokens - ) - payload['stream'] = True - yield from self._process_stream(payload) - else: - raise - except Exception as e: - raise LLMClientError(LLMError( - message=f"Streaming request failed: {e!s}", - provider=self.provider_name - )) from e def _process_stream( self, payload: dict @@ -672,20 +784,7 @@ def _process_stream( request, timeout=120, context=SSL_CONTEXT ) except urllib.error.HTTPError as e: - error_body = e.read().decode('utf-8') - try: - error_data = json.loads(error_body) - error_msg = error_data.get( - 'error', {} - ).get('message', str(e)) - except json.JSONDecodeError: - error_msg = error_body or str(e) - raise LLMClientError(LLMError( - message=error_msg, - code=str(e.code), - provider=self.provider_name, - retryable=e.code in (429, 500, 502, 503, 504) - )) + raise LLMClientError(self._http_error_to_llm_error(e)) except urllib.error.URLError as e: raise LLMClientError(LLMError( message=f"Connection error: {e.reason}", @@ -723,6 +822,8 @@ def _read_openai_stream( finish_reason = None model_name = self._model usage = Usage() + reasoning_seen = False + reasoning_tokens = 0 while True: line_bytes = response.readline() @@ -753,6 +854,8 @@ def _read_openai_stream( output_tokens=u.get('completion_tokens', 0), total_tokens=u.get('total_tokens', 0) ) + details = u.get('completion_tokens_details') or {} + reasoning_tokens = details.get('reasoning_tokens') or 0 if 'model' in data: model_name = data['model'] @@ -767,6 +870,11 @@ def _read_openai_stream( if choice.get('finish_reason'): finish_reason = choice['finish_reason'] + # Reasoning content is tracked (for budget-exhaustion + # detection) but not yielded as visible content. + if delta.get('reasoning_content'): + reasoning_seen = True + # Text content text_chunk = delta.get('content') if text_chunk: @@ -820,12 +928,21 @@ def _read_openai_stream( ) if not content and not tool_calls: + if finish_reason == 'length' and ( + reasoning_seen or reasoning_tokens > 0 + ): + # A reasoning model exhausted its output budget on + # internal reasoning before producing any content. + self._raise_reasoning_budget_error() raise LLMClientError(LLMError( message='No response content returned from API', provider=self.provider_name, retryable=False )) + if reasoning_seen or reasoning_tokens > 0: + self._remember('reasoning') + yield LLMResponse( content=content, tool_calls=tool_calls, @@ -881,6 +998,7 @@ def tool_call_entry(output_index, item_id): usage = Usage() resp_status = '' resp_incomplete = {} + reasoning_seen = False while True: line_bytes = response.readline() @@ -938,7 +1056,11 @@ def tool_call_entry(output_index, item_id): data.get('output_index'), data.get('item_id')) tc['arguments'] += data.get('delta', '') - elif event_type == 'response.completed': + elif event_type in ('response.reasoning_text.delta', + 'response.reasoning_summary_text.delta'): + reasoning_seen = True + + elif event_type in ('response.completed', 'response.incomplete'): resp = data.get('response', {}) u = resp.get('usage', {}) usage = Usage( @@ -946,6 +1068,9 @@ def tool_call_entry(output_index, item_id): output_tokens=u.get('output_tokens', 0), total_tokens=u.get('total_tokens', 0) ) + details = u.get('output_tokens_details') or {} + reasoning_seen = reasoning_seen or ( + details.get('reasoning_tokens') or 0) > 0 model_name = resp.get('model', model_name) resp_status = resp.get('status', '') resp_incomplete = resp.get('incomplete_details', {}) @@ -981,12 +1106,17 @@ def tool_call_entry(output_index, item_id): stop_reason = StopReason.UNKNOWN if not content and not tool_calls: + if stop_reason == StopReason.MAX_TOKENS and reasoning_seen: + self._raise_reasoning_budget_error() raise LLMClientError(LLMError( message='No response content returned from API', provider=self.provider_name, retryable=False )) + if reasoning_seen: + self._remember('reasoning') + yield LLMResponse( content=content, tool_calls=tool_calls, diff --git a/web/pgadmin/llm/tests/test_openai_stream.py b/web/pgadmin/llm/tests/test_openai_stream.py index 59ffd27480d..c2249ce8cc9 100644 --- a/web/pgadmin/llm/tests/test_openai_stream.py +++ b/web/pgadmin/llm/tests/test_openai_stream.py @@ -278,3 +278,344 @@ def runTest(self): actual = [(tc.id, tc.name, tc.arguments) for tc in result.tool_calls] self.assertEqual(actual, self.expected) + + +class OpenAICompatibleServerTestCase(BaseTestGenerator): + """Check adaptive requests using scripted transport, without API calls.""" + + scenarios = [(name, dict(check=name)) for name in ( + 'compatible_default', 'token_fallback', 'token_cache', + 'unrelated_error', 'structured_only', 'responses_cache', + 'retry_chain', 'retry_cap', 'reasoning_retry', 'reasoning_cache', + 'reasoning_success', 'reasoning_null', 'reasoning_exhausted', + 'partial_stream', 'malformed_errors', 'hostnames', + 'reasoning_content', 'responses_rejection', 'three_send_cap', + 'http_fallback', 'responses_reasoning', 'debug_once', + 'responses_model_param')] + + def runTest(self): + """RunTest.""" + from pgadmin.llm.providers import openai as module + module._adaptive_cache.clear() + try: + getattr(self, '_check_' + self.check)(False) + module._adaptive_cache.clear() + if self.check not in ('partial_stream', 'malformed_errors', + 'hostnames'): + getattr(self, '_check_' + self.check)(True) + finally: + module._adaptive_cache.clear() + + def _client(self, url='https://gateway.example/v1', model='test-model'): + """Client.""" + return OpenAIClient(api_key='test-key', api_url=url, model=model) + + def _error(self, param='max_tokens', code='unsupported_parameter', + status=400, message='Unsupported parameter'): + """Error.""" + import io + import urllib.error + from pgadmin.llm.client import LLMClientError + body = json.dumps({'error': dict(param=param, code=code, + message=message)}).encode() + return LLMClientError(self._client()._http_error_to_llm_error( + urllib.error.HTTPError('https://gateway.example', status, + 'error', {}, io.BytesIO(body)))) + + def _ok(self, reasoning=None, content='hello', finish='stop'): + """Ok.""" + return dict(choices=[dict(message=dict(content=content), + finish_reason=finish)], + usage=dict(completion_tokens_details=dict( + reasoning_tokens=reasoning))) + + def _run(self, client, script, streaming, budget=500): + """Run.""" + from pgadmin.llm.models import Message + calls = [] + + def request(payload): + """Request.""" + calls.append((client._api_url, payload.copy())) + item = script[len(calls) - 1] + if isinstance(item, Exception): + raise item + return item + + def stream(payload): + """Stream.""" + item = request(payload) + if client._use_responses_api: + events = [] + for output in item.get('output', []): + for part in output.get('content', []): + events.append(_sse(dict( + type='response.output_text.delta', + delta=part['text']))) + events.append(_sse(dict(type='response.completed', + response=item))) + yield from client._read_responses_stream(_FakeStream(events)) + else: + choice = item['choices'][0] + yield from client._read_openai_stream(_FakeStream([ + _sse(dict(choices=[dict( + delta=choice['message'], + finish_reason=choice['finish_reason'])], + usage=item.get('usage', {})))])) + + client._make_request = request + client._process_stream = stream + self.calls = calls + messages = [Message.user('hi')] + if streaming: + return list(client.chat_stream(messages, max_tokens=budget)) + return client.chat(messages, max_tokens=budget) + + def _budgets(self): + """Budgets.""" + return [next(p[k] for k in ('max_tokens', 'max_completion_tokens', + 'max_output_tokens') if k in p) + for _, p in self.calls] + + def _check_compatible_default(self, streaming): + """Check compatible default.""" + self._run(self._client(), [self._ok()], streaming) + self.assertIn('max_tokens', self.calls[0][1]) + + def _check_token_fallback(self, streaming): + """Check token fallback.""" + self._run(self._client(), [self._error(), self._ok()], streaming) + self.assertIn('max_tokens', self.calls[0][1]) + self.assertIn('max_completion_tokens', self.calls[1][1]) + + def _check_token_cache(self, streaming): + """Check token cache.""" + self._check_token_fallback(streaming) + self._run(self._client(), [self._ok()], streaming) + self.assertIn('max_completion_tokens', self.calls[0][1]) + self._run(self._client(model='other'), [self._ok()], streaming) + self.assertIn('max_tokens', self.calls[0][1]) + self._run(self._client(url='https://other.example/v1'), + [self._ok()], streaming) + self.assertIn('max_tokens', self.calls[0][1]) + + def _check_unrelated_error(self, streaming): + """Check unrelated error.""" + from pgadmin.llm.client import LLMClientError + error = self._error(param=None, message='not a chat model') + with self.assertRaises(LLMClientError): + self._run(self._client(), [error], streaming) + self.assertEqual(len(self.calls), 1) + + def _check_structured_only(self, streaming): + """Check structured only.""" + from pgadmin.llm.client import LLMClientError + for error in (self._error(param='other'), + self._error(status=500), + self._error(code='other')): + with self.assertRaises(LLMClientError): + self._run(self._client(), [error], streaming) + self.assertEqual(len(self.calls), 1) + + def _route_error(self): + """Route error.""" + return self._error( + param=None, status=404, + message='This model is only supported in v1/responses') + + def _responses_ok(self): + """Responses ok.""" + return dict(status='completed', output=[dict(type='message', + content=[dict(type='output_text', text='hello')])]) + + def _check_responses_cache(self, streaming): + """Check responses cache.""" + self._run(self._client(), [self._route_error(), + self._responses_ok()], streaming) + self.assertTrue(self.calls[-1][0].endswith('/responses')) + self._run(self._client(), [self._responses_ok()], streaming) + self.assertTrue(self.calls[0][0].endswith('/responses')) + + def _check_retry_chain(self, streaming): + """Check retry chain.""" + self._run(self._client(), [self._error(), self._route_error(), + self._responses_ok()], streaming) + self.assertEqual(len(self.calls), 3) + + def _check_retry_cap(self, streaming): + """Check retry cap.""" + from pgadmin.llm.client import LLMClientError + with self.assertRaises(LLMClientError): + self._run(self._client(), [ + self._error(), self._error()], streaming) + self.assertEqual(len(self.calls), 2) + + def _exhausted(self): + """Exhausted.""" + return self._ok(reasoning=10, content='', finish='length') + + def _check_reasoning_retry(self, streaming): + """Check reasoning retry.""" + self._run(self._client(), [self._error(), self._exhausted(), + self._ok()], streaming) + self.assertEqual(self._budgets(), [500, 500, 1500]) + + def _check_reasoning_cache(self, streaming): + """Check reasoning cache.""" + self._check_reasoning_retry(streaming) + self._run(self._client(), [self._ok()], streaming) + self.assertEqual(self._budgets(), [1500]) + + def _check_reasoning_success(self, streaming): + """Check reasoning success.""" + self._run(self._client(), [self._ok(reasoning=10)], streaming) + self._run(self._client(), [self._ok()], streaming) + self.assertEqual(self._budgets(), [1500]) + + def _check_reasoning_null(self, streaming): + """Check reasoning null.""" + self._run(self._client(), [self._ok(reasoning=None)], streaming) + self.assertEqual(self._budgets(), [500]) + + def _check_reasoning_exhausted(self, streaming): + """Check reasoning exhausted.""" + from pgadmin.llm.client import LLMClientError + with self.assertRaises(LLMClientError) as caught: + self._run(self._client(), [self._exhausted(), + self._exhausted()], streaming) + self.assertEqual(caught.exception.error.code, 'reasoning_budget') + self.assertEqual(self._budgets(), [500, 1500]) + + def _check_partial_stream(self, streaming): + """Check partial stream.""" + from pgadmin.llm.client import LLMClientError + from pgadmin.llm.models import Message + client = self._client() + calls = [] + + def partial(payload): + """Partial.""" + calls.append(payload) + yield 'hello' + raise self._error() + + client._process_stream = partial + response = client.chat_stream([Message.user('hi')]) + self.assertEqual(next(response), 'hello') + with self.assertRaises(LLMClientError): + next(response) + self.assertEqual(len(calls), 1) + + def _check_malformed_errors(self, streaming): + """Check malformed errors.""" + import io + import urllib.error + for body in (b'[]', b'{"error": []}', b'{"error": null}', + b'not json'): + error = self._client()._http_error_to_llm_error( + urllib.error.HTTPError('https://gateway.example', 400, + 'error', {}, io.BytesIO(body))) + self.assertTrue(error.message) + self.assertEqual(error.http_status, 400) + + def _check_hostnames(self, streaming): + """Check hostnames.""" + self.assertTrue(OpenAIClient( + api_url='https://api.openai.com.evil.example/v1').is_available()) + self.assertFalse(OpenAIClient( + api_url='https://api.openai.com:443/v1').is_available()) + + def _check_reasoning_content(self, streaming): + """Check reasoning content.""" + response = self._ok(content='', finish='length') + response['choices'][0]['message']['reasoning_content'] = 'thinking' + self._run(self._client(), [response, self._ok()], streaming) + self.assertEqual(self._budgets(), [500, 1500]) + + def _check_responses_rejection(self, streaming): + """Check responses rejection.""" + from pgadmin.llm.client import LLMClientError + for error in ( + self._error(param='other', + message='Model only supported in v1/responses'), + self._error( + param=None, status=500, + message='Model only supported in v1/responses')): + with self.assertRaises(LLMClientError): + self._run(self._client(), [error], streaming) + self.assertEqual(len(self.calls), 1) + + def _check_three_send_cap(self, streaming): + """Check three send cap.""" + from pgadmin.llm.client import LLMClientError + with self.assertRaises(LLMClientError): + self._run(self._client(), [self._error(), self._exhausted(), + self._route_error()], streaming) + self.assertEqual(len(self.calls), 3) + + def _check_http_fallback(self, streaming): + """Exercise the real request and error parser on both transports.""" + import io + import urllib.error + from unittest.mock import patch + from pgadmin.llm.models import Message + client = self._client() + sent = [] + + class Response(io.BytesIO): + pass + + def urlopen(request, **kwargs): + """Urlopen.""" + payload = json.loads(request.data) + sent.append(payload) + if len(sent) == 1: + body = json.dumps({'error': dict( + param='max_tokens', code='unsupported_parameter', + message='Unsupported')}) + raise urllib.error.HTTPError( + request.full_url, 400, 'Bad Request', {}, + io.BytesIO(body.encode())) + if streaming: + body = _sse(dict(choices=[dict(delta=dict(content='hello'), + finish_reason='stop')])) + '\n' + else: + body = json.dumps(self._ok()) + return Response(body.encode()) + + with patch('pgadmin.llm.providers.openai.urlopen_no_redirect', + side_effect=urlopen): + if streaming: + list(client.chat_stream([Message.user('hi')])) + else: + client.chat([Message.user('hi')]) + self.assertIn('max_tokens', sent[0]) + self.assertIn('max_completion_tokens', sent[1]) + + def _check_responses_reasoning(self, streaming): + """Reasoning budgets also work after a Responses API reroute.""" + exhausted = dict( + status='incomplete', output=[], + incomplete_details=dict(reason='max_output_tokens'), + usage=dict(output_tokens_details=dict(reasoning_tokens=10))) + self._run(self._client(), [self._route_error(), exhausted, + self._responses_ok()], streaming) + self.assertEqual(self._budgets(), [500, 500, 1500]) + self._run(self._client(), [self._responses_ok()], streaming) + self.assertEqual(self._budgets(), [1500]) + + def _check_debug_once(self, streaming): + """Check debug once.""" + from unittest.mock import patch + with patch('pgadmin.llm.providers.openai._logger.debug') as debug: + self._check_token_cache(streaming) + debug.assert_called_once() + + def _check_responses_model_param(self, streaming): + """Accept the official Responses-only model error shape.""" + self._run(self._client(), [self._error( + param='model', code=None, status=404, + message='This model is only supported in v1/responses ' + 'and not in v1/chat/completions.'), + self._responses_ok()], streaming) + self.assertTrue(self.calls[-1][0].endswith('/responses'))