Added ability to save record into dynamodb

Refactored code
Updated logic for getting suggestions from stream

app.py

@@ -1,21 +1,43 @@
 from openai import OpenAI
 import gradio as gr
+import json
+from bot_actions import functions_dictionary, save_record
 import os
  
+CSS ="""
+.contain { display: flex; flex-direction: column; }
+.gradio-container { height: 100vh !important; }
+#chatbot { flex-grow: 1; overflow: auto;}
+"""
+
 client = OpenAI(api_key=os.environ['openAIToken'])
 assistantId = os.environ['assistantId']
 initial_message = os.environ['initialMessage']
 
-# client = OpenAI(api_key="sk-proj-4GPZmEnOFsZx0fZeQ9B3T3BlbkFJkXHmGCuW0rGmK0RAavtY")
-# assistantId = "asst_Rce5MptxfyEwp4NjImSDvXGD"
-# initial_message = "Hello. Let's start defining new facility. Also provide options."
+def handle_requires_action(data):
+  tool_outputs = []
+  for tool in data.required_action.submit_tool_outputs.tool_calls:
+    function_name = tool.function.name
+    function_args = json.loads(tool.function.arguments)
+    print(function_name)
+    print(function_args)
+    if tool.function.name == "save_record":
+      try:
+        result = functions_dictionary[tool.function.name](**function_args)
+      except Exception as e:
+        print(e)
+      tool_outputs.append({"tool_call_id": tool.id, "output": result})
+    
+  # Submit all tool_outputs at the same time
+  return tool_outputs
+
 
-def createThread(sessionStorage):
+def create_thread_openai(sessionStorage):
   streaming_thread = client.beta.threads.create()
-  sessionStorage.append(streaming_thread.id)
+  sessionStorage["threadId"] = streaming_thread.id
   return sessionStorage
 
-def addMessage(text, threadId):
+def add_message_to_openai(text, threadId):
   print("User message: ", text)
   return client.beta.threads.messages.create(
     thread_id=threadId,
@@ -23,13 +45,10 @@ def addMessage(text, threadId):
     content=text
   )
 
-list_of_suggestions = []
-latest_suggestions = ""
-
-def handle_suggestions(string_of_suggestions):
+def transform_suggestions_into_list(string_of_suggestions):
   local_message = None
   parts = string_of_suggestions.split('#s#')
-  list_of_suggestions = [x.strip('"') for x in parts[0].strip('][').split('", ') if x.strip('"')]
+  list_of_suggestions = json.loads(parts[0])
   list_of_suggestions = [ x for x in list_of_suggestions if x]
   if len(parts) > 1:
     local_message = parts[1]
@@ -40,52 +59,84 @@ def create_suggestions_list(suggestions):
   update_hide = [gr.update(visible=False, value="") for _ in range(6-len(suggestions))]
   return update_show + update_hide
 
-btn_list = []
-list_of_suggestions = ""
-
-CSS ="""
-.contain { display: flex; flex-direction: column; }
-.gradio-container { height: 100vh !important; }
-#chatbot { flex-grow: 1; overflow: auto;}
-"""
+def process_text_chunk(text, list_of_suggestions, string_of_suggestions, is_loading_suggestions):
+  print(text, end="", flush=True)
+  local_message = None
+  if "[" in text:
+    is_loading_suggestions = True
+
+  if is_loading_suggestions != True:
+    local_message = text
+  else:
+    string_of_suggestions = string_of_suggestions + text
+    if "#s#" in string_of_suggestions:
+      is_loading_suggestions = False
+      list_of_suggestions, local_message = transform_suggestions_into_list(string_of_suggestions)
+      string_of_suggestions = ""
+    elif "]" in string_of_suggestions and "]#" not in string_of_suggestions and not string_of_suggestions.endswith("]"):
+      is_loading_suggestions = False
+      local_message = string_of_suggestions
+      string_of_suggestions = ""
+  return local_message, list_of_suggestions, string_of_suggestions, is_loading_suggestions
 
-def initiate_chatting(storage):
-  threadId = storage[0]
-  addMessage(initial_message, threadId)
-  global list_of_suggestions
+def handle_events(threadId, chat_history, storage):
   list_of_suggestions = []
   string_of_suggestions = ""
   is_loading_suggestions = False
-  is_it_first_response = True
-  response_message = ""
-  with client.beta.threads.runs.stream(
-    thread_id=threadId,
-    assistant_id=assistantId,
-  ) as stream:
-    for text in stream.text_deltas:
-      local_message = None
-      if "[" in text and is_it_first_response:
-        is_loading_suggestions = True
-
-      is_it_first_response = False
-
-      if is_loading_suggestions != True:
-        local_message = text
-      else:
-        string_of_suggestions = string_of_suggestions + text
-        if "#s#" in string_of_suggestions:
-          is_loading_suggestions = False
-          list_of_suggestions, local_message = handle_suggestions(string_of_suggestions)
-      if local_message is not None:
-        response_message += local_message 
-    stream.until_done()
-  
-  return [[None, response_message]]
-
-def createDemo():
+  try:
+    with client.beta.threads.runs.stream(
+      thread_id=threadId,
+      assistant_id=assistantId
+    ) as stream:
+      for event in stream:
+        if event.event == "thread.message.delta" and event.data.delta.content:
+          text = event.data.delta.content[0].text.value
+          local_message, list_of_suggestions, string_of_suggestions, is_loading_suggestions = process_text_chunk(text, list_of_suggestions, string_of_suggestions, is_loading_suggestions)
+          if local_message is not None:
+            chat_history[-1][1] += local_message 
+            yield [ chat_history,  storage]
+        if event.event == 'thread.run.requires_action':
+          tool_outputs = handle_requires_action(event.data)
+          with client.beta.threads.runs.submit_tool_outputs_stream(
+            thread_id=stream.current_run.thread_id,
+            run_id=event.data.id,
+            tool_outputs=tool_outputs,
+          ) as stream:
+            for text in stream.text_deltas:
+              local_message, list_of_suggestions, string_of_suggestions, is_loading_suggestions = process_text_chunk(text, list_of_suggestions, string_of_suggestions, is_loading_suggestions)
+              if local_message is not None:
+                chat_history[-1][1] += local_message
+                yield [chat_history, storage]
+      stream.until_done()
+      print("")
+      storage["list_of_suggestions"] = list_of_suggestions
+      return [chat_history, storage]
+  except Exception as e:
+    print(e)
+    chat_history[-1][1] = "Error occured during processing your message. Please try again"
+    yield [chat_history, storage]
+
+def initiate_chatting(chat_history, storage):
+  threadId = storage["threadId"]
+  chat_history = [[None, ""]]
+  add_message_to_openai(initial_message, threadId)
+  for response in handle_events(threadId, chat_history, storage):
+    yield response 
+
+def respond_on_user_msg(chat_history, storage):
+  message = chat_history[-1][0]
+  threadId = storage["threadId"]
+  print("Responding for threadId: ", threadId)
+  chat_history[-1][1] = ""
+  add_message_to_openai(message, threadId)
+  for response in handle_events(threadId, chat_history, storage):
+    yield response
+
+def create_application():
   with gr.Blocks(css=CSS, fill_height=True) as demo:
-    storage = gr.State([])
+    storage = gr.State({})
     chatbot = gr.Chatbot(label="Facility managment bot", elem_id="chatbot")
+    btn_list = []
     with gr.Row():
       for i in range(6):
           btn = gr.Button(visible=False)
@@ -95,54 +146,11 @@ def createDemo():
       clear_btn = gr.ClearButton(chatbot)
       btn = gr.Button("Submit")
 
-
     def user(user_message, history):
         return "", history + [[user_message, None]]
 
-    def respond(chat_history, storage):
-      global btn_list
-      message = chat_history[-1][0]
-      threadId = storage[0]
-      print("Responding for threadId: ", threadId)
-      chat_history[-1][1] = ""
-      addMessage(message, threadId)
-      global list_of_suggestions
-      list_of_suggestions = []
-      string_of_suggestions = ""
-      is_loading_suggestions = False
-      is_it_first_response = True
-      try:
-        with client.beta.threads.runs.stream(
-          thread_id=threadId,
-          assistant_id=assistantId,
-        ) as stream:
-          for text in stream.text_deltas:
-            print(text, end="")
-            local_message = None
-            if "[" in text and is_it_first_response:
-              is_loading_suggestions = True
-
-            is_it_first_response = False
-
-            if is_loading_suggestions != True:
-              local_message = text
-            else:
-              string_of_suggestions = string_of_suggestions + text
-              if "#s#" in string_of_suggestions:
-                is_loading_suggestions = False
-                list_of_suggestions, local_message = handle_suggestions(string_of_suggestions)
-            if local_message is not None:
-              chat_history[-1][1] += local_message 
-              yield {chatbot: chat_history}
-          stream.until_done()
-        print("")
-      except:
-        list_of_suggestions = []
-        chat_history[-1][1] = "Error occured during processing your message. Please try again"
-        yield {chatbot: chat_history}
-
-    def update_suggestions():
-      global list_of_suggestions
+    def update_suggestions(storage):
+      list_of_suggestions = storage['list_of_suggestions'] or []
       btn_list = create_suggestions_list(list_of_suggestions)
       return btn_list
     
@@ -158,8 +166,8 @@ def createDemo():
       return message_box
 
     add_user_message_flow = [user, [msg,chatbot],  [msg,chatbot]]
-    chat_response_flow = [respond, [chatbot, storage], [chatbot]]
-    update_suggestions_flow = [update_suggestions, None, btn_list]
+    chat_response_flow = [respond_on_user_msg, [chatbot, storage], [chatbot, storage]]
+    update_suggestions_flow = [update_suggestions, storage, btn_list]
     hide_suggestions_flow = [hide_suggestions, None, btn_list]
     disable_msg_flow = [disable_msg, None, msg] 
     enable_msg_flow = [enable_msg, None, msg]
@@ -185,12 +193,12 @@ def createDemo():
               ).then(*update_suggestions_flow
               ).then(*enable_msg_flow)
     
-    demo.load(createThread, inputs=storage, outputs=storage
-        ).then(initiate_chatting, inputs=storage, outputs=chatbot
+    demo.load(create_thread_openai, inputs=storage, outputs=storage
+        ).then(initiate_chatting, inputs=[chatbot, storage], outputs=[chatbot, storage]
         ).then(*update_suggestions_flow
         ).then(*enable_msg_flow)
   return demo
 
 if __name__ == "__main__":
-  demo = createDemo()
+  demo = create_application()
   demo.launch()

bot_actions.py

@@ -0,0 +1,11 @@
+import requests
+import time
+
+def save_record(arg_json):
+  r = requests.post("https://ogtk5k73z8.execute-api.us-east-1.amazonaws.com/Development/facilities", json={'pk': round(time.time()), 'json': arg_json})
+  print(r.status_code, r.reason)
+  return "Done"
+
+functions_dictionary = {
+  "save_record": save_record
+}