Update app.py

app.py

@@ -5,9 +5,14 @@ from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
 from fastapi import FastAPI
 from pydantic import BaseModel
-from threading import Thread
+from threading import Thread, Event
 import uvicorn
 
+# ----------------------------
+# STOP EVENT
+# ----------------------------
+stop_event = Event()
+
 # ----------------------------
 # Model
 # ----------------------------
@@ -31,8 +36,9 @@ llm("warmup", max_tokens=1)
 # System Prompt
 # ----------------------------
 SYSTEM_PROMPT = """
-You are a senior software engineer and coding teacher.
-Explain code step by step, fix bugs, and keep explanations clear.
+You are an advanced AI assistant. 
+Answer questions clearly and concisely. 
+You can handle multi-turn conversations and provide detailed responses if needed.
 """
 
 # ----------------------------
@@ -44,24 +50,55 @@ def generate_response(message, history):
 
     prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
     for h in history:
-        if isinstance(h, dict) and "role" in h and "message" in h:
+        if isinstance(h, dict):
             role = h["role"]
             msg = h["message"]
-            if role == "user":
-                prompt += f"<|im_start|>user\n{msg}<|im_end|>\n"
-            else:
-                prompt += f"<|im_start|>assistant\n{msg}<|im_end|>\n"
-        elif isinstance(h, (list, tuple)) and len(h) >= 2:
-            u, a = h[0], h[1]
+            prompt += f"<|im_start|>{role}\n{msg}<|im_end|>\n"
+        elif isinstance(h, (list, tuple)) and len(h) == 2:
+            u, a = h
             prompt += f"<|im_start|>user\n{u}<|im_end|>\n<|im_start|>assistant\n{a}<|im_end|>\n"
 
     prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
 
     output = ""
     for token in llm(prompt, max_tokens=2048, temperature=0.2, top_p=0.9, repeat_penalty=1.1, stream=True):
+        if stop_event.is_set():
+            stop_event.clear()
+            break
         output += token["choices"][0]["text"]
         yield output
 
+# ----------------------------
+# COPY / DOWNLOAD / EDIT FUNCTIONS
+# ----------------------------
+def copy_last(history):
+    if history:
+        last = history[-1]
+        if isinstance(last, dict):
+            return last["message"]
+        return last[1]
+    return ""
+
+def download_history(history):
+    file_path = "chat_history.txt"
+    text = ""
+    for h in history:
+        if isinstance(h, dict):
+            text += f'{h["role"]}: {h["message"]}\n\n'
+        else:
+            text += f'User: {h[0]}\nAI: {h[1]}\n\n'
+    with open(file_path, "w", encoding="utf-8") as f:
+        f.write(text)
+    return file_path
+
+def edit_message(history, index, new_text):
+    if index < len(history):
+        if isinstance(history[index], dict):
+            history[index]["message"] = new_text
+        else:
+            history[index] = (new_text, history[index][1])
+    return history
+
 # ----------------------------
 # FastAPI API
 # ----------------------------
@@ -87,22 +124,31 @@ def chat_endpoint(request: ChatRequest):
 # Gradio UI
 # ----------------------------
 with gr.Blocks(theme=gr.Theme.from_hub("JackismyShephard/ultimate-rvc-theme")) as demo:
-    gr.HTML("<h2 style='text-align:center; color:white;'>Code Explainer AI</h2>")
-
-    chatbot = gr.ChatInterface(
-        fn=generate_response,
-        chatbot=gr.Chatbot(height=600),
-        textbox=gr.Textbox(placeholder="Paste code or ask for explanation...", container=False)
-    )
+    gr.HTML("<h2 style='text-align:center; color:white;'>Advanced AI Chat</h2>")
+
+    chatbot = gr.Chatbot(height=600)
+    msg = gr.Textbox(placeholder="Type your question...", container=False)
+    state = gr.State([])
+
+    send_btn = gr.Button("Send")
+    stop_btn = gr.Button("🛑 Stop")
+    copy_btn = gr.Button("📋 Copy Last")
+    download_btn = gr.File(label="📥 Download Chat")
+    edit_index = gr.Number(label="Edit Message Index")
+    edit_text = gr.Textbox(label="New Text")
+    edit_btn = gr.Button("✏️ Edit")
+
+    # ----------------------------
+    # Button callbacks
+    # ----------------------------
+    send_btn.click(generate_response, [msg, state], [chatbot, state])
+    stop_btn.click(lambda: stop_event.set())
+    copy_btn.click(copy_last, inputs=[state], outputs=None)
+    download_btn.click(download_history, inputs=[state], outputs=download_btn)
+    edit_btn.click(edit_message, inputs=[state, edit_index, edit_text], outputs=[chatbot])
 
-    # Rounded corners for main container
     demo.css = """
-    .gradio-container {
-        border-radius: 25px !important;
-        max-width: 600px !important;
-        margin: auto !important;
-        overflow: hidden;
-    }
+    .gradio-container { border-radius: 25px !important; max-width: 600px !important; margin:auto !important; overflow:hidden; }
     .message.user { border-radius: 18px 18px 4px 18px !important; }
     .message.bot { border-radius: 18px 18px 18px 4px !important; }
     """