Add new provider classes for Together AI, Fireworks AI, and Grok with automatic model profile selection (#1842)

Author: DouweM

docs/api/providers.md

@@ -19,3 +19,9 @@
 ::: pydantic_ai.providers.cohere
 
 ::: pydantic_ai.providers.mistral
+
+::: pydantic_ai.providers.fireworks
+
+::: pydantic_ai.providers.grok
+
+::: pydantic_ai.providers.together

docs/models/openai.md

@@ -336,16 +336,16 @@ agent = Agent(model)
 ### Grok (xAI)
 
 Go to [xAI API Console](https://console.x.ai/) and create an API key.
-Once you have the API key, you can use it with the `OpenAIProvider`:
+Once you have the API key, you can use it with the `GrokProvider`:
 
 ```python
 from pydantic_ai import Agent
 from pydantic_ai.models.openai import OpenAIModel
-from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_ai.providers.grok import GrokProvider
 
 model = OpenAIModel(
     'grok-2-1212',
-    provider=OpenAIProvider(base_url='https://api.x.ai/v1', api_key='your-xai-api-key'),
+    provider=GrokProvider(api_key='your-xai-api-key'),
 )
 agent = Agent(model)
 ...
@@ -375,19 +375,16 @@ agent = Agent(model)
 ### Fireworks AI
 
 Go to [Fireworks.AI](https://fireworks.ai/) and create an API key in your account settings.
-Once you have the API key, you can use it with the `OpenAIProvider`:
+Once you have the API key, you can use it with the `FireworksProvider`:
 
 ```python
 from pydantic_ai import Agent
 from pydantic_ai.models.openai import OpenAIModel
-from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_ai.providers.fireworks import FireworksProvider
 
 model = OpenAIModel(
     'accounts/fireworks/models/qwq-32b',  # model library available at https://fireworks.ai/models
-    provider=OpenAIProvider(
-        base_url='https://api.fireworks.ai/inference/v1',
-        api_key='your-fireworks-api-key',
-    ),
+    provider=FireworksProvider(api_key='your-fireworks-api-key'),
 )
 agent = Agent(model)
 ...
@@ -396,19 +393,16 @@ agent = Agent(model)
 ### Together AI
 
 Go to [Together.ai](https://www.together.ai/) and create an API key in your account settings.
-Once you have the API key, you can use it with the `OpenAIProvider`:
+Once you have the API key, you can use it with the `TogetherProvider`:
 
 ```python
 from pydantic_ai import Agent
 from pydantic_ai.models.openai import OpenAIModel
-from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_ai.providers.together import TogetherProvider
 
 model = OpenAIModel(
     'meta-llama/Llama-3.3-70B-Instruct-Turbo-Free',  # model library available at https://www.together.ai/models
-    provider=OpenAIProvider(
-        base_url='https://api.together.xyz/v1',
-        api_key='your-together-api-key',
-    ),
+    provider=TogetherProvider(api_key='your-together-api-key'),
 )
 agent = Agent(model)
 ...

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -543,7 +543,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:
         from .cohere import CohereModel
 
         return CohereModel(model_name, provider=provider)
-    elif provider in ('deepseek', 'openai', 'azure', 'openrouter'):
+    elif provider in ('openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together'):
         from .openai import OpenAIModel
 
         return OpenAIModel(model_name, provider=provider)

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -170,7 +170,8 @@ def __init__(
         self,
         model_name: OpenAIModelName,
         *,
-        provider: Literal['openai', 'deepseek', 'azure', 'openrouter'] | Provider[AsyncOpenAI] = 'openai',
+        provider: Literal['openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together']
+        | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
         system_prompt_role: OpenAISystemPromptRole | None = None,
     ):
@@ -534,7 +535,8 @@ def __init__(
         self,
         model_name: OpenAIModelName,
         *,
-        provider: Literal['openai', 'deepseek', 'azure', 'openrouter'] | Provider[AsyncOpenAI] = 'openai',
+        provider: Literal['openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together']
+        | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
     ):
         """Initialize an OpenAI Responses model.

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -95,5 +95,17 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .cohere import CohereProvider
 
         return CohereProvider()
+    elif provider == 'grok':
+        from .grok import GrokProvider
+
+        return GrokProvider()
+    elif provider == 'fireworks':
+        from .fireworks import FireworksProvider
+
+        return FireworksProvider()
+    elif provider == 'together':
+        from .together import TogetherProvider
+
+        return TogetherProvider()
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')

pydantic_ai_slim/pydantic_ai/providers/fireworks.py

@@ -0,0 +1,99 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+from openai import AsyncOpenAI
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.google import google_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
+from pydantic_ai.profiles.qwen import qwen_model_profile
+from pydantic_ai.providers import Provider
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `openai` package to use the Fireworks AI provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class FireworksProvider(Provider[AsyncOpenAI]):
+    """Provider for Fireworks AI API."""
+
+    @property
+    def name(self) -> str:
+        return 'fireworks'
+
+    @property
+    def base_url(self) -> str:
+        return 'https://api.fireworks.ai/inference/v1'
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        prefix_to_profile = {
+            'llama': meta_model_profile,
+            'qwen': qwen_model_profile,
+            'deepseek': deepseek_model_profile,
+            'mistral': mistral_model_profile,
+            'gemma': google_model_profile,
+        }
+
+        prefix = 'accounts/fireworks/models/'
+
+        profile = None
+        if model_name.startswith(prefix):
+            model_name = model_name[len(prefix) :]
+            for provider, profile_func in prefix_to_profile.items():
+                if model_name.startswith(provider):
+                    profile = profile_func(model_name)
+                    break
+
+        # As the Fireworks API is OpenAI-compatible, let's assume we also need OpenAIJsonSchemaTransformer,
+        # unless json_schema_transformer is set explicitly
+        return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
+
+    @overload
+    def __init__(self) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str, http_client: AsyncHTTPClient) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        api_key = api_key or os.getenv('FIREWORKS_API_KEY')
+        if not api_key and openai_client is None:
+            raise UserError(
+                'Set the `FIREWORKS_API_KEY` environment variable or pass it via `FireworksProvider(api_key=...)`'
+                'to use the Fireworks AI provider.'
+            )
+
+        if openai_client is not None:
+            self._client = openai_client
+        elif http_client is not None:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+        else:
+            http_client = cached_async_http_client(provider='fireworks')
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)

pydantic_ai_slim/pydantic_ai/providers/grok.py

@@ -0,0 +1,82 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+from openai import AsyncOpenAI
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.grok import grok_model_profile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
+from pydantic_ai.providers import Provider
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `openai` package to use the Grok provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class GrokProvider(Provider[AsyncOpenAI]):
+    """Provider for Grok API."""
+
+    @property
+    def name(self) -> str:
+        return 'grok'
+
+    @property
+    def base_url(self) -> str:
+        return 'https://api.x.ai/v1'
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        profile = grok_model_profile(model_name)
+
+        # As the Grok API is OpenAI-compatible, let's assume we also need OpenAIJsonSchemaTransformer,
+        # unless json_schema_transformer is set explicitly.
+        # Also, Grok does not support strict tool definitions: https://github.com/pydantic/pydantic-ai/issues/1846
+        return OpenAIModelProfile(
+            json_schema_transformer=OpenAIJsonSchemaTransformer, openai_supports_strict_tool_definition=False
+        ).update(profile)
+
+    @overload
+    def __init__(self) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str, http_client: AsyncHTTPClient) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        api_key = api_key or os.getenv('GROK_API_KEY')
+        if not api_key and openai_client is None:
+            raise UserError(
+                'Set the `GROK_API_KEY` environment variable or pass it via `GrokProvider(api_key=...)`'
+                'to use the Grok provider.'
+            )
+
+        if openai_client is not None:
+            self._client = openai_client
+        elif http_client is not None:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+        else:
+            http_client = cached_async_http_client(provider='grok')
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)