diff --git a/docs/en_US/ai_tools.rst b/docs/en_US/ai_tools.rst index fdd6f438388..45baa0667de 100644 --- a/docs/en_US/ai_tools.rst +++ b/docs/en_US/ai_tools.rst @@ -26,7 +26,7 @@ Before using AI Reports, you must: **Note:** - * AI Reports using cloud providers (Anthropic, OpenAI) require an active internet connection. + * AI Reports using cloud providers (Anthropic, OpenAI, GLM) require an active internet connection. Local providers (Ollama, Docker Model Runner) do not require internet access. * API usage may incur costs depending on your LLM provider's pricing model. @@ -62,6 +62,13 @@ Select your preferred LLM provider from the dropdown: * **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. * **Model**: Select from available GPT models (e.g., gpt-4). +**GLM (Z.ai)** + Use GLM models from Z.ai through the provider's OpenAI-compatible API. + + * **API URL**: Custom API endpoint URL (leave empty for default: https://api.z.ai/api/paas/v4). Use https://api.z.ai/api/coding/paas/v4 with a Z.ai Coding Plan subscription. + * **API Key File**: Path to a file containing your Z.ai API key. 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. + * **Model**: Select from available GLM models (e.g., glm-4.7). There is no default; a model must be selected before the AI features can be used. + **Ollama** Use locally-hosted open-source models via Ollama. Requires a running Ollama instance. @@ -237,7 +244,7 @@ Troubleshooting You must establish a connection to the server or database before generating reports. **API Connection Errors** - * Verify your API key is correct (for Anthropic and OpenAI). + * Verify your API key is correct (for Anthropic, OpenAI, and GLM). * Check your internet connection (for cloud providers). * For Ollama, ensure the Ollama server is running and accessible. * For Docker Model Runner, ensure Docker Desktop 4.40+ is running with the model runner enabled. diff --git a/docs/en_US/preferences.rst b/docs/en_US/preferences.rst index 8342571ab5c..ea3a9b8bd46 100644 --- a/docs/en_US/preferences.rst +++ b/docs/en_US/preferences.rst @@ -43,7 +43,7 @@ in ``config.py``) for these preferences to be available. Use the fields on the *AI* panel to configure your LLM provider: * Use the *Default Provider* drop-down to select your LLM provider. Options include: - *Anthropic*, *OpenAI*, *Ollama*, or *Docker Model Runner*. + *Anthropic*, *OpenAI*, *GLM (Z.ai)*, *Ollama*, or *Docker Model Runner*. **Anthropic Settings:** @@ -81,6 +81,24 @@ Use the fields on the *AI* panel to configure your LLM provider: refresh button to fetch the latest available models from your configured endpoint. +**GLM (Z.ai) Settings:** + +* Use the *API URL* field to set a custom API endpoint URL. Leave empty to use + the default GLM API (``https://api.z.ai/api/paas/v4``). Use + ``https://api.z.ai/api/coding/paas/v4`` with a Z.ai Coding Plan + subscription. + +* Use the *API Key File* field to specify the path to a file containing your + GLM API key. 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. + +* Use the *Model* field to select from the available GLM models. Click the + refresh button to fetch the latest available models from your configured + endpoint. There is no default model, so one must be selected before the + AI features can be used. + **Ollama Settings:** * Use the *API URL* field to specify the Ollama server URL diff --git a/web/config.py b/web/config.py index 9cc01812ee8..5f9c309ce6a 100644 --- a/web/config.py +++ b/web/config.py @@ -1025,7 +1025,8 @@ # Default LLM Provider # Specifies which LLM provider to use by default when LLM_ENABLED is True. # Users can override this in their preferences. -# Valid values: 'anthropic', 'openai', 'ollama', 'docker', or '' (disabled) +# Valid values: 'anthropic', 'openai', 'glm', 'ollama', 'docker', +# or '' (disabled) DEFAULT_LLM_PROVIDER = '' # Anthropic Configuration @@ -1064,6 +1065,24 @@ # Examples: gpt-4o, gpt-4o-mini, gpt-4-turbo OPENAI_API_MODEL = '' +# GLM (Z.ai) Configuration +# URL for the GLM API endpoint. Leave empty to use the default +# (https://api.z.ai/api/paas/v4). Set this to +# https://api.z.ai/api/coding/paas/v4 for a Z.ai Coding Plan +GLM_API_URL = '' + +# Path to a file containing the GLM API key. The file should contain +# only the API key with no additional whitespace or formatting. +# Default: ~/.zai-api-key +GLM_API_KEY_FILE = '~/.zai-api-key' + +# The GLM model to use for AI features. There is no default: the models +# available differ between the Z.ai open platform, Coding Plan +# subscriptions and Zhipu's China platform, so one must be selected +# here or in the user's preferences. +# Examples: glm-4.6, glm-4.7, glm-5 +GLM_API_MODEL = '' + # Ollama Configuration # URL for the Ollama API endpoint. Leave empty to disable Ollama. # Typical value: http://localhost:11434 @@ -1105,6 +1124,7 @@ ALLOWED_LLM_API_URLS = [ 'https://api.anthropic.com:443', 'https://api.openai.com:443', + 'https://api.z.ai:443', # Loopback addresses on any port: covers Ollama (11434), Docker # Model Runner (12434), LiteLLM (4000), vLLM (8000), LM Studio # (1234), text-generation-webui (5000), and any self-hosted diff --git a/web/pgadmin/llm/README.md b/web/pgadmin/llm/README.md index caf7e39bada..0892d314565 100644 --- a/web/pgadmin/llm/README.md +++ b/web/pgadmin/llm/README.md @@ -13,15 +13,17 @@ This module provides AI/LLM functionality for pgAdmin, including database securi - **Anthropic Claude** (recommended) - **OpenAI GPT** +- **GLM (Z.ai)** - **Ollama** (local models) ## Configuration Configure LLM providers in `config.py`: -- `DEFAULT_LLM_PROVIDER`: Set to 'anthropic', 'openai', or 'ollama' +- `DEFAULT_LLM_PROVIDER`: Set to 'anthropic', 'openai', 'glm', or 'ollama' - `ANTHROPIC_API_KEY_FILE`: Path to file containing Anthropic API key - `OPENAI_API_KEY_FILE`: Path to file containing OpenAI API key +- `GLM_API_KEY_FILE`: Path to file containing GLM API key - `OLLAMA_API_URL`: URL for Ollama server (e.g., 'http://localhost:11434') If API keys are not found, the LLM features will be gracefully disabled. diff --git a/web/pgadmin/llm/__init__.py b/web/pgadmin/llm/__init__.py index bf035b59b52..58f8d355ab3 100644 --- a/web/pgadmin/llm/__init__.py +++ b/web/pgadmin/llm/__init__.py @@ -40,7 +40,7 @@ MODULE_NAME = 'llm' # Valid LLM providers -LLM_PROVIDERS = ['anthropic', 'openai', 'ollama', 'docker'] +LLM_PROVIDERS = ['anthropic', 'openai', 'glm', 'ollama', 'docker'] class LLMModule(PgAdminModule): @@ -71,6 +71,7 @@ def register_preferences(self): {'label': gettext('None (Disabled)'), 'value': ''}, {'label': gettext('Anthropic'), 'value': 'anthropic'}, {'label': gettext('OpenAI'), 'value': 'openai'}, + {'label': gettext('GLM (Z.ai)'), 'value': 'glm'}, {'label': gettext('Ollama'), 'value': 'ollama'}, {'label': gettext('Docker Model Runner'), 'value': 'docker'}, ] @@ -238,6 +239,66 @@ def register_preferences(self): } ) + # GLM Settings + # Get defaults from config + glm_url_default = getattr(config, 'GLM_API_URL', '') + glm_key_file_default = getattr(config, 'GLM_API_KEY_FILE', '') + glm_model_default = getattr(config, 'GLM_API_MODEL', '') + + self.glm_api_url = self.preference.register( + 'glm', 'glm_api_url', + gettext("API URL"), 'text', + glm_url_default, + category_label=gettext('GLM (Z.ai)'), + help_str=gettext( + 'URL for the GLM API endpoint. Leave empty to use ' + 'the default (https://api.z.ai/api/paas/v4). Use ' + 'https://api.z.ai/api/coding/paas/v4 for a Coding Plan ' + 'subscription, or set a custom URL to use another ' + 'GLM-compatible API provider.' + ) + ) + + self.glm_api_key_file = self.preference.register( + 'glm', 'glm_api_key_file', + gettext("API Key File"), 'text', + glm_key_file_default, + category_label=gettext('GLM (Z.ai)'), + help_str=gettext( + 'Path to a file containing your GLM API key. ' + 'This path must be on the server hosting pgAdmin, ' + 'e.g. inside the container when using Docker. ' + 'The file should contain only the API key.' + ) + ) + + glm_model_options = [] + + self.glm_api_model = self.preference.register( + 'glm', 'glm_api_model', + gettext("Model"), 'options', + glm_model_default, + category_label=gettext('GLM (Z.ai)'), + options=glm_model_options, + help_str=gettext( + 'The GLM model to use. Models are loaded dynamically ' + 'from your API key. You can also type a custom model name. ' + 'A model must be selected; there is no default.' + ), + control_props={ + 'allowClear': True, + 'creatable': True, + 'tags': True, + 'placeholder': gettext('Select or type a model name...'), + 'optionsUrl': 'llm.models_glm', + 'optionsRefreshUrl': 'llm.refresh_models_glm', + 'refreshDepNames': { + 'api_url': 'glm_api_url', + 'api_key_file': 'glm_api_key_file' + } + } + ) + # Ollama Settings # Get defaults from config ollama_url_default = getattr(config, 'OLLAMA_API_URL', '') @@ -334,10 +395,12 @@ def get_exposed_url_endpoints(self): return [ 'llm.models_anthropic', 'llm.models_openai', + 'llm.models_glm', 'llm.models_ollama', 'llm.models_docker', 'llm.refresh_models_anthropic', 'llm.refresh_models_openai', + 'llm.refresh_models_glm', 'llm.refresh_models_ollama', 'llm.refresh_models_docker', 'llm.status', @@ -378,7 +441,7 @@ def get_llm_status(): """ from pgadmin.llm.utils import ( is_llm_enabled, is_llm_enabled_system, get_default_provider, - get_anthropic_model, get_openai_model, get_ollama_model, + get_anthropic_model, get_openai_model, get_glm_model, get_ollama_model, get_docker_model ) @@ -388,6 +451,8 @@ def get_llm_status(): model = get_anthropic_model() elif provider == 'openai': model = get_openai_model() + elif provider == 'glm': + model = get_glm_model() elif provider == 'ollama': model = get_ollama_model() elif provider == 'docker': @@ -617,6 +682,108 @@ def refresh_openai_models(): ) +@blueprint.route("/models/glm", methods=["GET"], endpoint='models_glm') +@pga_login_required +def get_glm_models(): + """ + Fetch available GLM models. + Returns models from the OpenAI-compatible models endpoint. + """ + from pgadmin.llm.utils import get_glm_api_key, get_glm_api_url + + api_key = get_glm_api_key() + api_url = get_glm_api_url() + if not api_key and not api_url: + return make_json_response( + data={'models': [], 'error': 'No API key configured'}, + status=200 + ) + + try: + models = _fetch_glm_models(api_key, api_url) + return make_json_response(data={'models': models}, status=200) + except LLMApiError as e: + return make_json_response( + data={'models': [], 'error': str(e)}, + status=200 + ) + except Exception: + return make_json_response( + data={'models': [], + 'error': 'Failed to fetch GLM models. ' + 'Check the API key and URL configuration.'}, + status=200 + ) + + +@blueprint.route( + "/models/glm/refresh", + methods=["POST"], + endpoint='refresh_models_glm' +) +@pga_login_required +def refresh_glm_models(): + """ + Fetch available GLM models using a provided API key file path + and/or custom API URL. + Used by the preferences refresh button to load models before saving. + """ + from pgadmin.llm.utils import ( + read_api_key_file, validate_api_key_path, validate_api_url + ) + + data = request.get_json(force=True, silent=True) or {} + api_key_file = data.get('api_key_file', '') + api_url = data.get('api_url', '') + + if api_url and not validate_api_url(api_url): + return make_json_response( + data={'models': [], + 'error': 'API URL is not in the allowed list. ' + 'Contact your administrator to update ' + 'ALLOWED_LLM_API_URLS in the server ' + 'configuration.'}, + status=200 + ) + + api_key = None + if api_key_file: + safe_path = validate_api_key_path(api_key_file) + if safe_path is None: + return make_json_response( + data={'models': [], + 'error': 'API key file path is not permitted. ' + 'The file must be within your private ' + 'user storage directory; shared storage ' + 'and other locations are not allowed for ' + 'security reasons.'}, + status=200 + ) + api_key = read_api_key_file(safe_path) + + if not api_key and not api_url: + return make_json_response( + data={'models': [], 'error': 'No API key or custom URL provided'}, + status=200 + ) + + try: + models = _fetch_glm_models(api_key, api_url) + return make_json_response(data={'models': models}, status=200) + except LLMApiError as e: + return make_json_response( + data={'models': [], 'error': str(e)}, + status=200 + ) + except Exception: + return make_json_response( + data={'models': [], + 'error': 'Failed to fetch GLM models. ' + 'Check the API key file and URL.'}, + status=200 + ) + + @blueprint.route("/models/ollama", methods=["GET"], endpoint='models_ollama') @pga_login_required def get_ollama_models(): @@ -924,6 +1091,19 @@ def _fetch_openai_models(api_key, api_url=''): return models +def _fetch_glm_models(api_key, api_url=''): + """ + Fetch models from GLM (Z.ai) API. + Returns a list of model options with label and value. + """ + from pgadmin.llm.providers.glm import DEFAULT_API_BASE_URL + + return _fetch_openai_models( + api_key=api_key, + api_url=api_url or DEFAULT_API_BASE_URL + ) + + def _fetch_ollama_models(api_url): """ Fetch models from Ollama API. diff --git a/web/pgadmin/llm/client.py b/web/pgadmin/llm/client.py index 3c3dc823858..002f886c62c 100644 --- a/web/pgadmin/llm/client.py +++ b/web/pgadmin/llm/client.py @@ -154,9 +154,9 @@ def get_llm_client( Get an LLM client instance for the specified or default provider. Args: - provider: Optional provider name ('anthropic', 'openai', 'ollama', - 'docker'). If not specified, uses the configured default - provider. + provider: Optional provider name ('anthropic', 'openai', 'glm', + 'ollama', 'docker'). If not specified, uses the configured + default provider. model: Optional model name to use. If not specified, uses the configured default model for the provider. @@ -171,6 +171,7 @@ def get_llm_client( get_default_provider, get_anthropic_api_url, get_anthropic_api_key, get_anthropic_model, get_openai_api_url, get_openai_api_key, get_openai_model, + get_glm_api_url, get_glm_api_key, get_glm_model, get_ollama_api_url, get_ollama_model, get_docker_api_url, get_docker_model, is_pref_api_url_rejected, @@ -243,6 +244,34 @@ def _rejected_key_file_error(prov): api_key=api_key, model=model_name, api_url=api_url ) + elif provider == 'glm': + from pgadmin.llm.providers.glm import GLMClient + if is_pref_api_url_rejected('glm_api_url'): + raise _rejected_url_error('glm') + if is_pref_api_key_path_rejected('glm_api_key_file'): + raise _rejected_key_file_error('glm') + api_key = get_glm_api_key() + api_url = get_glm_api_url() + if not api_key and not api_url: + raise LLMClientError(LLMError( + message="GLM API key not configured", + provider="glm" + )) + model_name = model or get_glm_model() + if not model_name: + # GLM has no default model: the catalogue differs across the + # Z.ai open platform, Coding Plan subscriptions and Zhipu's + # China platform. Say so, rather than letting an unset model + # reach the API and come back as an opaque 404. + raise LLMClientError(LLMError( + message="No GLM model selected. Choose one in " + "Preferences > AI > GLM (Z.ai).", + provider="glm" + )) + return GLMClient( + api_key=api_key, model=model_name, api_url=api_url + ) + elif provider == 'ollama': from pgadmin.llm.providers.ollama import OllamaClient if is_pref_api_url_rejected('ollama_api_url'): diff --git a/web/pgadmin/llm/providers/__init__.py b/web/pgadmin/llm/providers/__init__.py index 5d469266865..9c7e54c8fc9 100644 --- a/web/pgadmin/llm/providers/__init__.py +++ b/web/pgadmin/llm/providers/__init__.py @@ -12,5 +12,6 @@ from pgadmin.llm.providers.anthropic import AnthropicClient from pgadmin.llm.providers.openai import OpenAIClient from pgadmin.llm.providers.ollama import OllamaClient +from pgadmin.llm.providers.glm import GLMClient -__all__ = ['AnthropicClient', 'OpenAIClient', 'OllamaClient'] +__all__ = ['AnthropicClient', 'OpenAIClient', 'OllamaClient', 'GLMClient'] diff --git a/web/pgadmin/llm/providers/glm.py b/web/pgadmin/llm/providers/glm.py new file mode 100644 index 00000000000..43fb12e777f --- /dev/null +++ b/web/pgadmin/llm/providers/glm.py @@ -0,0 +1,108 @@ +########################################################################## +# +# pgAdmin 4 - PostgreSQL Tools +# +# Copyright (C) 2013 - 2026, The pgAdmin Development Team +# This software is released under the PostgreSQL Licence +# +########################################################################## + +"""GLM (Z.ai) LLM client implementation.""" + +from typing import Optional +from urllib.parse import urlparse + +from pgadmin.llm.providers.openai import OpenAIClient + + +# Default API base URL. This is the Z.ai open platform endpoint. Users on a +# Coding Plan subscription or on Zhipu's China platform set their own URL in +# preferences (https://api.z.ai/api/coding/paas/v4 and +# https://open.bigmodel.cn/api/paas/v4 respectively). +DEFAULT_API_BASE_URL = 'https://api.z.ai/api/paas/v4' + +# Hosts that always require an API key. Any other host is assumed to be a +# self-hosted or proxied GLM-compatible endpoint, which may not need one. +AUTHENTICATED_HOSTS = ('api.z.ai') + + +class GLMClient(OpenAIClient): + """ + GLM API client. + + Implements the LLMClient interface for Z.ai GLM models using the + OpenAI-compatible API format. + + Deliberately has no default model: the models on offer differ between + the Z.ai open platform, a Coding Plan subscription and Zhipu's China + platform, so guessing one yields an opaque 404 rather than a useful + error. Callers must supply a model. + """ + + DEFAULT_API_BASE_URL = DEFAULT_API_BASE_URL + + def __init__(self, api_key: Optional[str] = None, + model: Optional[str] = None, + api_url: Optional[str] = None): + """ + Initialize the GLM client. + + Args: + api_key: The Z.ai API key. + model: The model name. Required; there is no default. + api_url: Optional custom API base URL. Defaults to + https://api.z.ai/api/paas/v4. + """ + super().__init__( + api_key=api_key, + model=model, + api_url=api_url or DEFAULT_API_BASE_URL + ) + # OpenAIClient falls back to its own default model ('gpt-4o') when + # none is given. That would silently send an OpenAI model name to + # Z.ai, so clear it: is_available() then reports the client as + # unusable and get_llm_client() raises a message naming the + # preference to set. + self._model = model or '' + + @property + def provider_name(self) -> str: + return 'glm' + + def is_available(self) -> bool: + """Check if the client is properly configured.""" + if not self._model: + return False + + host = (urlparse(self._base_url).hostname or '').lower() + if host in AUTHENTICATED_HOSTS: + return bool(self._api_key) + return True + + def _build_chat_payload(self, messages, tools, system_prompt, + max_tokens) -> dict: + """ + Build payload for the Chat Completions API. + + Z.ai documents ``max_tokens``. It ignores OpenAI's newer + ``max_completion_tokens`` outright, which silently removes the + output limit rather than failing. + """ + payload = super()._build_chat_payload( + messages, tools, system_prompt, max_tokens + ) + payload['max_tokens'] = payload.pop( + 'max_completion_tokens', max_tokens + ) + return payload + + def _should_use_responses_api(self, error) -> bool: + """ + Never fall back to the Responses API. + + Z.ai exposes no ``/responses`` endpoint. Without this, an unrelated + error whose text happens to match the inherited heuristic would + switch the client onto a path that cannot work, masking the real + failure. + """ + return False diff --git a/web/pgadmin/llm/tests/test_api_url_validation.py b/web/pgadmin/llm/tests/test_api_url_validation.py index 2229dbe0c16..1000e9ae41d 100644 --- a/web/pgadmin/llm/tests/test_api_url_validation.py +++ b/web/pgadmin/llm/tests/test_api_url_validation.py @@ -30,16 +30,40 @@ class ValidateApiUrlTestCase(BaseTestGenerator): allowed_urls=[ 'https://api.anthropic.com:443', 'https://api.openai.com:443', + 'https://api.z.ai:443', + 'https://open.bigmodel.cn:443', 'http://localhost:11434', 'http://localhost:12434', ], expected=True, )), + ('Allowed GLM URL', dict( + url='https://api.z.ai/api/paas/v4/models', + allowed_urls=[ + 'https://api.anthropic.com:443', + 'https://api.openai.com:443', + 'https://api.z.ai:443', + 'https://open.bigmodel.cn:443', + 'http://localhost:11434', + 'http://localhost:12434', + ], + expected=True, + )), + # Allowlist entries match on scheme+host+port, so one api.z.ai + # entry has to cover the Coding Plan path as well. + ('Allowed GLM Coding Plan URL', dict( + url='https://api.z.ai/api/coding/paas/v4/models', + allowed_urls=[ + 'https://api.z.ai:443', + ], + expected=True, + )), ('Allowed Ollama URL', dict( url='http://localhost:11434/api/tags', allowed_urls=[ 'https://api.anthropic.com:443', 'https://api.openai.com:443', + 'https://api.z.ai:443', 'http://localhost:11434', 'http://localhost:12434', ], @@ -50,6 +74,7 @@ class ValidateApiUrlTestCase(BaseTestGenerator): allowed_urls=[ 'https://api.anthropic.com:443', 'https://api.openai.com:443', + 'https://api.z.ai:443', 'http://localhost:11434', 'http://localhost:12434', ], @@ -273,6 +298,16 @@ class GetApiUrlResolutionTestCase(BaseTestGenerator): ], expect='', )), + ('GLM: rejected pref returns "" (no fallback)', dict( + getter='get_glm_api_url', + pref_value='http://evil.com:9999/', + config_attr='GLM_API_URL', + config_value='https://api.z.ai/api/paas/v4', + allowed_urls=[ + 'https://api.z.ai:443', + ], + expect='', + )), ('Docker: rejected pref returns "" (no fallback)', dict( getter='get_docker_api_url', pref_value='http://192.168.1.1:12434/', @@ -394,6 +429,15 @@ class GetApiKeyResolutionTestCase(BaseTestGenerator): read_admin_returns='sk-admin-fallback', expect=None, )), + ('GLM: rejected pref path returns None', dict( + getter='get_glm_api_key', + pref_value='/tmp/malicious-zai.key', + config_attr='GLM_API_KEY_FILE', + config_value='/etc/pgadmin/admin.key', + path_valid=False, + read_admin_returns='sk-admin-fallback', + expect=None, + )), ('Anthropic: allowed pref path returns user key', dict( getter='get_anthropic_api_key', pref_value='/home/u/private/anthropic.key', @@ -506,6 +550,11 @@ class GetLLMClientRejectedKeyFileTestCase(BaseTestGenerator): pref_key='anthropic_api_key_file', pref_value='/tmp/leak.key', )), + ('GLM rejected key file', dict( + provider='glm', + pref_key='glm_api_key_file', + pref_value='/tmp/leak-zai.key', + )), ] def setUp(self): @@ -550,6 +599,12 @@ class GetLLMClientRejectedUrlTestCase(BaseTestGenerator): pref_value='http://10.0.0.1:8080/', allowed_urls=['https://api.anthropic.com:443'], )), + ('GLM rejected URL surfaces clear error', dict( + provider='glm', + pref_name='glm_api_url', + pref_value='http://10.0.0.1:8080/', + allowed_urls=['https://api.z.ai:443'], + )), ('Ollama rejected URL surfaces clear error', dict( provider='ollama', pref_name='ollama_api_url', diff --git a/web/pgadmin/llm/tests/test_llm_status.py b/web/pgadmin/llm/tests/test_llm_status.py index 9aa145c6caa..ceea4c25aaa 100644 --- a/web/pgadmin/llm/tests/test_llm_status.py +++ b/web/pgadmin/llm/tests/test_llm_status.py @@ -34,6 +34,12 @@ class LLMStatusTestCase(BaseTestGenerator): expected_enabled=True, provider_name='openai' )), + ('LLM Status - GLM Enabled', dict( + url='/llm/status', + provider_enabled=True, + expected_enabled=True, + provider_name='glm' + )), ('LLM Status - Ollama Enabled', dict( url='/llm/status', provider_enabled=True, diff --git a/web/pgadmin/llm/tests/test_model_refresh_endpoints.py b/web/pgadmin/llm/tests/test_model_refresh_endpoints.py index ead5df0da40..3499432af7c 100644 --- a/web/pgadmin/llm/tests/test_model_refresh_endpoints.py +++ b/web/pgadmin/llm/tests/test_model_refresh_endpoints.py @@ -154,6 +154,77 @@ def runTest(self): p.stop() +class RefreshGLMEndpointTestCase(BaseTestGenerator): + """Integration tests for GLM refresh endpoint error handling""" + + scenarios = [ + ('Bad api_key_file returns generic error', dict( + post_data={'api_key_file': '/etc/passwd', 'api_url': ''}, + mock_validate_path_return=None, + mock_validate_url_return='skip', + expect_error_contains='not permitted', + )), + ('Disallowed api_url returns rejection', dict( + post_data={ + 'api_key_file': '', + 'api_url': 'http://169.254.169.254/' + }, + mock_validate_path_return='skip', + mock_validate_url_return=False, + expect_error_contains='not in the allowed list', + )), + ] + + def setUp(self): + pass + + def runTest(self): + patches = [ + patch( + 'pgadmin.authenticate.mfa.utils.mfa_required', + lambda f: f + ), + ] + + if self.mock_validate_path_return != 'skip': + patches.append(patch( + 'pgadmin.llm.utils.validate_api_key_path', + return_value=self.mock_validate_path_return + )) + + if hasattr(self, 'mock_validate_url_return') and \ + self.mock_validate_url_return != 'skip': + patches.append(patch( + 'pgadmin.llm.utils.validate_api_url', + return_value=self.mock_validate_url_return + )) + + for p in patches: + p.start() + + try: + response = self.tester.post( + '/llm/models/glm/refresh', + data=json.dumps(self.post_data), + content_type='application/json', + follow_redirects=True + ) + + self.assertEqual(response.status_code, 200) + data = json.loads(response.data) + self.assertIn('error', data['data']) + self.assertIn( + self.expect_error_contains, + data['data']['error'] + ) + self.assertNotIn( + 'root:', data['data'].get('error', '') + ) + finally: + for p in patches: + p.stop() + + class RefreshOllamaEndpointTestCase(BaseTestGenerator): """Integration tests for Ollama refresh endpoint error handling""" @@ -322,6 +393,47 @@ def runTest(self): self.assertEqual(data['data']['models'], _FAKE_MODELS) +class RefreshGLMHappyPathTestCase(BaseTestGenerator): + """Happy-path test: GLM refresh returns models when input valid.""" + + scenarios = [ + ('Valid key file and URL returns models', dict( + post_data={ + 'api_key_file': '/tmp/fake-key', + 'api_url': 'https://api.z.ai/api/paas/v4', + }, + )), + ] + + def setUp(self): + pass + + def runTest(self): + with patch( + 'pgadmin.authenticate.mfa.utils.mfa_required', lambda f: f + ), patch( + 'pgadmin.llm.utils.validate_api_url', return_value=True + ), patch( + 'pgadmin.llm.utils.validate_api_key_path', + return_value='/tmp/fake-key' + ), patch( + 'pgadmin.llm.utils.read_api_key_file', return_value='sk-test' + ), patch( + 'pgadmin.llm._fetch_glm_models', return_value=_FAKE_MODELS + ): + response = self.tester.post( + '/llm/models/glm/refresh', + data=json.dumps(self.post_data), + content_type='application/json', + follow_redirects=True + ) + + self.assertEqual(response.status_code, 200) + data = json.loads(response.data) + self.assertNotIn('error', data['data']) + self.assertEqual(data['data']['models'], _FAKE_MODELS) + + class RefreshOllamaHappyPathTestCase(BaseTestGenerator): """Happy-path test: Ollama refresh returns models when input valid.""" diff --git a/web/pgadmin/llm/utils.py b/web/pgadmin/llm/utils.py index c89b6c4f8c1..616607076fd 100644 --- a/web/pgadmin/llm/utils.py +++ b/web/pgadmin/llm/utils.py @@ -452,6 +452,40 @@ def get_openai_model(): return config.OPENAI_API_MODEL or '' +def get_glm_api_url(): + """ + Get the GLM API URL. See :func:`get_anthropic_api_url` for + the resolution and rejection rules. + """ + return _resolve_pref_url('glm_api_url', config.GLM_API_URL) + + +def get_glm_api_key(): + """ + Get the GLM API key. See :func:`_resolve_pref_key_file` for + resolution and rejection rules. + """ + return _resolve_pref_key_file( + 'glm_api_key_file', config.GLM_API_KEY_FILE + ) + + +def get_glm_model(): + """ + Get the GLM model to use. + + Checks user preferences first, then falls back to system configuration. + + Returns: + The model name string, or empty string if not configured. + """ + pref_model = _get_preference_value('glm_api_model') + if pref_model: + return pref_model + + return config.GLM_API_MODEL or '' + + def get_ollama_api_url(): """ Get the Ollama API URL. See :func:`get_anthropic_api_url` for @@ -513,7 +547,7 @@ def get_default_provider(): Returns None if disabled at system level or user preference is empty. Returns: - The provider name ('anthropic', 'openai', 'ollama', 'docker') + The provider name ('anthropic', 'openai', 'glm', 'ollama', 'docker') or None if disabled. """ # Check master switch first - cannot be overridden by user @@ -521,7 +555,7 @@ def get_default_provider(): return None # Valid provider values - valid_providers = {'anthropic', 'openai', 'ollama', 'docker'} + valid_providers = {'anthropic', 'openai', 'glm', 'ollama', 'docker'} # Get preference value (includes config default if not set by user) try: @@ -607,6 +641,10 @@ def get_llm_config(): 'api_key': str or None, 'model': str }, + 'glm': { + 'api_key': str or None, + 'model': str + }, 'ollama': { 'api_url': str, 'model': str @@ -630,6 +668,11 @@ def get_llm_config(): 'api_key': get_openai_api_key(), 'model': get_openai_model() }, + 'glm': { + 'api_url': get_glm_api_url(), + 'api_key': get_glm_api_key(), + 'model': get_glm_model() + }, 'ollama': { 'api_url': get_ollama_api_url(), 'model': get_ollama_model() diff --git a/web/pgadmin/preferences/__init__.py b/web/pgadmin/preferences/__init__.py index 66152b31ad2..09a938758c2 100644 --- a/web/pgadmin/preferences/__init__.py +++ b/web/pgadmin/preferences/__init__.py @@ -235,6 +235,7 @@ def save(): _provider_map = { 'anthropic_api_key_file': 'anthropic', 'openai_api_key_file': 'openai', + 'glm_api_key_file': 'glm', 'ollama_api_url': 'ollama', 'docker_api_url': 'docker', } @@ -333,6 +334,7 @@ def update(): _provider_map = { 'anthropic_api_key_file': 'anthropic', 'openai_api_key_file': 'openai', + 'glm_api_key_file': 'glm', 'ollama_api_url': 'ollama', 'docker_api_url': 'docker', }