Address review comments

Author: Motta Kin

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -699,6 +699,7 @@ def infer_model(model: Model | KnownModelName | str) -> Model:  # noqa: C901
         'together',
         'heroku',
         'github',
+        'litellm',
     ):
         from .openai import OpenAIModel
 

pydantic_ai_slim/pydantic_ai/providers/litellm.py

@@ -5,10 +5,21 @@
 from httpx import AsyncClient as AsyncHTTPClient
 from openai import AsyncOpenAI
 
+from pydantic_ai.models import cached_async_http_client
 from pydantic_ai.profiles import ModelProfile
+from pydantic_ai.profiles.amazon import amazon_model_profile
+from pydantic_ai.profiles.anthropic import anthropic_model_profile
+from pydantic_ai.profiles.cohere import cohere_model_profile
+from pydantic_ai.profiles.deepseek import deepseek_model_profile
+from pydantic_ai.profiles.google import google_model_profile
+from pydantic_ai.profiles.grok import grok_model_profile
+from pydantic_ai.profiles.groq import groq_model_profile
+from pydantic_ai.profiles.meta import meta_model_profile
+from pydantic_ai.profiles.mistral import mistral_model_profile
+from pydantic_ai.profiles.moonshotai import moonshotai_model_profile
 from pydantic_ai.profiles.openai import OpenAIJsonSchemaTransformer, OpenAIModelProfile, openai_model_profile
+from pydantic_ai.profiles.qwen import qwen_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
@@ -28,18 +39,45 @@ def name(self) -> str:
 
     @property
     def base_url(self) -> str:
-        return self._base_url
+        return str(self.client.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,
+        # Map provider prefixes to their profile functions
+        provider_to_profile = {
+            'anthropic': anthropic_model_profile,
+            'openai': openai_model_profile,
+            'google': google_model_profile,
+            'mistralai': mistral_model_profile,
+            'mistral': mistral_model_profile,
+            'cohere': cohere_model_profile,
+            'amazon': amazon_model_profile,
+            'bedrock': amazon_model_profile,
+            'meta-llama': meta_model_profile,
+            'meta': meta_model_profile,
+            'groq': groq_model_profile,
+            'deepseek': deepseek_model_profile,
+            'moonshotai': moonshotai_model_profile,
+            'x-ai': grok_model_profile,
+            'qwen': qwen_model_profile,
+        }
+
+        profile = None
+
+        # Check if model name contains a provider prefix (e.g., "anthropic/claude-3")
+        if '/' in model_name:
+            provider_prefix, model_suffix = model_name.split('/', 1)
+            if provider_prefix in provider_to_profile:
+                profile = provider_to_profile[provider_prefix](model_suffix)
+
+        # If no profile found, default to OpenAI profile
+        if profile is None:
+            profile = openai_model_profile(model_name)
+
+        # As LiteLLMProvider is used with OpenAIModel, which uses OpenAIJsonSchemaTransformer,
         # we maintain that behavior
         return OpenAIModelProfile(json_schema_transformer=OpenAIJsonSchemaTransformer).update(profile)
 
@@ -85,36 +123,20 @@ def __init__(
         """
         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'
+        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
+                base_url=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
+                base_url=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

@@ -18,7 +18,7 @@
 
 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.base_url == 'https://custom.litellm.com/v1/'
     assert provider.client.api_key == 'test-key'
 
 
@@ -30,7 +30,7 @@ def test_init_with_custom_llm_provider():
 def test_init_without_api_key():
     provider = LiteLLMProvider()
     assert provider.name == 'litellm'
-    assert provider.base_url == 'https://api.litellm.ai/v1'
+    assert provider.base_url == 'https://api.litellm.ai/v1/'
     assert provider.client.api_key == 'litellm-placeholder'
 
 
@@ -68,24 +68,77 @@ class MockProfile:
 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={})
+    # Create mocks for all profile functions
+    from dataclasses import dataclass
 
-    # 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',
-    ]
+    @dataclass
+    class MockProfile:
+        max_tokens: int = 4096
+        supports_streaming: bool = True
 
-    for model in test_models:
+    # Mock all profile functions
+    mock_profiles = {
+        'openai': mocker.patch('pydantic_ai.providers.litellm.openai_model_profile', return_value=MockProfile()),
+        'anthropic': mocker.patch('pydantic_ai.providers.litellm.anthropic_model_profile', return_value=MockProfile()),
+        'google': mocker.patch('pydantic_ai.providers.litellm.google_model_profile', return_value=MockProfile()),
+        'meta': mocker.patch('pydantic_ai.providers.litellm.meta_model_profile', return_value=MockProfile()),
+        'mistral': mocker.patch('pydantic_ai.providers.litellm.mistral_model_profile', return_value=MockProfile()),
+        'cohere': mocker.patch('pydantic_ai.providers.litellm.cohere_model_profile', return_value=MockProfile()),
+        'amazon': mocker.patch('pydantic_ai.providers.litellm.amazon_model_profile', return_value=MockProfile()),
+        'deepseek': mocker.patch('pydantic_ai.providers.litellm.deepseek_model_profile', return_value=MockProfile()),
+        'groq': mocker.patch('pydantic_ai.providers.litellm.groq_model_profile', return_value=MockProfile()),
+        'grok': mocker.patch('pydantic_ai.providers.litellm.grok_model_profile', return_value=MockProfile()),
+        'moonshotai': mocker.patch(
+            'pydantic_ai.providers.litellm.moonshotai_model_profile', return_value=MockProfile()
+        ),
+        'qwen': mocker.patch('pydantic_ai.providers.litellm.qwen_model_profile', return_value=MockProfile()),
+    }
+
+    # Test models without provider prefix (should use openai profile)
+    models_without_prefix = ['gpt-4', 'claude-3-sonnet', 'gemini-pro', 'llama2-70b']
+
+    for model in models_without_prefix:
         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)
+    # Verify openai_model_profile was called for each model without prefix
+    assert mock_profiles['openai'].call_count == len(models_without_prefix)
+
+    # Reset all call counts
+    for mock in mock_profiles.values():
+        mock.reset_mock()
+
+    # Test all provider prefixes
+    test_cases = [
+        ('anthropic/claude-3-haiku', 'anthropic', 'claude-3-haiku'),
+        ('openai/gpt-4-turbo', 'openai', 'gpt-4-turbo'),
+        ('google/gemini-1.5-pro', 'google', 'gemini-1.5-pro'),
+        ('mistralai/mistral-large', 'mistral', 'mistral-large'),
+        ('mistral/mistral-7b', 'mistral', 'mistral-7b'),
+        ('cohere/command-r', 'cohere', 'command-r'),
+        ('amazon/titan-text', 'amazon', 'titan-text'),
+        ('bedrock/claude-v2', 'amazon', 'claude-v2'),
+        ('meta-llama/llama-3-8b', 'meta', 'llama-3-8b'),
+        ('meta/llama-2-70b', 'meta', 'llama-2-70b'),
+        ('groq/llama3-70b', 'groq', 'llama3-70b'),
+        ('deepseek/deepseek-coder', 'deepseek', 'deepseek-coder'),
+        ('moonshotai/moonshot-v1', 'moonshotai', 'moonshot-v1'),
+        ('x-ai/grok-beta', 'grok', 'grok-beta'),
+        ('qwen/qwen-72b', 'qwen', 'qwen-72b'),
+    ]
+
+    for model_name, expected_profile, expected_suffix in test_cases:
+        profile = provider.model_profile(model_name)
+        assert isinstance(profile, OpenAIModelProfile)
+        assert profile.json_schema_transformer == OpenAIJsonSchemaTransformer
+        # Verify the correct profile function was called with the correct suffix
+        mock_profiles[expected_profile].assert_called_with(expected_suffix)
+        mock_profiles[expected_profile].reset_mock()
+
+    # Test unknown provider prefix (should fall back to openai)
+    provider.model_profile('unknown-provider/some-model')
+    mock_profiles['openai'].assert_called_once_with('unknown-provider/some-model')
 
 
 async def test_cached_http_client_usage(mocker: MockerFixture):