feat(client): enhance InferenceGatewayClient with support for multiple providers and message handling

Also add health check method.

docs(README): update documentation to include health check and new client usage examples
test: add unit tests for client initialization, health check, and message serialization

Signed-off-by: Eden Reich <eden.reich@gmail.com>

Author: edenreich

README.md

@@ -8,6 +8,7 @@ An SDK written in Python for the [Inference Gateway](https://github.com/edenreic
     - [Creating a Client](#creating-a-client)
     - [Listing Models](#listing-models)
     - [Generating Content](#generating-content)
+    - [Health Check](#health-check)
   - [License](#license)
 
 ## Installation
@@ -21,17 +22,12 @@ pip install inference-gateway
 ### Creating a Client
 
 ```python
-from inference_gateway.client import InferenceGatewayClient
+from inference_gateway.client import InferenceGatewayClient, Provider
 
+client = InferenceGatewayClient("http://localhost:8080")
 
-if __name__ == "__main__":
-    client = InferenceGatewayClient("http://localhost:8080")
-
-    models = client.list_models()
-    print("Available models:", models)
-
-    response = client.generate_content("providerName", "modelName", "your prompt here")
-    print("Generated content:", response["Response"]["Content"])
+# With authentication token(optional)
+client = InferenceGatewayClient("http://localhost:8080", token="your-token")
 ```
 
 ### Listing Models
@@ -48,8 +44,28 @@ print("Available models:", models)
 To generate content using a model, use the generate_content method:
 
 ```python
-response = client.generate_content("providerName", "modelName", "your prompt here")
-print("Generated content:", response["Response"]["Content"])
+from inference_gateway.client import Provider, Role, Message
+
+messages = [
+    Message(Role.SYSTEM, "You are a helpful assistant"),
+    Message(Role.USER, "Hello!"),
+]
+
+response = client.generate_content(
+    provider=Provider.OPENAI,
+    model="gpt-4",
+    messages=messages
+)
+print("Assistant:", response["choices"][0]["message"]["content"])
+```
+
+### Health Check
+
+To check the health of the API, use the health_check method:
+
+```python
+is_healthy = client.health_check()
+print("API Status:", "Healthy" if is_healthy else "Unhealthy")
 ```
 
 ## License

Taskfile.yml

@@ -7,6 +7,11 @@ tasks:
     cmds:
       - curl -o openapi.yaml https://raw.githubusercontent.com/inference-gateway/inference-gateway/refs/heads/main/openapi.yaml
 
+  test:
+    desc: Run tests
+    cmds:
+      - pytest tests/
+
   clean:
     desc: Clean up
     cmds:

inference_gateway/client.py

@@ -1,17 +1,92 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Dict, Optional
 import requests
 
 
+class Provider(str, Enum):
+    """Supported LLM providers"""
+    OLLAMA = "ollama"
+    GROQ = "groq"
+    OPENAI = "openai"
+    GOOGLE = "google"
+    CLOUDFLARE = "cloudflare"
+    COHERE = "cohere"
+
+
+class Role(str, Enum):
+    """Message role types"""
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+
+
+@dataclass
+class Message:
+    role: Role
+    content: str
+
+    def to_dict(self) -> Dict[str, str]:
+        """Convert message to dictionary format with string values"""
+        return {
+            "role": self.role.value,
+            "content": self.content
+        }
+
+
+class Model:
+    """Represents an LLM model"""
+
+    def __init__(self, id: str, object: str, owned_by: str, created: int):
+        self.id = id
+        self.object = object
+        self.owned_by = owned_by
+        self.created = created
+
+
+class ProviderModels:
+    """Groups models by provider"""
+
+    def __init__(self, provider: Provider, models: List[Model]):
+        self.provider = provider
+        self.models = models
+
+
 class InferenceGatewayClient:
-    def __init__(self, base_url):
-        self.base_url = base_url
+    """Client for interacting with the Inference Gateway API"""
+
+    def __init__(self, base_url: str, token: Optional[str] = None):
+        """Initialize the client with base URL and optional auth token"""
+        self.base_url = base_url.rstrip('/')
+        self.session = requests.Session()
+        if token:
+            self.session.headers.update({"Authorization": f"Bearer {token}"})
 
-    def list_models(self):
-        response = requests.get(f"{self.base_url}/llms")
+    def list_models(self) -> List[ProviderModels]:
+        """List all available language models"""
+        response = self.session.get(f"{self.base_url}/llms")
         response.raise_for_status()
         return response.json()
 
-    def generate_content(self, provider, model, prompt):
-        payload = {"modelName": model, "prompt": prompt}
-        response = requests.post(f"{self.base_url}/llms/{provider}/generate", json=payload)
+    def generate_content(
+        self,
+        provider: Provider,
+        model: str,
+        messages: List[Message]
+    ) -> Dict:
+        payload = {
+            "model": model,
+            "messages": [msg.to_dict() for msg in messages]
+        }
+
+        response = self.session.post(
+            f"{self.base_url}/llms/{provider.value}/generate",
+            json=payload
+        )
         response.raise_for_status()
         return response.json()
+
+    def health_check(self) -> bool:
+        """Check if the API is healthy"""
+        response = self.session.get(f"{self.base_url}/health")
+        return response.status_code == 200

tests/test_client.py

@@ -1,32 +1,102 @@
-import unittest
-from unittest.mock import patch, Mock
-from inference_gateway.client import InferenceGatewayClient
+import pytest
+from unittest.mock import Mock, patch
+from inference_gateway.client import InferenceGatewayClient, Provider, Role, Message
 
 
-class TestInferenceGatewayClient(unittest.TestCase):
-    def setUp(self):
-        self.client = InferenceGatewayClient("http://localhost:8080")
+@pytest.fixture
+def client():
+    """Create a test client instance"""
+    return InferenceGatewayClient("http://test-api")
 
-    @patch("inference_gateway.client.requests.get")
-    def test_list_models(self, mock_get):
-        mock_response = Mock()
-        mock_response.json.return_value = {"models": ["model1", "model2"]}
-        mock_response.raise_for_status = Mock()
-        mock_get.return_value = mock_response
 
-        models = self.client.list_models()
-        self.assertEqual(models, {"models": ["model1", "model2"]})
+@pytest.fixture
+def mock_response():
+    """Create a mock response"""
+    mock = Mock()
+    mock.status_code = 200
+    mock.json.return_value = {"response": "test"}
+    return mock
 
-    @patch("inference_gateway.client.requests.post")
-    def test_generate_content(self, mock_post):
-        mock_response = Mock()
-        mock_response.json.return_value = {"Response": {"Content": "generated content"}}
-        mock_response.raise_for_status = Mock()
-        mock_post.return_value = mock_response
 
-        response = self.client.generate_content("provider", "model", "prompt")
-        self.assertEqual(response, {"Response": {"Content": "generated content"}})
+def test_client_initialization():
+    """Test client initialization with and without token"""
+    client = InferenceGatewayClient("http://test-api")
+    assert client.base_url == "http://test-api"
+    assert "Authorization" not in client.session.headers
 
+    client_with_token = InferenceGatewayClient(
+        "http://test-api", token="test-token")
+    assert "Authorization" in client_with_token.session.headers
+    assert client_with_token.session.headers["Authorization"] == "Bearer test-token"
 
-if __name__ == "__main__":
-    unittest.main()
+
+@patch("requests.Session.get")
+def test_list_models(mock_get, client, mock_response):
+    """Test listing available models"""
+    mock_get.return_value = mock_response
+    response = client.list_models()
+
+    mock_get.assert_called_once_with("http://test-api/llms")
+    assert response == {"response": "test"}
+
+
+@patch("requests.Session.post")
+def test_generate_content(mock_post, client, mock_response):
+    """Test content generation"""
+    messages = [
+        Message(Role.SYSTEM, "You are a helpful assistant"),
+        Message(Role.USER, "Hello!")
+    ]
+
+    mock_post.return_value = mock_response
+    response = client.generate_content(Provider.OPENAI, "gpt-4", messages)
+
+    mock_post.assert_called_once_with(
+        "http://test-api/llms/openai/generate",
+        json={
+            "model": "gpt-4",
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant"},
+                {"role": "user", "content": "Hello!"}
+            ]
+        }
+    )
+    assert response == {"response": "test"}
+
+
+@patch("requests.Session.get")
+def test_health_check(mock_get, client):
+    """Test health check endpoint"""
+    mock_response = Mock()
+    mock_response.status_code = 200
+    mock_get.return_value = mock_response
+
+    assert client.health_check() is True
+    mock_get.assert_called_once_with("http://test-api/health")
+
+    # Test unhealthy response
+    mock_response.status_code = 500
+    assert client.health_check() is False
+
+
+def test_message_to_dict():
+    """Test Message class serialization"""
+    message = Message(Role.USER, "Hello!")
+    assert message.to_dict() == {"role": "user", "content": "Hello!"}
+
+
+def test_provider_enum():
+    """Test Provider enum values"""
+    assert Provider.OPENAI == "openai"
+    assert Provider.OLLAMA == "ollama"
+    assert Provider.GROQ == "groq"
+    assert Provider.GOOGLE == "google"
+    assert Provider.CLOUDFLARE == "cloudflare"
+    assert Provider.COHERE == "cohere"
+
+
+def test_role_enum():
+    """Test Role enum values"""
+    assert Role.SYSTEM == "system"
+    assert Role.USER == "user"
+    assert Role.ASSISTANT == "assistant"