Update app.py via AI Editor

app.py

@@ -46,8 +46,9 @@ def get_session_state(session_id):
         SESSION_DATA[session_id] = {
             "messages": [],
             "uploads": [],
-            "openai_file_ids": [],
-            "created": datetime.datetime.utcnow().isoformat()
+            "created": datetime.datetime.utcnow().isoformat(),
+            "streaming": False,
+            "stream_buffer": "",
         }
     return SESSION_DATA[session_id]
 
@@ -153,7 +154,8 @@ def right_main_static():
                 html.Div(id="error-message", style={"color": "#bb2124", "marginTop": "0.5rem"}),
             ])
         ], className="mt-3"),
-        dcc.Loading(id="loading", type="default", fullscreen=False, style={"position": "absolute", "top": "5%", "left": "50%"})
+        dcc.Loading(id="loading", type="default", fullscreen=False, style={"position": "absolute", "top": "5%", "left": "50%"}),
+        dcc.Interval(id="stream-interval", interval=400, n_intervals=0, disabled=True, max_intervals=1000)
     ], style={"padding": "1rem", "backgroundColor": "#fff", "height": "100vh", "overflowY": "auto"})
 
 app.layout = html.Div([
@@ -165,25 +167,8 @@ app.layout = html.Div([
     ], style={"display": "flex"})
 ])
 
-def _upload_file_to_openai(file_path, purpose="assistants"):
-    try:
-        with open(file_path, 'rb') as f:
-            res = openai.File.create(
-                file=f,
-                purpose=purpose
-            )
-        logger.info(f"Uploaded file to OpenAI: {res.id}")
-        return res.id
-    except Exception as e:
-        logger.error(f"Failed to upload file to OpenAI: {e}")
-        return None
-
-def _get_openai_file_ids(session_state):
-    return session_state.get("openai_file_ids", [])
-
 def _is_supported_doc(filename):
     ext = os.path.splitext(filename)[1].lower()
-    # OpenAI supports: txt, pdf, docx, md for assistants file search
     return ext in [".txt", ".pdf", ".md", ".docx"]
 
 @app.callback(
@@ -203,14 +188,17 @@ def assign_session_id(_):
     Output("chat-history-list", "children"),
     Output("chat-window", "children"),
     Output("error-message", "children"),
+    Output("stream-interval", "disabled"),
+    Output("stream-interval", "n_intervals"),
     Input("session-id", "data"),
     Input("send-btn", "n_clicks"),
     Input("file-upload", "contents"),
     State("file-upload", "filename"),
     State("user-input", "value"),
+    State("stream-interval", "n_intervals"),
     prevent_initial_call=False
 )
-def main_callback(session_id, send_clicks, file_contents, file_names, user_input):
+def main_callback(session_id, send_clicks, file_contents, file_names, user_input, stream_n):
     trigger = callback_context.triggered[0]['prop_id'].split('.')[0] if callback_context.triggered else ""
     if not session_id:
         session_id = get_session_id()
@@ -219,6 +207,7 @@ def main_callback(session_id, send_clicks, file_contents, file_names, user_input
         load_session_state(session_id)
         state = get_session_state(session_id)
         error = ""
+        start_streaming = False
 
         if trigger == "file-upload" and file_contents and file_names:
             uploads = []
@@ -235,64 +224,106 @@ def main_callback(session_id, send_clicks, file_contents, file_names, user_input
                 with open(fp, "wb") as f:
                     f.write(base64.b64decode(data))
                 uploads.append({"name": fname, "is_img": is_img, "path": fp})
-                # If document is supported, upload to OpenAI
-                if _is_supported_doc(fname):
-                    file_id = _upload_file_to_openai(fp)
-                    if file_id:
-                        if "openai_file_ids" not in state:
-                            state["openai_file_ids"] = []
-                        state["openai_file_ids"].append(file_id)
             state["uploads"].extend(uploads)
             save_session_state(session_id)
             logger.info(f"Session {session_id}: Uploaded files {[u['name'] for u in uploads]}")
 
         if trigger == "send-btn" and user_input and user_input.strip():
             state["messages"].append({"role": "user", "content": user_input})
-            try:
-                file_ids = _get_openai_file_ids(state)
-                system_prompt = load_system_prompt()
-                messages = [
-                    {"role": "system", "content": system_prompt},
-                ]
-                for m in state["messages"]:
-                    messages.append({"role": m["role"], "content": m["content"]})
-                if file_ids:
-                    response = openai.ChatCompletion.create(
-                        model="gpt-3.5-turbo-1106",
-                        messages=messages,
-                        tools=[{"type": "file_search"}],
-                        tool_choice="file_search",
-                        file_ids=file_ids,
-                        max_tokens=700,
-                        temperature=0.2,
-                    )
-                else:
+            state["streaming"] = True
+            state["stream_buffer"] = ""
+            save_session_state(session_id)
+
+            def run_stream(session_id, messages):
+                try:
+                    system_prompt = load_system_prompt()
+                    msg_list = [{"role": "system", "content": system_prompt}]
+                    for m in messages:
+                        msg_list.append({"role": m["role"], "content": m["content"]})
                     response = openai.ChatCompletion.create(
                         model="gpt-3.5-turbo",
-                        messages=messages,
+                        messages=msg_list,
                         max_tokens=700,
                         temperature=0.2,
+                        stream=True,
                     )
-                reply = response.choices[0].message.content
-                state["messages"].append({"role": "assistant", "content": reply})
-                logger.info(f"Session {session_id}: User: {user_input} | Assistant: {reply}")
-                error = ""
-            except Exception as e:
-                error = f"Error: {e}"
-                logger.error(f"Session {session_id}: {error}")
-            save_session_state(session_id)
+                    reply = ""
+                    for chunk in response:
+                        delta = chunk["choices"][0]["delta"]
+                        content = delta.get("content", "")
+                        if content:
+                            reply += content
+                            # Update buffer in session state
+                            session_lock = get_session_lock(session_id)
+                            with session_lock:
+                                load_session_state(session_id)
+                                state = get_session_state(session_id)
+                                state["stream_buffer"] = reply
+                                save_session_state(session_id)
+                    # Finalize message
+                    session_lock = get_session_lock(session_id)
+                    with session_lock:
+                        load_session_state(session_id)
+                        state = get_session_state(session_id)
+                        state["messages"].append({"role": "assistant", "content": reply})
+                        state["stream_buffer"] = ""
+                        state["streaming"] = False
+                        save_session_state(session_id)
+                    logger.info(f"Session {session_id}: User: {user_input} | Assistant: {reply}")
+                except Exception as e:
+                    session_lock = get_session_lock(session_id)
+                    with session_lock:
+                        load_session_state(session_id)
+                        state = get_session_state(session_id)
+                        state["streaming"] = False
+                        state["stream_buffer"] = ""
+                        save_session_state(session_id)
+                    logger.error(f"Session {session_id}: Streaming error: {e}")
+
+            threading.Thread(target=run_stream, args=(session_id, list(state["messages"])), daemon=True).start()
+            start_streaming = True
 
         chat_history = state.get("messages", [])
         uploads = state.get("uploads", [])
         upload_cards = [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in uploads]
         chat_items = [html.Li(html.Span((msg['role'] + ": " + msg['content'])[:40] + ("..." if len(msg['content']) > 40 else ""), style={"fontSize": "0.92rem"})) for msg in chat_history[-6:]]
         chat_cards = []
-        for msg in chat_history:
+        for i, msg in enumerate(chat_history):
+            if msg['role'] == "user":
+                chat_cards.append(chat_message_card(msg['content'], is_user=True))
+            elif msg['role'] == "assistant":
+                chat_cards.append(chat_message_card(msg['content'], is_user=False))
+        if state.get("streaming", False):
+            # Add a partial assistant message at the end
+            if state.get("stream_buffer", ""):
+                chat_cards.append(chat_message_card(state["stream_buffer"], is_user=False))
+            return upload_cards, chat_items, chat_cards, error, False, 0 if not start_streaming else False, 0
+        return upload_cards, chat_items, chat_cards, error, (not state.get("streaming", False)), 0
+
+@app.callback(
+    Output("chat-window", "children"),
+    Output("stream-interval", "disabled"),
+    Input("stream-interval", "n_intervals"),
+    State("session-id", "data"),
+    prevent_initial_call=True
+)
+def poll_stream(n_intervals, session_id):
+    session_lock = get_session_lock(session_id)
+    with session_lock:
+        load_session_state(session_id)
+        state = get_session_state(session_id)
+        chat_history = state.get("messages", [])
+        chat_cards = []
+        for i, msg in enumerate(chat_history):
             if msg['role'] == "user":
                 chat_cards.append(chat_message_card(msg['content'], is_user=True))
             elif msg['role'] == "assistant":
                 chat_cards.append(chat_message_card(msg['content'], is_user=False))
-        return upload_cards, chat_items, chat_cards, error
+        if state.get("streaming", False):
+            if state.get("stream_buffer", ""):
+                chat_cards.append(chat_message_card(state["stream_buffer"], is_user=False))
+            return chat_cards, False
+        return chat_cards, True
 
 @app_flask.after_request
 def set_session_cookie(resp):