Integration with Nebius AI Studio added

aisuite/providers/nebius_provider.py

@@ -0,0 +1,31 @@
+import os
+from aisuite.provider import Provider
+from openai import Client
+
+
+BASE_URL = "https://api.studio.nebius.ai/v1"
+
+
+class NebiusProvider(Provider):
+    def __init__(self, **config):
+        """
+        Initialize the Nebius AI Studio provider with the given configuration.
+        Pass the entire configuration dictionary to the OpenAI client constructor.
+        """
+        # Ensure API key is provided either in config or via environment variable
+        config.setdefault("api_key", os.getenv("NEBIUS_API_KEY"))
+        if not config["api_key"]:
+            raise ValueError(
+                "Nebius AI Studio API key is missing. Please provide it in the config or set the NEBIUS_API_KEY environment variable. You can get your API key at https://studio.nebius.ai/settings/api-keys"
+            )
+
+        config["base_url"] = BASE_URL
+        # Pass the entire config to the OpenAI client constructor
+        self.client = Client(**config)
+
+    def chat_completions_create(self, model, messages, **kwargs):
+        return self.client.chat.completions.create(
+            model=model,
+            messages=messages,
+            **kwargs  # Pass any additional arguments to the Nebius API
+        )

examples/QnA_with_pdf.ipynb

@@ -102,7 +102,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import aisuite as ai\n",
     "client = ai.Client()\n",
     "messages = [\n",
     "    {\"role\": \"system\", \"content\": \"You are a helpful assistant. Answer the question only based on the below text.\"},\n",
@@ -180,6 +179,25 @@
     "print(response.choices[0].message.content)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Yes, the price of organic avocados is higher than non-organic avocados. According to the text, the average price of organic avocados is generally 35-40% higher than conventional avocados.\n"
+     ]
+    }
+   ],
+   "source": [
+    "nebius_model = \"nebius:meta-llama/Meta-Llama-3.1-8B-Instruct-fast\"\n",
+    "response = client.chat.completions.create(model=nebius_model, messages=messages, top_p=0.01)\n",
+    "print(response.choices[0].message.content)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

examples/client.ipynb

@@ -61,6 +61,7 @@
     "    'AWS_ACCESS_KEY_ID': 'xxx',\n",
     "    'AWS_SECRET_ACCESS_KEY': 'xxx',\n",
     "    'ANTHROPIC_API_KEY': 'xxx',\n",
+    "    'NEBIUS_API_KEY': 'xxx',\n",
     "}\n",
     "\n",
     "# Configure environment\n",
@@ -208,6 +209,18 @@
     "print(response.choices[0].message.content)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f38d033a-a580-4239-9176-27f3d53e7fe1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nebius_model = \"nebius:Qwen/Qwen2.5-1.5B-Instruct\"\n",
+    "response = client.chat.completions.create(model=nebius_model, messages=messages, top_p=0.01)\n",
+    "print(response.choices[0].message.content)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -254,4 +267,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

tests/client/test_client.py

@@ -33,6 +33,9 @@ def provider_configs():
         "fireworks": {
             "api_key": "fireworks-api-key",
         },
+        "nebius": {
+            "api_key": "nebius-api-key",
+        },
         "watsonx": {
             "service_url": "https://watsonx-service-url.com",
             "api_key": "watsonx-api-key",
@@ -84,6 +87,11 @@ def provider_configs():
             "fireworks",
             "fireworks-model",
         ),
+        (
+            "aisuite.providers.nebius_provider.NebiusProvider.chat_completions_create",
+            "nebius",
+            "nebius-model",
+        ),
         (
             "aisuite.providers.watsonx_provider.WatsonxProvider.chat_completions_create",
             "watsonx",

tests/providers/test_nebius_provider.py

@@ -0,0 +1,45 @@
+import pytest
+from unittest.mock import patch, MagicMock
+
+from aisuite.providers.nebius_provider import NebiusProvider
+
+
+@pytest.fixture(autouse=True)
+def set_api_key_env_var(monkeypatch):
+    """Fixture to set environment variables for tests."""
+    monkeypatch.setenv("NEBIUS_API_KEY", "test-api-key")
+
+
+def test_nebius_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.75
+    response_text_content = "mocked-text-response-from-model"
+
+    provider = NebiusProvider()
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message = MagicMock()
+    mock_response.choices[0].message.content = response_text_content
+
+    with patch.object(
+        provider.client.chat.completions,
+        "create",
+        return_value=mock_response,
+    ) as mock_create:
+        response = provider.chat_completions_create(
+            messages=message_history,
+            model=selected_model,
+            temperature=chosen_temperature,
+        )
+
+        mock_create.assert_called_with(
+            messages=message_history,
+            model=selected_model,
+            temperature=chosen_temperature,
+        )
+
+        assert response.choices[0].message.content == response_text_content