diff --git a/api/apps/restful_apis/aimlapi_api.py b/api/apps/restful_apis/aimlapi_api.py index c1874bab19..3bd2bfbeca 100644 --- a/api/apps/restful_apis/aimlapi_api.py +++ b/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: diff --git a/common/aimlapi_utils.py b/common/aimlapi_utils.py new file mode 100644 index 0000000000..0d653008ad --- /dev/null +++ b/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 diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py index 1783b858ee..11274bb948 100644 --- a/rag/llm/chat_model.py +++ b/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) diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py index 14ef9557b1..c668186e10 100644 --- a/rag/llm/cv_model.py +++ b/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) diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index 005053366f..c182631d50 100644 --- a/rag/llm/embedding_model.py +++ b/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) diff --git a/rag/llm/model_meta.py b/rag/llm/model_meta.py index 6a91d4b6ea..cb518a3f36 100644 --- a/rag/llm/model_meta.py +++ b/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): diff --git a/test/unit_test/common/test_aimlapi_utils.py b/test/unit_test/common/test_aimlapi_utils.py new file mode 100644 index 0000000000..c307e6fa2f --- /dev/null +++ b/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}