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.

api/apps/restful_apis/aimlapi_api.py

@@ -25,6 +25,7 @@ API key is returned to the browser, to be filled into the provider form.
 Configuration (all via environment):
   AIMLAPI_APP_URL                  agent-auth host (default https://app.aimlapi.com)
   AIMLAPI_PARTNER_ID               partner id (defaults to RAGFlow's; override to change)
+  AIMLAPI_SOURCE                   client identifier sent on every request (default agent/ragflow)
   AIMLAPI_PARTNER_NAME             display name shown on the consent page
   AIMLAPI_VERIFICATION_BASE_URL    base of the consent page (default https://aimlapi.com)
   AIMLAPI_REQUESTED_USD_LIMIT_MINOR requested spend cap in USD minor units (default 1000)
@@ -34,11 +35,13 @@ import asyncio
 import json
 import logging
 import os
+from urllib.parse import urlencode
 
 import requests
 
 from api.apps import login_required, current_user
 from api.utils.api_utils import get_data_error_result, get_json_result, get_request_json, server_error_response, validate_request
+from common.aimlapi_utils import attribution_headers, partner_id as _partner_id, source as _source
 from rag.utils.redis_conn import REDIS_CONN
 
 LOGGER = logging.getLogger(__name__)
@@ -53,10 +56,6 @@ def _app_url() -> str:
     return os.environ.get("AIMLAPI_APP_URL", "https://app.aimlapi.com").rstrip("/")
 
 
-def _partner_id() -> str:
-    return os.environ.get("AIMLAPI_PARTNER_ID", "part_yNkcOvbGLtgxWLjy4sRysaer").strip()
-
-
 def _partner_name() -> str:
     return os.environ.get("AIMLAPI_PARTNER_NAME", "RAGFlow").strip()
 
@@ -98,6 +97,7 @@ async def aimlapi_authorize_start():
             requests.post,
             f"{_app_url()}/v3/agent-auth/authorizations",
             json=payload,
+            headers=attribution_headers(),
             timeout=_HTTP_TIMEOUT,
         )
     except Exception as e:
@@ -131,7 +131,10 @@ async def aimlapi_authorize_start():
     LOGGER.info("AIMLAPI authorize start ok: request_id=%s", request_id)
 
     # Rebuild the consent URL from AIMLAPI_VERIFICATION_BASE_URL so an env override applies; the create response returns an absolute URL.
-    verification_uri = f"{_verification_base_url()}/agent/authorize?request={request_id}"
+    # `source` carries the same client identifier as the request headers: if the user signs up while consenting, that happens in a browser
+    # flow on AIMLAPI's side, where a request header cannot reach — the query parameter is the only way to attribute the new account.
+    query = urlencode({"request": request_id, "source": _source()})
+    verification_uri = f"{_verification_base_url()}/agent/authorize?{query}"
     return get_json_result(
         data={
             "request_id": request_id,
@@ -170,6 +173,7 @@ async def aimlapi_authorize_poll():
             requests.post,
             f"{_app_url()}/v3/agent-auth/token",
             json=payload,
+            headers=attribution_headers(),
             timeout=_HTTP_TIMEOUT,
         )
     except Exception as e:

common/aimlapi_utils.py

@@ -0,0 +1,49 @@
+#
+#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+"""Shared request attribution for the aimlapi.com provider.
+
+Every call RAGFlow makes to AIMLAPI — chat, embeddings, vision, the model
+catalog and the agent-authorization endpoints — carries the same two headers:
+they tell AIMLAPI which client the traffic comes from and which integration it
+belongs to. The same client identifier also rides on the consent URL, where a
+header cannot reach (the account is created in a browser AIMLAPI controls).
+Defaults are RAGFlow's own; both are overridable via environment:
+
+  AIMLAPI_SOURCE       client identifier (default ``agent/ragflow``)
+  AIMLAPI_PARTNER_ID   integration id (defaults to RAGFlow's)
+"""
+
+import os
+
+DEFAULT_SOURCE = "agent/ragflow"
+DEFAULT_PARTNER_ID = "part_yNkcOvbGLtgxWLjy4sRysaer"
+
+
+def source() -> str:
+    return os.environ.get("AIMLAPI_SOURCE", DEFAULT_SOURCE).strip()
+
+
+def partner_id() -> str:
+    return os.environ.get("AIMLAPI_PARTNER_ID", DEFAULT_PARTNER_ID).strip()
+
+
+def attribution_headers() -> dict:
+    """Headers to send on every request to aimlapi.com."""
+    headers = {"X-AIMLAPI-Source": source()}
+    pid = partner_id()
+    if pid:
+        headers["X-AIMLAPI-Partner-ID"] = pid
+    return headers

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):

test/unit_test/common/test_aimlapi_utils.py

@@ -0,0 +1,32 @@
+from common.aimlapi_utils import DEFAULT_PARTNER_ID, DEFAULT_SOURCE, attribution_headers, partner_id, source
+
+
+def test_defaults(monkeypatch):
+    monkeypatch.delenv("AIMLAPI_SOURCE", raising=False)
+    monkeypatch.delenv("AIMLAPI_PARTNER_ID", raising=False)
+
+    assert source() == DEFAULT_SOURCE
+    assert attribution_headers() == {
+        "X-AIMLAPI-Source": DEFAULT_SOURCE,
+        "X-AIMLAPI-Partner-ID": DEFAULT_PARTNER_ID,
+    }
+
+
+def test_environment_overrides(monkeypatch):
+    monkeypatch.setenv("AIMLAPI_SOURCE", " agent/custom ")
+    monkeypatch.setenv("AIMLAPI_PARTNER_ID", " part_custom ")
+
+    assert source() == "agent/custom"
+    assert partner_id() == "part_custom"
+    assert attribution_headers() == {
+        "X-AIMLAPI-Source": "agent/custom",
+        "X-AIMLAPI-Partner-ID": "part_custom",
+    }
+
+
+def test_blank_partner_id_is_omitted(monkeypatch):
+    monkeypatch.delenv("AIMLAPI_SOURCE", raising=False)
+    monkeypatch.setenv("AIMLAPI_PARTNER_ID", "  ")
+
+    headers = attribution_headers()
+    assert headers == {"X-AIMLAPI-Source": DEFAULT_SOURCE}