My initial chatbot

app.py

@@ -62,65 +62,111 @@ css = """
 }
 """
 
-
-# gradio chatbot main function
+# initialize the chat model and sets up the Gradio interface
 def registry(name, token, examples=None, **kwargs):
   client = genai.Client(api_key=token)
-  chat = client.chats.create(
+  chat_locks = {}  # Dictionary to hold locks for each user's chat
+  chat_sessions = {} # Dictionary to hold each user chat
+  
+
+  # create a new chat session with the specified model and configuration
+  def create_chat():
+    return client.chats.create(
       model=name,
       config=types.GenerateContentConfig(
-          system_instruction=system_instruction,
-          temperature=0.5,
+        system_instruction=system_instruction,
+        temperature=0.5,
       ),
-  )
+    )
+
 
+  # send a user message to Gemini, streams the response back to the chatbot
+  # and updates the history
+  def stream_response(history, text, img, request: gr.Request):
+      user_id = request.client.host
+      if user_id not in chat_locks:
+        chat_locks[user_id] = Lock()
+        chat_sessions[user_id] = create_chat()
+
+      lock = chat_locks[user_id]
+      chat = chat_sessions[user_id]
 
-  # send user message to Gemini and append the answer at the chatbot history
-  def llm_response(history, text, img):
-    if not img:
-      response = chat.send_message(text)
-      history += [(None, response.text)]
-      return history, "", None
-    if img:
       try:
-        img = Image.open(img)
-        response = chat.send_message([text, img])
-        history += [(None, response.text)]
-        return history, "", None
+        with lock:
+          if not img:
+            response_stream = chat.send_message_stream(text)
+          else:
+            try:
+              img = Image.open(img)
+              response_stream = chat.send_message_stream([text, img])
+            except Exception as e:
+              print(f"Error processing image: {str(e)}")
+              return
+
+          partial_message = ""
+          for chunk in response_stream:
+            if chunk.text:
+              partial_message += chunk.text
+              history[-1] = (history[-1][0], partial_message)
+              yield history
+              time.sleep(0.02)  # Small delay to ensure smooth streaming
+
       except Exception as e:
-        print("something wrong happened at the llm_response call")
-        print(str(e))
-        return history, "", None
+        print(f"Error in stream_response: {str(e)}")
+        return
 
-  # instantiate the gradio chatbot
   print("Building the gradio app...")
   with gr.Blocks(css=css) as app:
-      gr.HTML(TITLE)
-      gr.HTML(SUBTITLE)
-      with gr.Row():
-        # define the input fields - image input
-        image_box = gr.Image(type="filepath")
-        chatbot = gr.Chatbot(
-          scale=2,
-          height=500,
-          container=True
-        )
-
-      # define the input fields - text input
-      text_box = gr.Textbox(
-        placeholder="Type your message and press enter and optionally upload an image",
-        container=False,
+    gr.HTML(TITLE)
+    gr.HTML(SUBTITLE)
+    
+    with gr.Row():
+      image_box = gr.Image(type="filepath")
+      chatbot = gr.Chatbot(
+        scale=2,
+        height=500,
+        container=True
       )
 
-      # define the main chatbot logic
-      btn = gr.Button("Send")
-      clicked = btn.click(query_message,
-              [chatbot,text_box,image_box],
-              [chatbot]
-              ).then(llm_response,
-                  [chatbot,text_box,image_box],
-                  [chatbot, text_box, image_box]
-              )
+    text_box = gr.Textbox(
+      placeholder="Type your message and press enter and optionally upload an image",
+      container=False,
+    )
+
+    btn = gr.Button("Send")
+    
+    # Update the event handlers to use streaming
+    btn.click(
+      fn=query_message,
+      inputs=[chatbot, text_box, image_box],
+      outputs=[chatbot],
+    ).then(
+      fn=stream_response,
+      inputs=[chatbot, text_box, image_box],
+      outputs=[chatbot],
+      api_name="stream_response"
+    ).then(
+      fn=lambda: (None, ""),  # Clear the image and text inputs after sending
+      inputs=None,
+      outputs=[image_box, text_box],
+    )
+
+    # Add enter key handler
+    text_box.submit(
+      fn=query_message,
+      inputs=[chatbot, text_box, image_box],
+      outputs=[chatbot],
+    ).then(
+      fn=stream_response,
+      inputs=[chatbot, text_box, image_box],
+      outputs=[chatbot],
+      api_name="stream_response"
+    ).then(
+      fn=lambda: (None, ""),  # Clear the image and text inputs after sending
+      inputs=None,
+      outputs=[image_box, text_box],
+    )
+
   return app
 
 if __name__ == "__main__":