fix: don't always use ollama provider (#18)

square · Aug 29, 2024 · 061695f · 061695f
1 parent f959ccd
commit 061695f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/exchange/providers/ollama.py b/src/exchange/providers/ollama.py
@@ -41,7 +41,7 @@ def __init__(self, client: httpx.Client) -> None:
 
     @classmethod
     def from_env(cls: Type["OllamaProvider"]) -> "OllamaProvider":
-        url = os.environ.get("OLLAMA_HOST", OLLAMA_HOST)
+        url = os.environ["OLLAMA_HOST"]
         client = httpx.Client(
             base_url=url,
             headers={"Content-Type": "application/json"},

diff --git a/tests/providers/test_ollama.py b/tests/providers/test_ollama.py
@@ -3,12 +3,13 @@
 
 import pytest
 from exchange import Message, Text
-from exchange.providers.ollama import OllamaProvider
+from exchange.providers.ollama import OLLAMA_HOST, OllamaProvider
 
 
 @pytest.fixture
 @patch.dict(os.environ, {})
 def ollama_provider():
+    os.environ["OLLAMA_HOST"] = OLLAMA_HOST
     return OllamaProvider.from_env()
 
 
@@ -45,6 +46,7 @@ def test_ollama_completion(mock_error, mock_warning, mock_sleep, mock_post, olla
 
 @pytest.mark.integration
 def test_ollama_integration():
+    os.environ["OLLAMA_HOST"] = OLLAMA_HOST
     provider = OllamaProvider.from_env()
     model = "llama2"  # specify a valid model
     system = "You are a helpful assistant."