Updates OpenAI plugin

src/dispatch/plugins/dispatch_openai/config.py

@@ -1,31 +1,19 @@
 from pydantic import Field, SecretStr
-from typing import List, TypeVar
 
 from dispatch.config import BaseConfigurationModel
 
-Stop = TypeVar("Stop", str, List)
-
 
 class OpenAIConfiguration(BaseConfigurationModel):
     """OpenAI configuration description."""
 
     api_key: SecretStr = Field(title="API Key", description="Your secret OpenAI API key.")
-    model: str = Field("text-davinci-003", title="Model", description="")
-    max_tokens: int = Field(
-        50,
-        title="Max Tokens",
-        description="The maximum number of tokens to generate in the completion.",
-    )
-    temperature: float = Field(
-        1, title="Temperature", description="What sampling temperature to use, between 0 and 2."
-    )
-    n: int = Field(
-        1,
-        title="Number of completions (n)",
-        description="How many completions to generate for each prompt.",
+    model: str = Field(
+        "gpt-3.5-turbo",
+        title="Model",
+        description="Available models can be found at https://platform.openai.com/docs/models",
     )
-    stop: Stop = Field(
-        None,
-        title="Stop",
-        description="Up to 4 sequences where the API will stop generating further tokens. The returned text will not contain the stop sequence.",
+    system_message: str = Field(
+        "You are a helpful assistant.",
+        title="System Message",
+        description="The system message to help set the behavior of the assistant.",
     )

src/dispatch/plugins/dispatch_openai/plugin.py

@@ -7,8 +7,7 @@
 """
 import logging
 
-import openai
-from openai import util
+from openai import OpenAI
 
 from dispatch.decorators import apply, counter, timer
 from dispatch.plugins import dispatch_openai as openai_plugin
@@ -35,19 +34,24 @@ def __init__(self):
         self.configuration_schema = OpenAIConfiguration
 
     def completion(self, prompt: str) -> dict:
-        openai.api_key = self.api_key
+        client = OpenAI(api_key=self.api_key)
 
         try:
-            response = openai.Completion.create(
-                max_tokens=self.max_tokens,
+            completion = client.chat.completions.create(
                 model=self.model,
-                n=self.n,
-                prompt=prompt,
-                stop=self.stop,
-                temperature=self.temperature,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": self.system_message,
+                    },
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    },
+                ],
             )
         except Exception as e:
             logger.error(e)
             raise
 
-        return util.convert_to_dict(response)
+        return completion.choices[0].message