drift

Signed-off-by: Adrian Cole <[email protected]>
square · Sep 23, 2024 · 9aee5f3 · 9aee5f3
1 parent 7cb44d1
commit 9aee5f3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/providers/openai/test_ollama.py b/tests/providers/openai/test_ollama.py
@@ -1,5 +1,6 @@
 from typing import Tuple
 
+import os
 import pytest
 
 from exchange import Text
@@ -26,7 +27,7 @@ def test_ollama_completion_integration():
 
 def ollama_complete() -> Tuple[Message, Usage]:
     provider = OllamaProvider.from_env()
-    model = OLLAMA_MODEL
+    model = os.getenv("OLLAMA_MODEL", OLLAMA_MODEL)
     system = "You are a helpful assistant."
     messages = [Message.user("Hello")]
     return provider.complete(model=model, system=system, messages=messages, tools=None)