From 13e7e9141de944ef9dead9dc2723c993f9acc20f Mon Sep 17 00:00:00 2001 From: VonLan233 Date: Wed, 9 Sep 2026 17:18:48 +0800 Subject: [PATCH] feat: add OpenCode Go and Zen providers --- astrbot/core/config/default.py | 24 +++ astrbot/core/provider/manager.py | 8 + .../sources/openai_responses_source.py | 4 + .../core/provider/sources/openai_source.py | 3 + .../provider/sources/opencode_go_source.py | 201 ++++++++++++++++++ .../provider/sources/opencode_zen_source.py | 25 +++ tests/test_opencode_source.py | 186 ++++++++++++++++ 7 files changed, 451 insertions(+) create mode 100644 astrbot/core/provider/sources/opencode_go_source.py create mode 100644 astrbot/core/provider/sources/opencode_zen_source.py create mode 100644 tests/test_opencode_source.py diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 8ce60ed7d0..593784bac2 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -1315,6 +1315,30 @@ "proxy": "", "custom_headers": {}, }, + "OpenCode Go": { + "id": "opencode-go", + "provider": "opencode-go", + "type": "opencode_go_chat_completion", + "provider_type": "chat_completion", + "enable": True, + "key": [], + "api_base": "https://opencode.ai/zen/go/v1", + "timeout": 120, + "proxy": "", + "custom_headers": {}, + }, + "OpenCode Zen": { + "id": "opencode", + "provider": "opencode", + "type": "opencode_zen_chat_completion", + "provider_type": "chat_completion", + "enable": True, + "key": [], + "api_base": "https://opencode.ai/zen/v1", + "timeout": 120, + "proxy": "", + "custom_headers": {}, + }, "Google Gemini": { "id": "google_gemini", "provider": "google", diff --git a/astrbot/core/provider/manager.py b/astrbot/core/provider/manager.py index 60044fb863..da80a253ae 100644 --- a/astrbot/core/provider/manager.py +++ b/astrbot/core/provider/manager.py @@ -425,6 +425,14 @@ def dynamic_import_provider(self, type: str) -> None: from .sources.openai_source import ( ProviderOpenAIOfficial as ProviderOpenAIOfficial, ) + case "opencode_go_chat_completion": + from .sources.opencode_go_source import ( + ProviderOpenCodeGo as ProviderOpenCodeGo, + ) + case "opencode_zen_chat_completion": + from .sources.opencode_zen_source import ( + ProviderOpenCodeZen as ProviderOpenCodeZen, + ) case "openai_responses": from .sources.openai_responses_source import ( ProviderOpenAIResponses as ProviderOpenAIResponses, diff --git a/astrbot/core/provider/sources/openai_responses_source.py b/astrbot/core/provider/sources/openai_responses_source.py index c5cb9bdb82..aa8cf7a1aa 100644 --- a/astrbot/core/provider/sources/openai_responses_source.py +++ b/astrbot/core/provider/sources/openai_responses_source.py @@ -241,6 +241,7 @@ async def _prepare_chat_payload( tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None, model: str | None = None, extra_user_content_parts: list[ContentPart] | None = None, + extra_headers: dict[str, str] | None = None, **kwargs: Any, ) -> tuple[dict, list[dict]]: """Build a stateless Responses API payload and replayable context. @@ -254,6 +255,7 @@ async def _prepare_chat_payload( tool_calls_result: Function calls and their returned outputs. model: Optional per-request model override. extra_user_content_parts: Additional user content blocks. + extra_headers: HTTP headers applied only to this request. **kwargs: Reserved provider request arguments. Returns: @@ -291,6 +293,8 @@ async def _prepare_chat_payload( } if system_prompt: payloads["instructions"] = system_prompt + if extra_headers: + payloads["extra_headers"] = extra_headers return payloads, context_query diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index f7870b7137..b38d39b096 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -956,6 +956,7 @@ async def _prepare_chat_payload( tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None, model: str | None = None, extra_user_content_parts: list[ContentPart] | None = None, + extra_headers: dict[str, str] | None = None, **kwargs, ) -> tuple: """准备聊天所需的有效载荷和上下文""" @@ -993,6 +994,8 @@ async def _prepare_chat_payload( model = model or self.get_model() payloads = {"messages": context_query, "model": model} + if extra_headers: + payloads["extra_headers"] = extra_headers self._finally_convert_payload(payloads) diff --git a/astrbot/core/provider/sources/opencode_go_source.py b/astrbot/core/provider/sources/opencode_go_source.py new file mode 100644 index 0000000000..cca1f8857d --- /dev/null +++ b/astrbot/core/provider/sources/opencode_go_source.py @@ -0,0 +1,201 @@ +import hashlib +from collections.abc import AsyncGenerator +from typing import Literal +from uuid import uuid4 + +from astrbot import __version__ +from astrbot.api.provider import Provider +from astrbot.core.agent.message import ContentPart, Message +from astrbot.core.agent.tool import ToolSet +from astrbot.core.provider.entities import LLMResponse, ToolCallsResult + +from ..register import register_provider_adapter +from .openai_responses_source import ProviderOpenAIResponses +from .openai_source import ProviderOpenAIOfficial + +OPENCODE_GO_API_BASE = "https://opencode.ai/zen/go/v1" +OPENCODE_GO_MODEL_PREFIX = "opencode-go/" +OPENCODE_GO_DEFAULT_MODEL = "kimi-k2.6" +OPENCODE_GO_RESPONSES_MODELS = { + "gpt-5.6-luna", + "grok-4.6", + "muse-spark-1.2-contributor", + "muse-spark-1.3-contributor", +} +OPENCODE_GO_MESSAGES_ONLY_MODELS = {"minimax-m2.5", "minimax-m2.7"} + + +@register_provider_adapter( + "opencode_go_chat_completion", + "OpenCode Go Subscription Provider Adapter", +) +class ProviderOpenCodeGo(Provider): + API_BASE = OPENCODE_GO_API_BASE + MODEL_PREFIX = OPENCODE_GO_MODEL_PREFIX + DEFAULT_MODEL = OPENCODE_GO_DEFAULT_MODEL + PROVIDER_NAME = "OpenCode Go" + RESPONSES_MODELS = OPENCODE_GO_RESPONSES_MODELS + RESPONSES_MODEL_PREFIXES: tuple[str, ...] = () + UNSUPPORTED_MODEL_ENDPOINTS = dict.fromkeys( + OPENCODE_GO_MESSAGES_ONLY_MODELS, "/v1/messages" + ) + UNSUPPORTED_MODEL_PREFIX_ENDPOINTS: tuple[tuple[str, str], ...] = () + + def __init__(self, provider_config: dict, provider_settings: dict) -> None: + super().__init__(provider_config, provider_settings) + self.api_base = provider_config.get("api_base", self.API_BASE).rstrip("/") + self.timeout = provider_config.get("timeout", 120) + if isinstance(self.timeout, str): + self.timeout = int(self.timeout) + + model = self._to_api_model(provider_config.get("model", self.DEFAULT_MODEL)) + self.set_model(model) + self._fallback_session_id = uuid4().hex + + uses_responses = model in self.RESPONSES_MODELS or model.startswith( + self.RESPONSES_MODEL_PREFIXES + ) + delegate_class = ( + ProviderOpenAIResponses if uses_responses else ProviderOpenAIOfficial + ) + self.openai_provider = delegate_class( + self._build_delegate_config(model=model), + provider_settings, + ) + + def _build_delegate_config(self, *, model: str) -> dict: + config = dict(self.provider_config) + config["api_base"] = self.api_base + config["model"] = model + custom_headers = config.get("custom_headers") + custom_headers = ( + dict(custom_headers) if isinstance(custom_headers, dict) else {} + ) + custom_headers.setdefault("User-Agent", f"AstrBot-Coding-Agent/{__version__}") + config["custom_headers"] = custom_headers + return config + + @classmethod + def _to_api_model(cls, model: str | None) -> str: + resolved_model = (model or cls.DEFAULT_MODEL).strip() + if resolved_model.startswith(cls.MODEL_PREFIX): + return resolved_model.removeprefix(cls.MODEL_PREFIX) + return resolved_model + + @classmethod + def _to_provider_model(cls, model: str) -> str: + return f"{cls.MODEL_PREFIX}{cls._to_api_model(model)}" + + @classmethod + def _unsupported_endpoint(cls, api_model: str) -> str | None: + if endpoint := cls.UNSUPPORTED_MODEL_ENDPOINTS.get(api_model): + return endpoint + for prefix, endpoint in cls.UNSUPPORTED_MODEL_PREFIX_ENDPOINTS: + if api_model.startswith(prefix): + return endpoint + return None + + @classmethod + def _ensure_supported_model(cls, model: str | None) -> str: + api_model = cls._to_api_model(model) + if endpoint := cls._unsupported_endpoint(api_model): + raise ValueError( + f"{cls.PROVIDER_NAME} model {cls.MODEL_PREFIX}{api_model} uses " + f"unsupported endpoint {endpoint}." + ) + return api_model + + def _resolve_model(self, model: str | None = None) -> str: + return self._ensure_supported_model(model or self.get_model()) + + def get_current_key(self) -> str: + return self.openai_provider.get_current_key() + + def get_keys(self) -> list[str]: + return self.openai_provider.get_keys() + + def set_key(self, key: str) -> None: + self.openai_provider.set_key(key) + + async def get_models(self) -> list[str]: + models = await self.openai_provider.get_models() + api_models = [] + for model in models: + api_model = self._to_api_model(model) + if api_model and not self._unsupported_endpoint(api_model): + api_models.append(api_model) + return sorted(api_models) + + async def text_chat( + self, + prompt: str | None = None, + session_id: str | None = None, + image_urls: list[str] | None = None, + audio_urls: list[str] | None = None, + func_tool: ToolSet | None = None, + contexts: list[Message] | list[dict] | None = None, + system_prompt: str | None = None, + tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None, + model: str | None = None, + extra_user_content_parts: list[ContentPart] | None = None, + tool_choice: Literal["auto", "required"] = "auto", + **kwargs, + ) -> LLMResponse: + extra_headers = dict(kwargs.pop("extra_headers", {}) or {}) + extra_headers["x-opencode-session"] = hashlib.sha256( + str(session_id or self._fallback_session_id).encode() + ).hexdigest() + return await self.openai_provider.text_chat( + prompt=prompt, + session_id=session_id, + image_urls=image_urls, + audio_urls=audio_urls, + func_tool=func_tool, + contexts=contexts, + system_prompt=system_prompt, + tool_calls_result=tool_calls_result, + model=self._resolve_model(model), + extra_user_content_parts=extra_user_content_parts, + tool_choice=tool_choice, + extra_headers=extra_headers, + **kwargs, + ) + + async def text_chat_stream( + self, + prompt: str | None = None, + session_id: str | None = None, + image_urls: list[str] | None = None, + audio_urls: list[str] | None = None, + func_tool: ToolSet | None = None, + contexts: list[Message] | list[dict] | None = None, + system_prompt: str | None = None, + tool_calls_result: ToolCallsResult | list[ToolCallsResult] | None = None, + model: str | None = None, + extra_user_content_parts: list[ContentPart] | None = None, + tool_choice: Literal["auto", "required"] = "auto", + **kwargs, + ) -> AsyncGenerator[LLMResponse, None]: + extra_headers = dict(kwargs.pop("extra_headers", {}) or {}) + extra_headers["x-opencode-session"] = hashlib.sha256( + str(session_id or self._fallback_session_id).encode() + ).hexdigest() + async for response in self.openai_provider.text_chat_stream( + prompt=prompt, + session_id=session_id, + image_urls=image_urls, + audio_urls=audio_urls, + func_tool=func_tool, + contexts=contexts, + system_prompt=system_prompt, + tool_calls_result=tool_calls_result, + model=self._resolve_model(model), + extra_user_content_parts=extra_user_content_parts, + tool_choice=tool_choice, + extra_headers=extra_headers, + **kwargs, + ): + yield response + + async def terminate(self) -> None: + await self.openai_provider.terminate() diff --git a/astrbot/core/provider/sources/opencode_zen_source.py b/astrbot/core/provider/sources/opencode_zen_source.py new file mode 100644 index 0000000000..90fde06471 --- /dev/null +++ b/astrbot/core/provider/sources/opencode_zen_source.py @@ -0,0 +1,25 @@ +from ..register import register_provider_adapter +from .opencode_go_source import ProviderOpenCodeGo + +OPENCODE_ZEN_API_BASE = "https://opencode.ai/zen/v1" +OPENCODE_ZEN_MODEL_PREFIX = "opencode/" +OPENCODE_ZEN_DEFAULT_MODEL = "kimi-k2.6" + + +@register_provider_adapter( + "opencode_zen_chat_completion", + "OpenCode Zen Provider Adapter", +) +class ProviderOpenCodeZen(ProviderOpenCodeGo): + API_BASE = OPENCODE_ZEN_API_BASE + MODEL_PREFIX = OPENCODE_ZEN_MODEL_PREFIX + DEFAULT_MODEL = OPENCODE_ZEN_DEFAULT_MODEL + PROVIDER_NAME = "OpenCode Zen" + RESPONSES_MODELS = set() + RESPONSES_MODEL_PREFIXES = ("gpt-", "grok-", "muse-spark-") + UNSUPPORTED_MODEL_ENDPOINTS = {} + UNSUPPORTED_MODEL_PREFIX_ENDPOINTS = ( + ("claude-", "/v1/messages"), + ("qwen", "/v1/messages"), + ("gemini-", "/v1/models/{model}"), + ) diff --git a/tests/test_opencode_source.py b/tests/test_opencode_source.py new file mode 100644 index 0000000000..e3c15a6a7f --- /dev/null +++ b/tests/test_opencode_source.py @@ -0,0 +1,186 @@ +import hashlib +from types import SimpleNamespace + +import pytest + +from astrbot.core.config.default import CONFIG_METADATA_2 +from astrbot.core.provider.sources.openai_responses_source import ( + ProviderOpenAIResponses, +) +from astrbot.core.provider.sources.openai_source import ProviderOpenAIOfficial +from astrbot.core.provider.sources.opencode_go_source import ProviderOpenCodeGo +from astrbot.core.provider.sources.opencode_zen_source import ProviderOpenCodeZen + + +class _DelegateStub: + def __init__(self, models: list[str] | None = None) -> None: + self.models = models or [] + self.chat_kwargs = None + self.stream_kwargs = None + + async def get_models(self) -> list[str]: + return self.models + + async def text_chat(self, **kwargs): + self.chat_kwargs = kwargs + return SimpleNamespace(role="assistant") + + async def text_chat_stream(self, **kwargs): + self.stream_kwargs = kwargs + yield SimpleNamespace(role="assistant") + + +def _unit_provider( + provider_class: type[ProviderOpenCodeGo], + delegate: _DelegateStub, +) -> ProviderOpenCodeGo: + """Build an OpenCode provider without creating an HTTP client. + + Args: + provider_class: OpenCode adapter class to instantiate. + delegate: Stub receiving delegated calls. + + Returns: + Partially initialized provider suitable for unit tests. + """ + provider = provider_class.__new__(provider_class) + provider.openai_provider = delegate + provider.model_name = provider_class.DEFAULT_MODEL + provider._fallback_session_id = "fallback-session" + return provider + + +def test_opencode_templates_are_available() -> None: + templates = CONFIG_METADATA_2["provider_group"]["metadata"]["provider"][ + "config_template" + ] + + assert templates["OpenCode Go"]["type"] == "opencode_go_chat_completion" + assert templates["OpenCode Zen"]["type"] == "opencode_zen_chat_completion" + + +@pytest.mark.asyncio +async def test_opencode_go_prefixes_models_and_filters_messages_endpoint() -> None: + provider = _unit_provider( + ProviderOpenCodeGo, + _DelegateStub( + [ + "kimi-k2.6", + "opencode-go/gpt-5.6-luna", + "muse-spark-1.3-contributor", + "minimax-m2.7", + ] + ), + ) + + assert await provider.get_models() == [ + "gpt-5.6-luna", + "kimi-k2.6", + "muse-spark-1.3-contributor", + ] + + +@pytest.mark.asyncio +async def test_opencode_zen_supports_responses_and_rejects_other_endpoints() -> None: + provider = _unit_provider( + ProviderOpenCodeZen, + _DelegateStub( + [ + "kimi-k2.6", + "muse-spark-1.3-contributor-free", + "claude-sonnet-4-6", + "gemini-3.1-pro", + ] + ), + ) + + assert await provider.get_models() == [ + "kimi-k2.6", + "muse-spark-1.3-contributor-free", + ] + with pytest.raises(ValueError, match="/v1/messages"): + provider._resolve_model("opencode/claude-sonnet-4-6") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("provider_class", "model", "delegate_class"), + [ + (ProviderOpenCodeGo, "opencode-go/kimi-k2.6", ProviderOpenAIOfficial), + ( + ProviderOpenCodeGo, + "opencode-go/muse-spark-1.3-contributor", + ProviderOpenAIResponses, + ), + ( + ProviderOpenCodeZen, + "opencode/muse-spark-1.3-contributor-free", + ProviderOpenAIResponses, + ), + ], +) +async def test_opencode_selects_delegate_for_model_endpoint( + provider_class, + model, + delegate_class, +) -> None: + provider = provider_class({"model": model, "key": ["test-key"]}, {}) + try: + assert isinstance(provider.openai_provider, delegate_class) + finally: + await provider.terminate() + + +@pytest.mark.asyncio +async def test_opencode_sends_coding_agent_and_hashed_session_headers() -> None: + delegate = _DelegateStub() + provider = _unit_provider(ProviderOpenCodeGo, delegate) + provider.provider_config = {"custom_headers": {"x-source": "configured"}} + provider.api_base = "https://example.test/v1" + + config = provider._build_delegate_config(model="kimi-k2.6") + await provider.text_chat( + prompt="hello", + session_id="telegram:user:123", + extra_headers={"x-request": "present"}, + ) + + assert config["custom_headers"]["User-Agent"].startswith("AstrBot-Coding-Agent/") + assert config["custom_headers"]["x-source"] == "configured" + assert delegate.chat_kwargs["extra_headers"] == { + "x-request": "present", + "x-opencode-session": hashlib.sha256(b"telegram:user:123").hexdigest(), + } + + +@pytest.mark.asyncio +async def test_opencode_stream_uses_stable_fallback_session() -> None: + delegate = _DelegateStub() + provider = _unit_provider(ProviderOpenCodeGo, delegate) + + responses = [item async for item in provider.text_chat_stream(prompt="hello")] + + assert responses + assert delegate.stream_kwargs["extra_headers"]["x-opencode-session"] == ( + hashlib.sha256(b"fallback-session").hexdigest() + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_class", + [ProviderOpenAIOfficial, ProviderOpenAIResponses], +) +async def test_openai_payload_forwards_per_request_headers(provider_class) -> None: + provider = provider_class( + {"id": "test", "model": "test-model", "key": ["test-key"]}, + {}, + ) + try: + payload, _ = await provider._prepare_chat_payload( + "hello", + extra_headers={"x-opencode-session": "session-hash"}, + ) + assert payload["extra_headers"] == {"x-opencode-session": "session-hash"} + finally: + await provider.terminate()