From b65c8dcfe0c5249d06b772052963df6094296f7b Mon Sep 17 00:00:00 2001 From: jxxghp Date: Tue, 26 May 2026 08:20:02 +0800 Subject: [PATCH] feat: support llm user agent --- app/agent/__init__.py | 6 +++ app/agent/llm/helper.py | 61 ++++++++++++++++++++++++++++--- app/agent/llm/provider.py | 54 ++++++++++++++++++++++++--- app/api/endpoints/llm.py | 13 +++++++ app/core/config.py | 2 + app/schemas/event.py | 3 +- tests/test_agent_tokens_events.py | 2 + tests/test_llm_helper_testcall.py | 44 +++++++++++++++++++++- tests/test_system_llm_test.py | 28 +++++++++++++- 9 files changed, 199 insertions(+), 14 deletions(-) diff --git a/app/agent/__init__.py b/app/agent/__init__.py index 0ed59b2a..43f9b946 100644 --- a/app/agent/__init__.py +++ b/app/agent/__init__.py @@ -469,6 +469,7 @@ class MoviePilotAgent: api_key=settings.LLM_API_KEY, base_url=settings.LLM_BASE_URL, base_url_preset=settings.LLM_BASE_URL_PRESET, + user_agent=getattr(settings, "LLM_USER_AGENT", None), thinking_level=None, ) selected_event = await eventmanager.async_send_event( @@ -497,6 +498,10 @@ class MoviePilotAgent: self._clean_optional_text(self._get_event_value(resolved_data, "base_url_preset")) or settings.LLM_BASE_URL_PRESET ) + user_agent = ( + self._clean_optional_text(self._get_event_value(resolved_data, "user_agent")) + or getattr(settings, "LLM_USER_AGENT", None) + ) thinking_level = self._clean_optional_text( self._get_event_value(resolved_data, "thinking_level") ) @@ -522,6 +527,7 @@ class MoviePilotAgent: "api_key": api_key, "base_url": base_url, "base_url_preset": base_url_preset, + "user_agent": user_agent, "thinking_level": thinking_level, } return self._llm_runtime_config diff --git a/app/agent/llm/helper.py b/app/agent/llm/helper.py index 45f02332..06115982 100644 --- a/app/agent/llm/helper.py +++ b/app/agent/llm/helper.py @@ -602,6 +602,7 @@ class LLMHelper: model_name: str | None, api_key: str | None = None, base_url: str | None = None, + user_agent: str | None = None, ) -> dict[str, Any]: """ 在 provider 目录不可用时回退到旧的直接构造逻辑。 @@ -625,12 +626,36 @@ class LLMHelper: "model_id": model_name, "api_key": api_key_value, "base_url": base_url_value, - "default_headers": None, + "default_headers": LLMHelper._build_openai_default_headers( + None, + user_agent=user_agent, + ), "use_responses_api": None, "model_record": None, "model_metadata": None, } + @staticmethod + def _build_openai_default_headers( + default_headers: dict[str, str] | None = None, + user_agent: str | None = None, + ) -> dict[str, str] | None: + """ + 合并 OpenAI 兼容接口默认请求头。 + + :param default_headers: provider 运行时已解析的默认请求头 + :param user_agent: 用户配置的 User-Agent，非空时写入标准请求头 + :return: 可传给 OpenAI SDK 的请求头字典 + """ + headers = dict(default_headers or {}) + normalized_user_agent = str(user_agent or "").strip() + if normalized_user_agent: + for key in list(headers.keys()): + if key.lower() == "user-agent": + headers.pop(key) + headers["User-Agent"] = normalized_user_agent + return headers or None + @classmethod def _resolve_thinking_level( cls, @@ -675,6 +700,7 @@ class LLMHelper: api_key: str | None = None, base_url: str | None = None, base_url_preset: str | None = None, + user_agent: str | None = None, ): """ 获取LLM实例 @@ -688,6 +714,7 @@ class LLMHelper: :param api_key: API Key。未显式传入时使用当前配置项 LLM_API_KEY。对于某些提供商（如 DeepSeek），可能需要同时提供 base_url。 :param base_url: API Base URL。未显式传入时使用当前配置项 LLM_BASE_URL。 :param base_url_preset: Base URL 预设。未显式传入时使用当前配置项 LLM_BASE_URL_PRESET。 + :param user_agent: OpenAI兼容接口请求 User-Agent。未显式传入时使用配置项 LLM_USER_AGENT。 :return: LLM实例 """ provider_name = str(provider if provider is not None else settings.LLM_PROVIDER).lower() @@ -697,6 +724,9 @@ class LLMHelper: base_url_preset_value = ( base_url_preset if base_url_preset is not None else settings.LLM_BASE_URL_PRESET ) + user_agent_value = ( + user_agent if user_agent is not None else getattr(settings, "LLM_USER_AGENT", None) + ) normalized_thinking_level = cls._resolve_thinking_level( thinking_level=thinking_level, ) @@ -711,6 +741,7 @@ class LLMHelper: api_key=api_key_value, base_url=base_url_value, base_url_preset_id=base_url_preset_value, + user_agent=user_agent_value, ) except Exception as err: logger.debug(f"LLM provider 目录不可用，回退到旧运行时逻辑: {err}") @@ -719,8 +750,13 @@ class LLMHelper: model_name=model_name, api_key=api_key_value, base_url=base_url_value, + user_agent=user_agent_value, ) model_name = runtime.get("model_id") or model_name + default_headers = cls._build_openai_default_headers( + runtime.get("default_headers"), + user_agent=user_agent_value, + ) thinking_kwargs = cls._build_thinking_kwargs( provider=provider_name, model=model_name, @@ -776,7 +812,7 @@ class LLMHelper: streaming=streaming, stream_usage=True, anthropic_proxy=settings.PROXY_HOST, - default_headers=runtime.get("default_headers"), + default_headers=default_headers, **thinking_kwargs, ) else: @@ -797,7 +833,7 @@ class LLMHelper: streaming=streaming, stream_usage=True, openai_proxy=settings.PROXY_HOST, - default_headers=runtime.get("default_headers"), + default_headers=default_headers, use_responses_api=runtime.get("use_responses_api"), **thinking_kwargs, ) @@ -873,6 +909,7 @@ class LLMHelper: api_key: str | None = None, base_url: str | None = None, base_url_preset: str | None = None, + user_agent: str | None = None, ) -> dict: """ 使用当前已保存配置执行一次最小 LLM 调用。 @@ -888,6 +925,7 @@ class LLMHelper: api_key=api_key, base_url=base_url, base_url_preset=base_url_preset, + user_agent=user_agent, ) try: response = await asyncio.wait_for(llm.ainvoke(prompt), timeout=timeout) @@ -918,6 +956,7 @@ class LLMHelper: api_key: str | None = None, base_url: str | None = None, base_url_preset: str | None = None, + user_agent: str | None = None, force_refresh: bool = False, ) -> List[dict[str, Any]]: """ @@ -935,6 +974,7 @@ class LLMHelper: api_key=api_key, base_url=base_url, base_url_preset_id=base_url_preset, + user_agent=user_agent, force_refresh=force_refresh, ) except Exception as err: @@ -963,6 +1003,7 @@ class LLMHelper: provider, api_key or "", model_list_base_url, + user_agent=user_agent, ) ] @@ -997,7 +1038,10 @@ class LLMHelper: @staticmethod async def _get_openai_compatible_models( - provider: str, api_key: str, base_url: str = None + provider: str, + api_key: str, + base_url: str = None, + user_agent: str | None = None, ) -> List[str]: """获取OpenAI兼容模型列表""" try: @@ -1006,7 +1050,14 @@ class LLMHelper: if provider == "deepseek": base_url = base_url or "https://api.deepseek.com" - client = AsyncOpenAI(api_key=api_key, base_url=base_url) + client = AsyncOpenAI( + api_key=api_key, + base_url=base_url, + default_headers=LLMHelper._build_openai_default_headers( + None, + user_agent=user_agent, + ), + ) models = await client.models.list() await client.close() return [model.id for model in models.data] diff --git a/app/agent/llm/provider.py b/app/agent/llm/provider.py index 963165b2..4e2e7e6c 100644 --- a/app/agent/llm/provider.py +++ b/app/agent/llm/provider.py @@ -1166,6 +1166,23 @@ class LLMProviderManager(metaclass=Singleton): return None return value.rstrip("/") + @staticmethod + def _merge_user_agent_header( + default_headers: Optional[dict[str, str]], + user_agent: Optional[str], + ) -> Optional[dict[str, str]]: + """ + 合并用户配置的 OpenAI 兼容接口 User-Agent 请求头。 + """ + headers = dict(default_headers or {}) + normalized_user_agent = str(user_agent or "").strip() + if normalized_user_agent: + for key in list(headers.keys()): + if key.lower() == "user-agent": + headers.pop(key) + headers["User-Agent"] = normalized_user_agent + return headers or None + @classmethod def _default_base_url_for_provider(cls, spec: ProviderSpec) -> Optional[str]: """获取 provider 的默认 Base URL。""" @@ -1825,6 +1842,7 @@ class LLMProviderManager(metaclass=Singleton): api_key: Optional[str] = None, base_url: Optional[str] = None, base_url_preset_id: Optional[str] = None, + user_agent: Optional[str] = None, force_refresh: bool = False, ) -> list[dict[str, Any]]: """返回标准化后的模型目录。""" @@ -1854,6 +1872,7 @@ class LLMProviderManager(metaclass=Singleton): api_key=api_key, base_url=base_url, base_url_preset_id=base_url_preset_id, + user_agent=user_agent, ) if resolved_model_list_strategy == "google": @@ -1877,7 +1896,10 @@ class LLMProviderManager(metaclass=Singleton): runtime["base_url"], base_url_preset_id=base_url_preset_id, ), - default_headers=runtime.get("default_headers"), + default_headers=self._merge_user_agent_header( + runtime.get("default_headers"), + user_agent, + ), ) if resolved_model_list_strategy == "anthropic_compatible": @@ -1905,7 +1927,10 @@ class LLMProviderManager(metaclass=Singleton): runtime["base_url"], base_url_preset_id=base_url_preset_id, ), - default_headers=runtime.get("default_headers"), + default_headers=self._merge_user_agent_header( + runtime.get("default_headers"), + user_agent, + ), ) async def resolve_model_metadata( @@ -2398,6 +2423,7 @@ class LLMProviderManager(metaclass=Singleton): api_key: Optional[str] = None, base_url: Optional[str] = None, base_url_preset_id: Optional[str] = None, + user_agent: Optional[str] = None, ) -> dict[str, Any]: """ 解析 provider 运行时参数。 @@ -2428,6 +2454,7 @@ class LLMProviderManager(metaclass=Singleton): api_key=api_key, base_url=base_url, base_url_preset_id=normalized_base_url_preset_id, + user_agent=user_agent, ) if item["id"] == model ), @@ -2470,7 +2497,10 @@ class LLMProviderManager(metaclass=Singleton): "runtime": "chatgpt", "api_key": auth["access_token"], "base_url": self._CHATGPT_CODEX_BASE_URL, - "default_headers": headers, + "default_headers": self._merge_user_agent_header( + headers, + user_agent, + ), "use_responses_api": True, "auth_mode": "oauth", } @@ -2484,6 +2514,10 @@ class LLMProviderManager(metaclass=Singleton): "api_key": normalized_api_key, "base_url": normalized_base_url or self._default_base_url_for_provider(spec), + "default_headers": self._merge_user_agent_header( + None, + user_agent, + ), "auth_mode": "api_key", } ) @@ -2508,9 +2542,12 @@ class LLMProviderManager(metaclass=Singleton): else "github_copilot", "api_key": token, "base_url": "https://api.githubcopilot.com", - "default_headers": self._copilot_headers( - token, - include_auth=transport == "anthropic", + "default_headers": self._merge_user_agent_header( + self._copilot_headers( + token, + include_auth=transport == "anthropic", + ), + user_agent, ), "auth_mode": "oauth" if auth else "api_key", } @@ -2543,6 +2580,10 @@ class LLMProviderManager(metaclass=Singleton): "base_url": self._normalize_base_url_for_anthropic( effective_base_url ), + "default_headers": self._merge_user_agent_header( + None, + user_agent, + ), "auth_mode": "api_key", } ) @@ -2557,6 +2598,7 @@ class LLMProviderManager(metaclass=Singleton): { "api_key": normalized_api_key, "base_url": effective_base_url, + "default_headers": self._merge_user_agent_header(None, user_agent), "auth_mode": "api_key", } ) diff --git a/app/api/endpoints/llm.py b/app/api/endpoints/llm.py index 2c1b1bd9..de6df61d 100644 --- a/app/api/endpoints/llm.py +++ b/app/api/endpoints/llm.py @@ -24,6 +24,10 @@ router = APIRouter() class LlmTestRequest(BaseModel): + """ + LLM 测试调用请求参数。 + """ + enabled: Optional[bool] = None provider: Optional[str] = None model: Optional[str] = None @@ -31,9 +35,14 @@ class LlmTestRequest(BaseModel): api_key: Optional[str] = None base_url: Optional[str] = None base_url_preset: Optional[str] = None + user_agent: Optional[str] = None class LlmProviderAuthStartRequest(BaseModel): + """ + LLM 提供商授权启动请求参数。 + """ + provider: str method: str @@ -67,6 +76,7 @@ async def get_llm_models( api_key: Optional[str] = None, base_url: Optional[str] = None, base_url_preset: Optional[str] = None, + user_agent: Optional[str] = None, force_refresh: Optional[bool] = False, _: User = Depends(get_current_active_user_async), ): @@ -80,6 +90,7 @@ async def get_llm_models( api_key=api_key, base_url=base_url, base_url_preset=base_url_preset, + user_agent=user_agent, force_refresh=bool(force_refresh), ) return schemas.Response( @@ -238,6 +249,7 @@ async def llm_test( api_key=settings.LLM_API_KEY, base_url=settings.LLM_BASE_URL, base_url_preset=settings.LLM_BASE_URL_PRESET, + user_agent=getattr(settings, "LLM_USER_AGENT", None), ) if not payload.provider: @@ -269,6 +281,7 @@ async def llm_test( api_key=payload.api_key, base_url=payload.base_url, base_url_preset=payload.base_url_preset, + user_agent=payload.user_agent, ) if not result.get("reply_preview"): return schemas.Response( diff --git a/app/core/config.py b/app/core/config.py index 43df19f2..d0c51d67 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -564,6 +564,8 @@ class ConfigModel(BaseModel): LLM_BASE_URL: Optional[str] = "https://api.deepseek.com" # LLM Base URL 预设标识，用于区分同一 Base URL 下的不同模型目录 LLM_BASE_URL_PRESET: Optional[str] = None + # LLM OpenAI兼容接口请求User-Agent + LLM_USER_AGENT: Optional[str] = None # LLM最大上下文Token数量（K） LLM_MAX_CONTEXT_TOKENS: int = 64 # LLM温度参数 diff --git a/app/schemas/event.py b/app/schemas/event.py index 201c82de..3da6761b 100644 --- a/app/schemas/event.py +++ b/app/schemas/event.py @@ -69,7 +69,7 @@ class AgentLLMProviderEventData(ChainEventData): Agent LLM 供应商选择事件数据。 事件发出方会带入当前系统配置作为默认值；插件可覆盖 provider、base_url、 - api_key、model 等字段，并通过 selected_provider_id 标记本次选择，方便 + api_key、model、user_agent 等字段，并通过 selected_provider_id 标记本次选择，方便 后续用量事件精确回写到同一个配额条目。 """ @@ -78,6 +78,7 @@ class AgentLLMProviderEventData(ChainEventData): api_key: Optional[str] = Field(default=None, description="API Key") model: Optional[str] = Field(default=None, description="模型名称") base_url_preset: Optional[str] = Field(default=None, description="Base URL 预设ID") + user_agent: Optional[str] = Field(default=None, description="OpenAI兼容接口User-Agent") thinking_level: Optional[str] = Field(default=None, description="思考模式级别") selected_provider_id: Optional[str] = Field(default=None, description="插件侧供应商ID") selected_provider_name: Optional[str] = Field(default=None, description="插件侧供应商名称") diff --git a/tests/test_agent_tokens_events.py b/tests/test_agent_tokens_events.py index e1c15882..3da7067f 100644 --- a/tests/test_agent_tokens_events.py +++ b/tests/test_agent_tokens_events.py @@ -79,6 +79,7 @@ class AgentTokensEventsTest(unittest.IsolatedAsyncioTestCase): data.api_key = "sk-agent-token" data.model = "free-model" data.base_url_preset = None + data.user_agent = "AgentTokens-UA/1.0" data.selected_provider_id = "provider-1" data.selected_provider_name = "Free Provider" data.source = "AgentTokens" @@ -105,6 +106,7 @@ class AgentTokensEventsTest(unittest.IsolatedAsyncioTestCase): api_key="sk-agent-token", base_url="https://tokens.example.com/v1", base_url_preset=None, + user_agent="AgentTokens-UA/1.0", thinking_level=None, ) self.assertEqual("provider-1", agent._llm_provider_selection["selected_provider_id"]) diff --git a/tests/test_llm_helper_testcall.py b/tests/test_llm_helper_testcall.py index 555e7346..98b2da01 100644 --- a/tests/test_llm_helper_testcall.py +++ b/tests/test_llm_helper_testcall.py @@ -118,6 +118,10 @@ def _build_fake_openai_modules(chat_openai_cls=_FakeChatOpenAIForPatch): }, base_module +_ORIGINAL_STUBBED_MODULES = { + name: sys.modules.get(name) + for name in ("app.core.config", "app.log") +} sys.modules.pop("app.agent.llm.helper", None) _stub_module( "app.core.config", @@ -127,6 +131,7 @@ _stub_module( LLM_API_KEY="global-key", LLM_BASE_URL="https://global.example.com", LLM_BASE_URL_PRESET=None, + LLM_USER_AGENT=None, LLM_THINKING_LEVEL=None, LLM_TEMPERATURE=0.1, LLM_MAX_CONTEXT_TOKENS=64, @@ -140,6 +145,11 @@ spec = importlib.util.spec_from_file_location("test_llm_module", module_path) llm_module = importlib.util.module_from_spec(spec) assert spec and spec.loader spec.loader.exec_module(llm_module) +for _module_name, _module in _ORIGINAL_STUBBED_MODULES.items(): + if _module is None: + sys.modules.pop(_module_name, None) + else: + sys.modules[_module_name] = _module class LlmHelperTestCallTest(unittest.TestCase): @@ -177,6 +187,7 @@ class LlmHelperTestCallTest(unittest.TestCase): api_key="sk-test", base_url="https://api.deepseek.com", base_url_preset="deepseek-default", + user_agent=None, ) self.assertEqual(result["provider"], "deepseek") self.assertEqual(result["model"], "deepseek-chat") @@ -436,7 +447,7 @@ class LlmHelperTestCallTest(unittest.TestCase): "model_id": kwargs["model"], "api_key": kwargs["api_key"], "base_url": kwargs["base_url"], - "default_headers": None, + "default_headers": {"X-Test": "1"}, "use_responses_api": None, "model_record": None, "model_metadata": None, @@ -477,6 +488,7 @@ class LlmHelperTestCallTest(unittest.TestCase): "api_key": "updated-key", "base_url": "https://updated.example.com/v1", "base_url_preset_id": "updated-preset", + "user_agent": None, }, ) self.assertEqual(len(llm_calls), 1) @@ -486,6 +498,36 @@ class LlmHelperTestCallTest(unittest.TestCase): llm_calls[0].get("base_url"), "https://updated.example.com/v1", ) + self.assertEqual(llm_calls[0].get("default_headers"), {"X-Test": "1"}) + + def test_get_llm_passes_user_agent_as_openai_default_header(self): + calls = [] + + class _FakeChatOpenAI: + def __init__(self, **kwargs): + calls.append(kwargs) + self.model = kwargs["model"] + self.profile = None + + with patch.dict( + sys.modules, + {"langchain_openai": SimpleNamespace(ChatOpenAI=_FakeChatOpenAI)}, + ): + asyncio.run( + llm_module.LLMHelper.get_llm( + provider="openai", + model="gpt-5-mini", + api_key="sk-test", + base_url="https://api.example.com/v1", + user_agent="MoviePilot-Test/1.0", + ) + ) + + self.assertEqual(len(calls), 1) + self.assertEqual( + calls[0].get("default_headers"), + {"User-Agent": "MoviePilot-Test/1.0"}, + ) def test_get_llm_keeps_openai_patch_global_without_model_marker(self): class _FakeProviderManager: diff --git a/tests/test_system_llm_test.py b/tests/test_system_llm_test.py index 7f75aa3f..bccc400b 100644 --- a/tests/test_system_llm_test.py +++ b/tests/test_system_llm_test.py @@ -5,7 +5,12 @@ from types import ModuleType from unittest.mock import AsyncMock, patch +_ORIGINAL_STUBBED_MODULES = {} + + def _stub_module(name: str, **attrs): + if name not in _ORIGINAL_STUBBED_MODULES: + _ORIGINAL_STUBBED_MODULES[name] = sys.modules.get(name) module = sys.modules.get(name) if module is None: module = ModuleType(name) @@ -36,7 +41,15 @@ _stub_module("app.helper.sites", SitesHelper=_Dummy) _stub_module("app.chain.mediaserver", MediaServerChain=_Dummy) _stub_module("app.chain.search", SearchChain=_Dummy) _stub_module("app.chain.system", SystemChain=_Dummy) -_stub_module("app.core.event", eventmanager=_Dummy()) +_stub_module( + "app.agent.llm", + LLMHelper=_Dummy, + LLMProviderManager=_Dummy, + LLMTestError=_DummyError, + LLMTestTimeout=_DummyError, + render_auth_result_html=lambda success, message: message, +) +_stub_module("app.core.event", eventmanager=_Dummy(), Event=_Dummy, EventManager=_Dummy) _stub_module("app.core.metainfo", MetaInfo=_Dummy) _stub_module("app.core.module", ModuleManager=_Dummy) _stub_module( @@ -79,6 +92,12 @@ _stub_module("version", APP_VERSION="test") from app.api.endpoints import llm as system_endpoint +for _module_name, _module in _ORIGINAL_STUBBED_MODULES.items(): + if _module is None: + sys.modules.pop(_module_name, None) + else: + sys.modules[_module_name] = _module + class LlmTestEndpointTest(unittest.TestCase): def test_llm_test_requires_ai_agent_enabled(self): @@ -126,6 +145,8 @@ class LlmTestEndpointTest(unittest.TestCase): system_endpoint.settings, "LLM_BASE_URL", "https://api.deepseek.com" ), patch.object( system_endpoint.settings, "LLM_BASE_URL_PRESET", "deepseek-default" + ), patch.object( + system_endpoint.settings, "LLM_USER_AGENT", "MoviePilot-Test/1.0", create=True ), patch.object( system_endpoint.LLMHelper, "test_current_settings", @@ -141,6 +162,7 @@ class LlmTestEndpointTest(unittest.TestCase): api_key="sk-test", base_url="https://api.deepseek.com", base_url_preset="deepseek-default", + user_agent="MoviePilot-Test/1.0", ) self.assertTrue(resp.success) self.assertEqual(resp.data["provider"], "deepseek") @@ -165,6 +187,7 @@ class LlmTestEndpointTest(unittest.TestCase): api_key="sk-live", base_url="https://example.com/v1", base_url_preset="openai-default", + user_agent="MoviePilot-Custom/1.0", ) with patch.object(system_endpoint.settings, "AI_AGENT_ENABLE", False), patch.object( @@ -188,6 +211,7 @@ class LlmTestEndpointTest(unittest.TestCase): api_key="sk-live", base_url="https://example.com/v1", base_url_preset="openai-default", + user_agent="MoviePilot-Custom/1.0", ) self.assertTrue(resp.success) self.assertEqual(resp.data["provider"], "openai") @@ -209,6 +233,7 @@ class LlmTestEndpointTest(unittest.TestCase): api_key="sk-live", base_url="https://api.deepseek.com", base_url_preset="deepseek-default", + user_agent=None, ) with patch.object(system_endpoint.settings, "AI_AGENT_ENABLE", False), patch.object( @@ -226,6 +251,7 @@ class LlmTestEndpointTest(unittest.TestCase): api_key="sk-live", base_url="https://api.deepseek.com", base_url_preset="deepseek-default", + user_agent=None, ) self.assertTrue(resp.success)