Update app.py

app.py

@@ -5,7 +5,8 @@ import ssl
 from contextlib import closing
 from typing import Optional, Tuple
 import datetime
-
+import promptlayer
+promptlayer.api_key = os.environ.get("PROMPTLAYER_KEY")
 import boto3
 import gradio as gr
 import requests
@@ -18,7 +19,8 @@ from langchain import ConversationChain, LLMChain
 
 from langchain.agents import load_tools, initialize_agent
 from langchain.chains.conversation.memory import ConversationBufferMemory
-from langchain.llms import OpenAI
+# from langchain.llms import OpenAI
+from promptlayer.langchain.llms import OpenAI
 from threading import Lock
 
 # Console to variable
@@ -49,7 +51,7 @@ tmdb_bearer_token = os.environ["TMDB_BEARER_TOKEN"]
 
 TOOLS_LIST = ['serpapi', 'wolfram-alpha', 'pal-math', 'pal-colored-objects', 'news-api']  #'google-search','news-api','tmdb-api','open-meteo-api'
 TOOLS_DEFAULT_LIST = ['serpapi', 'wolfram-alpha', 'pal-math', 'pal-colored-objects', 'news-api']
-BUG_FOUND_MSG = "Congratulations, you've found a bug in this application!"
+BUG_FOUND_MSG = "Error in the return response. Please try again."
 # AUTH_ERR_MSG = "Please paste your OpenAI key from openai.com to use this application. It is not necessary to hit a button or key after pasting it."
 AUTH_ERR_MSG = "Please paste your OpenAI key from openai.com to use this application. "
 MAX_TOKENS = 2048
@@ -262,29 +264,53 @@ def load_chain(tools_list, llm):
     return chain, express_chain, memory
 
 
-async def set_chain_state_api_key(api_key):
+# async def set_chain_state_api_key(api_key):
+async def set_openai_key(api_key):
     # Set the API key for chain_state
     chain_state.api_key = api_key
 
-async def set_express_chain_state_api_key(api_key):
-    # Set the API key for express_chain_state
-    express_chain_state.api_key = api_key
+# async def set_express_chain_state_api_key(api_key):
+#     # Set the API key for express_chain_state
+#     express_chain_state.api_key = api_key
 
-async def set_llm_state_api_key(api_key):
+# async def set_llm_state_api_key(api_key):
     # Set the API key for llm_state
-    llm_state.api_key = api_key
+#     llm_state.api_key = api_key
 
-async def set_embeddings_state_api_key(api_key):
+# async def set_embeddings_state_api_key(api_key):
     # Set the API key for embeddings_state
-    embeddings_state.api_key = api_key
+#     embeddings_state.api_key = api_key
 
-async def set_qa_chain_state_api_key(api_key):
+# async def set_qa_chain_state_api_key(api_key):
     # Set the API key for qa_chain_state
-    qa_chain_state.api_key = api_key
+#     qa_chain_state.api_key = api_key
 
-async def set_memory_state_api_key(api_key):
+# async def set_memory_state_api_key(api_key):
     # Set the API key for memory_state
-    memory_state.api_key = api_key
+#     memory_state.api_key = api_key
+
+    if api_key and api_key.startswith("sk-") and len(api_key) > 50:
+        os.environ["OPENAI_API_KEY"] = api_key
+        print("\n\n ++++++++++++++ Setting OpenAI API key ++++++++++++++ \n\n")
+        print(str(datetime.datetime.now()) + ": Before OpenAI, OPENAI_API_KEY length: " + str(
+            len(os.environ["OPENAI_API_KEY"])))
+        llm = OpenAI(temperature=0, max_tokens=MAX_TOKENS)
+        print(str(datetime.datetime.now()) + ": After OpenAI, OPENAI_API_KEY length: " + str(
+            len(os.environ["OPENAI_API_KEY"])))
+        chain, express_chain, memory = load_chain(TOOLS_DEFAULT_LIST, llm)
+
+        # Pertains to question answering functionality
+        embeddings = OpenAIEmbeddings()
+        qa_chain = load_qa_chain(OpenAI(temperature=0), chain_type="stuff")
+
+        print(str(datetime.datetime.now()) + ": After load_chain, OPENAI_API_KEY length: " + str(
+            len(os.environ["OPENAI_API_KEY"])))
+        os.environ["OPENAI_API_KEY"] = ""
+        return chain, express_chain, llm, embeddings, qa_chain, memory
+    return None, None, None, None, None, None
+
+PROMPTLAYER_API_BASE = "https://api.promptlayer.com"
+
 
 async def set_openai_api_key(api_key):
     """
@@ -924,7 +950,7 @@ openai_api_key_textbox.change(set_openai_api_key,
                                   inputs=[openai_api_key_textbox],
                                   outputs=None, _js="(api_key) => localStorage.setItem('open_api_key', api_key)")
 
-    openai_api_key_textbox.change(lambda api_key: asyncio.run(set_openai_api_key(api_key)),
+    openai_api_key_textbox.change(set_openai_api_key,
                               inputs=[openai_api_key_textbox],
                               outputs=[chain_state, express_chain_state, llm_state, embeddings_state,
                                        qa_chain_state, memory_state])