Handling of session states

app.py

@@ -23,50 +23,47 @@ import json
 import plotly.graph_objects as go
 import time
 
-#os.environ["OPENAI_API_KEY"] = "..."  # Replace with your key
+os.environ["OPENAI_API_KEY"] = "sk-..."  # Replace with your key
+assistant_id = "asst_7OC3NTeyCjEZrApdLRklplE7"
+
+client = openai.OpenAI()
+assistant = client.beta.assistants.retrieve(assistant_id)
 
 #######################################
 # TOOLS SETUP
 #######################################
 
-def update_map_state(latitude, longitude, zoom):
+def update_map_state(map_metadata_state, latitude, longitude, zoom):
     """OpenAI tool to update map in-app
     """
-    session_state[map_state] = {
-        "latitude": latitude,
-        "longitude": longitude,
-        "zoom": zoom,
-    }
-    print(session_state[map_state])
+
+    map_metadata_state["latitude"] = latitude
+    map_metadata_state["longitude"] = longitude
+    map_metadata_state["zoom"] = zoom
+
+    print(map_metadata_state)
+
     return "Map updated"
 
-def add_markers_state(latitudes, longitudes, labels):
+def add_markers_state(map_metadata_state, latitudes, longitudes, labels):
     """OpenAI tool to update markers in-app
     """
-    session_state[markers_state] = {
-        "lat": latitudes,
-        "lon": longitudes,
-        "text": labels,
-    }
+
+    map_metadata_state["lat"] = latitudes
+    map_metadata_state["lon"] = longitudes
+    map_metadata_state["text"] = labels
+
+    print(map_metadata_state)
+
     return "Markers added"
 
 tool_to_function = {
     "update_map": update_map_state,
-    "add_markers": add_markers_state,
+    "add_markers": add_markers_state
 }
 
 ## Helpers
 
-def get_assistant_id():
-    return session_state[assistant_state].id
-
-def get_thread_id():
-    return session_state[thread_state].id
-
-
-def get_run_id():
-    return session_state[last_openai_run_state].id
-
 def submit_message(assistant_id, thread_id, user_message):
     client.beta.threads.messages.create(
         thread_id=thread_id, role="user", content=user_message
@@ -85,43 +82,17 @@ def get_run_info(run_id, thread_id):
     return run
 
 #######################################
-# SESSION SETUP
+# INITIAL DATA FOR MAP
 #######################################
 
-client = openai.OpenAI()
-assistant_id = "asst_7OC3NTeyCjEZrApdLRklplE7"
-
-session_state = {}
-
-assistant_state = "assistant"
-thread_state = "thread"
-conversation_state = "conversation"
-last_openai_run_state = "last_openai_run"
-map_state = "map"
-markers_state = "markers"
-
-if (assistant_state not in session_state) or (thread_state not in session_state):
-    session_state[assistant_state] = client.beta.assistants.retrieve(assistant_id)
-    session_state[thread_state] = client.beta.threads.create()
-
-if conversation_state not in session_state:
-    session_state[conversation_state] = []
-
-if last_openai_run_state not in session_state:
-    session_state[last_openai_run_state] = None
-
-if map_state not in session_state:
-    session_state[map_state] = {
+map_metadata = {
         "latitude": 48.85,
         "longitude": 2.35,
         "zoom": 12,
-    }
-
-if markers_state not in session_state:
-    session_state[markers_state] = {
         "lat": [],
         "lon": [],
         "text": [],
+
     }
 
 fig = go.Figure(go.Scattermapbox())
@@ -131,29 +102,28 @@ fig.update_layout(
         hovermode='closest',
         mapbox=dict(
             center=go.layout.mapbox.Center(
-                lat=session_state[map_state]["latitude"],
-                lon=session_state[map_state]["longitude"]
+                lat=map_metadata["latitude"],
+                lon=map_metadata["longitude"]
             ),
-            zoom=session_state[map_state]["zoom"]
+            zoom=map_metadata["zoom"]
         ),
 )
 
-def respond(message, chat_history):
+def respond(message, chat_history, thread, map_metadata_state):
 
-    print(chat_history)
+    if thread is None:
+        thread = client.beta.threads.create()
 
-    run = submit_message(get_assistant_id(), get_thread_id(), message)
+    print(map_metadata_state)
 
-    session_state[last_openai_run_state] = run
-
-    print(run)
+    run = submit_message(assistant.id, thread.id, message)
 
     completed = False
 
     # Polling
     while not completed:
 
-        run = get_run_info(get_run_id(), get_thread_id())
+        run = get_run_info(run.id, thread.id)
 
         if run.status == "requires_action":
 
@@ -165,9 +135,9 @@ def respond(message, chat_history):
                 f_name = f.name
                 f_args = json.loads(f.arguments)
 
-                print(f"Launching function {f_name} with args {f_args}")
+                #print(f"Launching function {f_name} with args {f_args}")
 
-                tool_result = tool_to_function[f_name](**f_args)
+                tool_result = tool_to_function[f_name](map_metadata_state,**f_args)
 
                 tools_output.append(
                         {
@@ -176,11 +146,11 @@ def respond(message, chat_history):
                         }
                     )
 
-            print(f"Will submit {tools_output}")
+            #print(f"Will submit {tools_output}")
 
             client.beta.threads.runs.submit_tool_outputs(
-                    thread_id=get_thread_id(),
-                    run_id=get_run_id(),
+                    thread_id=thread.id,
+                    run_id=run.id,
                     tool_outputs=tools_output,
             )
 
@@ -191,12 +161,11 @@ def respond(message, chat_history):
         else:
             time.sleep(0.1)
 
-    session_state[conversation_state] = [
+    dialog = [
         [m.role, m.content[0].text.value]
-        for m in client.beta.threads.messages.list(get_thread_id(), order="asc").data
+        for m in client.beta.threads.messages.list(thread.id, order="asc").data
     ]
 
-    dialog = session_state[conversation_state]
     formatted_dialog = []
     for i in range(int(len(dialog)/2)):
       formatted_dialog.append([dialog[i*2][1],dialog[i*2+1][1]])
@@ -204,17 +173,18 @@ def respond(message, chat_history):
 
     chat_history = formatted_dialog
 
-    fig = None
+    print(formatted_dialog)
 
-    if session_state[markers_state] is None:
+    fig = None
 
+    if len(map_metadata_state["lat"])==0:
       fig = go.Figure(go.Scattermapbox())
 
     else :
       fig = go.Figure(go.Scattermapbox(
-            customdata=session_state[markers_state]["text"],
-            lat=session_state[markers_state]["lat"],
-            lon=session_state[markers_state]["lon"],
+            customdata=map_metadata_state["text"],
+            lat=map_metadata_state["lat"],
+            lon=map_metadata_state["lon"],
             mode='markers',
             marker=go.scattermapbox.Marker(
                 size=18
@@ -228,34 +198,38 @@ def respond(message, chat_history):
         hovermode='closest',
         mapbox=dict(
             center=go.layout.mapbox.Center(
-                lat=session_state[map_state]["latitude"],
-                lon=session_state[map_state]["longitude"]
+                lat=map_metadata_state["latitude"],
+                lon=map_metadata_state["longitude"]
             ),
             zoom=12
         ),
     )
 
-    return "", chat_history, fig
+    return "", chat_history, fig, thread, map_metadata_state
 
 with gr.Blocks(title="OpenAI assistant Wanderlust") as demo:
 
     gr.Markdown("# OpenAI assistant Wanderlust")
 
+    thread = gr.State()
+
+    map_metadata_state = gr.State(map_metadata)
+
     with gr.Column():
         with gr.Row():
 
           chatbot = gr.Chatbot()
           map = gr.Plot(fig)
 
-    msg = gr.Textbox("Move the map to Brussels and locate the best places for waffles")
+    msg = gr.Textbox("Move the map to Brussels and add markers on shops with the best waffles")
 
     with gr.Column():
         with gr.Row():
           submit = gr.Button("Submit")
           clear = gr.ClearButton([msg, chatbot])
 
-    msg.submit(respond, [msg, chatbot], [msg, chatbot, map])
-    submit.click(respond, [msg, chatbot], [msg, chatbot, map])
+    msg.submit(respond, [msg, chatbot, thread, map_metadata_state], [msg, chatbot, map, thread, map_metadata_state])
+    submit.click(respond, [msg, chatbot, thread, map_metadata_state], [msg, chatbot, map, thread, map_metadata_state])
 
     gr.Markdown(
         """