feat: Add XAI provider for text generation

packages/image-generation/README.md

@@ -41,7 +41,7 @@ uv add "celeste-ai[image-generation]"
 
 <img src="https://www.google.com/s2/favicons?domain=openai.com&sz=64" width="64" height="64" alt="OpenAI" title="OpenAI">
 <img src="https://www.google.com/s2/favicons?domain=google.com&sz=64" width="64" height="64" alt="Google" title="Google">
-<img src="https://www.google.com/s2/favicons?domain=seed.bytedance.com&sz=64" width="64" height="64" alt="ByteDance" title="ByteDance">
+<img src="https://www.google.com/s2/favicons?domain=byteplus.com&sz=64" width="64" height="64" alt="ByteDance" title="ByteDance">
 
 
 **Missing a provider?** [Request it](https://github.com/withceleste/celeste-python/issues/new) – ⚡ **we ship fast**.

packages/text-generation/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "celeste-text-generation"
-version = "0.2.9"
+version = "0.2.10"
 description = "Text generation package for Celeste AI. Unified interface for all providers"
 authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}]
 readme = "README.md"

packages/text-generation/src/celeste_text_generation/models.py

@@ -8,11 +8,13 @@
 from celeste_text_generation.providers.google.models import MODELS as GOOGLE_MODELS
 from celeste_text_generation.providers.mistral.models import MODELS as MISTRAL_MODELS
 from celeste_text_generation.providers.openai.models import MODELS as OPENAI_MODELS
+from celeste_text_generation.providers.xai.models import MODELS as XAI_MODELS
 
 MODELS: list[Model] = [
     *ANTHROPIC_MODELS,
     *COHERE_MODELS,
     *GOOGLE_MODELS,
     *MISTRAL_MODELS,
     *OPENAI_MODELS,
+    *XAI_MODELS,
 ]

packages/text-generation/src/celeste_text_generation/providers/__init__.py

@@ -23,13 +23,17 @@ def _get_providers() -> list[tuple[Provider, type[Client]]]:
     from celeste_text_generation.providers.openai.client import (
         OpenAITextGenerationClient,
     )
+    from celeste_text_generation.providers.xai.client import (
+        XAITextGenerationClient,
+    )
 
     return [
         (Provider.ANTHROPIC, AnthropicTextGenerationClient),
         (Provider.COHERE, CohereTextGenerationClient),
         (Provider.GOOGLE, GoogleTextGenerationClient),
         (Provider.MISTRAL, MistralTextGenerationClient),
         (Provider.OPENAI, OpenAITextGenerationClient),
+        (Provider.XAI, XAITextGenerationClient),
     ]
 
 

packages/text-generation/src/celeste_text_generation/providers/mistral/config.py

@@ -3,7 +3,7 @@
 # HTTP Configuration
 BASE_URL = "https://api.mistral.ai"
 ENDPOINT = "/v1/chat/completions"
-STREAM_ENDPOINT = ENDPOINT  # Same endpoint
+STREAM_ENDPOINT = ENDPOINT
 
 # Authentication
 AUTH_HEADER_NAME = "Authorization"

packages/text-generation/src/celeste_text_generation/providers/xai/__init__.py

@@ -0,0 +1,7 @@
+"""XAI provider for text generation."""
+
+from .client import XAITextGenerationClient
+from .models import MODELS
+from .streaming import XAITextGenerationStream
+
+__all__ = ["MODELS", "XAITextGenerationClient", "XAITextGenerationStream"]

packages/text-generation/src/celeste_text_generation/providers/xai/client.py

@@ -0,0 +1,142 @@
+"""XAI client implementation for text generation."""
+
+from collections.abc import AsyncIterator
+from typing import Any, Unpack
+
+import httpx
+from pydantic import BaseModel
+
+from celeste.mime_types import ApplicationMimeType
+from celeste.parameters import ParameterMapper
+from celeste_text_generation.client import TextGenerationClient
+from celeste_text_generation.io import (
+    TextGenerationFinishReason,
+    TextGenerationInput,
+    TextGenerationUsage,
+)
+from celeste_text_generation.parameters import TextGenerationParameters
+
+from . import config
+from .parameters import XAI_PARAMETER_MAPPERS
+from .streaming import XAITextGenerationStream
+
+
+class XAITextGenerationClient(TextGenerationClient):
+    """XAI client for text generation."""
+
+    @classmethod
+    def parameter_mappers(cls) -> list[ParameterMapper]:
+        return XAI_PARAMETER_MAPPERS
+
+    def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]:
+        """Initialize request from XAI messages array format."""
+        messages = [
+            {
+                "role": "user",
+                "content": inputs.prompt,
+            }
+        ]
+
+        return {"messages": messages}
+
+    def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage:
+        """Parse usage from response."""
+        usage_data = response_data.get("usage", {})
+        prompt_tokens_details = usage_data.get("prompt_tokens_details", {})
+        completion_tokens_details = usage_data.get("completion_tokens_details", {})
+
+        return TextGenerationUsage(
+            input_tokens=usage_data.get("prompt_tokens"),
+            output_tokens=usage_data.get("completion_tokens"),
+            total_tokens=usage_data.get("total_tokens"),
+            cached_tokens=prompt_tokens_details.get("cached_tokens"),
+            reasoning_tokens=completion_tokens_details.get("reasoning_tokens"),
+            billed_tokens=None,
+        )
+
+    def _parse_content(
+        self,
+        response_data: dict[str, Any],
+        **parameters: Unpack[TextGenerationParameters],
+    ) -> str | BaseModel:
+        """Parse content from response."""
+        choices = response_data.get("choices", [])
+        if not choices:
+            msg = "No choices in response"
+            raise ValueError(msg)
+
+        message = choices[0].get("message", {})
+        content = message.get("content") or ""
+
+        return self._transform_output(content, **parameters)
+
+    def _parse_finish_reason(
+        self, response_data: dict[str, Any]
+    ) -> TextGenerationFinishReason | None:
+        """Parse finish reason from response."""
+        choices = response_data.get("choices", [])
+        if not choices:
+            return None
+
+        choice = choices[0]
+        finish_reason = choice.get("finish_reason")
+
+        if not finish_reason:
+            return None
+
+        return TextGenerationFinishReason(reason=finish_reason)
+
+    def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]:
+        """Build metadata dictionary from response data."""
+        # Filter content field before calling super
+        content_fields = {"choices"}
+        filtered_data = {
+            k: v for k, v in response_data.items() if k not in content_fields
+        }
+        return super()._build_metadata(filtered_data)
+
+    async def _make_request(
+        self,
+        request_body: dict[str, Any],
+        **parameters: Unpack[TextGenerationParameters],
+    ) -> httpx.Response:
+        """Make HTTP request(s) and return response object."""
+        request_body["model"] = self.model.id
+
+        headers = {
+            config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}",
+            "Content-Type": ApplicationMimeType.JSON,
+        }
+
+        return await self.http_client.post(
+            f"{config.BASE_URL}{config.ENDPOINT}",
+            headers=headers,
+            json_body=request_body,
+        )
+
+    def _stream_class(self) -> type[XAITextGenerationStream]:
+        """Return the Stream class for this client."""
+        return XAITextGenerationStream
+
+    def _make_stream_request(
+        self,
+        request_body: dict[str, Any],
+        **parameters: Unpack[TextGenerationParameters],
+    ) -> AsyncIterator[dict[str, Any]]:
+        """Make HTTP streaming request and return async iterator of events."""
+        request_body["model"] = self.model.id
+        request_body["stream"] = True
+
+        headers = {
+            config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}",
+            "Content-Type": ApplicationMimeType.JSON,
+        }
+
+        return self.http_client.stream_post(
+            f"{config.BASE_URL}{config.STREAM_ENDPOINT}",
+            headers=headers,
+            json_body=request_body,
+        )
+
+
+__all__ = ["XAITextGenerationClient"]

packages/text-generation/src/celeste_text_generation/providers/xai/config.py

@@ -0,0 +1,10 @@
+"""XAI provider configuration for text generation."""
+
+# HTTP Configuration
+BASE_URL = "https://api.x.ai/v1"
+ENDPOINT = "/chat/completions"
+STREAM_ENDPOINT = ENDPOINT
+
+# Authentication
+AUTH_HEADER_NAME = "Authorization"
+AUTH_HEADER_PREFIX = "Bearer "

packages/text-generation/src/celeste_text_generation/providers/xai/models.py

@@ -0,0 +1,76 @@
+"""XAI models for text generation."""
+
+from celeste import Model, Provider
+from celeste.constraints import Choice, Range, Schema
+from celeste.core import Parameter
+from celeste_text_generation.parameters import TextGenerationParameter
+
+MODELS: list[Model] = [
+    Model(
+        id="grok-4-1-fast-reasoning",
+        provider=Provider.XAI,
+        display_name="Grok 4.1 Fast Reasoning",
+        streaming=True,
+        parameter_constraints={
+            Parameter.TEMPERATURE: Range(min=0.0, max=2.0),
+            Parameter.MAX_TOKENS: Range(min=1, max=30000),
+            TextGenerationParameter.OUTPUT_SCHEMA: Schema(),
+        },
+    ),
+    Model(
+        id="grok-4-1-fast-non-reasoning",
+        provider=Provider.XAI,
+        display_name="Grok 4.1 Fast Non-Reasoning",
+        streaming=True,
+        parameter_constraints={
+            Parameter.TEMPERATURE: Range(min=0.0, max=2.0),
+            Parameter.MAX_TOKENS: Range(min=1, max=30000),
+            TextGenerationParameter.OUTPUT_SCHEMA: Schema(),
+        },
+    ),
+    Model(
+        id="grok-4-fast-reasoning",
+        provider=Provider.XAI,
+        display_name="Grok 4 Fast Reasoning",
+        streaming=True,
+        parameter_constraints={
+            Parameter.TEMPERATURE: Range(min=0.0, max=2.0),
+            Parameter.MAX_TOKENS: Range(min=1, max=30000),
+            TextGenerationParameter.OUTPUT_SCHEMA: Schema(),
+        },
+    ),
+    Model(
+        id="grok-4-fast-non-reasoning",
+        provider=Provider.XAI,
+        display_name="Grok 4 Fast Non-Reasoning",
+        streaming=True,
+        parameter_constraints={
+            Parameter.TEMPERATURE: Range(min=0.0, max=2.0),
+            Parameter.MAX_TOKENS: Range(min=1, max=30000),
+            TextGenerationParameter.OUTPUT_SCHEMA: Schema(),
+        },
+    ),
+    Model(
+        id="grok-4-0709",
+        provider=Provider.XAI,
+        display_name="Grok 4",
+        streaming=True,
+        parameter_constraints={
+            Parameter.TEMPERATURE: Range(min=0.0, max=2.0),
+            Parameter.MAX_TOKENS: Range(min=1, max=64000),
+            TextGenerationParameter.OUTPUT_SCHEMA: Schema(),
+        },
+    ),
+    Model(
+        id="grok-3-mini",
+        provider=Provider.XAI,
+        display_name="Grok 3 Mini",
+        streaming=True,
+        parameter_constraints={
+            Parameter.TEMPERATURE: Range(min=0.0, max=2.0),
+            Parameter.MAX_TOKENS: Range(min=1, max=16000),
+            TextGenerationParameter.THINKING_LEVEL: Choice(options=["low", "high"]),
+            TextGenerationParameter.OUTPUT_SCHEMA: Schema(),
+        },
+    ),
+]