Update llm_monitors to work with openai>=1.0.0

gustavocidornelas · whoseoyster · commit d8483da1fa2c · 2023-11-27T00:22:17.000Z
diff --git a/examples/monitoring/quickstart/llms/openai_llm_monitor.ipynb b/examples/monitoring/quickstart/llms/openai_llm_monitor.ipynb
@@ -29,7 +29,12 @@
    "outputs": [],
    "source": [
     "import os\n",
+    "import openai\n",
+    "\n",
+    "# OpenAI env variable\n",
+    "os.environ[\"OPENAI_API_KEY\"] = \"YOUR_OPENAI_API_KEY_HERE\"\n",
     "\n",
+    "# Openlayer env variables\n",
     "os.environ[\"OPENLAYER_API_KEY\"] = \"YOUR_OPENLAYER_API_KEY_HERE\"\n",
     "os.environ[\"OPENLAYER_PROJECT_NAME\"] = \"Your Openlayer Project Name Here\" "
    ]
@@ -49,9 +54,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import openlayer\n",
     "from openlayer import llm_monitors\n",
     "\n",
-    "openai_monitor = llm_monitors.OpenAIMonitor(publish=True)  # with publish=True, every row gets published to Openlayer automatically"
+    "# If you're using `openai>=1.0.0`:\n",
+    "openai_client = openai.OpenAI()\n",
+    "openai_monitor = llm_monitors.OpenAIMonitor(client=openai_client, publish=True)  # with publish=True, every row gets published to Openlayer automatically\n",
+    "\n",
+    "# Otherwise, use:\n",
+    "# openai_monitor = llm_monitors.OpenAIMonitor(publish=True)  # with publish=True, every row gets published to Openlayer automatically"
    ]
   },
   {
@@ -83,45 +94,25 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fb648049-00bd-447c-8feb-ecf794d45ba7",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import openai\n",
-    "\n",
-    "openai.api_key = \"YOUR_OPENAI_API_KEY_HERE\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7d1e8c74-c667-44a7-b8ea-8f5b83049ea7",
+   "id": "e00c1c79",
    "metadata": {},
    "outputs": [],
    "source": [
-    "completion = openai.ChatCompletion.create(\n",
-    "  model=\"gpt-3.5-turbo\",\n",
+    "# If you're using `openai>=1.0.0`:\n",
+    "completion = openai_client.chat.completions.create(model=\"gpt-3.5-turbo\",\n",
     "  messages=[\n",
     "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
     "    {\"role\": \"user\", \"content\": \"How are you doing today?\"}\n",
-    "  ]\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "58dd9603",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "completion = openai.ChatCompletion.create(\n",
-    "  model=\"gpt-3.5-turbo\",\n",
-    "  messages=[\n",
-    "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
-    "    {\"role\": \"user\", \"content\": \"Is Python strongly typed?\"}\n",
-    "  ]\n",
-    ")"
+    "  ])\n",
+    "\n",
+    "# Othwewise, use:\n",
+    "# completion = openai.ChatCompletion.create(\n",
+    "#   model=\"gpt-3.5-turbo\",\n",
+    "#   messages=[\n",
+    "#     {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
+    "#     {\"role\": \"user\", \"content\": \"How are you doing today?\"}\n",
+    "#   ]\n",
+    "# )"
    ]
   },
   {
diff --git a/openlayer/llm_monitors.py b/openlayer/llm_monitors.py
@@ -9,7 +9,7 @@
 
 import openlayer
 
-from . import utils
+from . import tasks, utils
 
 logger = logging.getLogger(__name__)
 
@@ -23,6 +23,8 @@ class OpenAIMonitor:
         Whether to publish the data to Openlayer as soon as it is available. If True,
         the Openlayer credentials must be provided (either as keyword arguments or as
         environment variables).
+    client : openai.api_client.Client, optional
+        The OpenAI client. It is required if you are using openai>=1.0.0.
     openlayer_api_key : str, optional
         The Openlayer API key. If not provided, it is read from the environment
         variable `OPENLAYER_API_KEY`. This is required if `publish` is set to True.
@@ -44,22 +46,32 @@ class OpenAIMonitor:
     1. Set the environment variables:
 
     .. code-block:: bash
+        export OPENAI_API_KEY=<your-openai-api-key>
 
-        export OPENLAYER_API_KEY=<your-api-key>
+        export OPENLAYER_API_KEY=<your-openlayer-api-key>
         export OPENLAYER_PROJECT_NAME=<your-project-name>
 
     2. Instantiate the monitor:
 
+    ** If you are using openai<1.0.0 **
     >>> from opemlayer import llm_monitors
     >>>
     >>> monitor = llm_monitors.OpenAIMonitor(publish=True)
 
+    ** If you are using openai>=1.0.0 **
+    >>> from opemlayer import llm_monitors
+    >>> from openai import OpenAI
+    >>>
+    >>> openai_client = OpenAI()
+    >>> monitor = llm_monitors.OpenAIMonitor(publish=True, client=openai_client)
+
     3. Start monitoring:
 
     >>> monitor.start_monitoring()
 
     From this point onwards, you can continue making requests to your model normally:
 
+    ** If you are using openai<1.0.0 **
     >>> openai.ChatCompletion.create(
     >>>     model="gpt-3.5-turbo",
     >>>     messages=[
@@ -68,6 +80,15 @@ class OpenAIMonitor:
     >>>     ],
     >>> )
 
+    ** If you are using openai>=1.0.0 **
+    >>> openai_client.chat.completions.create(
+    >>>     model="gpt-3.5-turbo",
+    >>>     messages=[
+    >>>         {"role": "system", "content": "You are a helpful assistant."},
+    >>>         {"role": "user", "content": "How are you doing today?"}
+    >>>     ],
+    >>> )
+
     Your data is automatically being published to your Openlayer project!
 
     If you no longer want to monitor your model, you can stop monitoring by calling:
@@ -79,6 +100,7 @@ class OpenAIMonitor:
     def __init__(
         self,
         publish: bool = False,
+        client=None,
         openlayer_api_key: Optional[str] = None,
         openlayer_project_name: Optional[str] = None,
         openlayer_inference_pipeline_name: Optional[str] = None,
@@ -97,6 +119,13 @@ def __init__(
         self._load_inference_pipeline()
 
         # OpenAI setup
+        self.openai_version = openai.__version__
+        if self.openai_version.split(".")[0] == "1" and client is None:
+            raise ValueError(
+                "You must provide the OpenAI client for as the kwarg `client` for"
+                " openai>=1.0.0."
+            )
+        self.openai_client = client
         self.create_chat_completion: callable = None
         self.create_completion: callable = None
         self.modified_create_chat_completion: callable = None
@@ -148,7 +177,9 @@ def _load_inference_pipeline(self) -> None:
                 client = openlayer.OpenlayerClient(
                     api_key=self.openlayer_api_key,
                 )
-                project = client.load_project(name=self.openlayer_project_name)
+                project = client.create_project(
+                    name=self.openlayer_project_name, task_type=tasks.TaskType.LLM
+                )
                 if self.openlayer_inference_pipeline_name:
                     inference_pipeline = project.load_inference_pipeline(
                         name=self.openlayer_inference_pipeline_name
@@ -160,8 +191,14 @@ def _load_inference_pipeline(self) -> None:
 
     def _initialize_openai(self) -> None:
         """Initializes the OpenAI attributes."""
-        self.create_chat_completion = openai.ChatCompletion.create
-        self.create_completion = openai.Completion.create
+        if self.openai_version.split(".")[0] == "0":
+            openai_api_key = utils.get_env_variable("OPENAI_API_KEY")
+            openai.api_key = openai_api_key
+            self.create_chat_completion = openai.ChatCompletion.create
+            self.create_completion = openai.Completion.create
+        else:
+            self.create_chat_completion = self.openai_client.chat.completions.create
+            self.create_completion = self.openai_client.completions.create
         self.modified_create_chat_completion = (
             self._get_modified_create_chat_completion()
         )
@@ -178,7 +215,7 @@ def modified_create_chat_completion(*args, **kwargs) -> str:
             try:
                 input_data = self._format_user_messages(kwargs["messages"])
                 output_data = response.choices[0].message.content.strip()
-                num_of_tokens = response.usage["total_tokens"]
+                num_of_tokens = response.usage.total_tokens
 
                 self._append_row_to_df(
                     input_data=input_data,
@@ -211,7 +248,7 @@ def modified_create_completion(*args, **kwargs):
 
                 for input_data, choices in zip(prompts, choices_splits):
                     output_data = choices[0].text.strip()
-                    num_of_tokens = int(response.usage["total_tokens"] / len(prompts))
+                    num_of_tokens = int(response.usage.total_tokens / len(prompts))
 
                     self._append_row_to_df(
                         input_data=input_data,
@@ -313,8 +350,14 @@ def start_monitoring(self) -> None:
 
     def _overwrite_completion_methods(self) -> None:
         """Overwrites OpenAI's completion methods with the modified versions."""
-        openai.ChatCompletion.create = self.modified_create_chat_completion
-        openai.Completion.create = self.modified_create_completion
+        if self.openai_version.split(".")[0] == "0":
+            openai.ChatCompletion.create = self.modified_create_chat_completion
+            openai.Completion.create = self.modified_create_completion
+        else:
+            self.openai_client.chat.completions.create = (
+                self.modified_create_chat_completion
+            )
+            self.openai_client.completions.create = self.modified_create_completion
 
     def stop_monitoring(self):
         """Switches monitoring for OpenAI LLMs off.
@@ -335,8 +378,12 @@ def stop_monitoring(self):
 
     def _restore_completion_methods(self) -> None:
         """Restores OpenAI's completion methods to their original versions."""
-        openai.ChatCompletion.create = self.create_chat_completion
-        openai.Completion.create = self.create_completion
+        if self.openai_version.split(".")[0] == "0":
+            openai.ChatCompletion.create = self.create_chat_completion
+            openai.Completion.create = self.create_completion
+        else:
+            self.openai_client.chat.completions.create = self.create_chat_completion
+            self.openai_client.completions.create = self.create_completion
 
     def publish_batch_data(self):
         """Manually publish the accumulated data to Openlayer when automatic publishing
