feat: add extra_headers support to all client methods (#193)

* feat: add extra_headers support to all client methods (#178)

Add `extra_headers: dict[str, str] | None = None` alongside `extra_body`
on all client methods, enabling users to pass provider-specific HTTP
headers (e.g., Anthropic's 1M context beta header).

- Extend `_json_headers()` to accept and merge extra headers
- Add `_merge_headers()` static helper for edge cases (WebSocket, multipart)
- Thread `extra_headers` through `_predict()` → `_make_request()` and
  `_stream()` → `_make_stream_request()` across all providers
- Update all 3 templates and the template contract test

Closes #178

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: improve extra_headers consistency across all providers

- Make _merge_headers return a new dict instead of mutating input
- Add extra_headers to base _make_stream_request stub signature
- Fix OpenAI videos multipart sending wrong Content-Type
- Add extra_headers to Google _get_interaction and download_content
- Simplify BFL header construction to use _json_headers(extra_headers)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: Kamilbenkirane

src/celeste/client.py

@@ -53,9 +53,27 @@ def http_client(self) -> HTTPClient:
         """HTTP client with connection pooling for this provider."""
         ...
 
-    def _json_headers(self) -> dict[str, str]:
+    def _json_headers(
+        self, extra_headers: dict[str, str] | None = None
+    ) -> dict[str, str]:
         """Build standard JSON request headers with auth."""
-        return {**self.auth.get_headers(), "Content-Type": ApplicationMimeType.JSON}
+        headers = {**self.auth.get_headers(), "Content-Type": ApplicationMimeType.JSON}
+        if extra_headers:
+            headers.update(extra_headers)
+        return headers
+
+    @staticmethod
+    def _merge_headers(
+        headers: dict[str, str],
+        extra_headers: dict[str, str] | None = None,
+    ) -> dict[str, str]:
+        """Merge user-provided extra headers into provider headers.
+
+        User-provided headers take precedence over provider defaults.
+        """
+        if extra_headers:
+            return {**headers, **extra_headers}
+        return headers
 
     @staticmethod
     def _deep_merge(
@@ -160,6 +178,7 @@ async def _predict(
         *,
         endpoint: str | None = None,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[Params],  # type: ignore[misc]
     ) -> Out:
         """Generic prediction - called by operation methods.
@@ -168,6 +187,7 @@ async def _predict(
             inputs: Operation-specific input object.
             endpoint: Optional endpoint path (e.g., "/generations").
             extra_body: Additional parameters to merge into the request body.
+            extra_headers: Additional headers to merge into the request headers.
             **parameters: Operation-specific keyword arguments.
 
         Returns:
@@ -176,7 +196,7 @@ async def _predict(
         inputs, parameters = self._validate_artifacts(inputs, **parameters)
         request_body = self._build_request(inputs, extra_body=extra_body, **parameters)
         response_data = await self._make_request(
-            request_body, endpoint=endpoint, **parameters
+            request_body, endpoint=endpoint, extra_headers=extra_headers, **parameters
         )
         content = self._parse_content(response_data, **parameters)
         content = self._transform_output(content, **parameters)
@@ -195,6 +215,7 @@ def _stream(
         endpoint: str | None = None,
         base_url: str | None = None,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[Params],  # type: ignore[misc]
     ) -> Stream[Out, Params, Chunk]:
         """Generic streaming - called by operation methods.
@@ -206,6 +227,7 @@ def _stream(
             inputs: Operation-specific input object.
             stream_class: The Stream class to instantiate.
             extra_body: Additional parameters to merge into the request body.
+            extra_headers: Additional headers to merge into the request headers.
             **parameters: Operation-specific keyword arguments.
 
         Returns:
@@ -222,7 +244,11 @@ def _stream(
             inputs, extra_body=extra_body, streaming=True, **parameters
         )
         sse_iterator = self._make_stream_request(
-            request_body, endpoint=endpoint, base_url=base_url, **parameters
+            request_body,
+            endpoint=endpoint,
+            base_url=base_url,
+            extra_headers=extra_headers,
+            **parameters,
         )
         return stream_class(
             sse_iterator,
@@ -287,6 +313,7 @@ async def _make_request(
         request_body: dict[str, Any],
         *,
         endpoint: str | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[Params],  # type: ignore[misc]
     ) -> dict[str, Any]:
         """Make HTTP request(s) and return response data."""
@@ -295,6 +322,9 @@ async def _make_request(
     def _make_stream_request(
         self,
         request_body: dict[str, Any],
+        *,
+        endpoint: str | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[Params],  # type: ignore[misc]
     ) -> AsyncIterator[dict[str, Any]]:
         """Make HTTP streaming request and return async iterator of events."""

src/celeste/modalities/audio/client.py

@@ -49,6 +49,7 @@ def speak(
         text: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[AudioParameters],
     ) -> AudioStream:
         """Stream speech generation."""
@@ -57,6 +58,7 @@ def speak(
             inputs,
             stream_class=self._client._stream_class(),
             extra_body=extra_body,
+            extra_headers=extra_headers,
             **parameters,
         )
 
@@ -72,12 +74,13 @@ def speak(
         text: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[AudioParameters],
     ) -> AudioOutput:
         """Blocking speech generation."""
         inputs = AudioInput(text=text)
         return async_to_sync(self._client._predict)(
-            inputs, extra_body=extra_body, **parameters
+            inputs, extra_body=extra_body, extra_headers=extra_headers, **parameters
         )
 
     @property
@@ -97,6 +100,7 @@ def speak(
         text: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[AudioParameters],
     ) -> AudioStream:
         """Sync streaming speech generation.
@@ -110,7 +114,9 @@ def speak(
             stream.output.content.save("output.mp3")
         """
         # Return same stream as async version - __iter__/__next__ handle sync iteration
-        return self._client.stream.speak(text, extra_body=extra_body, **parameters)
+        return self._client.stream.speak(
+            text, extra_body=extra_body, extra_headers=extra_headers, **parameters
+        )
 
 
 __all__ = [

src/celeste/modalities/embeddings/client.py

@@ -38,13 +38,15 @@ async def embed(
         text: str | list[str],
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[EmbeddingsParameters],
     ) -> EmbeddingsOutput:
         """Generate embeddings from text.
 
         Args:
             text: Text to embed. Single string or list of strings.
             extra_body: Additional provider-specific fields to merge into request.
+            extra_headers: Additional HTTP headers to include in the request.
             **parameters: Embedding parameters (e.g., dimensions).
 
         Returns:
@@ -53,7 +55,9 @@ async def embed(
             - list[list[float]] if text was a list
         """
         inputs = EmbeddingsInput(text=text)
-        output = await self._predict(inputs, extra_body=extra_body, **parameters)
+        output = await self._predict(
+            inputs, extra_body=extra_body, extra_headers=extra_headers, **parameters
+        )
 
         # If single text input, unwrap from batch format to single embedding
         if (
@@ -83,11 +87,12 @@ def embed(
         text: str | list[str],
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[EmbeddingsParameters],
     ) -> EmbeddingsOutput:
         """Blocking embeddings generation."""
         return async_to_sync(self._client.embed)(
-            text, extra_body=extra_body, **parameters
+            text, extra_body=extra_body, extra_headers=extra_headers, **parameters
         )
 
 

src/celeste/modalities/images/client.py

@@ -53,6 +53,7 @@ def generate(
         prompt: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[ImageParameters],
     ) -> ImagesStream:
         """Stream image generation."""
@@ -61,6 +62,7 @@ def generate(
             inputs,
             stream_class=self._client._stream_class(),
             extra_body=extra_body,
+            extra_headers=extra_headers,
             **parameters,
         )
 
@@ -70,6 +72,7 @@ def edit(
         prompt: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[ImageParameters],
     ) -> ImagesStream:
         """Stream image editing."""
@@ -78,6 +81,7 @@ def edit(
             inputs,
             stream_class=self._client._stream_class(),
             extra_body=extra_body,
+            extra_headers=extra_headers,
             **parameters,
         )
 
@@ -96,6 +100,7 @@ def generate(
         prompt: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[ImageParameters],
     ) -> ImageOutput:
         """Blocking image generation.
@@ -106,7 +111,7 @@ def generate(
         """
         inputs = ImageInput(prompt=prompt)
         return async_to_sync(self._client._predict)(
-            inputs, extra_body=extra_body, **parameters
+            inputs, extra_body=extra_body, extra_headers=extra_headers, **parameters
         )
 
     def edit(
@@ -115,6 +120,7 @@ def edit(
         prompt: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[ImageParameters],
     ) -> ImageOutput:
         """Blocking image edit.
@@ -125,7 +131,7 @@ def edit(
         """
         inputs = ImageInput(prompt=prompt, image=image)
         return async_to_sync(self._client._predict)(
-            inputs, extra_body=extra_body, **parameters
+            inputs, extra_body=extra_body, extra_headers=extra_headers, **parameters
         )
 
     @property
@@ -145,6 +151,7 @@ def generate(
         prompt: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[ImageParameters],
     ) -> ImagesStream:
         """Sync streaming image generation.
@@ -158,14 +165,17 @@ def generate(
             print(stream.output.usage)
         """
         # Return same stream as async version - __iter__/__next__ handle sync iteration
-        return self._client.stream.generate(prompt, extra_body=extra_body, **parameters)
+        return self._client.stream.generate(
+            prompt, extra_body=extra_body, extra_headers=extra_headers, **parameters
+        )
 
     def edit(
         self,
         image: ImageArtifact,
         prompt: str,
         *,
         extra_body: dict[str, Any] | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[ImageParameters],
     ) -> ImagesStream:
         """Sync streaming image editing.
@@ -179,7 +189,11 @@ def edit(
             print(stream.output.usage)
         """
         return self._client.stream.edit(
-            image, prompt, extra_body=extra_body, **parameters
+            image,
+            prompt,
+            extra_body=extra_body,
+            extra_headers=extra_headers,
+            **parameters,
         )
 
 

src/celeste/modalities/images/providers/byteplus/client.py

@@ -149,6 +149,7 @@ async def _make_request(
         request_body: dict[str, Any],
         *,
         endpoint: str | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[ImageParameters],
     ) -> dict[str, Any]:
         """Make HTTP request with parameter validation."""
@@ -164,7 +165,7 @@ async def _make_request(
             raise ConstraintViolationError(msg)
 
         return await super()._make_request(
-            request_body, endpoint=endpoint, **parameters
+            request_body, endpoint=endpoint, extra_headers=extra_headers, **parameters
         )
 
     def _stream_class(self) -> type[ImagesStream]:

src/celeste/modalities/images/providers/google/client.py

@@ -101,11 +101,13 @@ async def _make_request(
         request_body: dict[str, Any],
         *,
         endpoint: str | None = None,
+        extra_headers: dict[str, str] | None = None,
         **parameters: Unpack[ImageParameters],
     ) -> dict[str, Any]:
         return await self._strategy._make_request(  # type: ignore[union-attr]
             request_body,
             endpoint=endpoint,
+            extra_headers=extra_headers,
             **parameters,
         )