Add Watsonx provider with tests using the python SDK

Author: DaxConverge

.env.sample

@@ -25,3 +25,9 @@ FIREWORKS_API_KEY=
 
 # Together AI
 TOGETHER_API_KEY=
+
+# WatsonX
+WATSONX_SERVICE_URL=
+WATSONX_API_KEY=
+WATSONX_PROJECT_ID=
+

README.md

@@ -7,7 +7,7 @@ Simple, unified interface to multiple Generative AI providers.
 `aisuite` makes it easy for developers to use multiple LLM through a standardized interface. Using an interface similar to OpenAI's, `aisuite` makes it easy to interact with the most popular LLMs and compare the results. It is a thin wrapper around python client libraries, and allows creators to seamlessly swap out and test responses from different LLM providers without changing their code. Today, the library is primarily focussed on chat completions. We will expand it cover more use cases in near future.
 
 Currently supported providers are -
-OpenAI, Anthropic, Azure, Google, AWS, Groq, Mistral, HuggingFace and Ollama.
+OpenAI, Anthropic, Azure, Google, AWS, Groq, Mistral, HuggingFace, Ollama and Watsonx.
 To maximize stability, `aisuite` uses either the HTTP endpoint or the SDK for making calls to the provider.
 
 ## Installation

aisuite/providers/watsonx_provider.py

@@ -0,0 +1,33 @@
+from aisuite.provider import Provider
+import os
+from ibm_watsonx_ai import Credentials
+from ibm_watsonx_ai.foundation_models import ModelInference
+from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
+
+DEFAULT_TEMPERATURE = 0.7
+
+
+class WatsonxProvider(Provider):
+ def __init__(self, **config):
+ self.service_url = config.get("service_url") or os.getenv("WATSONX_SERVICE_URL")
+ self.api_key = config.get("api_key") or os.getenv("WATSONX_API_KEY")
+ self.project_id = config.get("project_id") or os.getenv("WATSONX_PROJECT_ID")
+
+ if not self.service_url or not self.api_key or not self.project_id:
+ raise EnvironmentError(
+ "Missing one or more required WatsonX environment variables: "
+ "WATSONX_SERVICE_URL, WATSONX_API_KEY, WATSONX_PROJECT_ID. "
+ "Please refer to the setup guide: /guides/watsonx.md."
+ )
+
+ def chat_completions_create(self, model, messages, **kwargs):
+ model = ModelInference(
+ model_id=model,
+ params={
+ GenParams.TEMPERATURE: kwargs.get("temperature", DEFAULT_TEMPERATURE),
+ },
+ credentials=Credentials(api_key=self.api_key, url=self.service_url),
+ project_id=self.project_id,
+ )
+
+ return model.chat(prompt=messages, **kwargs)

guides/watsonx.md

@@ -0,0 +1,83 @@
+# Watsonx with `aisuite`
+
+A a step-by-step guide to set up Watsonx with the `aisuite` library, enabling you to use IBM Watsonx's powerful AI models for various tasks.
+
+## Setup Instructions
+
+### Step 1: Create a Watsonx Account
+
+1. Visit [IBM Watsonx](https://www.ibm.com/watsonx).
+2. Sign up for a new account or log in with your existing IBM credentials.
+3. Once logged in, navigate to the **Watsonx Dashboard**.
+
+---
+
+### Step 2: Obtain API Credentials
+
+1. **Generate an API Key**:
+ - Go to the **API Keys** section in your Watsonx account settings.
+ - Click on **Create API Key**.
+ - Provide a name for your API key (e.g., `MyWatsonxKey`).
+ - Click **Generate**, then download or copy the API key. **Keep this key secure!**
+
+2. **Locate the Service URL**:
+ - Go to the **Endpoints** section in the Watsonx dashboard.
+ - Find the URL corresponding to your service and note it. This is your `WATSONX_SERVICE_URL`.
+
+3. **Get the Project ID**:
+ - Navigate to the **Projects** tab in the dashboard.
+ - Select the project you want to use.
+ - Copy the **Project ID**. This will serve as your `WATSONX_PROJECT_ID`.
+
+---
+
+### Step 3: Set Environment Variables
+
+To simplify authentication, set the following environment variables:
+
+Run the following commands in your terminal:
+
+```bash
+export WATSONX_API_KEY="your-watsonx-api-key"
+export WATSONX_SERVICE_URL="your-watsonx-service-url"
+export WATSONX_PROJECT_ID="your-watsonx-project-id"
+```
+
+
+## Create a Chat Completion
+
+Install the `ibm-watsonx-ai` Python client:
+
+Example with pip:
+
+```shell
+pip install ibm-watsonx-ai
+```
+
+Example with poetry:
+
+```shell
+poetry add ibm-watsonx-ai
+```
+
+In your code:
+
+```python
+import aisuite as ai
+client = ai.Client()
+
+provider = "watsonx"
+model_id = "meta-llama/llama-3-70b-instruct"
+
+messages = [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Tell me a joke."},
+]
+
+response = client.chat.completions.create(
+ model=f"{provider}:{model_id}",
+ messages=messages,
+)
+
+print(response.choices[0].message.content)
+```

poetry.lock

@@ -13,6 +13,7 @@ vertexai = { version = "^1.63.0", optional = true }
 groq = { version = "^0.9.0", optional = true }
 mistralai = { version = "^1.0.3", optional = true }
 openai = { version = "^1.35.8", optional = true }
+ibm-watsonx-ai = { version = "^1.1.16", optional = true }
 
 # Optional dependencies for different providers
 [tool.poetry.extras]
@@ -25,7 +26,8 @@ huggingface = []
 mistral = ["mistralai"]
 ollama = []
 openai = ["openai"]
-all = ["anthropic", "aws", "google", "groq", "mistral", "openai"] # To install all providers
+watsonx = ["ibm-watsonx-ai"]
+all = ["anthropic", "aws", "google", "groq", "mistral", "openai", "watsonx"] # To install all providers
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.2.2"
@@ -44,6 +46,7 @@ chromadb = "^0.5.4"
 sentence-transformers = "^3.0.1"
 datasets = "^2.20.0"
 vertexai = "^1.63.0"
+ibm-watsonx-ai = "^1.1.16"
 
 [build-system]
 requires = ["poetry-core"]

pyproject.toml

@@ -4,6 +4,7 @@
 
 
 class TestClient(unittest.TestCase):
+
  @patch("aisuite.providers.mistral_provider.MistralProvider.chat_completions_create")
  @patch("aisuite.providers.groq_provider.GroqProvider.chat_completions_create")
  @patch("aisuite.providers.openai_provider.OpenaiProvider.chat_completions_create")
@@ -16,6 +17,7 @@ class TestClient(unittest.TestCase):
  @patch(
  "aisuite.providers.fireworks_provider.FireworksProvider.chat_completions_create"
  )
+ @patch("aisuite.providers.watsonx_provider.WatsonxProvider.chat_completions_create")
  def test_client_chat_completions(
  self,
  mock_fireworks,
@@ -26,6 +28,7 @@ def test_client_chat_completions(
  mock_openai,
  mock_groq,
  mock_mistral,
+ mock_watsonx,
  ):
  # Mock responses from providers
  mock_openai.return_value = "OpenAI Response"
@@ -36,6 +39,7 @@ def test_client_chat_completions(
  mock_mistral.return_value = "Mistral Response"
  mock_google.return_value = "Google Response"
  mock_fireworks.return_value = "Fireworks Response"
+ mock_watsonx.return_value = "Watsonx Response"
 
  # Provider configurations
  provider_configs = {
@@ -64,6 +68,11 @@ def test_client_chat_completions(
  "fireworks": {
  "api_key": "fireworks-api-key",
  },
+ "watsonx": {
+ "service_url": "https://watsonx-service-url.com",
+ "api_key": "watsonx-api-key",
+ "project_id": "watsonx-project-id",
+ },
  }
 
  # Initialize the client
@@ -134,6 +143,14 @@ def test_client_chat_completions(
  self.assertEqual(fireworks_response, "Fireworks Response")
  mock_fireworks.assert_called_once()
 
+ # Test Watsonx model
+ watsonx_model = "watsonx" + ":" + "watsonx-model"
+ watsonx_response = client.chat.completions.create(
+ watsonx_model, messages=messages
+ )
+ self.assertEqual(watsonx_response, "Watsonx Response")
+ mock_watsonx.assert_called_once()
+
  # Test that new instances of Completion are not created each time we make an inference call.
  compl_instance = client.chat.completions
  next_compl_instance = client.chat.completions

tests/client/test_client.py

@@ -0,0 +1,63 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from ibm_watsonx_ai import Credentials
+from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
+
+from aisuite.providers.watsonx_provider import WatsonxProvider
+
+
+@pytest.fixture(autouse=True)
+def set_api_key_env_var(monkeypatch):
+ """Fixture to set environment variables for tests."""
+ monkeypatch.setenv("WATSONX_SERVICE_URL", "https://watsonx-service-url.com")
+ monkeypatch.setenv("WATSONX_API_KEY", "test-api-key")
+ monkeypatch.setenv("WATSONX_PROJECT_ID", "test-project-id")
+
+
+def test_watsonx_provider():
+ """High-level test that the provider is initialized and chat completions are requested successfully."""
+
+ user_greeting = "Hello!"
+ message_history = [{"role": "user", "content": user_greeting}]
+ selected_model = "our-favorite-model"
+ chosen_temperature = 0.7
+ response_text_content = "mocked-text-response-from-model"
+
+ provider = WatsonxProvider()
+ mock_response = MagicMock()
+ mock_response.choices = [MagicMock()]
+ mock_response.choices[0].message = MagicMock()
+ mock_response.choices[0].message.content = response_text_content
+
+ with patch(
+ "aisuite.providers.watsonx_provider.ModelInference"
+ ) as mock_model_inference:
+ mock_model = MagicMock()
+ mock_model_inference.return_value = mock_model
+ mock_model.chat.return_value = mock_response
+
+ response = provider.chat_completions_create(
+ messages=message_history,
+ model=selected_model,
+ temperature=chosen_temperature,
+ )
+
+ # Assert that ModelInference was called with correct arguments.
+ mock_model_inference.assert_called_once()
+ args, kwargs = mock_model_inference.call_args
+ assert kwargs["model_id"] == selected_model
+ assert kwargs["params"] == {GenParams.TEMPERATURE: chosen_temperature}
+
+ # Assert that the credentials have the correct API key and service URL.
+ credentials = kwargs["credentials"]
+ assert credentials.api_key == provider.api_key
+ assert credentials.url == provider.service_url
+
+ # Assert that chat was called with correct history and temperature.
+ mock_model.chat.assert_called_once_with(
+ prompt=message_history,
+ temperature=chosen_temperature,
+ )
+
+ assert response.choices[0].message.content == response_text_content