fix(aimlapi): send attribution headers on every aimlapi.com request (#17491)

### Summary

The **aimlapi.com** provider added in #17311 does not identify itself on
any of its outgoing requests, so its traffic cannot be attributed to the
integration. This PR adds the two headers AIMLAPI expects —
`X-AIMLAPI-Source` and `X-AIMLAPI-Partner-ID` — to every request the
provider makes.

rag/llm/chat_model.py

@@ -31,6 +31,7 @@ import openai
 from openai import AsyncOpenAI, OpenAI
 from enum import StrEnum
 
+from common.aimlapi_utils import attribution_headers
 from common.misc_utils import thread_pool_exec
 from common.llm_request_context import current_llm_user
 from common.token_utils import num_tokens_from_string, total_token_count_from_response, usage_from_response
@@ -1560,6 +1561,9 @@ class AIMLAPIChat(Base):
     def __init__(self, key, model_name, base_url="", **kwargs):
         base_url = base_url or os.environ.get("AIMLAPI_API_URL", "https://api.aimlapi.com/v1")
         super().__init__(key, model_name, base_url, **kwargs)
+        headers = attribution_headers()
+        self.client = self.client.with_options(default_headers=headers)
+        self.async_client = self.async_client.with_options(default_headers=headers)
         logging.info("[aimlapi.com] Chat initialized with model %s", model_name)
 
 

rag/llm/cv_model.py

@@ -31,6 +31,7 @@ import requests
 from openai import OpenAI, AsyncOpenAI
 from openai.lib.azure import AzureOpenAI, AsyncAzureOpenAI
 
+from common.aimlapi_utils import attribution_headers
 from common.token_utils import num_tokens_from_string, total_token_count_from_response
 from rag.nlp import is_english
 from rag.prompts.generator import vision_llm_describe_prompt
@@ -1404,6 +1405,9 @@ class AIMLAPICV(GptV4):
     def __init__(self, key, model_name, lang="Chinese", base_url="", **kwargs):
         base_url = base_url or os.environ.get("AIMLAPI_API_URL", "https://api.aimlapi.com/v1")
         super().__init__(key, model_name, lang=lang, base_url=base_url, **kwargs)
+        headers = attribution_headers()
+        self.client = self.client.with_options(default_headers=headers)
+        self.async_client = self.async_client.with_options(default_headers=headers)
         logging.info("[aimlapi.com] CV initialized with model %s", model_name)
 
 

rag/llm/embedding_model.py

@@ -29,6 +29,7 @@ from openai import OpenAI
 from zai import ZhipuAiClient
 
 from common import settings
+from common.aimlapi_utils import attribution_headers
 from common.exceptions import ModelException
 from common.token_utils import num_tokens_from_string, truncate, total_token_count_from_response
 from rag.llm.key_utils import _normalize_replicate_key
@@ -362,6 +363,7 @@ class AIMLAPIEmbed(OpenAIEmbed):
     def __init__(self, key, model_name="text-embedding-3-small", base_url=""):
         base_url = base_url or os.environ.get("AIMLAPI_API_URL", "https://api.aimlapi.com/v1")
         super().__init__(key, model_name, base_url)
+        self.client = self.client.with_options(default_headers=attribution_headers())
         logging.info("[aimlapi.com] Embedding initialized with model %s", model_name)
 
 

rag/llm/model_meta.py

@@ -14,12 +14,14 @@
 #  limitations under the License.
 #
 import json
+import logging
 import aiohttp
 from abc import ABC
 from urllib.parse import urlparse
 from json.decoder import JSONDecodeError
 from typing import ClassVar
 
+from common.aimlapi_utils import attribution_headers
 from common.constants import LLMType
 
 
@@ -594,6 +596,23 @@ class AIMLAPI(Base):
         "-vl",
     )
 
+    # aiohttp defaults to a 5-minute total timeout, long enough for a stalled catalog
+    # request to hold the task; 60s matches the timeout the other providers here use.
+    _MODEL_LIST_TIMEOUT = aiohttp.ClientTimeout(total=60)
+
+    async def _get_raw_model_list(self):
+        url = self._get_model_list_url()
+        if not url:
+            logging.warning("[aimlapi.com] Model list skipped: no base URL configured")
+            return None
+        headers = {"Authorization": f"Bearer {self._get_api_key()}", **attribution_headers()}
+        async with aiohttp.ClientSession(timeout=self._MODEL_LIST_TIMEOUT) as session:
+            async with session.get(url, headers=headers) as resp:
+                if resp.status != 200:
+                    logging.warning("[aimlapi.com] Model list request to %s failed with HTTP %s", url, resp.status)
+                    return None
+                return await resp.json()
+
     def _format_model_list(self, raw_model_list):
         models = raw_model_list.get("data") if isinstance(raw_model_list, dict) else raw_model_list
         if not isinstance(models, list):