From e720fafff307ac06a74c4861e2c146b991439d49 Mon Sep 17 00:00:00 2001
From: Joshua Carroll <joshua.carroll@snowflake.com>
Date: Mon, 13 Nov 2023 14:30:44 -0800
Subject: [PATCH] Bump openai version to >=1 and fix up APIs where needed

---
 Chatbot.py                         | 12 ++++----
 app_test.py                        | 44 ++++++++++++++++++------------
 pages/5_Chat_with_user_feedback.py |  7 ++---
 requirements.txt                   |  2 +-
 4 files changed, 37 insertions(+), 28 deletions(-)

diff --git a/Chatbot.py b/Chatbot.py
index 960f50dce..b6fe7ee32 100644
--- a/Chatbot.py
+++ b/Chatbot.py
@@ -1,4 +1,4 @@
-import openai
+from openai import OpenAI
 import streamlit as st
 
 with st.sidebar:
@@ -20,10 +20,10 @@
         st.info("Please add your OpenAI API key to continue.")
         st.stop()
 
-    openai.api_key = openai_api_key
+    client = OpenAI(api_key=openai_api_key)
     st.session_state.messages.append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
-    response = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=st.session_state.messages)
-    msg = response.choices[0].message
-    st.session_state.messages.append(msg)
-    st.chat_message("assistant").write(msg.content)
+    response = client.chat.completions.create(model="gpt-3.5-turbo", messages=st.session_state.messages)
+    msg = response.choices[0].message.content
+    st.session_state.messages.append({"role": "assistant", "content": msg})
+    st.chat_message("assistant").write(msg)
diff --git a/app_test.py b/app_test.py
index 75cf76570..829f04856 100644
--- a/app_test.py
+++ b/app_test.py
@@ -1,21 +1,31 @@
+import datetime
 from unittest.mock import patch
 from streamlit.testing.v1 import AppTest
-from openai.openai_object import OpenAIObject
-
-
-# See https://github.com/openai/openai-python/issues/398
-def create_openai_object_sync(response: str, role: str = "assistant") -> OpenAIObject:
-    obj = OpenAIObject()
-    message = OpenAIObject()
-    content = OpenAIObject()
-    content.content = response
-    content.role = role
-    message.message = content
-    obj.choices = [message]
-    return obj
-
-
-@patch("openai.ChatCompletion.create")
+from openai.types.chat import ChatCompletionMessage
+from openai.types.chat.chat_completion import ChatCompletion, Choice
+
+
+# See https://github.com/openai/openai-python/issues/715#issuecomment-1809203346
+def create_chat_completion(response: str, role: str = "assistant") -> ChatCompletion:
+    return ChatCompletion(
+        id="foo",
+        model="gpt-3.5-turbo",
+        object="chat.completion",
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=ChatCompletionMessage(
+                    content=response,
+                    role=role,
+                ),
+            )
+        ],
+        created=int(datetime.datetime.now().timestamp()),
+    )
+
+
+@patch("openai.resources.chat.Completions.create")
 def test_Chatbot(openai_create):
     at = AppTest.from_file("Chatbot.py").run()
     assert not at.exception
@@ -23,7 +33,7 @@ def test_Chatbot(openai_create):
     assert at.info[0].value == "Please add your OpenAI API key to continue."
 
     JOKE = "Why did the chicken cross the road? To get to the other side."
-    openai_create.return_value = create_openai_object_sync(JOKE)
+    openai_create.return_value = create_chat_completion(JOKE)
     at.text_input(key="chatbot_api_key").set_value("sk-...")
     at.chat_input[0].set_value("Do you know any jokes?").run()
     print(at)
diff --git a/pages/5_Chat_with_user_feedback.py b/pages/5_Chat_with_user_feedback.py
index 23d1fca74..5f58f139c 100644
--- a/pages/5_Chat_with_user_feedback.py
+++ b/pages/5_Chat_with_user_feedback.py
@@ -1,4 +1,4 @@
-import openai
+from openai import OpenAI
 import streamlit as st
 from streamlit_feedback import streamlit_feedback
 import trubrics
@@ -34,9 +34,8 @@
     if not openai_api_key:
         st.info("Please add your OpenAI API key to continue.")
         st.stop()
-    else:
-        openai.api_key = openai_api_key
-    response = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=messages)
+    client = OpenAI(api_key=openai_api_key)
+    response = client.chat.completions.create(model="gpt-3.5-turbo", messages=messages)
     st.session_state["response"] = response.choices[0].message.content
     with st.chat_message("assistant"):
         messages.append({"role": "assistant", "content": st.session_state["response"]})
diff --git a/requirements.txt b/requirements.txt
index 9b90305d8..07fad9589 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
 streamlit>=1.28
 langchain>=0.0.217
-openai
+openai>=1.2
 duckduckgo-search
 anthropic>=0.3.0
 trubrics>=1.4.3