Update app.py via AI Editor

app.py

@@ -49,6 +49,7 @@ def get_session_state(session_id):
             "created": datetime.datetime.utcnow().isoformat(),
             "streaming": False,
             "stream_buffer": "",
+            "chat_histories": []
         }
     return SESSION_DATA[session_id]
 
@@ -149,7 +150,16 @@ def right_main_static():
                         wrap="soft",
                         maxLength=1000
                     ),
-                    dbc.Button("Send", id="send-btn", color="primary", className="mt-2", style={"float": "right", "minWidth": "100px"}),
+                    html.Div([
+                        dbc.Button("Send", id="send-btn", color="primary", className="mt-2 me-2", style={"minWidth": "100px"}),
+                        dbc.Button("New Chat", id="new-chat-btn", color="secondary", className="mt-2", style={"minWidth": "110px"}),
+                    ], style={"float": "right", "display": "flex", "gap": "0.5rem"}),
+                    dcc.Input(
+                        id="chat-name-input",
+                        placeholder="Enter chat name...",
+                        type="text",
+                        style={"width": "100%", "marginTop": "0.5rem", "display": "none"}
+                    ),
                 ], style={"marginTop": "1rem"}),
                 html.Div(id="error-message", style={"color": "#bb2124", "marginTop": "0.5rem"}),
             ])
@@ -190,15 +200,19 @@ def assign_session_id(_):
     Output("error-message", "children"),
     Output("stream-interval", "disabled"),
     Output("stream-interval", "n_intervals"),
+    Output("chat-name-input", "style"),
     Input("session-id", "data"),
     Input("send-btn", "n_clicks"),
     Input("file-upload", "contents"),
+    Input("new-chat-btn", "n_clicks"),
     State("file-upload", "filename"),
     State("user-input", "value"),
+    State("chat-name-input", "value"),
     State("stream-interval", "n_intervals"),
+    State("chat-name-input", "style"),
     prevent_initial_call=False
 )
-def main_callback(session_id, send_clicks, file_contents, file_names, user_input, stream_n):
+def main_callback(session_id, send_clicks, file_contents, new_chat_clicks, file_names, user_input, chat_name, stream_n, chat_name_style):
     trigger = callback_context.triggered[0]['prop_id'].split('.')[0] if callback_context.triggered else ""
     if not session_id:
         session_id = get_session_id()
@@ -209,6 +223,7 @@ def main_callback(session_id, send_clicks, file_contents, file_names, user_input
         error = ""
         start_streaming = False
 
+        # Handle File Upload
         if trigger == "file-upload" and file_contents and file_names:
             uploads = []
             if not isinstance(file_contents, list):
@@ -228,6 +243,7 @@ def main_callback(session_id, send_clicks, file_contents, file_names, user_input
             save_session_state(session_id)
             logger.info(f"Session {session_id}: Uploaded files {[u['name'] for u in uploads]}")
 
+        # Handle Send
         if trigger == "send-btn" and user_input and user_input.strip():
             state["messages"].append({"role": "user", "content": user_input})
             state["streaming"] = True
@@ -253,14 +269,12 @@ def main_callback(session_id, send_clicks, file_contents, file_names, user_input
                         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)
@@ -283,22 +297,89 @@ def main_callback(session_id, send_clicks, file_contents, file_names, user_input
             threading.Thread(target=run_stream, args=(session_id, list(state["messages"])), daemon=True).start()
             start_streaming = True
 
+        # Handle New Chat button logic
+        show_chat_name = chat_name_style if isinstance(chat_name_style, dict) else {}
+        if trigger == "new-chat-btn":
+            # If chat name input box is not yet visible, show it
+            if show_chat_name.get("display", "none") == "none":
+                show_chat_name["display"] = "block"
+                return (
+                    [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in state.get("uploads", [])],
+                    [
+                        html.Li(
+                            html.Span(
+                                chat["name"],
+                                style={"fontSize": "0.92rem"}
+                            )
+                        ) for chat in state.get("chat_histories", [])[-6:]
+                    ],
+                    [
+                        chat_message_card(msg['content'], is_user=(msg['role'] == "user"))
+                        for msg in state.get("messages", [])
+                    ] + (
+                        [chat_message_card(state["stream_buffer"], is_user=False)]
+                        if state.get("streaming", False) and state.get("stream_buffer", "") else []
+                    ),
+                    "",
+                    not state.get("streaming", False),
+                    0,
+                    show_chat_name
+                )
+            # If input box is visible and has a value, save chat history
+            else:
+                chat_dialog = list(state.get("messages", []))
+                if not chat_dialog:
+                    error = "Cannot save empty chat."
+                else:
+                    chat_title = chat_name if chat_name and chat_name.strip() else "Chat " + datetime.datetime.utcnow().strftime("%Y-%m-%d %H:%M")
+                    state.setdefault("chat_histories", []).append({
+                        "name": chat_title,
+                        "dialog": chat_dialog
+                    })
+                    state["messages"] = []
+                    state["stream_buffer"] = ""
+                    state["streaming"] = False
+                    save_session_state(session_id)
+                    logger.info(f"Session {session_id}: Saved chat history '{chat_title}'")
+                show_chat_name["display"] = "none"
+                return (
+                    [uploaded_file_card(os.path.basename(f["name"]), f["is_img"]) for f in state.get("uploads", [])],
+                    [
+                        html.Li(
+                            html.Span(
+                                chat["name"],
+                                style={"fontSize": "0.92rem", "fontWeight": "bold"}
+                            )
+                        ) for chat in state.get("chat_histories", [])[-6:]
+                    ],
+                    [],
+                    error,
+                    not state.get("streaming", False),
+                    0,
+                    show_chat_name
+                )
+
+        # Build Uploads, Chat History and Chat Window
         chat_history = state.get("messages", [])
         uploads = state.get("uploads", [])
+        chat_histories = state.get("chat_histories", [])
         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_history_items = [
+            html.Li(
+                html.Span(
+                    chat["name"],
+                    style={"fontSize": "0.92rem", "fontWeight": "bold" if i == len(chat_histories)-1 else "normal"}
+                )
+            ) for i, chat in enumerate(chat_histories[-6:])
+        ]
         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))
+        for msg in chat_history:
+            chat_cards.append(chat_message_card(msg['content'], is_user=(msg['role'] == "user")))
         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
+            return upload_cards, chat_history_items, chat_cards, error, False, 0, show_chat_name
+        return upload_cards, chat_history_items, chat_cards, error, (not state.get("streaming", False)), 0, show_chat_name
 
 @app.callback(
     Output("chat-window", "children"),
@@ -314,11 +395,8 @@ def poll_stream(n_intervals, 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))
+        for msg in chat_history:
+            chat_cards.append(chat_message_card(msg['content'], is_user=(msg['role'] == "user")))
         if state.get("streaming", False):
             if state.get("stream_buffer", ""):
                 chat_cards.append(chat_message_card(state["stream_buffer"], is_user=False))