From 071a6cea2aab6ed8f636d6db63763c3fd476283c Mon Sep 17 00:00:00 2001 From: Yijia-Xiao Date: Sat, 3 Oct 2026 22:44:32 +0000 Subject: [PATCH] refactor(llm): let a provider sign in with its own credential at a pinned endpoint - a provider entry can supply a credential, a token callable per request, and pin its base URL - such a provider never sends an API key, and refuses another backend_url or an explicit api_key --- tests/test_provider_registry.py | 56 ++++++++++++++++++++++ tradingagents/llm_clients/openai_client.py | 22 ++++++++- 2 files changed, 77 insertions(+), 1 deletion(-) diff --git a/tests/test_provider_registry.py b/tests/test_provider_registry.py index 36bd4fcbb..69e2b053b 100644 --- a/tests/test_provider_registry.py +++ b/tests/test_provider_registry.py @@ -58,3 +58,59 @@ def test_key_optionality(): assert OPENAI_COMPATIBLE_PROVIDERS["xai"].key_optional is False # OLLAMA_BASE_URL is the only base-URL env override. assert OPENAI_COMPATIBLE_PROVIDERS["ollama"].base_url_env == "OLLAMA_BASE_URL" + + +# --- a provider that signs in: its own credential, its own pinned endpoint --------- + +class _Recorder: + """Stands in for the chat class and records what the client built.""" + + def __init__(self, **kwargs): + self.kwargs = kwargs + + +def _signed_in_provider(monkeypatch, token="plan-token"): + from tradingagents.llm_clients import openai_client + + spec = openai_client.ProviderSpec( + base_url="https://api.openai.com/v1", pinned_base_url=True, use_responses_api=True, + credential=lambda: (lambda: token), chat_class=_Recorder, + ) + monkeypatch.setitem(openai_client.OPENAI_COMPATIBLE_PROVIDERS, "signed-in", spec) + return openai_client + + +@pytest.mark.unit +def test_a_signed_in_provider_sends_its_own_credential_never_an_api_key(monkeypatch): + openai_client = _signed_in_provider(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "sk-must-not-be-used") + llm = openai_client.OpenAIClient("gpt-x", provider="signed-in").get_llm() + assert callable(llm.kwargs["api_key"]) and llm.kwargs["api_key"]() == "plan-token" + assert llm.kwargs["base_url"] == "https://api.openai.com/v1" + + +@pytest.mark.unit +@pytest.mark.parametrize("client_kwargs", [ + {"base_url": "https://relay.example/v1"}, # backend_url / TRADINGAGENTS_LLM_BACKEND_URL + {"api_key": "sk-explicit"}, +]) +def test_a_signed_in_provider_refuses_another_endpoint_or_key(monkeypatch, client_kwargs): + openai_client = _signed_in_provider(monkeypatch) + base_url = client_kwargs.pop("base_url", None) + client = openai_client.OpenAIClient("gpt-x", base_url=base_url, provider="signed-in", **client_kwargs) + with pytest.raises(ValueError): + client.get_llm() + + +@pytest.mark.unit +def test_a_missing_sign_in_fails_before_any_request(monkeypatch): + from tradingagents.llm_clients import openai_client + + def not_signed_in(): + raise ValueError("Not signed in; run `tradingagents login chatgpt`.") + + spec = openai_client.ProviderSpec(base_url="https://api.openai.com/v1", pinned_base_url=True, + credential=not_signed_in, chat_class=_Recorder) + monkeypatch.setitem(openai_client.OPENAI_COMPATIBLE_PROVIDERS, "signed-in", spec) + with pytest.raises(ValueError, match="login chatgpt"): + openai_client.OpenAIClient("gpt-x", provider="signed-in").get_llm() diff --git a/tradingagents/llm_clients/openai_client.py b/tradingagents/llm_clients/openai_client.py index b1d7b65d6..56dc50845 100644 --- a/tradingagents/llm_clients/openai_client.py +++ b/tradingagents/llm_clients/openai_client.py @@ -1,5 +1,6 @@ import os import re +from collections.abc import Callable from dataclasses import dataclass from typing import Any from urllib.parse import urlparse @@ -204,6 +205,11 @@ class ProviderSpec: placeholder_key: str = "EMPTY" # sent when no key is available (keyless local servers) require_base_url: bool = False # error if no base_url is resolved (generic endpoint) use_responses_api: bool = False # native OpenAI Responses API + # A provider the user signs in to: ``credential()`` returns the api_key callable + # (a fresh token per request) or raises when nobody is signed in. Its token goes + # only to its pinned base_url, and never alongside an API key. + credential: Callable[[], Callable[[], str]] | None = None + pinned_base_url: bool = False # refuse any other endpoint for this provider # Single source of truth for the OpenAI-compatible provider family. Dual-region @@ -234,6 +240,14 @@ OPENAI_COMPATIBLE_PROVIDERS: dict[str, ProviderSpec] = { } +def check_endpoint(provider: str, base_url: str | None) -> None: + """Raise ValueError when a provider pinned to its own endpoint is pointed elsewhere.""" + spec = OPENAI_COMPATIBLE_PROVIDERS.get(provider.lower()) + if spec is not None and spec.pinned_base_url and base_url and base_url != spec.base_url: + raise ValueError(f"Provider '{provider}' only connects to {spec.base_url}; " + "unset backend_url / TRADINGAGENTS_LLM_BACKEND_URL for it.") + + def is_openai_compatible(provider: str) -> bool: """Whether ``provider`` is served by the OpenAI-compatible registry.""" return provider.lower() in OPENAI_COMPATIBLE_PROVIDERS @@ -289,6 +303,7 @@ class OpenAIClient(BaseLLMClient): # OLLAMA_BASE_URL) > provider default. None means use the SDK default. env_base_url = os.environ.get(spec.base_url_env) if spec.base_url_env else None base_url = self.base_url or env_base_url or spec.base_url + check_endpoint(self.provider, base_url) if spec.require_base_url and not base_url: raise ValueError( f"Provider '{self.provider}' requires a base_url. Set it via " @@ -301,9 +316,14 @@ class OpenAIClient(BaseLLMClient): # API key: required unless key_optional; keyless local servers get a # placeholder. The env-var name is the single source in api_key_env. + # A signed-in provider sends its own credential, and only that. api_key_env = get_api_key_env(self.provider) api_key = os.environ.get(api_key_env) if api_key_env else None - if api_key: + if spec.credential is not None: + if "api_key" in self.kwargs: + raise ValueError(f"Provider '{self.provider}' uses its sign-in, not an api_key.") + llm_kwargs["api_key"] = spec.credential() + elif api_key: llm_kwargs["api_key"] = api_key elif spec.key_optional: llm_kwargs["api_key"] = spec.placeholder_key