Merge pull request #16 from andrewyng/add-openai-tests

Add OpenAI test file

examples/multi_fm_client.ipynb

@@ -14,15 +14,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 22,
    "id": "initial_id",
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-07-04T15:30:02.064319Z",
      "start_time": "2024-07-04T15:30:02.051986Z"
     }
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "import sys\n",
     "sys.path.append('../../aimodels')\n",
@@ -236,11 +247,26 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 23,
    "id": "611210a4dc92845f",
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Why did the pirate go to the seafood restaurant? \n",
+      "Because he heard they had some great fish tales! Arrr!\n"
+     ]
+    }
+   ],
+   "source": [
+    "openai_gpt35 = \"openai:gpt-3.5-turbo\"\n",
+    "\n",
+    "response = client.chat.completions.create(model=openai_gpt35, messages=messages, temperature=0.75)\n",
+    "\n",
+    "print(response.choices[0].message.content)"
+   ]
   }
  ],
  "metadata": {

tests/providers/test_openai_interface.py

@@ -0,0 +1,42 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from aimodels.providers.openai_interface import OpenAIInterface
+
+
+@pytest.fixture(autouse=True)
+def set_api_key_env_var(monkeypatch):
+    """Set environment variables for tests."""
+    monkeypatch.setenv("OPENAI_API_KEY", "test-api-key")
+
+
+def test_openai_interface():
+    """Test chat completions."""
+
+    user_greeting = "Hello Hello!"
+    message_history = [{"role": "user", "content": user_greeting}]
+    selected_model = "our-favorite-model"
+    chosen_temperature = 0.9
+    response_text_content = "mocked-text-response-from-model"
+
+    interface = OpenAIInterface()
+    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(
+        interface.openai_client.chat.completions, "create", return_value=mock_response
+    ) as mock_create:
+        response = interface.chat_completion_create(
+            messages=message_history,
+            model=selected_model,
+            temperature=chosen_temperature,
+        )
+
+        mock_create.assert_called_with(
+            model=selected_model,
+            messages=message_history,
+            temperature=chosen_temperature,
+        )
+
+        assert response.choices[0].message.content == response_text_content