Add LiteLLM provider for OpenAI API compatible models

Author: Motta Kin

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -225,6 +225,7 @@ def __init__(
             'openrouter',
             'together',
             'vercel',
+            'litellm',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -252,6 +253,7 @@ def __init__(
             'openrouter',
             'together',
             'vercel',
+            'litellm',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,
@@ -278,6 +280,7 @@ def __init__(
             'openrouter',
             'together',
             'vercel',
+            'litellm',
         ]
         | Provider[AsyncOpenAI] = 'openai',
         profile: ModelProfileSpec | None = None,

pydantic_ai_slim/pydantic_ai/providers/__init__.py

@@ -135,6 +135,10 @@ def infer_provider_class(provider: str) -> type[Provider[Any]]:  # noqa: C901
         from .github import GitHubProvider
 
         return GitHubProvider
+    elif provider == 'litellm':
+        from .litellm import LiteLLMProvider
+
+        return LiteLLMProvider
     else:  # pragma: no cover
         raise ValueError(f'Unknown provider: {provider}')
 

pydantic_ai_slim/pydantic_ai/providers/litellm.py

@@ -0,0 +1,120 @@
+from __future__ import annotations as _annotations
+
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+from openai import AsyncOpenAI
+
+from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
+from pydantic_ai.providers import Provider
+from pydantic_ai_slim.pydantic_ai.models import cached_async_http_client
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `openai` package to use the LiteLLM provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class LiteLLMProvider(Provider[AsyncOpenAI]):
+    """Provider for LiteLLM API."""
+
+    @property
+    def name(self) -> str:
+        return 'litellm'
+
+    @property
+    def base_url(self) -> str:
+        return self._base_url
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    def model_profile(self, model_name: str) -> ModelProfile | None:
+        # For LiteLLM, we use a basic OpenAI profile since it's OpenAI-compatible
+        # Users can override this with their own profile if needed
+        profile = openai_model_profile(model_name)
+
+        # As LiteLLMProvider is used with OpenAIModel, which used to use OpenAIJsonSchemaTransformer,
+        # we maintain that behavior
+        return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
+
+    @overload
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        custom_llm_provider: str | None = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        custom_llm_provider: str | None = None,
+        http_client: AsyncHTTPClient,
+    ) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        custom_llm_provider: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        """Initialize a LiteLLM provider.
+
+        Args:
+            api_key: API key for the model provider. If None, LiteLLM will try to get it from environment variables.
+            api_base: Base URL for the model provider. Use this for custom endpoints or self-hosted models.
+            custom_llm_provider: Custom LLM provider name for LiteLLM. Use this if LiteLLM can't auto-detect the provider.
+            openai_client: Pre-configured OpenAI client. If provided, other parameters are ignored.
+            http_client: Custom HTTP client to use.
+        """
+        if openai_client is not None:
+            self._client = openai_client
+            self._base_url = str(openai_client.base_url)
+            return
+
+        # Set up LiteLLM configuration
+        if api_key:
+            # Store API key in LiteLLM's global config if needed
+            # LiteLLM will handle provider-specific API key names
+            pass
+
+        if custom_llm_provider:
+            # LiteLLM can auto-detect most providers, but this allows override
+            pass
+
+        # Use api_base if provided, otherwise use a generic base URL
+        # LiteLLM doesn't actually use this URL - it routes internally
+        self._base_url = api_base or 'https://api.litellm.ai/v1'
+
+        # Create OpenAI client that will be used with LiteLLM's completion function
+        # The actual API calls will be intercepted and routed through LiteLLM
+        if http_client is not None:
+            self._client = AsyncOpenAI(
+                base_url=self._base_url, api_key=api_key or 'litellm-placeholder', http_client=http_client
+            )
+        else:
+            http_client = cached_async_http_client(provider='litellm')
+            self._client = AsyncOpenAI(
+                base_url=self._base_url, api_key=api_key or 'litellm-placeholder', http_client=http_client
+            )
+
+        # Store configuration for LiteLLM
+        self._api_key = api_key
+        self._api_base = api_base
+        self._custom_llm_provider = custom_llm_provider

tests/providers/test_litellm.py

@@ -0,0 +1,113 @@
+import httpx
+import pytest
+from pytest_mock import MockerFixture
+
+from ..conftest import try_import
+
+with try_import() as imports_successful:
+    from openai import AsyncOpenAI
+
+    from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile
+    from pydantic_ai.providers.litellm import LiteLLMProvider
+
+pytestmark = [
+    pytest.mark.skipif(not imports_successful(), reason='OpenAI client not installed'),
+    pytest.mark.anyio,
+]
+
+
+def test_init_with_api_config():
+    provider = LiteLLMProvider(api_key='test-key', api_base='https://custom.litellm.com/v1')
+    assert provider.base_url == 'https://custom.litellm.com/v1'
+    assert provider.client.api_key == 'test-key'
+
+
+def test_init_with_custom_llm_provider():
+    provider = LiteLLMProvider(api_key='test-key', custom_llm_provider='anthropic')
+    assert provider.name == 'litellm'
+
+
+def test_init_without_api_key():
+    provider = LiteLLMProvider()
+    assert provider.name == 'litellm'
+    assert provider.base_url == 'https://api.litellm.ai/v1'
+    assert provider.client.api_key == 'litellm-placeholder'
+
+
+def test_init_with_openai_client():
+    openai_client = AsyncOpenAI(api_key='custom-key', base_url='https://custom.openai.com/v1')
+    provider = LiteLLMProvider(openai_client=openai_client)
+    assert provider.client == openai_client
+    assert provider.base_url == 'https://custom.openai.com/v1/'
+
+
+def test_model_profile_returns_openai_compatible_profile(mocker: MockerFixture):
+    provider = LiteLLMProvider(api_key='test-key')
+
+    # Create a proper mock profile object that can be updated
+    from dataclasses import dataclass
+
+    @dataclass
+    class MockProfile:
+        max_tokens: int = 4096
+        supports_streaming: bool = True
+
+    mock_profile = MockProfile()
+    mock_openai_profile = mocker.patch('pydantic_ai.providers.litellm.openai_model_profile', return_value=mock_profile)
+
+    profile = provider.model_profile('gpt-3.5-turbo')
+
+    # Verify openai_model_profile was called with the model name
+    mock_openai_profile.assert_called_once_with('gpt-3.5-turbo')
+
+    # Verify the returned profile is an OpenAIModelProfile with OpenAIJsonSchemaTransformer
+    assert isinstance(profile, OpenAIModelProfile)
+    assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+
+def test_model_profile_with_different_models(mocker: MockerFixture):
+    provider = LiteLLMProvider(api_key='test-key')
+
+    mock_openai_profile = mocker.patch('pydantic_ai.providers.litellm.openai_model_profile', return_value={})
+
+    # Test with different model formats that LiteLLM supports
+    test_models = [
+        'gpt-4',
+        'claude-3-sonnet-20240229',
+        'gemini-pro',
+        'llama2-70b-chat',
+        'anthropic/claude-3-haiku-20240307',
+    ]
+
+    for model in test_models:
+        profile = provider.model_profile(model)
+        assert isinstance(profile, OpenAIModelProfile)
+        assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+
+    # Verify openai_model_profile was called for each model
+    assert mock_openai_profile.call_count == len(test_models)
+
+
+async def test_cached_http_client_usage(mocker: MockerFixture):
+    # Create a real AsyncClient for the mock
+    async with httpx.AsyncClient() as mock_cached_client:
+        mock_cached_http_client_func = mocker.patch(
+            'pydantic_ai.providers.litellm.cached_async_http_client', return_value=mock_cached_client
+        )
+
+        provider = LiteLLMProvider(api_key='test-key')
+
+        # Verify cached_async_http_client was called with 'litellm' provider
+        mock_cached_http_client_func.assert_called_once_with(provider='litellm')
+
+        # Verify the client was created
+        assert isinstance(provider.client, AsyncOpenAI)
+
+
+async def test_init_with_http_client_overrides_cached():
+    async with httpx.AsyncClient() as custom_client:
+        provider = LiteLLMProvider(api_key='test-key', http_client=custom_client)
+
+        # Verify the provider was created successfully with custom client
+        assert isinstance(provider.client, AsyncOpenAI)
+        assert provider.client.api_key == 'test-key'