From c71a535f16ab09f2f6baa9c98053771645c19d4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torbj=C3=B8rn=20Lindahl?= Date: Mon, 22 Dec 2025 23:13:29 +0100 Subject: [PATCH] added opencode zen as provider --- .env.example | 25 ++- conf/zen_models.json | 186 +++++++++++++++ docs/configuration.md | 8 + providers/__init__.py | 2 + providers/registries/zen.py | 35 +++ providers/registry.py | 2 + providers/shared/provider_type.py | 1 + providers/zen.py | 141 ++++++++++++ pyproject.toml | 2 +- server.py | 26 ++- tests/conftest.py | 13 +- tests/test_zen_model_registry.py | 166 ++++++++++++++ tests/test_zen_provider.py | 361 ++++++++++++++++++++++++++++++ tools/listmodels.py | 1 + 14 files changed, 956 insertions(+), 13 deletions(-) create mode 100644 conf/zen_models.json create mode 100644 providers/registries/zen.py create mode 100644 providers/zen.py create mode 100644 tests/test_zen_model_registry.py create mode 100644 tests/test_zen_provider.py diff --git a/.env.example b/.env.example index 2d29bdb..9e2a6da 100644 --- a/.env.example +++ b/.env.example @@ -3,11 +3,12 @@ # API Keys - At least one is required # -# IMPORTANT: Choose ONE approach: -# - Native APIs (Gemini/OpenAI/XAI) for direct access -# - DIAL for unified enterprise access -# - OpenRouter for unified cloud access -# Having multiple unified providers creates ambiguity about which serves each model. + # IMPORTANT: Choose ONE approach: + # - Native APIs (Gemini/OpenAI/XAI) for direct access + # - DIAL for unified enterprise access + # - OpenCode Zen for curated coding-focused models + # - OpenRouter for unified cloud access + # Having multiple unified providers creates ambiguity about which serves each model. # # Option 1: Use native APIs (recommended for direct access) # Get your Gemini API key from: https://makersuite.google.com/app/apikey @@ -29,11 +30,15 @@ AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/ # Get your X.AI API key from: https://console.x.ai/ XAI_API_KEY=your_xai_api_key_here -# Get your DIAL API key and configure host URL -# DIAL provides unified access to multiple AI models through a single API -DIAL_API_KEY=your_dial_api_key_here -# DIAL_API_HOST=https://core.dialx.ai # Optional: Base URL without /openai suffix (auto-appended) -# DIAL_API_VERSION=2025-01-01-preview # Optional: API version header for DIAL requests + # Get your DIAL API key and configure host URL + # DIAL provides unified access to multiple AI models through a single API + DIAL_API_KEY=your_dial_api_key_here + # DIAL_API_HOST=https://core.dialx.ai # Optional: Base URL without /openai suffix (auto-appended) + # DIAL_API_VERSION=2025-01-01-preview # Optional: API version header for DIAL requests + + # Get your OpenCode Zen API key from: https://opencode.ai/auth + # Zen provides curated, tested models optimized for coding agents + ZEN_API_KEY=your_zen_api_key_here # Option 2: Use OpenRouter for access to multiple models through one API # Get your OpenRouter API key from: https://openrouter.ai/ diff --git a/conf/zen_models.json b/conf/zen_models.json new file mode 100644 index 0000000..54c26b6 --- /dev/null +++ b/conf/zen_models.json @@ -0,0 +1,186 @@ +{ + "_README": { + "description": "Model metadata for OpenCode Zen curated models.", + "documentation": "https://opencode.ai/docs/zen", + "usage": "Models listed here are exposed through OpenCode Zen. Aliases are case-insensitive.", + "field_notes": "Matches providers/shared/model_capabilities.py.", + "field_descriptions": { + "model_name": "The model identifier as returned by Zen API (e.g., 'gpt-5.1-codex')", + "aliases": "Array of short names users can type instead of the full model name", + "context_window": "Total number of tokens the model can process (input + output combined)", + "max_output_tokens": "Maximum number of tokens the model can generate in a single response", + "supports_extended_thinking": "Whether the model supports extended reasoning tokens", + "supports_json_mode": "Whether the model can guarantee valid JSON output", + "supports_function_calling": "Whether the model supports function/tool calling", + "supports_images": "Whether the model can process images/visual input", + "max_image_size_mb": "Maximum total size in MB for all images combined", + "supports_temperature": "Whether the model accepts temperature parameter", + "temperature_constraint": "Type of temperature constraint or omit for default range", + "description": "Human-readable description of the model", + "intelligence_score": "1-20 human rating used for auto-mode model ordering", + "allow_code_generation": "Whether this model can generate working code" + } + }, + "models": [ + { + "model_name": "claude-opus-4-5", + "aliases": ["zen-opus", "zen-opus4.5", "zen-claude-opus"], + "context_window": 200000, + "max_output_tokens": 64000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": true, + "max_image_size_mb": 5.0, + "description": "Claude Opus 4.5 via OpenCode Zen - Anthropic's frontier reasoning model for complex software engineering", + "intelligence_score": 18, + "allow_code_generation": true + }, + { + "model_name": "claude-sonnet-4-5", + "aliases": ["zen-sonnet", "zen-sonnet4.5"], + "context_window": 200000, + "max_output_tokens": 64000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": true, + "max_image_size_mb": 5.0, + "description": "Claude Sonnet 4.5 via OpenCode Zen - Balanced performance for coding and general tasks", + "intelligence_score": 17, + "allow_code_generation": true + }, + { + "model_name": "claude-haiku-4-5", + "aliases": ["zen-haiku", "zen-haiku4.5"], + "context_window": 200000, + "max_output_tokens": 64000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": true, + "max_image_size_mb": 5.0, + "description": "Claude Haiku 4.5 via OpenCode Zen - Fast and efficient for coding tasks", + "intelligence_score": 16, + "allow_code_generation": true + }, + { + "model_name": "gpt-5.1-codex", + "aliases": ["zen-gpt-codex", "zen-codex"], + "context_window": 400000, + "max_output_tokens": 64000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": false, + "description": "GPT 5.1 Codex via OpenCode Zen - Specialized for code generation and understanding", + "intelligence_score": 17, + "allow_code_generation": true + }, + { + "model_name": "gpt-5.1", + "aliases": ["zen-gpt5.1"], + "context_window": 400000, + "max_output_tokens": 64000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": false, + "description": "GPT 5.1 via OpenCode Zen - Latest GPT model for general AI tasks", + "intelligence_score": 16, + "allow_code_generation": true + }, + { + "model_name": "gemini-3-pro", + "aliases": ["zen-gemini", "zen-gemini-pro"], + "context_window": 1000000, + "max_output_tokens": 64000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": true, + "max_image_size_mb": 10.0, + "description": "Gemini 3 Pro via OpenCode Zen - Google's multimodal model with large context", + "intelligence_score": 16, + "allow_code_generation": true + }, + { + "model_name": "glm-4.6", + "aliases": ["zen-glm"], + "context_window": 205000, + "max_output_tokens": 32000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": false, + "description": "GLM 4.6 via OpenCode Zen - High-performance model for coding and reasoning", + "intelligence_score": 15, + "allow_code_generation": true + }, + { + "model_name": "kimi-k2", + "aliases": ["zen-kimi"], + "context_window": 400000, + "max_output_tokens": 32000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": false, + "description": "Kimi K2 via OpenCode Zen - Advanced reasoning model", + "intelligence_score": 15, + "allow_code_generation": true + }, + { + "model_name": "qwen3-coder", + "aliases": ["zen-qwen", "zen-qwen-coder"], + "context_window": 480000, + "max_output_tokens": 32000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": false, + "description": "Qwen3 Coder via OpenCode Zen - Specialized coding model with large context", + "intelligence_score": 15, + "allow_code_generation": true + }, + { + "model_name": "grok-code", + "aliases": ["zen-grok"], + "context_window": 200000, + "max_output_tokens": 32000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": false, + "description": "Grok Code via OpenCode Zen - xAI's coding-focused model", + "intelligence_score": 14, + "allow_code_generation": true + }, + { + "model_name": "big-pickle", + "aliases": ["zen-pickle"], + "context_window": 200000, + "max_output_tokens": 32000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": false, + "description": "Big Pickle via OpenCode Zen - Stealth model for coding tasks", + "intelligence_score": 13, + "allow_code_generation": true + }, + { + "model_name": "gpt-5-nano", + "aliases": ["zen-nano"], + "context_window": 400000, + "max_output_tokens": 32000, + "supports_extended_thinking": false, + "supports_json_mode": true, + "supports_function_calling": true, + "supports_images": false, + "description": "GPT 5 Nano via OpenCode Zen - Lightweight GPT model", + "intelligence_score": 12, + "allow_code_generation": true + } + ] +} \ No newline at end of file diff --git a/docs/configuration.md b/docs/configuration.md index d084f2b..8783c1e 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -47,6 +47,14 @@ OPENROUTER_API_KEY=your_openrouter_api_key_here # If using OpenRouter, comment out native API keys above ``` +**Option 2.5: OpenCode Zen (Curated coding models)** +```env +# OpenCode Zen for tested and verified coding models +ZEN_API_KEY=your_zen_api_key_here +# Get from: https://opencode.ai/auth +# Provides access to models optimized for coding agents +``` + **Option 3: Custom API Endpoints (Local models)** ```env # For Ollama, vLLM, LM Studio, etc. diff --git a/providers/__init__.py b/providers/__init__.py index 8a499d6..b56f69d 100644 --- a/providers/__init__.py +++ b/providers/__init__.py @@ -8,6 +8,7 @@ from .openai_compatible import OpenAICompatibleProvider from .openrouter import OpenRouterProvider from .registry import ModelProviderRegistry from .shared import ModelCapabilities, ModelResponse +from .zen import ZenProvider __all__ = [ "ModelProvider", @@ -19,4 +20,5 @@ __all__ = [ "OpenAIModelProvider", "OpenAICompatibleProvider", "OpenRouterProvider", + "ZenProvider", ] diff --git a/providers/registries/zen.py b/providers/registries/zen.py new file mode 100644 index 0000000..a024cae --- /dev/null +++ b/providers/registries/zen.py @@ -0,0 +1,35 @@ +"""OpenCode Zen model registry for managing model configurations and aliases.""" + +from __future__ import annotations + +from ..shared import ModelCapabilities, ProviderType +from .base import CAPABILITY_FIELD_NAMES, CapabilityModelRegistry + + +class ZenModelRegistry(CapabilityModelRegistry): + """Capability registry backed by ``conf/zen_models.json``.""" + + def __init__(self, config_path: str | None = None) -> None: + super().__init__( + env_var_name="ZEN_MODELS_CONFIG_PATH", + default_filename="zen_models.json", + provider=ProviderType.ZEN, + friendly_prefix="OpenCode Zen ({model})", + config_path=config_path, + ) + + def _finalise_entry(self, entry: dict) -> tuple[ModelCapabilities, dict]: + provider_override = entry.get("provider") + if isinstance(provider_override, str): + entry_provider = ProviderType(provider_override.lower()) + elif isinstance(provider_override, ProviderType): + entry_provider = provider_override + else: + entry_provider = ProviderType.ZEN + + entry.setdefault("friendly_name", f"OpenCode Zen ({entry['model_name']})") + + filtered = {k: v for k, v in entry.items() if k in CAPABILITY_FIELD_NAMES} + filtered.setdefault("provider", entry_provider) + capability = ModelCapabilities(**filtered) + return capability, {} diff --git a/providers/registry.py b/providers/registry.py index cd28c42..9e9a2ed 100644 --- a/providers/registry.py +++ b/providers/registry.py @@ -40,6 +40,7 @@ class ModelProviderRegistry: ProviderType.OPENAI, # Direct OpenAI access ProviderType.AZURE, # Azure-hosted OpenAI deployments ProviderType.XAI, # Direct X.AI GROK access + ProviderType.ZEN, # OpenCode Zen curated models ProviderType.DIAL, # DIAL unified API access ProviderType.CUSTOM, # Local/self-hosted models ProviderType.OPENROUTER, # Catch-all for cloud models @@ -336,6 +337,7 @@ class ModelProviderRegistry: ProviderType.OPENAI: "OPENAI_API_KEY", ProviderType.AZURE: "AZURE_OPENAI_API_KEY", ProviderType.XAI: "XAI_API_KEY", + ProviderType.ZEN: "ZEN_API_KEY", ProviderType.OPENROUTER: "OPENROUTER_API_KEY", ProviderType.CUSTOM: "CUSTOM_API_KEY", # Can be empty for providers that don't need auth ProviderType.DIAL: "DIAL_API_KEY", diff --git a/providers/shared/provider_type.py b/providers/shared/provider_type.py index a1b3137..bb9a533 100644 --- a/providers/shared/provider_type.py +++ b/providers/shared/provider_type.py @@ -15,3 +15,4 @@ class ProviderType(Enum): OPENROUTER = "openrouter" CUSTOM = "custom" DIAL = "dial" + ZEN = "zen" diff --git a/providers/zen.py b/providers/zen.py new file mode 100644 index 0000000..463073c --- /dev/null +++ b/providers/zen.py @@ -0,0 +1,141 @@ +"""OpenCode Zen provider implementation.""" + +import logging + +from .openai_compatible import OpenAICompatibleProvider +from .registries.zen import ZenModelRegistry +from .shared import ( + ModelCapabilities, + ProviderType, +) + + +class ZenProvider(OpenAICompatibleProvider): + """Client for OpenCode Zen's curated model service. + + Role + Surface OpenCode Zen's tested and verified models through the same interface as + native providers so tools can reference Zen models without special cases. + + Characteristics + * Pulls model definitions from :class:`ZenModelRegistry` + (capabilities, metadata, pricing information) + * Reuses :class:`OpenAICompatibleProvider` infrastructure for request + execution so Zen endpoints behave like standard OpenAI-style APIs. + * Supports OpenCode Zen's curated list of coding-focused models. + """ + + FRIENDLY_NAME = "OpenCode Zen" + + # Model registry for managing configurations + _registry: ZenModelRegistry | None = None + + def __init__(self, api_key: str, **kwargs): + """Initialize OpenCode Zen provider. + + Args: + api_key: OpenCode Zen API key + **kwargs: Additional configuration + """ + base_url = "https://opencode.ai/zen/v1" + super().__init__(api_key, base_url=base_url, **kwargs) + + # Initialize model registry + if ZenProvider._registry is None: + ZenProvider._registry = ZenModelRegistry() + # Log loaded models only on first load + models = self._registry.list_models() + logging.info(f"OpenCode Zen loaded {len(models)} models") + + # ------------------------------------------------------------------ + # Capability surface + # ------------------------------------------------------------------ + + def _lookup_capabilities( + self, + canonical_name: str, + requested_name: str | None = None, + ) -> ModelCapabilities | None: + """Fetch Zen capabilities from the registry.""" + + capabilities = self._registry.get_capabilities(canonical_name) + if capabilities: + return capabilities + + # For unknown models, return None to let base class handle error + logging.debug("Model '%s' not found in Zen registry", canonical_name) + return None + + # ------------------------------------------------------------------ + # Provider identity + # ------------------------------------------------------------------ + + def get_provider_type(self) -> ProviderType: + """Identify this provider for restrictions and logging.""" + return ProviderType.ZEN + + # ------------------------------------------------------------------ + # Registry helpers + # ------------------------------------------------------------------ + + def list_models( + self, + *, + respect_restrictions: bool = True, + include_aliases: bool = True, + lowercase: bool = False, + unique: bool = False, + ) -> list[str]: + """Return formatted Zen model names, respecting restrictions.""" + + if not self._registry: + return [] + + from utils.model_restrictions import get_restriction_service + + restriction_service = get_restriction_service() if respect_restrictions else None + allowed_configs: dict[str, ModelCapabilities] = {} + + for model_name in self._registry.list_models(): + config = self._registry.resolve(model_name) + if not config: + continue + + if restriction_service: + if not restriction_service.is_allowed(self.get_provider_type(), model_name): + continue + + allowed_configs[model_name] = config + + if not allowed_configs: + return [] + + return ModelCapabilities.collect_model_names( + allowed_configs, + include_aliases=include_aliases, + lowercase=lowercase, + unique=unique, + ) + + def _resolve_model_name(self, model_name: str) -> str: + """Resolve aliases defined in the Zen registry.""" + + config = self._registry.resolve(model_name) + if config and config.model_name != model_name: + logging.debug("Resolved Zen model alias '%s' to '%s'", model_name, config.model_name) + return config.model_name + + return model_name + + def get_all_model_capabilities(self) -> dict[str, ModelCapabilities]: + """Expose registry-backed Zen capabilities.""" + + if not self._registry: + return {} + + capabilities: dict[str, ModelCapabilities] = {} + for model_name in self._registry.list_models(): + config = self._registry.resolve(model_name) + if config: + capabilities[model_name] = config + return capabilities diff --git a/pyproject.toml b/pyproject.toml index c60506d..a386a72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ name = "pal-mcp-server" version = "9.8.2" description = "AI-powered MCP server with multiple model providers" -requires-python = ">=3.9" +requires-python = ">=3.10" dependencies = [ "mcp>=1.0.0", "google-genai>=1.19.0", diff --git a/server.py b/server.py index 74f7ed8..9a04174 100644 --- a/server.py +++ b/server.py @@ -400,11 +400,13 @@ def configure_providers(): from providers.openrouter import OpenRouterProvider from providers.shared import ProviderType from providers.xai import XAIModelProvider + from providers.zen import ZenProvider from utils.model_restrictions import get_restriction_service valid_providers = [] has_native_apis = False has_openrouter = False + has_zen = False has_custom = False # Check for Gemini API key @@ -475,6 +477,19 @@ def configure_providers(): else: logger.debug("OpenRouter API key is placeholder value") + # Check for OpenCode Zen API key + zen_key = get_env("ZEN_API_KEY") + logger.debug(f"OpenCode Zen key check: key={'[PRESENT]' if zen_key else '[MISSING]'}") + if zen_key and zen_key != "your_zen_api_key_here": + valid_providers.append("OpenCode Zen") + has_zen = True + logger.info("OpenCode Zen API key found - Curated models available via Zen") + else: + if not zen_key: + logger.debug("OpenCode Zen API key not found in environment") + else: + logger.debug("OpenCode Zen API key is placeholder value") + # Check for custom API endpoint (Ollama, vLLM, etc.) custom_url = get_env("CUSTOM_API_URL") if custom_url: @@ -530,7 +545,13 @@ def configure_providers(): registered_providers.append(ProviderType.CUSTOM.value) logger.debug(f"Registered provider: {ProviderType.CUSTOM.value}") - # 3. OpenRouter last (catch-all for everything else) + # 3. OpenCode Zen + if has_zen: + ModelProviderRegistry.register_provider(ProviderType.ZEN, ZenProvider) + registered_providers.append(ProviderType.ZEN.value) + logger.debug(f"Registered provider: {ProviderType.ZEN.value}") + + # 4. OpenRouter last (catch-all for everything else) if has_openrouter: ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider) registered_providers.append(ProviderType.OPENROUTER.value) @@ -548,6 +569,7 @@ def configure_providers(): "- OPENAI_API_KEY for OpenAI models\n" "- XAI_API_KEY for X.AI GROK models\n" "- DIAL_API_KEY for DIAL models\n" + "- ZEN_API_KEY for OpenCode Zen (curated models)\n" "- OPENROUTER_API_KEY for OpenRouter (multiple models)\n" "- CUSTOM_API_URL for local models (Ollama, vLLM, etc.)" ) @@ -558,6 +580,8 @@ def configure_providers(): priority_info = [] if has_native_apis: priority_info.append("Native APIs (Gemini, OpenAI)") + if has_zen: + priority_info.append("OpenCode Zen") if has_custom: priority_info.append("Custom endpoints") if has_openrouter: diff --git a/tests/conftest.py b/tests/conftest.py index 6772f26..2e08d58 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -82,9 +82,20 @@ def project_path(tmp_path): return test_dir +@pytest.fixture +def zen_provider(): + """ + Provides a Zen provider instance for testing. + Uses dummy API key for isolated testing. + """ + from providers.zen import ZenProvider + + return ZenProvider(api_key="test-zen-key") + + def _set_dummy_keys_if_missing(): """Set dummy API keys only when they are completely absent.""" - for var in ("GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY"): + for var in ("GEMINI_API_KEY", "OPENAI_API_KEY", "XAI_API_KEY", "ZEN_API_KEY"): if not os.environ.get(var): os.environ[var] = "dummy-key-for-tests" diff --git a/tests/test_zen_model_registry.py b/tests/test_zen_model_registry.py new file mode 100644 index 0000000..b8bdaa4 --- /dev/null +++ b/tests/test_zen_model_registry.py @@ -0,0 +1,166 @@ +"""Tests for OpenCode Zen model registry functionality.""" + +import json +import os +import tempfile +from unittest.mock import patch + +from providers.registries.zen import ZenModelRegistry +from providers.shared import ProviderType + + +class TestZenModelRegistry: + """Test cases for Zen model registry.""" + + def test_registry_initialization(self): + """Test registry initializes with default config.""" + registry = ZenModelRegistry() + + # Should load models from default location + assert len(registry.list_models()) > 0 + assert len(registry.list_aliases()) > 0 + + # Should include our configured models + assert "claude-sonnet-4-5" in registry.list_models() + assert "gpt-5.1-codex" in registry.list_models() + + def test_custom_config_path(self): + """Test registry with custom config path.""" + # Create temporary config + config_data = { + "models": [ + { + "model_name": "test/zen-model-1", + "aliases": ["zen-test1", "zt1"], + "context_window": 4096, + "max_output_tokens": 2048, + "intelligence_score": 15, + } + ] + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(config_data, f) + temp_path = f.name + + try: + registry = ZenModelRegistry(config_path=temp_path) + assert len(registry.list_models()) == 1 + assert "test/zen-model-1" in registry.list_models() + assert "zen-test1" in registry.list_aliases() + assert "zt1" in registry.list_aliases() + finally: + os.unlink(temp_path) + + def test_get_capabilities(self): + """Test capability retrieval.""" + registry = ZenModelRegistry() + + # Test getting capabilities for a known model + caps = registry.get_capabilities("claude-sonnet-4-5") + assert caps is not None + assert caps.provider == ProviderType.ZEN + assert caps.model_name == "claude-sonnet-4-5" + assert caps.friendly_name == "OpenCode Zen (claude-sonnet-4-5)" + assert caps.context_window == 200000 + assert caps.intelligence_score == 17 + + # Test getting capabilities for unknown model + caps = registry.get_capabilities("unknown-model") + assert caps is None + + def test_resolve_model(self): + """Test model resolution with aliases.""" + registry = ZenModelRegistry() + + # Test resolving a direct model name + config = registry.resolve("claude-sonnet-4-5") + assert config is not None + assert config.model_name == "claude-sonnet-4-5" + + # Test resolving an alias + config = registry.resolve("zen-sonnet") + assert config is not None + assert config.model_name == "claude-sonnet-4-5" + + # Test resolving unknown model + config = registry.resolve("unknown-model") + assert config is None + + def test_list_aliases(self): + """Test alias listing.""" + registry = ZenModelRegistry() + + aliases = registry.list_aliases() + assert isinstance(aliases, list) + assert len(aliases) > 0 + + # Should include our configured aliases + assert "zen-sonnet" in aliases + assert "zen-codex" in aliases + assert "zen-gemini" in aliases + + def test_environment_config_path(self): + """Test registry respects environment variable for config path.""" + config_data = { + "models": [ + { + "model_name": "env/test-model", + "aliases": ["env-test"], + "context_window": 8192, + "max_output_tokens": 4096, + "intelligence_score": 10, + } + ] + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(config_data, f) + temp_path = f.name + + try: + with patch.dict(os.environ, {"ZEN_MODELS_CONFIG_PATH": temp_path}): + registry = ZenModelRegistry() + assert "env/test-model" in registry.list_models() + assert "env-test" in registry.list_aliases() + finally: + os.unlink(temp_path) + + def test_malformed_config(self): + """Test registry handles malformed config gracefully.""" + malformed_config = { + "models": [ + { + "model_name": "test/bad-model", + # Missing required fields + } + ] + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(malformed_config, f) + temp_path = f.name + + try: + registry = ZenModelRegistry(config_path=temp_path) + # Should still initialize but model may not load properly + # This tests error handling in config loading + registry.list_models() # Test that this doesn't crash + # May or may not include the malformed model depending on validation + finally: + os.unlink(temp_path) + + def test_empty_config(self): + """Test registry with empty config.""" + empty_config = {"models": []} + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(empty_config, f) + temp_path = f.name + + try: + registry = ZenModelRegistry(config_path=temp_path) + assert len(registry.list_models()) == 0 + assert len(registry.list_aliases()) == 0 + finally: + os.unlink(temp_path) diff --git a/tests/test_zen_provider.py b/tests/test_zen_provider.py new file mode 100644 index 0000000..b1456e5 --- /dev/null +++ b/tests/test_zen_provider.py @@ -0,0 +1,361 @@ +"""Tests for OpenCode Zen provider.""" + +import os +from unittest.mock import Mock, patch + +import pytest + +from providers.registry import ModelProviderRegistry +from providers.shared import ProviderType +from providers.zen import ZenProvider + + +class TestZenProvider: + """Test cases for OpenCode Zen provider.""" + + def test_provider_initialization(self): + """Test Zen provider initialization.""" + provider = ZenProvider(api_key="test-key") + assert provider.api_key == "test-key" + assert provider.base_url == "https://opencode.ai/zen/v1" + assert provider.FRIENDLY_NAME == "OpenCode Zen" + + def test_get_provider_type(self): + """Test provider type identification.""" + provider = ZenProvider(api_key="test-key") + assert provider.get_provider_type() == ProviderType.ZEN + + def test_model_validation(self): + """Test model validation.""" + provider = ZenProvider(api_key="test-key") + + # Zen accepts models that are in the registry + assert provider.validate_model_name("claude-sonnet-4-5") is True + assert provider.validate_model_name("gpt-5.1-codex") is True + + # Unknown models are rejected + assert provider.validate_model_name("unknown-model") is False + + def test_get_capabilities(self): + """Test capability generation.""" + provider = ZenProvider(api_key="test-key") + + # Test with a model in the registry + caps = provider.get_capabilities("claude-sonnet-4-5") + assert caps.provider == ProviderType.ZEN + assert caps.model_name == "claude-sonnet-4-5" + assert caps.friendly_name == "OpenCode Zen (claude-sonnet-4-5)" + + # Test with a model not in registry - should raise error + with pytest.raises(ValueError, match="Unsupported model 'unknown-model' for provider zen"): + provider.get_capabilities("unknown-model") + + def test_model_alias_resolution(self): + """Test model alias resolution.""" + provider = ZenProvider(api_key="test-key") + + # Test alias resolution + assert provider._resolve_model_name("zen-sonnet") == "claude-sonnet-4-5" + assert provider._resolve_model_name("zen-sonnet4.5") == "claude-sonnet-4-5" + assert provider._resolve_model_name("zen-codex") == "gpt-5.1-codex" + assert provider._resolve_model_name("zen-gpt-codex") == "gpt-5.1-codex" + + # Test case-insensitive + assert provider._resolve_model_name("ZEN-SONNET") == "claude-sonnet-4-5" + assert provider._resolve_model_name("Zen-Codex") == "gpt-5.1-codex" + + # Test direct model names (should pass through unchanged) + assert provider._resolve_model_name("claude-sonnet-4-5") == "claude-sonnet-4-5" + assert provider._resolve_model_name("gpt-5.1-codex") == "gpt-5.1-codex" + + # Test unknown models pass through + assert provider._resolve_model_name("unknown-model") == "unknown-model" + + def test_list_models(self): + """Test model listing with various options.""" + provider = ZenProvider(api_key="test-key") + + # Test basic model listing + models = provider.list_models() + assert isinstance(models, list) + assert len(models) > 0 + + # Should include our configured models + assert "claude-sonnet-4-5" in models + assert "gpt-5.1-codex" in models + + # Should include aliases + assert "zen-sonnet" in models + assert "zen-codex" in models + + def test_list_models_with_options(self): + """Test model listing with different options.""" + provider = ZenProvider(api_key="test-key") + + # Test without aliases + models_no_aliases = provider.list_models(include_aliases=False) + assert "zen-sonnet" not in models_no_aliases + assert "claude-sonnet-4-5" in models_no_aliases + + # Test lowercase + models_lower = provider.list_models(lowercase=True) + assert all(model == model.lower() for model in models_lower) + + def test_registry_capabilities(self): + """Test that registry capabilities are properly loaded.""" + provider = ZenProvider(api_key="test-key") + + # Test that we have a registry + assert provider._registry is not None + + # Test getting all capabilities + capabilities = provider.get_all_model_capabilities() + assert isinstance(capabilities, dict) + assert len(capabilities) > 0 + + # Should include our configured models + assert "claude-sonnet-4-5" in capabilities + assert "gpt-5.1-codex" in capabilities + + # Check capability structure + caps = capabilities["claude-sonnet-4-5"] + assert caps.provider == ProviderType.ZEN + assert caps.context_window == 200000 + assert caps.intelligence_score == 17 + + def test_model_capabilities_lookup(self): + """Test capability lookup for known and unknown models.""" + provider = ZenProvider(api_key="test-key") + + # Test known model + caps = provider._lookup_capabilities("claude-sonnet-4-5") + assert caps is not None + assert caps.provider == ProviderType.ZEN + + # Test unknown model returns None (base class handles error) + caps = provider._lookup_capabilities("unknown-zen-model") + assert caps is None + + def test_zen_registration(self): + """Test Zen can be registered and retrieved.""" + with patch.dict(os.environ, {"ZEN_API_KEY": "test-key"}): + # Clean up any existing registration + ModelProviderRegistry.unregister_provider(ProviderType.ZEN) + + # Register the provider + ModelProviderRegistry.register_provider(ProviderType.ZEN, ZenProvider) + + # Retrieve and verify + provider = ModelProviderRegistry.get_provider(ProviderType.ZEN) + assert provider is not None + assert isinstance(provider, ZenProvider) + + +class TestZenAutoMode: + """Test auto mode functionality when only Zen is configured.""" + + def setup_method(self): + """Store original state before each test.""" + self.registry = ModelProviderRegistry() + self._original_providers = self.registry._providers.copy() + self._original_initialized = self.registry._initialized_providers.copy() + + # Clear the registry state for this test + self.registry._providers.clear() + self.registry._initialized_providers.clear() + + self._original_env = {} + for key in ["ZEN_API_KEY", "GEMINI_API_KEY", "OPENAI_API_KEY", "DEFAULT_MODEL"]: + self._original_env[key] = os.environ.get(key) + + def teardown_method(self): + """Restore original state after each test.""" + self.registry._providers.clear() + self.registry._initialized_providers.clear() + self.registry._providers.update(self._original_providers) + self.registry._initialized_providers.update(self._original_initialized) + + for key, value in self._original_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + @pytest.mark.no_mock_provider + def test_zen_only_auto_mode(self): + """Test that auto mode works when only Zen is configured.""" + os.environ.pop("GEMINI_API_KEY", None) + os.environ.pop("OPENAI_API_KEY", None) + os.environ["ZEN_API_KEY"] = "test-zen-key" + os.environ["DEFAULT_MODEL"] = "auto" + + mock_registry = Mock() + model_names = [ + "claude-sonnet-4-5", + "claude-haiku-4-5", + "gpt-5.1-codex", + "gemini-3-pro", + "glm-4.6", + ] + mock_registry.list_models.return_value = model_names + + # Mock resolve to return a ModelCapabilities-like object for each model + def mock_resolve(model_name): + if model_name in model_names: + mock_config = Mock() + mock_config.provider = ProviderType.ZEN + mock_config.aliases = [] # Empty list of aliases + mock_config.get_effective_capability_rank = Mock(return_value=50) # Add ranking method + return mock_config + return None + + mock_registry.resolve.side_effect = mock_resolve + + ModelProviderRegistry.register_provider(ProviderType.ZEN, ZenProvider) + + provider = ModelProviderRegistry.get_provider(ProviderType.ZEN) + assert provider is not None, "Zen provider should be available with API key" + provider._registry = mock_registry + + available_models = ModelProviderRegistry.get_available_models(respect_restrictions=True) + + assert len(available_models) > 0, "Should find Zen models in auto mode" + assert all(provider_type == ProviderType.ZEN for provider_type in available_models.values()) + + for model in model_names: + assert model in available_models, f"Model {model} should be available" + + +class TestZenIntegration: + """Integration tests for Zen provider with server components.""" + + def test_zen_provider_in_server_init(self): + """Test that Zen provider is properly handled during server initialization.""" + # This test verifies that the server can handle Zen provider configuration + # without actual server startup + with patch.dict(os.environ, {"ZEN_API_KEY": "test-integration-key"}): + # Import server module to trigger provider setup + from providers.registry import ModelProviderRegistry + + # Verify Zen provider can be registered + ModelProviderRegistry.register_provider(ProviderType.ZEN, ZenProvider) + provider = ModelProviderRegistry.get_provider(ProviderType.ZEN) + assert provider is not None + assert isinstance(provider, ZenProvider) + + def test_zen_config_loading(self): + """Test that Zen configuration loads properly in integration context.""" + with patch.dict(os.environ, {"ZEN_API_KEY": "test-config-key"}): + from providers.registries.zen import ZenModelRegistry + + # Test registry loads configuration + registry = ZenModelRegistry() + models = registry.list_models() + aliases = registry.list_aliases() + + assert len(models) > 0, "Should load models from zen_models.json" + assert len(aliases) > 0, "Should load aliases from zen_models.json" + + # Verify specific models are loaded + assert "claude-sonnet-4-5" in models + assert "zen-sonnet" in aliases + + def test_zen_provider_priority(self): + """Test that Zen provider follows correct priority order.""" + # Zen should be prioritized after native APIs but before OpenRouter + from providers.registry import ModelProviderRegistry + + priority_order = ModelProviderRegistry.PROVIDER_PRIORITY_ORDER + zen_index = priority_order.index(ProviderType.ZEN) + openrouter_index = priority_order.index(ProviderType.OPENROUTER) + + # Zen should come before OpenRouter in priority + assert zen_index < openrouter_index, "Zen should have higher priority than OpenRouter" + + +class TestZenAPIMocking: + """Test API interactions with mocked OpenAI SDK.""" + + def test_chat_completion_mock(self): + """Test chat completion with mocked API response.""" + provider = ZenProvider(api_key="test-key") + + # Mock the OpenAI client and response + mock_response = Mock() + mock_response.choices = [Mock()] + mock_response.choices[0].message.content = "Mocked response from Zen" + mock_response.usage = Mock() + mock_response.usage.prompt_tokens = 10 + mock_response.usage.completion_tokens = 20 + + with patch.object(provider.client.chat.completions, "create", return_value=mock_response): + # Test the completion method - this will initialize the client + response = provider.complete( + model="claude-sonnet-4-5", messages=[{"role": "user", "content": "Hello"}], temperature=0.7 + ) + + assert response.content == "Mocked response from Zen" + + def test_streaming_completion_mock(self): + """Test streaming completion with mocked API.""" + provider = ZenProvider(api_key="test-key") + + # Mock streaming response + mock_chunk1 = Mock() + mock_chunk1.choices = [Mock()] + mock_chunk1.choices[0].delta.content = "Hello" + mock_chunk1.choices[0].finish_reason = None + + mock_chunk2 = Mock() + mock_chunk2.choices = [Mock()] + mock_chunk2.choices[0].delta.content = " world!" + mock_chunk2.choices[0].finish_reason = "stop" + + mock_stream = [mock_chunk1, mock_chunk2] + + # Access client to initialize it first + _ = provider.client + with patch.object(provider.client.chat.completions, "create", return_value=mock_stream): + # Test streaming completion + stream = provider.complete_stream( + model="gpt-5.1-codex", + messages=[{"role": "user", "content": "Say hello"}], + ) + + chunks = list(stream) + assert len(chunks) == 2 + assert chunks[0].content == "Hello" + assert chunks[1].content == " world!" + + def test_api_error_handling(self): + """Test error handling for API failures.""" + provider = ZenProvider(api_key="test-key") + + # Mock API error + from openai import APIError + + api_error = APIError("Mock API error", request=Mock(), body="error details") + + with patch.object(provider._client.chat.completions, "create", side_effect=api_error): + with pytest.raises(APIError): + provider.complete(model="claude-sonnet-4-5", messages=[{"role": "user", "content": "Test"}]) + + def test_invalid_model_error(self): + """Test error handling for invalid models.""" + provider = ZenProvider(api_key="test-key") + + with pytest.raises(ValueError, match="Unsupported model 'invalid-model' for provider zen"): + provider.get_capabilities("invalid-model") + + def test_authentication_error(self): + """Test handling of authentication errors.""" + provider = ZenProvider(api_key="invalid-key") + + # Mock authentication error + from openai import AuthenticationError + + auth_error = AuthenticationError("Invalid API key", request=Mock(), body="auth failed") + + with patch.object(provider._client.chat.completions, "create", side_effect=auth_error): + with pytest.raises(AuthenticationError): + provider.complete(model="claude-sonnet-4-5", messages=[{"role": "user", "content": "Test"}]) diff --git a/tools/listmodels.py b/tools/listmodels.py index 120afc1..1b9a98e 100644 --- a/tools/listmodels.py +++ b/tools/listmodels.py @@ -102,6 +102,7 @@ class ListModelsTool(BaseTool): ProviderType.OPENAI: {"name": "OpenAI", "env_key": "OPENAI_API_KEY"}, ProviderType.AZURE: {"name": "Azure OpenAI", "env_key": "AZURE_OPENAI_API_KEY"}, ProviderType.XAI: {"name": "X.AI (Grok)", "env_key": "XAI_API_KEY"}, + ProviderType.ZEN: {"name": "OpenCode Zen", "env_key": "ZEN_API_KEY"}, ProviderType.DIAL: {"name": "AI DIAL", "env_key": "DIAL_API_KEY"}, }