fix: fix validation for inference_key and agent_key auth

Author: dgellow

src/do_gradientai/_client.py

@@ -272,13 +272,13 @@ def default_headers(self) -> dict[str, str | Omit]:
 
     @override
     def _validate_headers(self, headers: Headers, custom_headers: Headers) -> None:
-        if self.api_key and headers.get("Authorization"):
+        if (self.api_key or self.agent_key or self.inference_key) and headers.get("Authorization"):
             return
         if isinstance(custom_headers.get("Authorization"), Omit):
             return
 
         raise TypeError(
-            '"Could not resolve authentication method. Expected the api_key to be set. Or for the `Authorization` headers to be explicitly omitted"'
+            '"Could not resolve authentication method. Expected api_key, agent_key, or inference_key to be set. Or for the `Authorization` headers to be explicitly omitted"'
         )
 
     def copy(
@@ -569,13 +569,13 @@ def default_headers(self) -> dict[str, str | Omit]:
 
     @override
     def _validate_headers(self, headers: Headers, custom_headers: Headers) -> None:
-        if self.api_key and headers.get("Authorization"):
+        if (self.api_key or self.agent_key or self.inference_key) and headers.get("Authorization"):
             return
         if isinstance(custom_headers.get("Authorization"), Omit):
             return
 
         raise TypeError(
-            '"Could not resolve authentication method. Expected the api_key to be set. Or for the `Authorization` headers to be explicitly omitted"'
+            '"Could not resolve authentication method. Expected api_key, agent_key, or inference_key to be set. Or for the `Authorization` headers to be explicitly omitted"'
         )
 
     def copy(

src/do_gradientai/resources/agents/chat/completions.py

@@ -470,6 +470,14 @@ def create(
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> CompletionCreateResponse | Stream[ChatCompletionChunk]:
+        # This method requires an agent_key to be set via client argument or environment variable
+        if not self._client.agent_key:
+            raise TypeError(
+                "Could not resolve authentication method. Expected agent_key to be set for chat completions."
+            )
+        headers = extra_headers or {}
+        headers = {"Authorization": f"Bearer {self._client.agent_key}", **headers}
+
         return self._post(
             "/chat/completions?agent=true"
             if self._client._base_url_overridden
@@ -501,7 +509,7 @@ def create(
                 else completion_create_params.CompletionCreateParamsNonStreaming,
             ),
             options=make_request_options(
-                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                extra_headers=headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
             ),
             cast_to=CompletionCreateResponse,
             stream=stream or False,
@@ -953,6 +961,14 @@ async def create(
         extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
     ) -> CompletionCreateResponse | AsyncStream[ChatCompletionChunk]:
+        # This method requires an agent_key to be set via client argument or environment variable
+        if not self._client.agent_key:
+            raise TypeError(
+                "Could not resolve authentication method. Expected agent_key to be set for chat completions."
+            )
+        headers = extra_headers or {}
+        headers = {"Authorization": f"Bearer {self._client.agent_key}", **headers}
+
         return await self._post(
             "/chat/completions?agent=true"
             if self._client._base_url_overridden
@@ -984,7 +1000,7 @@ async def create(
                 else completion_create_params.CompletionCreateParamsNonStreaming,
             ),
             options=make_request_options(
-                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                extra_headers=headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
             ),
             cast_to=CompletionCreateResponse,
             stream=stream or False,

src/do_gradientai/resources/chat/completions.py

@@ -464,7 +464,7 @@ def create(
         # This method requires an inference_key to be set via client argument or environment variable
         if not self._client.inference_key:
             raise TypeError(
-                "Could not resolve authentication method. Expected the inference_key to be set for chat completions."
+                "Could not resolve authentication method. Expected inference_key to be set for chat completions."
             )
         headers = extra_headers or {}
         headers = {"Authorization": f"Bearer {self._client.inference_key}", **headers}
@@ -946,7 +946,7 @@ async def create(
         # This method requires an inference_key to be set via client argument or environment variable
         if not hasattr(self._client, "inference_key") or not self._client.inference_key:
             raise TypeError(
-                "Could not resolve authentication method. Expected the inference_key to be set for chat completions."
+                "Could not resolve authentication method. Expected inference_key to be set for chat completions."
             )
         headers = extra_headers or {}
         headers = {"Authorization": f"Bearer {self._client.inference_key}", **headers}

tests/test_client.py

@@ -414,7 +414,7 @@ def test_validate_headers(self) -> None:
 
         with pytest.raises(
             TypeError,
-            match="Could not resolve authentication method. Expected the api_key to be set. Or for the `Authorization` headers to be explicitly omitted",
+            match="Could not resolve authentication method. Expected api_key, agent_key, or inference_key to be set. Or for the `Authorization` headers to be explicitly omitted",
         ):
             client2._build_request(FinalRequestOptions(method="get", url="/foo"))
 
@@ -1416,7 +1416,7 @@ def test_validate_headers(self) -> None:
 
         with pytest.raises(
             TypeError,
-            match="Could not resolve authentication method. Expected the api_key to be set. Or for the `Authorization` headers to be explicitly omitted",
+            match="Could not resolve authentication method. Expected api_key, agent_key, or inference_key to be set. Or for the `Authorization` headers to be explicitly omitted",
         ):
             client2._build_request(FinalRequestOptions(method="get", url="/foo"))