add tests for Ollama (#182)

Author: samuelcolvin

.github/workflows/ci.yml

@@ -83,6 +83,10 @@ jobs:
         with:
           enable-cache: true
 
+      - uses: pydantic/ollama-action@main
+        with:
+          model: qwen:0.5b
+
       - run: >
           uv run
           --python 3.12

pydantic_ai_slim/pydantic_ai/models/ollama.py

@@ -91,15 +91,13 @@ def __init__(
         self.model_name = model_name
         if openai_client is not None:
             assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
-            self.openai_model = OpenAIModel(model_name=model_name, openai_client=openai_client, http_client=http_client)
-        elif http_client is not None:
-            # API key is not required for ollama but a value is required to create the client
-            oai_client = AsyncOpenAI(base_url=base_url, api_key='ollama', http_client=http_client)
-            self.openai_model = OpenAIModel(model_name=model_name, openai_client=oai_client, http_client=http_client)
+            assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
+            self.openai_model = OpenAIModel(model_name=model_name, openai_client=openai_client)
         else:
             # API key is not required for ollama but a value is required to create the client
-            oai_client = AsyncOpenAI(base_url=base_url, api_key='ollama', http_client=cached_async_http_client())
-            self.openai_model = OpenAIModel(model_name=model_name, openai_client=oai_client, http_client=http_client)
+            http_client_ = http_client or cached_async_http_client()
+            oai_client = AsyncOpenAI(base_url=base_url, api_key='ollama', http_client=http_client_)
+            self.openai_model = OpenAIModel(model_name=model_name, openai_client=oai_client)
 
     async def agent_model(
         self,

tests/test_live.py

@@ -48,11 +48,18 @@ def groq(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
     return GroqModel('llama-3.1-70b-versatile', http_client=http_client)
 
 
+def ollama(http_client: httpx.AsyncClient, _tmp_path: Path) -> Model:
+    from pydantic_ai.models.ollama import OllamaModel
+
+    return OllamaModel('qwen:0.5b', http_client=http_client)
+
+
 params = [
     pytest.param(openai, id='openai'),
     pytest.param(gemini, id='gemini'),
     pytest.param(vertexai, id='vertexai'),
     pytest.param(groq, id='groq'),
+    pytest.param(ollama, id='ollama'),
 ]
 GetModel = Callable[[httpx.AsyncClient, Path], Model]
 
@@ -83,14 +90,18 @@ async def test_stream(http_client: httpx.AsyncClient, tmp_path: Path, get_model:
     assert 'paris' in data.lower()
     print('Stream cost:', result.cost())
     cost = result.cost()
-    assert cost.total_tokens is not None and cost.total_tokens > 0
+    if get_model.__name__ != 'ollama':
+        assert cost.total_tokens is not None and cost.total_tokens > 0
 
 
 class MyModel(BaseModel):
     city: str
 
 
-@pytest.mark.parametrize('get_model', params)
+structured_params = [p for p in params if p.id != 'ollama']
+
+
+@pytest.mark.parametrize('get_model', structured_params)
 async def test_structured(http_client: httpx.AsyncClient, tmp_path: Path, get_model: GetModel):
     agent = Agent(get_model(http_client, tmp_path), result_type=MyModel)
     result = await agent.run('What is the capital of the UK?')