From d579842d3427c698ef6273e5d450fd5224deecab Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 12 May 2026 13:19:36 -0500 Subject: [PATCH 1/3] feat: add `api_headers` parameter and callable `api_key` support (#190) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two improvements for custom authentication on OpenAI-compatible providers: 1. New `api_headers` parameter — accepts a dict of HTTP headers or a zero-argument callable returning one. Applied via `extra_headers` on every API call, enabling token refresh and other dynamic auth patterns. 2. Widen `api_key` to accept `Callable[[], str]` — passed straight through to the openai SDK, which natively supports callable API keys with per-request resolution. Both parameters are added to all OpenAI-compatible Chat functions (ChatOpenAI, ChatOpenAICompletions, ChatGithub, ChatGroq, ChatPerplexity, ChatDeepSeek, ChatOpenRouter, ChatHuggingFace, ChatPortkey, ChatLMStudio, ChatMistral, ChatCloudflare, ChatAzureOpenAI, ChatAzureOpenAICompletions). Closes #190 --- chatlas/__init__.py | 2 + chatlas/_api_headers.py | 36 ++++++ chatlas/_provider_cloudflare.py | 12 +- chatlas/_provider_deepseek.py | 12 +- chatlas/_provider_github.py | 16 ++- chatlas/_provider_groq.py | 12 +- chatlas/_provider_huggingface.py | 16 ++- chatlas/_provider_lmstudio.py | 21 +++- chatlas/_provider_mistral.py | 16 ++- chatlas/_provider_openai.py | 22 +++- chatlas/_provider_openai_azure.py | 28 +++-- chatlas/_provider_openai_completions.py | 26 +++- chatlas/_provider_openai_generic.py | 13 +- chatlas/_provider_openrouter.py | 12 +- chatlas/_provider_perplexity.py | 12 +- chatlas/_provider_portkey.py | 18 ++- tests/test_api_headers.py | 158 ++++++++++++++++++++++++ 17 files changed, 384 insertions(+), 48 deletions(-) create mode 100644 chatlas/_api_headers.py create mode 100644 tests/test_api_headers.py diff --git a/chatlas/__init__.py b/chatlas/__init__.py index 6f6073d5..0fdb6193 100644 --- a/chatlas/__init__.py +++ b/chatlas/__init__.py @@ -1,4 +1,5 @@ from . import types +from ._api_headers import ApiHeaders from ._auto import ChatAuto from ._batch_chat import ( batch_chat, @@ -46,6 +47,7 @@ __version__ = "0.0.0" # stub value for docs __all__ = ( + "ApiHeaders", "batch_chat", "batch_chat_completed", "batch_chat_structured", diff --git a/chatlas/_api_headers.py b/chatlas/_api_headers.py new file mode 100644 index 00000000..8671e54b --- /dev/null +++ b/chatlas/_api_headers.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from typing import Callable, Union + +ApiHeadersValue = dict[str, str] +"""A dict of HTTP header name-value pairs.""" + +ApiHeaders = Union[ApiHeadersValue, Callable[[], ApiHeadersValue]] +""" +Extra HTTP headers to include with every API request. + +Can be: + +* A **dict** of ``{header_name: header_value}`` — sent as-is on every request. +* A **zero-argument callable** returning such a dict — called on every + request, enabling token refresh and other dynamic auth patterns. +""" + + +def resolve_api_headers(api_headers: ApiHeaders | None) -> dict[str, str] | None: + """ + Resolve api_headers into a dict of HTTP headers (or None). + + Called at request time so that callables can return fresh values. + """ + if api_headers is None: + return None + + value = api_headers() if callable(api_headers) else api_headers + + if isinstance(value, dict): + return value + + raise TypeError( + f"api_headers must be (or return) a dict, got {type(value).__name__}" + ) diff --git a/chatlas/_provider_cloudflare.py b/chatlas/_provider_cloudflare.py index 41c350fb..cd78b8ed 100644 --- a/chatlas/_provider_cloudflare.py +++ b/chatlas/_provider_cloudflare.py @@ -1,8 +1,9 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider_openai_completions import OpenAICompletionsProvider @@ -18,7 +19,8 @@ def ChatCloudflare( account: Optional[str] = None, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, seed: Optional[int] | MISSING_TYPE = MISSING, kwargs: Optional["ChatClientArgs"] = None, ) -> Chat["SubmitInputArgs", ChatCompletion]: @@ -73,6 +75,11 @@ def ChatCloudflare( The API key to use for authentication. You generally should not supply this directly, but instead set the `CLOUDFLARE_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. seed Optional integer seed that ChatGPT uses to try and make output more reproducible. @@ -159,6 +166,7 @@ def ChatCloudflare( base_url=base_url, seed=seed, name="Cloudflare", + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, diff --git a/chatlas/_provider_deepseek.py b/chatlas/_provider_deepseek.py index bf33eecb..e3118a6f 100644 --- a/chatlas/_provider_deepseek.py +++ b/chatlas/_provider_deepseek.py @@ -1,8 +1,9 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional, cast +from typing import TYPE_CHECKING, Callable, Optional, cast +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider_openai_completions import OpenAICompletionsProvider @@ -19,7 +20,8 @@ def ChatDeepSeek( *, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, base_url: str = "https://api.deepseek.com", seed: Optional[int] | MISSING_TYPE = MISSING, kwargs: Optional["ChatClientArgs"] = None, @@ -67,6 +69,11 @@ def ChatDeepSeek( api_key The API key to use for authentication. You generally should not supply this directly, but instead set the `DEEPSEEK_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. base_url The base URL to the endpoint; the default uses DeepSeek's API. seed @@ -138,6 +145,7 @@ def ChatDeepSeek( seed=seed, preserve_thinking=True, name="DeepSeek", + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, diff --git a/chatlas/_provider_github.py b/chatlas/_provider_github.py index af97ba20..10084d23 100644 --- a/chatlas/_provider_github.py +++ b/chatlas/_provider_github.py @@ -1,10 +1,11 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional import requests +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider import ModelInfo @@ -20,7 +21,8 @@ def ChatGithub( *, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, base_url: str = "https://models.github.ai/inference/", seed: Optional[int] | MISSING_TYPE = MISSING, kwargs: Optional["ChatClientArgs"] = None, @@ -64,6 +66,11 @@ def ChatGithub( api_key The API key to use for authentication. You generally should not supply this directly, but instead set the `GITHUB_TOKEN` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. base_url The base URL to the endpoint; the default uses Github's API. seed @@ -134,6 +141,7 @@ def ChatGithub( base_url=base_url, seed=seed, name="GitHub", + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, @@ -141,8 +149,8 @@ def ChatGithub( class GitHubProvider(OpenAICompletionsProvider): - def __init__(self, base_url: str, **kwargs): - super().__init__(base_url=base_url, **kwargs) + def __init__(self, base_url: str, api_headers: Optional[ApiHeaders] = None, **kwargs): + super().__init__(base_url=base_url, api_headers=api_headers, **kwargs) self._base_url = base_url def list_models(self) -> list[ModelInfo]: diff --git a/chatlas/_provider_groq.py b/chatlas/_provider_groq.py index 140877b4..ebb5622e 100644 --- a/chatlas/_provider_groq.py +++ b/chatlas/_provider_groq.py @@ -1,8 +1,9 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider_openai_completions import OpenAICompletionsProvider @@ -17,7 +18,8 @@ def ChatGroq( *, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, base_url: str = "https://api.groq.com/openai/v1", seed: Optional[int] | MISSING_TYPE = MISSING, kwargs: Optional["ChatClientArgs"] = None, @@ -58,6 +60,11 @@ def ChatGroq( api_key The API key to use for authentication. You generally should not supply this directly, but instead set the `GROQ_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. base_url The base URL to the endpoint; the default uses Groq's API. seed @@ -128,6 +135,7 @@ def ChatGroq( base_url=base_url, seed=seed, name="Groq", + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, diff --git a/chatlas/_provider_huggingface.py b/chatlas/_provider_huggingface.py index 7f24b4f0..6b68a462 100644 --- a/chatlas/_provider_huggingface.py +++ b/chatlas/_provider_huggingface.py @@ -1,8 +1,9 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider_openai_completions import OpenAICompletionsProvider @@ -17,7 +18,8 @@ def ChatHuggingFace( *, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, kwargs: Optional["ChatClientArgs"] = None, ) -> Chat["SubmitInputArgs", ChatCompletion]: """ @@ -63,6 +65,11 @@ def ChatHuggingFace( The API key to use for authentication. You generally should not supply this directly, but instead set the `HUGGINGFACE_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. kwargs Additional arguments to pass to the underlying OpenAI client constructor. @@ -131,6 +138,7 @@ def ChatHuggingFace( provider=HuggingFaceProvider( api_key=api_key, model=model, + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, @@ -141,8 +149,9 @@ class HuggingFaceProvider(OpenAICompletionsProvider): def __init__( self, *, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, model: str, + api_headers: Optional[ApiHeaders] = None, kwargs: Optional["ChatClientArgs"] = None, ): # https://huggingface.co/docs/inference-providers/en/index?python-clients=requests#http--curl @@ -151,5 +160,6 @@ def __init__( model=model, api_key=api_key, base_url="https://router.huggingface.co/v1", + api_headers=api_headers, kwargs=kwargs, ) diff --git a/chatlas/_provider_lmstudio.py b/chatlas/_provider_lmstudio.py index 8efc3e03..dc0ca061 100644 --- a/chatlas/_provider_lmstudio.py +++ b/chatlas/_provider_lmstudio.py @@ -3,10 +3,11 @@ import os import re import urllib.request -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional import orjson +from ._api_headers import ApiHeaders from ._chat import Chat from ._provider import ModelInfo from ._provider_openai_completions import OpenAICompletionsProvider @@ -22,7 +23,8 @@ def ChatLMStudio( *, system_prompt: Optional[str] = None, base_url: str = "http://localhost:1234", - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, seed: int | None | MISSING_TYPE = MISSING, kwargs: Optional["ChatClientArgs"] = None, ) -> "Chat[SubmitInputArgs, ChatCompletion]": @@ -79,6 +81,11 @@ def ChatLMStudio( usage. If you're accessing an LM Studio instance behind a reverse proxy or secured endpoint that enforces bearer-token authentication, you can set the `LMSTUDIO_API_KEY` environment variable or provide a value here. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. seed Optional integer seed that helps to make output more reproducible. kwargs @@ -95,10 +102,12 @@ def ChatLMStudio( if api_key is None: api_key = os.getenv("LMSTUDIO_API_KEY", "") - if not has_lmstudio(base_url, api_key=api_key): + resolved_key = api_key() if callable(api_key) else api_key + + if not has_lmstudio(base_url, api_key=resolved_key): raise RuntimeError("Can't find locally running LM Studio.") - models = lmstudio_model_info(base_url, api_key=api_key) + models = lmstudio_model_info(base_url, api_key=resolved_key) model_ids = [m["id"] for m in models] if model is None: @@ -123,6 +132,7 @@ def ChatLMStudio( base_url=base_url, seed=seed, name="LM Studio", + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, @@ -130,13 +140,14 @@ def ChatLMStudio( class LMStudioProvider(OpenAICompletionsProvider): - def __init__(self, *, api_key, model, base_url, seed, name, kwargs): + def __init__(self, *, api_key: str | Callable[[], str], model, base_url, seed, name, api_headers=None, kwargs): super().__init__( api_key=api_key, model=model, base_url=f"{base_url}/v1", seed=seed, name=name, + api_headers=api_headers, kwargs=kwargs, ) self.base_url = base_url diff --git a/chatlas/_provider_mistral.py b/chatlas/_provider_mistral.py index f7826f17..6dbb4446 100644 --- a/chatlas/_provider_mistral.py +++ b/chatlas/_provider_mistral.py @@ -1,8 +1,9 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider_openai_completions import OpenAICompletionsProvider @@ -18,7 +19,8 @@ def ChatMistral( *, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, base_url: str = "https://api.mistral.ai/v1/", seed: int | None | MISSING_TYPE = MISSING, kwargs: Optional["ChatClientArgs"] = None, @@ -65,6 +67,11 @@ def ChatMistral( The API key to use for authentication. You generally should not supply this directly, but instead set the `MISTRAL_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. base_url The base URL to the endpoint; the default uses Mistral AI. seed @@ -130,6 +137,7 @@ def ChatMistral( model=model, base_url=base_url, seed=seed, + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, @@ -140,11 +148,12 @@ class MistralProvider(OpenAICompletionsProvider): def __init__( self, *, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, model: str, base_url: str = "https://api.mistral.ai/v1/", seed: Optional[int] = None, name: str = "Mistral", + api_headers: Optional[ApiHeaders] = None, kwargs: Optional["ChatClientArgs"] = None, ): super().__init__( @@ -153,6 +162,7 @@ def __init__( base_url=base_url, seed=seed, name=name, + api_headers=api_headers, kwargs=kwargs, ) diff --git a/chatlas/_provider_openai.py b/chatlas/_provider_openai.py index 40077444..cefc06d9 100644 --- a/chatlas/_provider_openai.py +++ b/chatlas/_provider_openai.py @@ -2,13 +2,14 @@ import base64 import warnings -from typing import TYPE_CHECKING, Literal, Optional, cast +from typing import TYPE_CHECKING, Callable, Literal, Optional, cast from urllib.parse import urlparse import orjson from openai.types.responses import Response, ResponseStreamEvent from pydantic import BaseModel +from ._api_headers import ApiHeaders from ._chat import Chat from ._content import ( Content, @@ -57,7 +58,8 @@ def ChatOpenAI( service_tier: Optional[ Literal["auto", "default", "flex", "scale", "priority"] ] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, kwargs: Optional["ChatClientArgs"] = None, ) -> Chat["SubmitInputArgs", Response]: """ @@ -115,6 +117,11 @@ def ChatOpenAI( The API key to use for authentication. You generally should not supply this directly, but instead set the `OPENAI_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. kwargs Additional arguments to pass to the `openai.OpenAI()` client constructor. @@ -187,6 +194,7 @@ def ChatOpenAI( api_key=api_key, model=model, base_url=base_url, + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, @@ -212,7 +220,10 @@ def chat_perform( kwargs: Optional["SubmitInputArgs"] = None, ): kwargs = self._chat_perform_args(stream, turns, tools, data_model, kwargs) - return self._client.responses.create(**kwargs) # type: ignore + return self._client.responses.create( # type: ignore + **kwargs, + extra_headers=self._get_extra_headers(), + ) async def chat_perform_async( self, @@ -224,7 +235,10 @@ async def chat_perform_async( kwargs: Optional["SubmitInputArgs"] = None, ): kwargs = self._chat_perform_args(stream, turns, tools, data_model, kwargs) - return await self._async_client.responses.create(**kwargs) # type: ignore + return await self._async_client.responses.create( # type: ignore + **kwargs, + extra_headers=self._get_extra_headers(), + ) def _chat_perform_args( self, diff --git a/chatlas/_provider_openai_azure.py b/chatlas/_provider_openai_azure.py index be5ad3d6..53d751a2 100644 --- a/chatlas/_provider_openai_azure.py +++ b/chatlas/_provider_openai_azure.py @@ -1,10 +1,11 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal, Optional +from typing import TYPE_CHECKING, Callable, Literal, Optional from openai import AsyncAzureOpenAI, AzureOpenAI from openai.types.chat import ChatCompletion +from ._api_headers import ApiHeaders from ._chat import Chat from ._provider_openai import OpenAIProvider from ._provider_openai_completions import OpenAICompletionsProvider @@ -27,7 +28,8 @@ def ChatAzureOpenAI( endpoint: str, deployment_id: str, api_version: str, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, system_prompt: Optional[str] = None, reasoning: "Optional[ReasoningEffort | Reasoning]" = None, service_tier: Optional[ @@ -72,6 +74,11 @@ def ChatAzureOpenAI( The API key to use for authentication. You generally should not supply this directly, but instead set the `AZURE_OPENAI_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. system_prompt A system prompt to set the behavior of the assistant. reasoning @@ -111,6 +118,7 @@ def ChatAzureOpenAI( deployment_id=deployment_id, api_version=api_version, api_key=api_key, + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, @@ -125,7 +133,8 @@ def __init__( endpoint: Optional[str] = None, deployment_id: str, api_version: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, name: str = "Azure/OpenAI", model: Optional[str] = "UnusedValue", kwargs: Optional["ChatAzureClientArgs"] = None, @@ -136,13 +145,14 @@ def __init__( # The OpenAI() constructor will fail if no API key is present. # However, a dummy value is fine -- AzureOpenAI() handles the auth. api_key=api_key or "not-used", + api_headers=api_headers, ) kwargs_full: "ChatAzureClientArgs" = { "azure_endpoint": endpoint, "azure_deployment": deployment_id, "api_version": api_version, - "api_key": api_key, + "api_key": api_key, # type: ignore[typeddict-item] # ChatAzureClientArgs is generated from AsyncAzureOpenAI which requires Awaitable; sync AzureOpenAI accepts plain Callable[[], str] **(kwargs or {}), } @@ -157,7 +167,8 @@ def ChatAzureOpenAICompletions( endpoint: str, deployment_id: str, api_version: str, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, system_prompt: Optional[str] = None, seed: int | None | MISSING_TYPE = MISSING, kwargs: Optional["ChatAzureClientArgs"] = None, @@ -178,6 +189,7 @@ def ChatAzureOpenAICompletions( deployment_id=deployment_id, api_version=api_version, api_key=api_key, + api_headers=api_headers, seed=seed, kwargs=kwargs, ), @@ -192,7 +204,8 @@ def __init__( endpoint: Optional[str] = None, deployment_id: str, api_version: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, seed: int | None = None, name: str = "Azure/OpenAI", model: Optional[str] = "UnusedValue", @@ -205,13 +218,14 @@ def __init__( # The OpenAI() constructor will fail if no API key is present. # However, a dummy value is fine -- AzureOpenAI() handles the auth. api_key=api_key or "not-used", + api_headers=api_headers, ) kwargs_full: "ChatAzureClientArgs" = { "azure_endpoint": endpoint, "azure_deployment": deployment_id, "api_version": api_version, - "api_key": api_key, + "api_key": api_key, # type: ignore[typeddict-item] # ChatAzureClientArgs is generated from AsyncAzureOpenAI which requires Awaitable; sync AzureOpenAI accepts plain Callable[[], str] **(kwargs or {}), } diff --git a/chatlas/_provider_openai_completions.py b/chatlas/_provider_openai_completions.py index b4d7ab1d..9c803cba 100644 --- a/chatlas/_provider_openai_completions.py +++ b/chatlas/_provider_openai_completions.py @@ -2,7 +2,7 @@ import base64 import warnings -from typing import TYPE_CHECKING, Any, Optional, cast +from typing import TYPE_CHECKING, Any, Callable, Optional, cast import orjson from openai.types.chat import ( @@ -16,6 +16,7 @@ ) from pydantic import BaseModel +from ._api_headers import ApiHeaders from ._chat import Chat from ._content import ( Content, @@ -59,7 +60,8 @@ def ChatOpenAICompletions( base_url: str = "https://api.openai.com/v1", system_prompt: Optional[str] = None, model: "Optional[ChatModel | str]" = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, seed: int | None | MISSING_TYPE = MISSING, preserve_thinking: bool = False, kwargs: Optional["ChatClientArgs"] = None, @@ -89,6 +91,11 @@ def ChatOpenAICompletions( The API key to use for authentication. You generally should not supply this directly, but instead set the `OPENAI_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. seed Optional seed for reproducible output. preserve_thinking @@ -111,6 +118,7 @@ def ChatOpenAICompletions( base_url=base_url, seed=seed, preserve_thinking=preserve_thinking, + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, @@ -128,12 +136,13 @@ class OpenAICompletionsProvider( def __init__( self, *, - api_key: str | None = None, + api_key: str | Callable[[], str] | None = None, model: str, base_url: str = "https://api.openai.com/v1", name: str = "OpenAI", seed: int | None = None, preserve_thinking: bool = False, + api_headers: Optional[ApiHeaders] = None, kwargs: Optional["ChatClientArgs"] = None, ): super().__init__( @@ -141,6 +150,7 @@ def __init__( model=model, base_url=base_url, name=name, + api_headers=api_headers, kwargs=kwargs, ) self._seed = seed @@ -156,7 +166,10 @@ def chat_perform( kwargs: Optional["SubmitInputArgs"] = None, ): kwargs = self._chat_perform_args(stream, turns, tools, data_model, kwargs) - return self._client.chat.completions.create(**kwargs) # type: ignore + return self._client.chat.completions.create( # type: ignore + **kwargs, + extra_headers=self._get_extra_headers(), + ) async def chat_perform_async( self, @@ -168,7 +181,10 @@ async def chat_perform_async( kwargs: Optional["SubmitInputArgs"] = None, ): kwargs = self._chat_perform_args(stream, turns, tools, data_model, kwargs) - return await self._async_client.chat.completions.create(**kwargs) # type: ignore + return await self._async_client.chat.completions.create( # type: ignore + **kwargs, + extra_headers=self._get_extra_headers(), + ) def _chat_perform_args( self, diff --git a/chatlas/_provider_openai_generic.py b/chatlas/_provider_openai_generic.py index 82120979..83a42a39 100644 --- a/chatlas/_provider_openai_generic.py +++ b/chatlas/_provider_openai_generic.py @@ -6,13 +6,14 @@ import tempfile from abc import abstractmethod from datetime import datetime -from typing import TYPE_CHECKING, Any, Generic, Literal, Optional +from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Optional import orjson from openai import AsyncOpenAI, OpenAI from openai.types.batch import Batch from pydantic import BaseModel +from ._api_headers import ApiHeaders, resolve_api_headers from ._content import Content, ContentImage, ContentImageRemote from ._provider import ( BatchStatus, @@ -72,16 +73,19 @@ class OpenAIAbstractProvider( def __init__( self, *, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, model: str, base_url: str = "https://api.openai.com/v1", name: str = "OpenAI", + api_headers: Optional[ApiHeaders] = None, kwargs: Optional["ChatClientArgs"] = None, ): super().__init__(name=name, model=model) + self._api_headers = api_headers + kwargs_full: "ChatClientArgs" = { - "api_key": api_key, + "api_key": api_key, # type: ignore[typeddict-item] # ChatClientArgs is generated from AsyncOpenAI which requires Awaitable; sync OpenAI accepts plain Callable[[], str] "base_url": base_url, **(kwargs or {}), } @@ -93,6 +97,9 @@ def __init__( self._client = OpenAI(**sync_kwargs) # type: ignore self._async_client = AsyncOpenAI(**async_kwargs) + def _get_extra_headers(self) -> dict[str, str] | None: + return resolve_api_headers(self._api_headers) + def list_models(self): models = self._client.models.list() diff --git a/chatlas/_provider_openrouter.py b/chatlas/_provider_openrouter.py index ffab69be..c3b41528 100644 --- a/chatlas/_provider_openrouter.py +++ b/chatlas/_provider_openrouter.py @@ -1,8 +1,9 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider_openai_completions import OpenAICompletionsProvider @@ -17,7 +18,8 @@ def ChatOpenRouter( *, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, base_url: str = "https://openrouter.ai/api/v1", seed: Optional[int] | MISSING_TYPE = MISSING, kwargs: Optional["ChatClientArgs"] = None, @@ -60,6 +62,11 @@ def ChatOpenRouter( api_key The API key to use for authentication. You generally should not supply this directly, but instead set the `OPENROUTER_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. base_url The base URL to the endpoint; the default uses OpenRouter's API. seed @@ -133,6 +140,7 @@ def ChatOpenRouter( seed=seed, name="OpenRouter", preserve_thinking=True, + api_headers=api_headers, kwargs=kwargs2, ), system_prompt=system_prompt, diff --git a/chatlas/_provider_perplexity.py b/chatlas/_provider_perplexity.py index 7153e32b..e5990dbd 100644 --- a/chatlas/_provider_perplexity.py +++ b/chatlas/_provider_perplexity.py @@ -1,8 +1,9 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider_openai_completions import OpenAICompletionsProvider @@ -17,7 +18,8 @@ def ChatPerplexity( *, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, base_url: str = "https://api.perplexity.ai/", seed: Optional[int] | MISSING_TYPE = MISSING, kwargs: Optional["ChatClientArgs"] = None, @@ -62,6 +64,11 @@ def ChatPerplexity( The API key to use for authentication. You generally should not supply this directly, but instead set the `PERPLEXITY_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. base_url The base URL to the endpoint; the default uses Perplexity's API. seed @@ -132,6 +139,7 @@ def ChatPerplexity( base_url=base_url, seed=seed, name="Perplexity", + api_headers=api_headers, kwargs=kwargs, ), system_prompt=system_prompt, diff --git a/chatlas/_provider_portkey.py b/chatlas/_provider_portkey.py index 44e51ea6..f831061a 100644 --- a/chatlas/_provider_portkey.py +++ b/chatlas/_provider_portkey.py @@ -1,8 +1,9 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional +from ._api_headers import ApiHeaders from ._chat import Chat from ._logging import log_model_default from ._provider_openai_completions import OpenAICompletionsProvider @@ -17,7 +18,8 @@ def ChatPortkey( *, system_prompt: Optional[str] = None, model: Optional[str] = None, - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, + api_headers: Optional[ApiHeaders] = None, virtual_key: Optional[str] = None, base_url: str = "https://api.portkey.ai/v1", kwargs: Optional["ChatClientArgs"] = None, @@ -62,6 +64,11 @@ def ChatPortkey( api_key The API key to use for authentication. You generally should not supply this directly, but instead set the `PORTKEY_API_KEY` environment variable. + api_headers + Extra HTTP headers to include with every API request. Can be a dict + of ``{header_name: header_value}`` pairs, or a zero-argument callable + returning such a dict. A callable is invoked on every request, + enabling dynamic auth patterns like token refresh. virtual_key An (optional) virtual identifier, storing the LLM provider's API key. See [documentation](https://portkey.ai/docs/product/ai-gateway/virtual-keys). @@ -101,6 +108,7 @@ def ChatPortkey( model=model, base_url=base_url, name="Portkey", + api_headers=api_headers, kwargs=kwargs2, ), system_prompt=system_prompt, @@ -109,13 +117,15 @@ def ChatPortkey( def add_default_headers( kwargs: "ChatClientArgs", - api_key: Optional[str] = None, + api_key: Optional[str | Callable[[], str]] = None, virtual_key: Optional[str] = None, ) -> "ChatClientArgs": headers = kwargs.get("default_headers", None) + # Callables cannot be serialised as header strings; only pass plain strings + api_key_header = api_key if isinstance(api_key, str) else None default_headers = drop_none( { - "x-portkey-api-key": api_key, + "x-portkey-api-key": api_key_header, "x-portkey-virtual-key": virtual_key, **(headers or {}), } diff --git a/tests/test_api_headers.py b/tests/test_api_headers.py new file mode 100644 index 00000000..6abb2336 --- /dev/null +++ b/tests/test_api_headers.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import pytest + +from chatlas._api_headers import ApiHeaders, resolve_api_headers + + +class TestResolveApiHeaders: + def test_none_returns_none(self): + assert resolve_api_headers(None) is None + + def test_dict_passed_through(self): + headers = {"Authorization": "Bearer token123", "X-Org-Id": "org-456"} + result = resolve_api_headers(headers) + assert result == headers + + def test_callable_returning_dict(self): + headers = {"Authorization": "Bearer refreshed"} + result = resolve_api_headers(lambda: headers) + assert result == headers + + def test_callable_called_each_time(self): + call_count = 0 + + def rotating_token() -> dict[str, str]: + nonlocal call_count + call_count += 1 + return {"Authorization": f"Bearer token-{call_count}"} + + assert resolve_api_headers(rotating_token) == { + "Authorization": "Bearer token-1" + } + assert resolve_api_headers(rotating_token) == { + "Authorization": "Bearer token-2" + } + + def test_callable_returning_wrong_type_raises(self): + with pytest.raises(TypeError, match="dict"): + resolve_api_headers(lambda: "not-a-dict") # type: ignore + + def test_wrong_type_raises(self): + with pytest.raises(TypeError, match="dict"): + resolve_api_headers("not-a-dict") # type: ignore + + +NO_PROXY: dict[str, str] = {"NO_PROXY": "*"} + + +class TestProviderApiHeaders: + """Test that api_headers flow through to the provider correctly.""" + + def test_openai_provider_stores_api_headers(self, monkeypatch): + from chatlas._provider_openai_completions import OpenAICompletionsProvider + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + headers = {"Authorization": "Bearer dynamic-key"} + provider = OpenAICompletionsProvider( + api_key="dummy", + model="gpt-4o", + api_headers=headers, + ) + assert provider._api_headers is not None + assert provider._get_extra_headers() == headers + + def test_openai_provider_callable_api_headers(self, monkeypatch): + from chatlas._provider_openai_completions import OpenAICompletionsProvider + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + provider = OpenAICompletionsProvider( + api_key="dummy", + model="gpt-4o", + api_headers=lambda: {"Authorization": "Bearer dynamic-key"}, + ) + assert provider._get_extra_headers() == { + "Authorization": "Bearer dynamic-key" + } + + def test_openai_provider_no_api_headers(self, monkeypatch): + from chatlas._provider_openai_completions import OpenAICompletionsProvider + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + provider = OpenAICompletionsProvider( + api_key="real-key", + model="gpt-4o", + ) + assert provider._api_headers is None + assert provider._get_extra_headers() is None + + +class TestChatFunctionApiHeaders: + """Test that Chat* functions accept api_headers.""" + + def test_chat_openai_completions_accepts_api_headers(self, monkeypatch): + from chatlas import ChatOpenAICompletions + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + chat = ChatOpenAICompletions( + model="gpt-4o", + api_key="dummy", + api_headers={"X-Custom": "value"}, + ) + assert chat.provider._api_headers is not None + + def test_chat_openai_accepts_api_headers(self, monkeypatch): + from chatlas import ChatOpenAI + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + chat = ChatOpenAI( + model="gpt-4o", + api_key="dummy", + api_headers={"Authorization": "Bearer my-token"}, + ) + assert chat.provider._api_headers is not None + + +class TestCallableApiKey: + """Test that api_key accepts callables (passed through to the openai SDK).""" + + def test_callable_api_key_accepted(self, monkeypatch): + from chatlas import ChatOpenAICompletions + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + chat = ChatOpenAICompletions( + model="gpt-4o", + api_key=lambda: "dynamic-key", + ) + assert chat.provider._client.api_key is not None + + def test_callable_api_key_responses_api(self, monkeypatch): + from chatlas import ChatOpenAI + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + chat = ChatOpenAI( + model="gpt-4o", + api_key=lambda: "dynamic-key", + ) + assert chat.provider._client.api_key is not None From 1535c9ac1ed599a331253c38860a498a76e6dba0 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 12 May 2026 15:12:54 -0500 Subject: [PATCH 2/3] fix: wrap sync callable api_key for AsyncOpenAI and clarify api_headers scope - Sync callable api_key is now wrapped with wrap_async() before passing to AsyncOpenAI/AsyncAzureOpenAI, which expect Callable[[], Awaitable[str]]. Without this, chat_async() would fail when api_key is a sync callable. - Clarified api_headers docstring: applies to chat API requests, not all API requests (list_models, batch, etc. use the client's default auth). --- chatlas/_provider_cloudflare.py | 2 +- chatlas/_provider_deepseek.py | 2 +- chatlas/_provider_github.py | 2 +- chatlas/_provider_groq.py | 2 +- chatlas/_provider_huggingface.py | 2 +- chatlas/_provider_lmstudio.py | 2 +- chatlas/_provider_mistral.py | 2 +- chatlas/_provider_openai.py | 2 +- chatlas/_provider_openai_azure.py | 50 +++++++++++++++++++------ chatlas/_provider_openai_completions.py | 2 +- chatlas/_provider_openai_generic.py | 24 ++++++++---- chatlas/_provider_openrouter.py | 2 +- chatlas/_provider_perplexity.py | 2 +- chatlas/_provider_portkey.py | 2 +- 14 files changed, 66 insertions(+), 32 deletions(-) diff --git a/chatlas/_provider_cloudflare.py b/chatlas/_provider_cloudflare.py index cd78b8ed..f8256be2 100644 --- a/chatlas/_provider_cloudflare.py +++ b/chatlas/_provider_cloudflare.py @@ -76,7 +76,7 @@ def ChatCloudflare( this directly, but instead set the `CLOUDFLARE_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_deepseek.py b/chatlas/_provider_deepseek.py index e3118a6f..474feb11 100644 --- a/chatlas/_provider_deepseek.py +++ b/chatlas/_provider_deepseek.py @@ -70,7 +70,7 @@ def ChatDeepSeek( The API key to use for authentication. You generally should not supply this directly, but instead set the `DEEPSEEK_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_github.py b/chatlas/_provider_github.py index 10084d23..19f9dd22 100644 --- a/chatlas/_provider_github.py +++ b/chatlas/_provider_github.py @@ -67,7 +67,7 @@ def ChatGithub( The API key to use for authentication. You generally should not supply this directly, but instead set the `GITHUB_TOKEN` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_groq.py b/chatlas/_provider_groq.py index ebb5622e..e84d31f9 100644 --- a/chatlas/_provider_groq.py +++ b/chatlas/_provider_groq.py @@ -61,7 +61,7 @@ def ChatGroq( The API key to use for authentication. You generally should not supply this directly, but instead set the `GROQ_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_huggingface.py b/chatlas/_provider_huggingface.py index 6b68a462..1481fbf4 100644 --- a/chatlas/_provider_huggingface.py +++ b/chatlas/_provider_huggingface.py @@ -66,7 +66,7 @@ def ChatHuggingFace( this directly, but instead set the `HUGGINGFACE_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_lmstudio.py b/chatlas/_provider_lmstudio.py index dc0ca061..306a3357 100644 --- a/chatlas/_provider_lmstudio.py +++ b/chatlas/_provider_lmstudio.py @@ -82,7 +82,7 @@ def ChatLMStudio( or secured endpoint that enforces bearer-token authentication, you can set the `LMSTUDIO_API_KEY` environment variable or provide a value here. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_mistral.py b/chatlas/_provider_mistral.py index 6dbb4446..b7a1c2ec 100644 --- a/chatlas/_provider_mistral.py +++ b/chatlas/_provider_mistral.py @@ -68,7 +68,7 @@ def ChatMistral( this directly, but instead set the `MISTRAL_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_openai.py b/chatlas/_provider_openai.py index cefc06d9..42ce12c9 100644 --- a/chatlas/_provider_openai.py +++ b/chatlas/_provider_openai.py @@ -118,7 +118,7 @@ def ChatOpenAI( this directly, but instead set the `OPENAI_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_openai_azure.py b/chatlas/_provider_openai_azure.py index 53d751a2..8f8c6b73 100644 --- a/chatlas/_provider_openai_azure.py +++ b/chatlas/_provider_openai_azure.py @@ -9,7 +9,13 @@ from ._chat import Chat from ._provider_openai import OpenAIProvider from ._provider_openai_completions import OpenAICompletionsProvider -from ._utils import MISSING, MISSING_TYPE, is_testing, split_http_client_kwargs +from ._utils import ( + MISSING, + MISSING_TYPE, + is_testing, + split_http_client_kwargs, + wrap_async, +) if TYPE_CHECKING: from openai.types.responses import Response @@ -75,7 +81,7 @@ def ChatAzureOpenAI( this directly, but instead set the `AZURE_OPENAI_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. @@ -148,18 +154,28 @@ def __init__( api_headers=api_headers, ) - kwargs_full: "ChatAzureClientArgs" = { + async_api_key = wrap_async(api_key) if callable(api_key) else api_key + + sync_full: "ChatAzureClientArgs" = { + "azure_endpoint": endpoint, + "azure_deployment": deployment_id, + "api_version": api_version, + "api_key": api_key, # type: ignore + **(kwargs or {}), + } + async_full: "ChatAzureClientArgs" = { "azure_endpoint": endpoint, "azure_deployment": deployment_id, "api_version": api_version, - "api_key": api_key, # type: ignore[typeddict-item] # ChatAzureClientArgs is generated from AsyncAzureOpenAI which requires Awaitable; sync AzureOpenAI accepts plain Callable[[], str] + "api_key": async_api_key, # type: ignore **(kwargs or {}), } - sync_kwargs, async_kwargs = split_http_client_kwargs(kwargs_full) + sync_full, _ = split_http_client_kwargs(sync_full) + _, async_full = split_http_client_kwargs(async_full) - self._client = AzureOpenAI(**sync_kwargs) # type: ignore - self._async_client = AsyncAzureOpenAI(**async_kwargs) # type: ignore + self._client = AzureOpenAI(**sync_full) # type: ignore + self._async_client = AsyncAzureOpenAI(**async_full) # type: ignore def ChatAzureOpenAICompletions( @@ -221,15 +237,25 @@ def __init__( api_headers=api_headers, ) - kwargs_full: "ChatAzureClientArgs" = { + async_api_key = wrap_async(api_key) if callable(api_key) else api_key + + sync_full: "ChatAzureClientArgs" = { + "azure_endpoint": endpoint, + "azure_deployment": deployment_id, + "api_version": api_version, + "api_key": api_key, # type: ignore + **(kwargs or {}), + } + async_full: "ChatAzureClientArgs" = { "azure_endpoint": endpoint, "azure_deployment": deployment_id, "api_version": api_version, - "api_key": api_key, # type: ignore[typeddict-item] # ChatAzureClientArgs is generated from AsyncAzureOpenAI which requires Awaitable; sync AzureOpenAI accepts plain Callable[[], str] + "api_key": async_api_key, # type: ignore **(kwargs or {}), } - sync_kwargs, async_kwargs = split_http_client_kwargs(kwargs_full) + sync_full, _ = split_http_client_kwargs(sync_full) + _, async_full = split_http_client_kwargs(async_full) - self._client = AzureOpenAI(**sync_kwargs) # type: ignore - self._async_client = AsyncAzureOpenAI(**async_kwargs) # type: ignore + self._client = AzureOpenAI(**sync_full) # type: ignore + self._async_client = AsyncAzureOpenAI(**async_full) # type: ignore diff --git a/chatlas/_provider_openai_completions.py b/chatlas/_provider_openai_completions.py index 9c803cba..bbd0fe5a 100644 --- a/chatlas/_provider_openai_completions.py +++ b/chatlas/_provider_openai_completions.py @@ -92,7 +92,7 @@ def ChatOpenAICompletions( this directly, but instead set the `OPENAI_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_openai_generic.py b/chatlas/_provider_openai_generic.py index 83a42a39..73017e14 100644 --- a/chatlas/_provider_openai_generic.py +++ b/chatlas/_provider_openai_generic.py @@ -27,7 +27,7 @@ from ._tokens import get_price_info from ._tools import Tool, ToolBuiltIn from ._turn import AssistantTurn, Turn, UserTurn, user_turn -from ._utils import split_http_client_kwargs +from ._utils import split_http_client_kwargs, wrap_async if TYPE_CHECKING: from .types.openai import ChatClientArgs @@ -84,18 +84,26 @@ def __init__( self._api_headers = api_headers - kwargs_full: "ChatClientArgs" = { - "api_key": api_key, # type: ignore[typeddict-item] # ChatClientArgs is generated from AsyncOpenAI which requires Awaitable; sync OpenAI accepts plain Callable[[], str] + # AsyncOpenAI expects Callable[[], Awaitable[str]] so wrap sync callables + async_api_key = wrap_async(api_key) if callable(api_key) else api_key + + sync_full: "ChatClientArgs" = { + "api_key": api_key, # type: ignore + "base_url": base_url, + **(kwargs or {}), + } + async_full: "ChatClientArgs" = { + "api_key": async_api_key, # type: ignore "base_url": base_url, **(kwargs or {}), } - # Avoid passing the wrong sync/async client to the OpenAI constructor. - sync_kwargs, async_kwargs = split_http_client_kwargs(kwargs_full) + # Avoid passing the wrong sync/async http_client to each constructor. + sync_full, _ = split_http_client_kwargs(sync_full) + _, async_full = split_http_client_kwargs(async_full) - # TODO: worth bringing in AsyncOpenAI types? - self._client = OpenAI(**sync_kwargs) # type: ignore - self._async_client = AsyncOpenAI(**async_kwargs) + self._client = OpenAI(**sync_full) # type: ignore + self._async_client = AsyncOpenAI(**async_full) # type: ignore def _get_extra_headers(self) -> dict[str, str] | None: return resolve_api_headers(self._api_headers) diff --git a/chatlas/_provider_openrouter.py b/chatlas/_provider_openrouter.py index c3b41528..bc2c6e27 100644 --- a/chatlas/_provider_openrouter.py +++ b/chatlas/_provider_openrouter.py @@ -63,7 +63,7 @@ def ChatOpenRouter( The API key to use for authentication. You generally should not supply this directly, but instead set the `OPENROUTER_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_perplexity.py b/chatlas/_provider_perplexity.py index e5990dbd..bb711c79 100644 --- a/chatlas/_provider_perplexity.py +++ b/chatlas/_provider_perplexity.py @@ -65,7 +65,7 @@ def ChatPerplexity( this directly, but instead set the `PERPLEXITY_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. diff --git a/chatlas/_provider_portkey.py b/chatlas/_provider_portkey.py index f831061a..ea38987f 100644 --- a/chatlas/_provider_portkey.py +++ b/chatlas/_provider_portkey.py @@ -65,7 +65,7 @@ def ChatPortkey( The API key to use for authentication. You generally should not supply this directly, but instead set the `PORTKEY_API_KEY` environment variable. api_headers - Extra HTTP headers to include with every API request. Can be a dict + Extra HTTP headers to include with every chat API request. Can be a dict of ``{header_name: header_value}`` pairs, or a zero-argument callable returning such a dict. A callable is invoked on every request, enabling dynamic auth patterns like token refresh. From a98262bce2049317c4e0739b6af7491df84457a8 Mon Sep 17 00:00:00 2001 From: Carson Date: Tue, 12 May 2026 15:23:22 -0500 Subject: [PATCH 3/3] fix: reject callable api_key in ChatPortkey and add async callable test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ChatPortkey needs to serialize the API key into the x-portkey-api-key header at construction time, so callable api_key is not supported — raise a clear TypeError directing users to api_headers instead. Also adds test coverage for callable api_key with the async client, verifying the sync→async wrapping works correctly. --- chatlas/_provider_portkey.py | 13 +++++++++---- tests/test_api_headers.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/chatlas/_provider_portkey.py b/chatlas/_provider_portkey.py index ea38987f..b3c738f6 100644 --- a/chatlas/_provider_portkey.py +++ b/chatlas/_provider_portkey.py @@ -96,6 +96,13 @@ def ChatPortkey( if api_key is None: api_key = os.getenv("PORTKEY_API_KEY") + if callable(api_key): + raise TypeError( + "ChatPortkey() does not support callable `api_key` because the key " + "must be sent as the `x-portkey-api-key` header at construction time. " + "For dynamic auth, use `api_headers` with a callable instead." + ) + kwargs2 = add_default_headers( kwargs or {}, api_key=api_key, @@ -117,15 +124,13 @@ def ChatPortkey( def add_default_headers( kwargs: "ChatClientArgs", - api_key: Optional[str | Callable[[], str]] = None, + api_key: Optional[str] = None, virtual_key: Optional[str] = None, ) -> "ChatClientArgs": headers = kwargs.get("default_headers", None) - # Callables cannot be serialised as header strings; only pass plain strings - api_key_header = api_key if isinstance(api_key, str) else None default_headers = drop_none( { - "x-portkey-api-key": api_key_header, + "x-portkey-api-key": api_key, "x-portkey-virtual-key": virtual_key, **(headers or {}), } diff --git a/tests/test_api_headers.py b/tests/test_api_headers.py index 6abb2336..0dacb08b 100644 --- a/tests/test_api_headers.py +++ b/tests/test_api_headers.py @@ -156,3 +156,35 @@ def test_callable_api_key_responses_api(self, monkeypatch): api_key=lambda: "dynamic-key", ) assert chat.provider._client.api_key is not None + + def test_callable_api_key_async_client_wrapped(self, monkeypatch): + import inspect + + from chatlas import ChatOpenAI + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + chat = ChatOpenAI( + model="gpt-4o", + api_key=lambda: "dynamic-key", + ) + # The SDK stores the callable on _api_key_provider internally. + # The async client's provider should be an async function (wrapped). + async_provider = chat.provider._async_client._api_key_provider + assert callable(async_provider) + assert inspect.iscoroutinefunction(async_provider) + + def test_portkey_callable_api_key_rejected(self, monkeypatch): + from chatlas import ChatPortkey + + for k in ("HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(k, raising=False) + monkeypatch.delenv(k.lower(), raising=False) + + with pytest.raises(TypeError, match="callable"): + ChatPortkey( + model="gpt-4o", + api_key=lambda: "dynamic-key", + )