Merge pull request #901 from Mirascope/fix-899

Fix 899

Author: willbakst

mirascope/llm/_call.py

@@ -50,24 +50,36 @@
 def _get_local_provider_call(
     provider: LocalProvider,
     client: Any | None,  # noqa: ANN401
+    is_async: bool,
 ) -> tuple[Callable, Any | None]:
     if provider == "ollama":
         from ..core.openai import openai_call
 
         if client:
             return openai_call, client
-        from openai import OpenAI
+        if is_async:
+            from openai import AsyncOpenAI
 
-        client = OpenAI(api_key="ollama", base_url="http://localhost:11434/v1")
+            client = AsyncOpenAI(api_key="ollama", base_url="http://localhost:11434/v1")
+        else:
+            from openai import OpenAI
+
+            client = OpenAI(api_key="ollama", base_url="http://localhost:11434/v1")
         return openai_call, client
     else:  # provider == "vllm"
         from ..core.openai import openai_call
 
         if client:
             return openai_call, client
-        from openai import OpenAI
 
-        client = OpenAI(api_key="ollama", base_url="http://localhost:8000/v1")
+        if is_async:
+            from openai import AsyncOpenAI
+
+            client = AsyncOpenAI(api_key="ollama", base_url="http://localhost:8000/v1")
+        else:
+            from openai import OpenAI
+
+            client = OpenAI(api_key="ollama", base_url="http://localhost:8000/v1")
         return openai_call, client
 
 
@@ -245,7 +257,9 @@ async def inner_async(
 
                 if effective_provider in get_args(LocalProvider):
                     provider_call, effective_client = _get_local_provider_call(
-                        cast(LocalProvider, effective_provider), effective_client
+                        cast(LocalProvider, effective_provider),
+                        effective_client,
+                        True,
                     )
                     effective_call_args["client"] = effective_client
                 else:
@@ -293,7 +307,9 @@ def inner(
 
                 if effective_provider in get_args(LocalProvider):
                     provider_call, effective_client = _get_local_provider_call(
-                        cast(LocalProvider, effective_provider), effective_client
+                        cast(LocalProvider, effective_provider),
+                        effective_client,
+                        False,
                     )
                     effective_call_args["client"] = effective_client
                 else:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mirascope"
-version = "1.21.0"
+version = "1.21.1"
 description = "LLM abstractions that aren't obstructions"
 readme = "README.md"
 license = { file = "LICENSE" }

tests/llm/test_call.py

@@ -3,7 +3,7 @@
 from unittest.mock import Mock, patch
 
 import pytest
-from openai import OpenAI
+from openai import AsyncOpenAI, OpenAI
 from pydantic import computed_field
 
 from mirascope.core.base import (
@@ -231,29 +231,43 @@ def test_get_provider_call_xai():
 
 def test_get_local_provider_call_ollama():
     with patch("mirascope.core.openai.openai_call", new="openai_ollama_mock"):
-        func, client = _get_local_provider_call("ollama", None)
+        func, client = _get_local_provider_call("ollama", None, False)
         assert func == "openai_ollama_mock"
         assert (
             isinstance(client, OpenAI)
             and client.api_key == "ollama"
             and str(client.base_url) == "http://localhost:11434/v1/"
         )
+        func, client = _get_local_provider_call("ollama", None, True)
+        assert func == "openai_ollama_mock"
+        assert (
+            isinstance(client, AsyncOpenAI)
+            and client.api_key == "ollama"
+            and str(client.base_url) == "http://localhost:11434/v1/"
+        )
         mock_client = Mock()
-        _, client = _get_local_provider_call("ollama", mock_client)
+        _, client = _get_local_provider_call("ollama", mock_client, False)
         assert client == mock_client
 
 
 def test_get_local_provider_call_vllm():
     with patch("mirascope.core.openai.openai_call", new="openai_vllm_mock"):
-        func, client = _get_local_provider_call("vllm", None)
+        func, client = _get_local_provider_call("vllm", None, False)
         assert func == "openai_vllm_mock"
         assert (
             isinstance(client, OpenAI)
             and client.api_key == "ollama"
             and str(client.base_url) == "http://localhost:8000/v1/"
         )
+        func, client = _get_local_provider_call("vllm", None, True)
+        assert func == "openai_vllm_mock"
+        assert (
+            isinstance(client, AsyncOpenAI)
+            and client.api_key == "ollama"
+            and str(client.base_url) == "http://localhost:8000/v1/"
+        )
         mock_client = Mock()
-        _, client = _get_local_provider_call("vllm", mock_client)
+        _, client = _get_local_provider_call("vllm", mock_client, False)
         assert client == mock_client