diff --git a/openhands-sdk/openhands/sdk/llm/llm.py b/openhands-sdk/openhands/sdk/llm/llm.py index 2dd49d7164..127715c863 100644 --- a/openhands-sdk/openhands/sdk/llm/llm.py +++ b/openhands-sdk/openhands/sdk/llm/llm.py @@ -344,6 +344,15 @@ class LLM(BaseModel, RetryMixin, NonNativeToolCallingMixin): json_schema_extra=field_meta(), ) + # AIMLAPI (https://aimlapi.com) partner attribution, sent only on calls + # routed through the ``aiml`` LiteLLM provider so AIMLAPI can attribute + # traffic to OpenHands (see AIMLAPI_PARTNER_ID in the CLI's aimlapi/config.py + # for the sibling implementation). + aimlapi_partner_id: str = Field( + default="part_uDVajKg3xPLrOdNdQetOtoGA", + json_schema_extra=field_meta(), + ) + num_retries: int = Field(default=5, ge=0, json_schema_extra=field_meta()) retry_multiplier: float = Field(default=8.0, ge=0, json_schema_extra=field_meta()) retry_min_wait: int = Field(default=8, ge=0, json_schema_extra=field_meta()) @@ -822,6 +831,21 @@ def _openrouter_headers(self) -> dict[str, str]: headers["X-Title"] = self.openrouter_app_name return headers + def _aiml_headers(self) -> dict[str, str]: + """Build AIMLAPI partner-attribution headers for per-call use. + + Unlike ``_openrouter_headers()``, gated on the parsed LiteLLM provider + being ``aiml`` — an unrecognized ``X-AIMLAPI-Partner-ID`` header would + just be dead weight on every other provider's request. + """ + if self._provider_info is None or self._provider_info.name != "aiml": + return {} + headers: dict[str, str] = {} + if self.aimlapi_partner_id: + headers["X-AIMLAPI-Partner-ID"] = self.aimlapi_partner_id + headers["X-AIMLAPI-Source"] = "agent" + return headers + def _aws_kwargs(self) -> dict[str, str]: """Build kwargs dict for AWS params to pass to litellm calls.""" kw: dict[str, str] = {} diff --git a/openhands-sdk/openhands/sdk/llm/options/common.py b/openhands-sdk/openhands/sdk/llm/options/common.py index d6b4602e37..30da86baf9 100644 --- a/openhands-sdk/openhands/sdk/llm/options/common.py +++ b/openhands-sdk/openhands/sdk/llm/options/common.py @@ -27,9 +27,10 @@ def apply_extra_headers(user_kwargs: dict[str, Any], llm: LLM) -> dict[str, Any] """Return a new dict with the LLM's extra headers applied. Propagates ``llm.extra_headers`` when the caller did not set its own, then - merges the OpenRouter attribution headers. These go through ``extra_headers`` - rather than ``os.environ`` so they don't leak across conversations in a - multi-tenant server (see issue #3138). User-supplied headers win. + merges the OpenRouter and AIMLAPI attribution headers. These go through + ``extra_headers`` rather than ``os.environ`` so they don't leak across + conversations in a multi-tenant server (see issue #3138). User-supplied + headers win. - Pure and deterministic; does not mutate inputs """ @@ -43,6 +44,11 @@ def apply_extra_headers(user_kwargs: dict[str, Any], llm: LLM) -> dict[str, Any] existing = out.get("extra_headers") or {} out["extra_headers"] = {**openrouter_headers, **existing} + aiml_headers = llm._aiml_headers() + if aiml_headers: + existing = out.get("extra_headers") or {} + out["extra_headers"] = {**aiml_headers, **existing} + return out diff --git a/tests/sdk/llm/test_chat_options.py b/tests/sdk/llm/test_chat_options.py index 7b8462ad74..0cbbe9ee08 100644 --- a/tests/sdk/llm/test_chat_options.py +++ b/tests/sdk/llm/test_chat_options.py @@ -29,6 +29,8 @@ class DummyLLM: _call_context: LLMCallContext = field(default_factory=LLMCallContext) openrouter_site_url: str = "" openrouter_app_name: str = "" + aimlapi_partner_id: str = "" + is_aiml_provider: bool = False def _openrouter_headers(self) -> dict[str, str]: headers: dict[str, str] = {} @@ -38,6 +40,15 @@ def _openrouter_headers(self) -> dict[str, str]: headers["X-Title"] = self.openrouter_app_name return headers + def _aiml_headers(self) -> dict[str, str]: + if not self.is_aiml_provider: + return {} + headers: dict[str, str] = {} + if self.aimlapi_partner_id: + headers["X-AIMLAPI-Partner-ID"] = self.aimlapi_partner_id + headers["X-AIMLAPI-Source"] = "agent" + return headers + def _model_name_for_capabilities(self) -> str: return self.model_canonical_name or self.model @@ -390,3 +401,28 @@ def test_chat_options_omits_openrouter_headers_when_unset(): llm = DummyLLM(model="gpt-4o") out = select_chat_options(llm, user_kwargs={}, has_tools=False) assert "extra_headers" not in out + + +def test_chat_options_injects_aimlapi_partner_headers_for_aiml_provider(): + """AIMLAPI partner attribution must flow per-call when routed via aiml.""" + llm = DummyLLM( + model="aiml/anthropic/claude-opus-5", + is_aiml_provider=True, + aimlapi_partner_id="part_uDVajKg3xPLrOdNdQetOtoGA", + ) + out = select_chat_options(llm, user_kwargs={}, has_tools=False) + assert out["extra_headers"]["X-AIMLAPI-Partner-ID"] == ( + "part_uDVajKg3xPLrOdNdQetOtoGA" + ) + assert out["extra_headers"]["X-AIMLAPI-Source"] == "agent" + + +def test_chat_options_omits_aimlapi_headers_for_non_aiml_provider(): + """The partner id must not leak onto unrelated providers as a stray header.""" + llm = DummyLLM( + model="anthropic/claude-opus-5", + is_aiml_provider=False, + aimlapi_partner_id="part_uDVajKg3xPLrOdNdQetOtoGA", + ) + out = select_chat_options(llm, user_kwargs={}, has_tools=False) + assert "extra_headers" not in out diff --git a/tests/sdk/llm/test_llm.py b/tests/sdk/llm/test_llm.py index f0c805ed21..1a27ea33cb 100644 --- a/tests/sdk/llm/test_llm.py +++ b/tests/sdk/llm/test_llm.py @@ -1910,4 +1910,27 @@ def test_max_output_tokens_not_capped_when_below_context_window( assert llm.effective_max_output_tokens == 8192 +def test_aiml_headers_present_for_aiml_provider(): + """AIMLAPI partner attribution is built for a model routed via ``aiml``.""" + llm = LLM( + model="aiml/anthropic/claude-opus-5", + api_key=SecretStr("test-key"), + usage_id="test-llm", + ) + assert llm._aiml_headers() == { + "X-AIMLAPI-Partner-ID": "part_uDVajKg3xPLrOdNdQetOtoGA", + "X-AIMLAPI-Source": "agent", + } + + +def test_aiml_headers_empty_for_non_aiml_provider(): + """The partner id must not leak onto an unrelated provider's request.""" + llm = LLM( + model="anthropic/claude-opus-5", + api_key=SecretStr("test-key"), + usage_id="test-llm", + ) + assert llm._aiml_headers() == {} + + # LLM Registry Tests