feat(llm): add public api_key property for in-place credential rotation

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,

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."""

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

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(

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

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

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"