ai: Enable gradio default stop button.

jarvis.py

@@ -61,10 +61,15 @@ class SessionWithID(requests.Session):
     def __init__(self):
         super().__init__()
         self.session_id = str(uuid.uuid4())
+        self.stop_event = asyncio.Event()
 
 def create_session():
     return SessionWithID()
 
+def ensure_stop_event(sess):
+    if not hasattr(sess, "stop_event"):
+        sess.stop_event = asyncio.Event()
+
 def get_available_items(items, marked):
     a = [i for i in items if i not in marked]
     random.shuffle(a)
@@ -199,7 +204,7 @@ def extract_file_content(fp):
             return f"{fp}: {e}"
 
 async def fetch_response_async(host, key, model, msgs, cfg, sid):
-    for t in [60, 80, 120, 240]:
+    for t in [1, 2]:
         try:
             async with httpx.AsyncClient(timeout=t) as client:
                 r = await client.post(host, json={"model": model, "messages": msgs, **cfg, "session_id": sid}, headers={"Authorization": f"Bearer {key}"})
@@ -223,6 +228,7 @@ async def chat_with_model_async(history, user_input, model_display, sess, custom
         return RESPONSES["RESPONSE_3"]
     if not hasattr(sess, "session_id"):
         sess.session_id = str(uuid.uuid4())
+        sess.stop_event = asyncio.Event()
     model_key = get_model_key(model_display)
     cfg = MODEL_CONFIG.get(model_key, DEFAULT_CONFIG)
     msgs = [{"role": "user", "content": u} for u, _ in history] + [{"role": "assistant", "content": a} for _, a in history if a]
@@ -250,54 +256,73 @@ async def chat_with_model_async(history, user_input, model_display, sess, custom
     return RESPONSES["RESPONSE_2"]
 
 async def respond_async(multi, history, model_display, sess, custom_prompt):
-    msg = {"text": multi.get("text", "").strip(), "files": multi.get("files", [])}
-    if not msg["text"] and not msg["files"]:
-        yield history, gr.MultimodalTextbox(value=None, interactive=True), sess
+    if not hasattr(sess, "stop_event"):
+        sess.stop_event = asyncio.Event()
+    sess.stop_event.clear()
+    msg_input = {"text": multi.get("text", "").strip(), "files": multi.get("files", [])}
+    if not msg_input["text"] and not msg_input["files"]:
+        yield history, gr.MultimodalTextbox(value="", interactive=True, submit_btn=True, stop_btn=False), sess
         return
     inp = ""
-    for f in msg["files"]:
+    for f in msg_input["files"]:
         if isinstance(f, dict):
             fp = f.get("data", f.get("name", ""))
         else:
             fp = f
         inp += f"{Path(fp).name}\n\n{extract_file_content(fp)}\n\n"
-    if msg["text"]:
-        inp += msg["text"]
-    history.append([inp, ""])
+    if msg_input["text"]:
+        inp += msg_input["text"]
+    history.append([inp, RESPONSES["RESPONSE_8"]])
+    yield history, gr.MultimodalTextbox(value="", interactive=False, submit_btn=False, stop_btn=True), sess
     ai = await chat_with_model_async(history, inp, model_display, sess, custom_prompt)
     history[-1][1] = ""
-    def to_str(d):
-        if isinstance(d, (str, int, float)):
-            return str(d)
-        if isinstance(d, bytes):
-            return d.decode("utf-8", errors="ignore")
-        if isinstance(d, (list, tuple)):
-            return "".join(map(to_str, d))
-        if isinstance(d, dict):
-            return json.dumps(d, ensure_ascii=False)
-        return repr(d)
-    for c in ai:
-        history[-1][1] += to_str(c)
-        await asyncio.sleep(0.0008)
-        yield history, gr.MultimodalTextbox(value=None, interactive=True), sess
+    buffer = []
+    last_update = asyncio.get_event_loop().time()
+    for char in ai:
+        if sess.stop_event.is_set():
+            history[-1][1] = RESPONSES["RESPONSE_1"]
+            yield history, gr.MultimodalTextbox(value="", interactive=True, submit_btn=True, stop_btn=False), sess
+            sess.stop_event.clear()
+            return
+        buffer.append(char)
+        current_time = asyncio.get_event_loop().time()
+        if len(buffer) >= 8 or (current_time - last_update) > 0.04:
+            history[-1][1] += "".join(buffer)
+            buffer.clear()
+            last_update = current_time
+            yield history, gr.MultimodalTextbox(value="", interactive=False, submit_btn=False, stop_btn=True), sess
+            await asyncio.sleep(0.016)
+    if buffer:
+        history[-1][1] += "".join(buffer)
+        yield history, gr.MultimodalTextbox(value="", interactive=False, submit_btn=False, stop_btn=True), sess
+    yield history, gr.MultimodalTextbox(value="", interactive=True, submit_btn=True, stop_btn=False), sess
 
 def change_model(new):
     visible = new != MODEL_CHOICES[0]
     default = SYSTEM_PROMPT_MAPPING.get(get_model_key(new), SYSTEM_PROMPT_DEFAULT)
     return [], create_session(), new, default, gr.update(value=default, visible=visible)
 
+def stop_response(history, sess):
+    if not hasattr(sess, "stop_event"):
+        sess.stop_event = asyncio.Event()
+    sess.stop_event.set()
+    if history:
+        history[-1][1] = RESPONSES["RESPONSE_1"]
+    new_session = create_session()
+    return history, gr.MultimodalTextbox(value="", interactive=True, submit_btn=True, stop_btn=False), new_session
+
 with gr.Blocks(fill_height=True, fill_width=True, title=AI_TYPES["AI_TYPE_4"], head=META_TAGS) as jarvis:
     user_history = gr.State([])
     user_session = gr.State(create_session())
     selected_model = gr.State(MODEL_CHOICES[0] if MODEL_CHOICES else "")
     custom_prompt_state = gr.State("")
     chatbot = gr.Chatbot(label=AI_TYPES["AI_TYPE_1"], show_copy_button=True, scale=1, elem_id=AI_TYPES["AI_TYPE_2"])
-    with gr.Row():
-        msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)
+    msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)
     with gr.Accordion(AI_TYPES["AI_TYPE_6"], open=False):
         model_dropdown = gr.Dropdown(show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0])
         system_prompt = gr.Textbox(label=AI_TYPES["AI_TYPE_7"], lines=2, interactive=True, visible=False)
     model_dropdown.change(fn=change_model, inputs=[model_dropdown], outputs=[user_history, user_session, selected_model, custom_prompt_state, system_prompt])
     system_prompt.change(fn=lambda x: x, inputs=[system_prompt], outputs=[custom_prompt_state])
     msg.submit(fn=respond_async, inputs=[msg, user_history, selected_model, user_session, custom_prompt_state], outputs=[chatbot, msg, user_session], api_name=INTERNAL_AI_GET_SERVER)
+    msg.stop(fn=stop_response, inputs=[user_history, user_session], outputs=[chatbot, msg, user_session])
 jarvis.launch(max_file_size="1mb")