diff --git a/nemoguardrails/llm/clients/base.py b/nemoguardrails/llm/clients/base.py index d70dcf8edd..d89a8b9620 100644 --- a/nemoguardrails/llm/clients/base.py +++ b/nemoguardrails/llm/clients/base.py @@ -137,6 +137,23 @@ def provider_name(self) -> Optional[str]: def provider_url(self) -> Optional[str]: return None + @property + def api_key(self) -> Optional[str]: + """The bearer token/API key used for the Authorization header.""" + return self._api_key + + @api_key.setter + def api_key(self, value: Optional[str]) -> None: + """Update the API key in place. + + Safe to call between requests on an in-flight client: headers are + rebuilt fresh per call in ``_build_headers()``, so this only affects + requests started after the assignment. Intended for callers that hold + one long-lived client/model across many short-lived credential + rotations (e.g. OAuth client-credentials tokens). + """ + self._api_key = value + def _error_context(self) -> ErrorContext: return ErrorContext( model_name=None, diff --git a/nemoguardrails/llm/models/instrumented.py b/nemoguardrails/llm/models/instrumented.py index 617d6671a9..b4fb9f37f1 100644 --- a/nemoguardrails/llm/models/instrumented.py +++ b/nemoguardrails/llm/models/instrumented.py @@ -93,6 +93,20 @@ def provider_name(self) -> Optional[str]: def provider_url(self) -> Optional[str]: return self._model.provider_url + @property + def api_key(self) -> Optional[str]: + """The wrapped model's bearer token/API key, if it exposes one. + + Raises ``AttributeError`` for wrapped models without one, so + ``hasattr(rails.llm, "api_key")`` still reports support correctly + through this decorator. + """ + return getattr(self._model, "api_key") + + @api_key.setter + def api_key(self, value: Optional[str]) -> None: + setattr(self._model, "api_key", value) + @property def wrapped_model(self) -> LLMModel: """Return the underlying model for direct access or re-instrumentation.""" diff --git a/nemoguardrails/llm/models/openai_chat.py b/nemoguardrails/llm/models/openai_chat.py index 3620a1b5ad..ee5604f6c2 100644 --- a/nemoguardrails/llm/models/openai_chat.py +++ b/nemoguardrails/llm/models/openai_chat.py @@ -73,6 +73,15 @@ def provider_name(self) -> str: def provider_url(self) -> Optional[str]: return self._client.provider_url + @property + def api_key(self) -> Optional[str]: + """The bearer token/API key used by the underlying HTTP client.""" + return self._client.api_key + + @api_key.setter + def api_key(self, value: Optional[str]) -> None: + self._client.api_key = value + def _enrich(self, exc: LLMClientError) -> LLMClientError: exc.provider_name = self._provider_name exc.model_name = self._model diff --git a/nemoguardrails/types.py b/nemoguardrails/types.py index 4ff3d61889..ec76448aa1 100644 --- a/nemoguardrails/types.py +++ b/nemoguardrails/types.py @@ -252,6 +252,10 @@ class LLMModel(Protocol): objects. Adapters convert ``ChatMessage`` to whatever their SDK expects. ``**kwargs`` are forwarded to the underlying SDK (e.g. temperature, max_tokens). + + Bearer-token backends (e.g. ``OpenAIChatModel``) additionally expose a + settable ``api_key`` property for in-place credential rotation; this is + not part of the protocol, so check with ``hasattr`` before relying on it. """ async def generate_async( diff --git a/tests/llm/clients/test_client_config.py b/tests/llm/clients/test_client_config.py index 4f3c48c6dc..c2fa3690be 100644 --- a/tests/llm/clients/test_client_config.py +++ b/tests/llm/clients/test_client_config.py @@ -101,6 +101,26 @@ async def test_stored(self): assert client._custom_query == {"api-version": "2024-02-01"} +class TestApiKey: + @pytest.mark.asyncio + async def test_getter_returns_constructor_value(self): + async with _make_client() as client: + assert client.api_key == "sk-test" + + @pytest.mark.asyncio + async def test_setter_updates_value(self): + async with _make_client() as client: + client.api_key = "sk-rotated" + assert client.api_key == "sk-rotated" + + @pytest.mark.asyncio + async def test_setter_reflected_in_next_request_headers(self): + async with _make_client() as client: + client.api_key = "sk-rotated" + headers = client._build_headers() + assert headers["Authorization"] == "Bearer sk-rotated" + + class TestHttpClientInjection: @pytest.mark.asyncio async def test_uses_injected_client(self): diff --git a/tests/llm/models/test_instrumented.py b/tests/llm/models/test_instrumented.py index 8ce80fab79..61da14468c 100644 --- a/tests/llm/models/test_instrumented.py +++ b/tests/llm/models/test_instrumented.py @@ -358,6 +358,25 @@ async def test_instrumentation_is_idempotent_and_does_not_own_model(span_exporte assert not hasattr(first, "aclose") +def test_api_key_delegates_to_wrapped_model_when_supported(): + class KeyedModel(RecordingModel): + api_key = "sk-initial" + + model = KeyedModel() + instrumented = InstrumentedLLMModel(model, metrics_enabled=True) + + assert instrumented.api_key == "sk-initial" + + instrumented.api_key = "sk-rotated" + assert model.api_key == "sk-rotated" + + +def test_api_key_hasattr_false_when_wrapped_model_lacks_it(): + instrumented = InstrumentedLLMModel(RecordingModel(), metrics_enabled=True) + + assert not hasattr(instrumented, "api_key") + + @pytest.mark.asyncio async def test_stream_cleanup_runs_outside_duration_metric(metric_reader): close_delay = 0.2 diff --git a/tests/llm/models/test_openai_chat.py b/tests/llm/models/test_openai_chat.py index 6dee45a881..c9304d0a78 100644 --- a/tests/llm/models/test_openai_chat.py +++ b/tests/llm/models/test_openai_chat.py @@ -872,3 +872,15 @@ def test_provider_url_delegated(self): mc.provider_url = "https://example.com/v1" m = _model(mc) assert m.provider_url == "https://example.com/v1" + + def test_api_key_getter_delegates_to_client(self): + mc = _mock_client() + mc.api_key = "sk-initial" + m = _model(mc) + assert m.api_key == "sk-initial" + + def test_api_key_setter_delegates_to_client(self): + mc = _mock_client() + m = _model(mc) + m.api_key = "sk-rotated" + assert mc.api_key == "sk-rotated"