diff --git a/core/config.py b/core/config.py index 02345339..a6c36f04 100644 --- a/core/config.py +++ b/core/config.py @@ -194,6 +194,7 @@ class ProvidersConfig(_Base): and adding the matching :class:`~core.providers.registry.ProviderSpec`. """ + aimlapi: ProviderConfig = Field(default_factory=ProviderConfig) custom: ProviderConfig = Field(default_factory=ProviderConfig) openrouter: ProviderConfig = Field(default_factory=ProviderConfig) forge: ProviderConfig = Field(default_factory=ProviderConfig) diff --git a/core/providers/catalog_service.py b/core/providers/catalog_service.py index e03c3035..734923d3 100644 --- a/core/providers/catalog_service.py +++ b/core/providers/catalog_service.py @@ -409,6 +409,12 @@ def _catalog_request(connection: ResolvedConnection) -> tuple[str, dict[str, str base = connection.api_base.rstrip("/") if connection.api_key: headers["Authorization"] = f"Bearer {connection.api_key}" + if connection.provider_name == "aimlapi": + # AI/ML API serves one directory for every endpoint family it hosts: + # 936 rows, of which only 353 are chat models — the rest are image, + # video, speech and embedding endpoints that a chat request rejects, + # and ids repeat across families. Ask for the chat surface only. + return f"{base}/models?type=openai%2Fchat-completions", headers return f"{base}/models", headers @@ -421,15 +427,21 @@ def _parse_model(value: dict[str, Any]) -> CatalogModel | None: top_provider = ( value.get("top_provider") if isinstance(value.get("top_provider"), dict) else {} ) + # Some gateways (AI/ML API) nest the descriptive fields one level down + # instead of publishing OpenRouter's flat keys. Read it after the flat + # names so a provider that has both keeps winning at the top level. + info = value.get("info") if isinstance(value.get("info"), dict) else {} context = ( value.get("context_length") or value.get("context_window") or value.get("max_input_tokens") + or info.get("contextLength") or fallback.context_window ) output = ( value.get("max_output_tokens") or top_provider.get("max_completion_tokens") + or info.get("outputMax") or fallback.max_output_tokens ) supported = value.get("supported_parameters", []) @@ -446,7 +458,12 @@ def _parse_model(value: dict[str, Any]) -> CatalogModel | None: ) return CatalogModel( id=model_id, - name=str(value.get("name") or value.get("display_name") or model_id), + name=str( + value.get("name") + or value.get("display_name") + or info.get("name") + or model_id + ), context_window=max(1, int(context)), max_output_tokens=max(1, int(output)), supported_parameters=supported_parameters, diff --git a/core/providers/openai_compat.py b/core/providers/openai_compat.py index f6301086..83f05f45 100644 --- a/core/providers/openai_compat.py +++ b/core/providers/openai_compat.py @@ -12,6 +12,7 @@ import uuid from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any +from urllib.parse import urlparse import json_repair from loguru import logger @@ -79,6 +80,17 @@ "HTTP-Referer": "https://github.com/HKUDS/DeepCode", "X-Title": "DeepCode", } +_DEFAULT_AIMLAPI_HEADERS = { + "HTTP-Referer": "https://github.com/HKUDS/DeepCode", + "X-Title": "DeepCode", + "X-AIMLAPI-Source": "agent/deepcode", + "X-AIMLAPI-Partner-ID": "part_CxcOejScJ3hI0O2cExWchiBy", +} +# Hosts that are actually AI/ML API. Attribution is keyed on the resolved +# request origin rather than on the selected template, so a user who repoints +# the ``aimlapi`` template at a proxy of their own does not hand a third party +# DeepCode's partner identity. +_AIMLAPI_HOST_SUFFIX = ".aimlapi.com" # Per-model thinking / reasoning quirks now live declaratively in # ``core.providers.model_compat`` (resolved via ``resolve_model_compat``); # this module only assembles requests from the resolved value. @@ -171,6 +183,22 @@ def _uses_requesty_attribution( return bool(api_base and "requesty" in api_base.lower()) +def _uses_aimlapi_attribution( + spec: "ProviderSpec | None", api_base: str | None +) -> bool: + """Apply DeepCode attribution headers to AI/ML API requests by default. + + Stricter than the two helpers above: a substring test would also fire for + a gateway that merely *fronts* AI/ML API, so the host of the resolved base + URL has to be ours. The template name only decides the case where no base + URL was resolved at all. + """ + if not api_base: + return bool(spec and spec.name == "aimlapi") + host = (urlparse(api_base.strip()).hostname or "").lower() + return host == "aimlapi.com" or host.endswith(_AIMLAPI_HOST_SUFFIX) + + _RESPONSES_FAILURE_THRESHOLD = 3 _RESPONSES_PROBE_INTERVAL_S = 300 # 5 minutes @@ -250,6 +278,8 @@ def __init__( default_headers.update(_DEFAULT_OPENROUTER_HEADERS) if _uses_requesty_attribution(spec, effective_base): default_headers.update(_DEFAULT_REQUESTY_HEADERS) + if _uses_aimlapi_attribution(spec, effective_base): + default_headers.update(_DEFAULT_AIMLAPI_HEADERS) if extra_headers: default_headers.update(extra_headers) diff --git a/core/providers/registry.py b/core/providers/registry.py index b483c3af..a34592f5 100644 --- a/core/providers/registry.py +++ b/core/providers/registry.py @@ -54,6 +54,21 @@ def label(self) -> str: PROVIDERS: tuple[ProviderSpec, ...] = ( + ProviderSpec( + name="aimlapi", + keywords=("aimlapi",), + env_key="AIMLAPI_API_KEY", + # The vendor writes its own name lowercase, with the TLD. + display_name="aimlapi.com", + backend="openai_compat", + is_gateway=True, + detect_by_base_keyword="aimlapi", + default_api_base="https://api.aimlapi.com/v1", + # ``vendor/model`` ids like the OpenRouter-style gateways below, not + # Forge's bare ids. ``cache_control`` markers are honoured and the + # gateway reports ``cached_tokens`` back in usage. + supports_prompt_caching=True, + ), ProviderSpec( name="custom", keywords=(), diff --git a/tests/test_aimlapi_provider.py b/tests/test_aimlapi_provider.py new file mode 100644 index 00000000..28ddf24b --- /dev/null +++ b/tests/test_aimlapi_provider.py @@ -0,0 +1,265 @@ +"""AI/ML API gateway registration and its attribution headers. + +The gateway is wired on the same generic ``openai_compat`` path as OpenRouter +and Requesty, so most of this file pins the registry entry against those two. + +The attribution block gets more attention than the registry entry because it +fails *silently*: the gateway accepts a malformed partner id with a 200 and +simply records the traffic as untagged, so a typo is invisible at runtime and +only a shape assertion here can catch it. The same reasoning covers the origin +check — headers that ride a repointed base URL leak DeepCode's identity to a +third party without any error to notice. +""" + +from __future__ import annotations + +import os +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from core.config import ( # noqa: E402 + DeepCodeConfig, + ProviderConfig, + ProvidersConfig, +) +from core.providers.catalog_service import _catalog_request, _parse_model # noqa: E402 +from core.providers.openai_compat import ( # noqa: E402 + _DEFAULT_AIMLAPI_HEADERS, + OpenAICompatProvider, + _uses_aimlapi_attribution, +) +from core.providers.profiles import ConnectionResolver # noqa: E402 +from core.providers.registry import ( # noqa: E402 + PROVIDERS, + find_by_model, + find_by_name, +) + +AIMLAPI = find_by_name("aimlapi") +OPENROUTER = find_by_name("openrouter") + +# apps/api gateway contract: /^part_[A-Za-z0-9]{1,64}$/ — no dashes, no +# underscores after the prefix. +PARTNER_ID_PATTERN = re.compile(r"^part_[A-Za-z0-9]{1,64}$") +# /, channel from a closed enum, client [a-z0-9-]{1,32}. +SOURCE_PATTERN = re.compile(r"^(?:web|agent|mcp)/[a-z0-9-]{1,32}$") + + +# ---- registry -------------------------------------------------------------- + + +def test_aimlapi_is_registered_as_a_gateway() -> None: + assert AIMLAPI is not None + assert AIMLAPI.name == "aimlapi" + assert AIMLAPI.is_gateway is True + assert AIMLAPI.backend == "openai_compat" + + +def test_aimlapi_display_name_is_the_vendor_spelling() -> None: + # The vendor writes its own name lowercase and with the TLD; the label is + # what every provider list renders, so it is pinned here. + assert AIMLAPI is not None + assert AIMLAPI.display_name == "aimlapi.com" + assert AIMLAPI.label == "aimlapi.com" + + +def test_aimlapi_provider_specific_endpoint() -> None: + assert AIMLAPI is not None + assert AIMLAPI.default_api_base == "https://api.aimlapi.com/v1" + assert AIMLAPI.env_key == "AIMLAPI_API_KEY" + assert AIMLAPI.detect_by_base_keyword == "aimlapi" + + +def test_aimlapi_does_not_borrow_openrouter_key_prefix() -> None: + # OpenRouter keys start with ``sk-or-``; AI/ML API keys do not, so the + # prefix heuristic must not be copied over. + assert AIMLAPI is not None + assert AIMLAPI.detect_by_key_prefix == "" + + +def test_aimlapi_mirrors_openrouter_generic_wiring() -> None: + """Unlike Forge, AI/ML API resolves ``vendor/model`` ids, and it honours + ``cache_control`` markers (verified against the live gateway: a repeated + prefix comes back with ``cached_tokens`` in usage).""" + assert AIMLAPI is not None and OPENROUTER is not None + assert AIMLAPI.strip_model_prefix is OPENROUTER.strip_model_prefix is False + assert AIMLAPI.supports_prompt_caching is True + assert AIMLAPI.is_local is False + assert AIMLAPI.is_oauth is False + + +def test_aimlapi_shares_provider_slash_model_naming() -> None: + # ``vendor/model`` slugs resolve to the owning vendor, exactly like + # OpenRouter -- AI/ML API adds no new namespace. + for model in ("openai/gpt-4o-mini", "anthropic/claude-sonnet-4.5"): + spec = find_by_model(model) + assert spec is not None + assert spec.name in {"openai", "anthropic"} + + +def test_providers_config_exposes_aimlapi() -> None: + """``config.py`` reads providers via ``getattr(..., spec.name)``, so a + missing field silently disables the provider everywhere.""" + assert hasattr(ProvidersConfig(), "aimlapi") + + +def test_aimlapi_is_listed_first() -> None: + """Fork-only placement. Every provider list DeepCode renders — the + Settings view, the connection resolver, the app-server provider payload — + iterates ``PROVIDERS`` in declaration order, so position 0 is the whole + mechanism. Drop this test together with the placement commit before + offering the provider upstream.""" + assert PROVIDERS[0].name == "aimlapi" + assert list(ProvidersConfig.model_fields)[0] == "aimlapi" + + +# ---- attribution ----------------------------------------------------------- + + +def test_partner_id_matches_the_gateway_contract() -> None: + """A malformed id is accepted with a 200 and earns nothing — the failure + mode is silence, so the shape is asserted rather than observed.""" + partner_id = _DEFAULT_AIMLAPI_HEADERS["X-AIMLAPI-Partner-ID"] + assert PARTNER_ID_PATTERN.match(partner_id), partner_id + + +def test_source_matches_the_channel_contract() -> None: + source = _DEFAULT_AIMLAPI_HEADERS["X-AIMLAPI-Source"] + assert SOURCE_PATTERN.match(source), source + + +def test_referer_and_title_identify_deepcode_not_the_gateway() -> None: + # OpenRouter convention: these name the calling application. + assert _DEFAULT_AIMLAPI_HEADERS["X-Title"] == "DeepCode" + assert "HKUDS/DeepCode" in _DEFAULT_AIMLAPI_HEADERS["HTTP-Referer"] + + +def test_attribution_is_scoped_to_our_own_origin() -> None: + assert _uses_aimlapi_attribution(AIMLAPI, "https://api.aimlapi.com/v1") + assert _uses_aimlapi_attribution(None, "https://api.aimlapi.com/v1") + # A proxy that merely fronts the same API is somebody else's origin. + assert not _uses_aimlapi_attribution(AIMLAPI, "https://proxy.example.com/v1") + # ...and a lookalike domain must not satisfy the suffix test. + assert not _uses_aimlapi_attribution(AIMLAPI, "https://api.aimlapi.com.evil.io/v1") + assert not _uses_aimlapi_attribution(AIMLAPI, "https://notaimlapi.com/v1") + + +def test_headers_reach_the_client_without_clobbering_the_caller() -> None: + provider = OpenAICompatProvider( + api_key="test-key", + spec=AIMLAPI, + extra_headers={"X-Title": "Mine", "X-Custom": "kept"}, + ) + sent = provider._client.default_headers + assert sent["X-AIMLAPI-Partner-ID"] == "part_CxcOejScJ3hI0O2cExWchiBy" + assert sent["X-AIMLAPI-Source"] == "agent/deepcode" + # Merge, never assign: the caller's own value wins on a clash and their + # unrelated headers survive. + assert sent["X-Title"] == "Mine" + assert sent["X-Custom"] == "kept" + + +def test_attribution_does_not_ride_other_providers() -> None: + for name in ("openrouter", "requesty", "forge", "openai", "deepseek"): + spec = find_by_name(name) + assert spec is not None + sent = OpenAICompatProvider( + api_key="test-key", spec=spec + )._client.default_headers + assert "X-AIMLAPI-Partner-ID" not in sent, name + assert "X-AIMLAPI-Source" not in sent, name + + +def test_the_shared_header_constant_is_never_mutated() -> None: + """Each client gets a fresh dict; a per-instance override must not leak + into the next connection built from the same template.""" + before = dict(_DEFAULT_AIMLAPI_HEADERS) + OpenAICompatProvider( + api_key="test-key", spec=AIMLAPI, extra_headers={"X-Title": "Mine"} + ) + assert _DEFAULT_AIMLAPI_HEADERS == before + second = OpenAICompatProvider(api_key="test-key", spec=AIMLAPI) + assert second._client.default_headers["X-Title"] == "DeepCode" + + +# ---- model discovery ------------------------------------------------------- + + +def _aimlapi_connection(): + config = DeepCodeConfig( + providers=ProvidersConfig(aimlapi=ProviderConfig(api_key="test-key")) + ) + return ConnectionResolver(config).resolve_connection("aimlapi") + + +def test_discovery_asks_for_the_chat_surface_only() -> None: + """The directory serves every endpoint family the gateway hosts — image, + video, speech and embedding rows included, with ids repeating across + families. Unfiltered, the model picker offers rows a chat request + rejects.""" + url, headers = _catalog_request(_aimlapi_connection()) + assert url == ("https://api.aimlapi.com/v1/models?type=openai%2Fchat-completions") + assert headers["Authorization"] == "Bearer test-key" + + +def test_other_providers_keep_the_plain_models_url() -> None: + config = DeepCodeConfig( + providers=ProvidersConfig(requesty=ProviderConfig(api_key="test-key")) + ) + connection = ConnectionResolver(config).resolve_connection("requesty") + url, _ = _catalog_request(connection) + assert url == "https://router.requesty.ai/v1/models" + + +def test_nested_metadata_is_read_when_the_flat_keys_are_absent() -> None: + """AI/ML API nests the numbers under ``info``. Without this the picker + shows a fallback 128k/8k for every model and the runtime budgets context + against a window the model does not have.""" + model = _parse_model( + { + "id": "alibaba/glm-5.2", + "type": "openai/chat-completions", + "info": { + "name": "GLM 5.2", + "contextLength": 1000000, + "outputMax": 131072, + }, + } + ) + assert model is not None + assert model.name == "GLM 5.2" + assert model.context_window == 1000000 + assert model.max_output_tokens == 131072 + + +def test_flat_keys_still_win_over_nested_ones() -> None: + model = _parse_model( + { + "id": "some/model", + "name": "Flat", + "context_length": 8192, + "max_output_tokens": 512, + "info": {"name": "Nested", "contextLength": 999, "outputMax": 99}, + } + ) + assert model is not None + assert (model.name, model.context_window, model.max_output_tokens) == ( + "Flat", + 8192, + 512, + ) + + +def test_provider_construction_never_mutates_the_environment(monkeypatch) -> None: + """The key stays on the instance; nothing ambient learns it.""" + assert AIMLAPI is not None + monkeypatch.delenv(AIMLAPI.env_key, raising=False) + before = dict(os.environ) + OpenAICompatProvider(api_key="fresh-key", spec=AIMLAPI) + assert dict(os.environ) == before + assert AIMLAPI.env_key not in os.environ