Support VertexAI models in infer_model (#89)

Author: samuelcolvin

pydantic_ai/models/__init__.py

@@ -15,6 +15,7 @@
 
 import httpx
 
+from ..exceptions import UserError
 from ..messages import Message, ModelAnyResponse, ModelStructuredResponse
 
 if TYPE_CHECKING:
@@ -46,6 +47,8 @@
     'groq:gemma-7b-it',
     'gemini-1.5-flash',
     'gemini-1.5-pro',
+    'vertexai:gemini-1.5-flash',
+    'vertexai:gemini-1.5-pro',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -245,9 +248,11 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .groq import GroqModel
 
         return GroqModel(model[5:])  # pyright: ignore[reportArgumentType]
-    else:
-        from ..exceptions import UserError
+    elif model.startswith('vertexai:'):
+        from .vertexai import VertexAIModel
 
+        return VertexAIModel(model[9:])  # pyright: ignore[reportArgumentType]
+    else:
         raise UserError(f'Unknown model: {model}')
 
 

pydantic_ai_examples/pydantic_model.py

@@ -24,6 +24,7 @@ class MyModel(BaseModel):
 
 
 model = cast(KnownModelName, os.getenv('PYDANTIC_AI_MODEL', 'openai:gpt-4o'))
+print(f'Using model: {model}')
 agent = Agent(model, result_type=MyModel)
 
 if __name__ == '__main__':

tests/models/test_model.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING
+
 import pytest
 
 from pydantic_ai import UserError
@@ -6,6 +8,19 @@
 from pydantic_ai.models.openai import OpenAIModel
 from tests.conftest import TestEnv
 
+if TYPE_CHECKING:
+    from pydantic_ai.models.vertexai import VertexAIModel
+
+    google_auth_installed = True
+
+else:
+    try:
+        from pydantic_ai.models.vertexai import VertexAIModel
+    except ImportError:
+        google_auth_installed = False
+    else:
+        google_auth_installed = True
+
 
 def test_infer_str_openai(env: TestEnv):
     env.set('OPENAI_API_KEY', 'via-env-var')
@@ -24,6 +39,13 @@ def test_infer_str_gemini(env: TestEnv):
     assert m.name() == 'gemini-1.5-flash'
 
 
+@pytest.mark.skipif(not google_auth_installed, reason='google-auth not installed')
+def test_infer_vertexai(env: TestEnv):
+    m = infer_model('vertexai:gemini-1.5-flash')
+    assert isinstance(m, VertexAIModel)
+    assert m.name() == 'vertexai:gemini-1.5-flash'
+
+
 def test_infer_str_unknown():
     with pytest.raises(UserError, match='Unknown model: foobar'):
         infer_model('foobar')  # pyright: ignore[reportArgumentType]