Update app.py

app.py

@@ -8,9 +8,6 @@ from pydantic import BaseModel
 from threading import Thread, Event
 import uvicorn
 
-# ----------------------------
-# STOP EVENT
-# ----------------------------
 stop_event = Event()
 
 # ----------------------------
@@ -32,22 +29,17 @@ llm = Llama(
 
 llm("warmup", max_tokens=1)
 
-# ----------------------------
-# System Prompt
-# ----------------------------
 SYSTEM_PROMPT = """
-You are an advanced AI assistant. 
-Answer questions clearly and concisely. 
-You can handle multi-turn conversations and provide detailed responses if needed.
+You are an advanced AI assistant.
+Answer questions clearly and concisely.
 """
 
 # ----------------------------
-# Chat Function
+# Chat function
 # ----------------------------
 def generate_response(message, history):
     yield "🤖 Thinking..."
-    time.sleep(0.5)
-
+    time.sleep(0.3)
     prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
     for h in history:
         if isinstance(h, dict):
@@ -59,7 +51,6 @@ def generate_response(message, history):
             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():
@@ -68,9 +59,6 @@ def generate_response(message, history):
         output += token["choices"][0]["text"]
         yield output
 
-# ----------------------------
-# COPY / DOWNLOAD / EDIT FUNCTIONS
-# ----------------------------
 def copy_last(history):
     if history:
         last = history[-1]
@@ -89,7 +77,7 @@ def download_history(history):
     file_path = "chat_history.txt"
     with open(file_path, "w", encoding="utf-8") as f:
         f.write(text)
-    return file_path
+    return file_path  # return path for gr.File output
 
 def edit_message(history, index, new_text):
     if index < len(history):
@@ -103,7 +91,6 @@ def edit_message(history, index, new_text):
 # FastAPI API
 # ----------------------------
 app = FastAPI()
-
 class ChatRequest(BaseModel):
     message: str
     history: list = []
@@ -115,7 +102,6 @@ def chat_endpoint(request: ChatRequest):
     for u, a in request.history:
         prompt += f"<|im_start|>user\n{u}<|im_end|>\n<|im_start|>assistant\n{a}<|im_end|>\n"
     prompt += f"<|im_start|>user\n{request.message}<|im_end|>\n<|im_start|>assistant\n"
-
     for token in llm(prompt, max_tokens=2048, temperature=0.2, top_p=0.9, repeat_penalty=1.1, stream=True):
         output += token["choices"][0]["text"]
     return {"response": output}
@@ -124,22 +110,21 @@ def chat_endpoint(request: ChatRequest):
 # Gradio UI
 # ----------------------------
 with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column(scale=1, min_width=800):
-            chatbot = gr.Chatbot(height=700)
-            msg = gr.Textbox(placeholder="Type your message...", container=False)
-            state = gr.State([])
-
-            with gr.Row():
-                send_btn = gr.Button("Send")
-                stop_btn = gr.Button("🛑 Stop")
-                copy_btn = gr.Button("📋 Copy Last")
-                download_btn = gr.File(label="📥 Download Chat")
-
-            with gr.Row():
-                edit_index = gr.Number(label="Edit Index")
-                edit_text = gr.Textbox(label="New Text")
-                edit_btn = gr.Button("✏️ Edit")
+    with gr.Column():
+        chatbot = gr.Chatbot(height=700)
+        state = gr.State([])
+        msg = gr.Textbox(placeholder="Type your message...", container=False)
+
+        with gr.Row():
+            send_btn = gr.Button("Send")
+            stop_btn = gr.Button("🛑 Stop")
+            copy_btn = gr.Button("📋 Copy Last")
+            download_file = gr.File(label="Download Chat", file_types=[".txt"])
+
+        with gr.Row():
+            edit_index = gr.Number(label="Edit Index")
+            edit_text = gr.Textbox(label="New Text")
+            edit_btn = gr.Button("✏️ Edit")
 
     # ----------------------------
     # Callbacks
@@ -147,7 +132,7 @@ with gr.Blocks() as demo:
     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.update(value=None)  # gr.File doesn't have click; we update value
+    download_file.output = download_history  # assign output function
     edit_btn.click(edit_message, inputs=[state, edit_index, edit_text], outputs=[chatbot])
 
     demo.css = """
@@ -165,9 +150,6 @@ with gr.Blocks() as demo:
     .message.bot { border-radius: 18px 18px 18px 4px !important; background:#1c1f2a !important; color:white !important;}
     """
 
-# ----------------------------
-# Run Gradio + FastAPI
-# ----------------------------
 def run_gradio():
     demo.launch(server_name="0.0.0.0", server_port=7860)