Add support for HerokuProvider (#1933)

Author: Kludex

docs/api/providers.md

@@ -18,10 +18,14 @@
 
 ::: pydantic_ai.providers.cohere
 
-::: pydantic_ai.providers.mistral
+::: pydantic_ai.providers.mistral.MistralProvider
 
-::: pydantic_ai.providers.fireworks
+::: pydantic_ai.providers.fireworks.FireworksProvider
 
-::: pydantic_ai.providers.grok
+::: pydantic_ai.providers.grok.GrokProvider
 
-::: pydantic_ai.providers.together
+::: pydantic_ai.providers.together.TogetherProvider
+
+::: pydantic_ai.providers.heroku.HerokuProvider
+
+::: pydantic_ai.providers.openrouter.OpenRouterProvider

docs/models/index.md

@@ -22,6 +22,7 @@ In addition, many providers are compatible with the OpenAI API, and can be used
 * [Fireworks AI](openai.md#fireworks-ai)
 * [Together AI](openai.md#together-ai)
 * [Azure AI Foundry](openai.md#azure-ai-foundry)
+* [Heroku](openai.md#heroku-ai)
 
 PydanticAI also comes with [`TestModel`](../api/models/test.md) and [`FunctionModel`](../api/models/function.md)
 for testing and development.

docs/models/openai.md

@@ -183,7 +183,7 @@ agent = Agent(model)
 ### DeepSeek
 
 To use the [DeepSeek](https://deepseek.com) provider, first create an API key by following the [Quick Start guide](https://api-docs.deepseek.com/).
-Once you have the API key, you can use it with the `DeepSeekProvider`:
+Once you have the API key, you can use it with the [`DeepSeekProvider`][pydantic_ai.providers.deepseek.DeepSeekProvider]:
 
 ```python
 from pydantic_ai import Agent
@@ -295,7 +295,8 @@ print(result.usage())
 
 ### Azure AI Foundry
 
-If you want to use [Azure AI Foundry](https://ai.azure.com/) as your provider, you can do so by using the `AzureProvider` class.
+If you want to use [Azure AI Foundry](https://ai.azure.com/) as your provider, you can do so by using the
+[`AzureProvider`][pydantic_ai.providers.azure.AzureProvider] class.
 
 ```python
 from pydantic_ai import Agent
@@ -318,7 +319,7 @@ agent = Agent(model)
 
 To use [OpenRouter](https://openrouter.ai), first create an API key at [openrouter.ai/keys](https://openrouter.ai/keys).
 
-Once you have the API key, you can use it with the `OpenRouterProvider`:
+Once you have the API key, you can use it with the [`OpenRouterProvider`][pydantic_ai.providers.openrouter.OpenRouterProvider]:
 
 ```python
 from pydantic_ai import Agent
@@ -336,7 +337,7 @@ 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 `GrokProvider`:
+Once you have the API key, you can use it with the [`GrokProvider`][pydantic_ai.providers.grok.GrokProvider]:
 
 ```python
 from pydantic_ai import Agent
@@ -375,7 +376,7 @@ 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 `FireworksProvider`:
+Once you have the API key, you can use it with the [`FireworksProvider`][pydantic_ai.providers.fireworks.FireworksProvider]:
 
 ```python
 from pydantic_ai import Agent
@@ -393,7 +394,7 @@ 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 `TogetherProvider`:
+Once you have the API key, you can use it with the [`TogetherProvider`][pydantic_ai.providers.together.TogetherProvider]:
 
 ```python
 from pydantic_ai import Agent
@@ -407,3 +408,27 @@ model = OpenAIModel(
 agent = Agent(model)
 ...
 ```
+
+### Heroku AI
+
+To use [Heroku AI](https://www.heroku.com/ai), you can use the [`HerokuProvider`][pydantic_ai.providers.heroku.HerokuProvider]:
+
+```python
+from pydantic_ai import Agent
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.heroku import HerokuProvider
+
+model = OpenAIModel(
+    'claude-3-7-sonnet',
+    provider=HerokuProvider(api_key='your-heroku-inference-key'),
+)
+agent = Agent(model)
+...
+```
+
+You can set the `HEROKU_INFERENCE_KEY` and `HEROKU_INFERENCE_URL` environment variables to set the API key and base URL, respectively:
+
+```bash
+export HEROKU_INFERENCE_KEY='your-heroku-inference-key'
+export HEROKU_INFERENCE_URL='https://us.inference.heroku.com'
+```

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -211,6 +211,11 @@
         'groq:llama-3.2-3b-preview',
         'groq:llama-3.2-11b-vision-preview',
         'groq:llama-3.2-90b-vision-preview',
+        'heroku:claude-3-5-haiku',
+        'heroku:claude-3-5-sonnet-latest',
+        'heroku:claude-3-7-sonnet',
+        'heroku:claude-4-sonnet',
+        'heroku:claude-3-haiku',
         'mistral:codestral-latest',
         'mistral:mistral-large-latest',
         'mistral:mistral-moderation-latest',
@@ -543,7 +548,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:
         from .cohere import CohereModel
 
         return CohereModel(model_name, provider=provider)
-    elif provider in ('openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together'):
+    elif provider in ('openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together', 'heroku'):
         from .openai import OpenAIModel
 
         return OpenAIModel(model_name, provider=provider)

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -228,7 +228,7 @@ async def _make_request(
 
         if gemini_labels := model_settings.get('gemini_labels'):
             if self._system == 'google-vertex':
-                request_data['labels'] = gemini_labels
+                request_data['labels'] = gemini_labels  # pragma: lax no cover
 
         headers = {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}
         url = f'/{self._model_name}:{"streamGenerateContent" if streamed else "generateContent"}'

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -170,7 +170,7 @@ def __init__(
         self,
         model_name: OpenAIModelName,
         *,
-        provider: Literal['openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together']
+        provider: Literal['openai', 'deepseek', 'azure', 'openrouter', 'grok', 'fireworks', 'together', 'heroku']
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
         system_prompt_role: OpenAISystemPromptRole | None = None,

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -48,7 +48,7 @@ def model_profile(self, model_name: str) -> ModelProfile | None:
         return None  # pragma: no cover
 
 
-def infer_provider(provider: str) -> Provider[Any]:
+def infer_provider(provider: str) -> Provider[Any]:  # noqa: C901
     """Infer the provider from the provider name."""
     if provider == 'openai':
         from .openai import OpenAIProvider
@@ -107,5 +107,9 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .together import TogetherProvider
 
         return TogetherProvider()
+    elif provider == 'heroku':
+        from .heroku import HerokuProvider
+
+        return HerokuProvider()
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')

pydantic_ai_slim/pydantic_ai/providers/google_vertex.py

@@ -50,7 +50,7 @@ def client(self) -> httpx.AsyncClient:
         return self._client
 
     def model_profile(self, model_name: str) -> ModelProfile | None:
-        return google_model_profile(model_name)
+        return google_model_profile(model_name)  # pragma: lax no cover
 
     @overload
     def __init__(

pydantic_ai_slim/pydantic_ai/providers/heroku.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.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 Heroku provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class HerokuProvider(Provider[AsyncOpenAI]):
+    """Provider for Heroku API."""
+
+    @property
+    def name(self) -> str:
+        return 'heroku'
+
+    @property
+    def base_url(self) -> str:
+        return str(self.client.base_url)
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        # As the Heroku API is OpenAI-compatible, let's assume we also need OpenAIJsonSchemaTransformer.
+        return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer)
+
+    @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,
+        *,
+        base_url: str | None = None,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        if openai_client is not None:
+            assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
+            assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
+            self._client = openai_client
+        else:
+            api_key = api_key or os.environ.get('HEROKU_INFERENCE_KEY')
+            if not api_key:
+                raise UserError(
+                    'Set the `HEROKU_INFERENCE_KEY` environment variable or pass it via `HerokuProvider(api_key=...)`'
+                    'to use the Heroku provider.'
+                )
+
+            base_url = base_url or os.environ.get('HEROKU_INFERENCE_URL', 'https://us.inference.heroku.com')
+            base_url = base_url.rstrip('/') + '/v1'
+
+            if http_client is not None:
+                self._client = AsyncOpenAI(api_key=api_key, http_client=http_client, base_url=base_url)
+            else:
+                http_client = cached_async_http_client(provider='heroku')
+                self._client = AsyncOpenAI(api_key=api_key, http_client=http_client, base_url=base_url)

tests/conftest.py

@@ -281,6 +281,11 @@ def openrouter_api_key() -> str:
     return os.getenv('OPENROUTER_API_KEY', 'mock-api-key')
 
 
+@pytest.fixture(scope='session')
+def heroku_inference_key() -> str:
+    return os.getenv('HEROKU_INFERENCE_KEY', 'mock-api-key')
+
+
 @pytest.fixture(scope='session')
 def bedrock_provider():
     try: