Address comments

Author: ParthSareen

ollama/_client.py

@@ -106,6 +106,9 @@ def __init__(
     )
 
 
+CONNECTION_ERROR_MESSAGE = 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download'
+
+
 class Client(BaseClient):
   def __init__(self, host: Optional[str] = None, **kwargs) -> None:
     super().__init__(httpx.Client, host, **kwargs)
@@ -118,7 +121,7 @@ def _request_raw(self, *args, **kwargs):
     except httpx.HTTPStatusError as e:
       raise ResponseError(e.response.text, e.response.status_code) from None
     except httpx.ConnectError:
-      raise ConnectionError('Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download') from None
+      raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None
 
   @overload
   def _request(
@@ -622,7 +625,7 @@ async def _request_raw(self, *args, **kwargs):
     except httpx.HTTPStatusError as e:
       raise ResponseError(e.response.text, e.response.status_code) from None
     except httpx.ConnectError:
-      raise ConnectionError('Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download') from None
+      raise ConnectionError(CONNECTION_ERROR_MESSAGE) from None
 
   @overload
   async def _request(

tests/test_client.py

@@ -9,7 +9,7 @@
 from pytest_httpserver import HTTPServer, URIPattern
 from werkzeug.wrappers import Request, Response
 
-from ollama._client import AsyncClient, Client, _copy_tools
+from ollama._client import CONNECTION_ERROR_MESSAGE, AsyncClient, Client, _copy_tools
 
 PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC'
 PNG_BYTES = base64.b64decode(PNG_BASE64)
@@ -1116,15 +1116,15 @@ def test_tool_validation():
 
 def test_client_connection_error():
   client = Client('http://localhost:1234')
-  with pytest.raises(ConnectionError) as exc_info:
+
+  with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE):
     client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}])
-  assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download'
-  with pytest.raises(ConnectionError) as exc_info:
+  with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE):
+    client.chat('model', messages=[{'role': 'user', 'content': 'prompt'}])
+  with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE):
     client.generate('model', 'prompt')
-  assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download'
-  with pytest.raises(ConnectionError) as exc_info:
+  with pytest.raises(ConnectionError, match=CONNECTION_ERROR_MESSAGE):
     client.show('model')
-  assert str(exc_info.value) == 'Failed to connect to Ollama. Please check that Ollama is downloaded, running and accessible. https://ollama.com/download'
 
 
 @pytest.mark.asyncio