Update app.py

app.py

@@ -5,11 +5,9 @@ from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
 from fastapi import FastAPI
 from pydantic import BaseModel
-from threading import Thread, Event
+from threading import Thread
 import uvicorn
 
-stop_event = Event()
-
 # ----------------------------
 # Model
 # ----------------------------
@@ -29,68 +27,47 @@ 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.
 """
 
 # ----------------------------
-# Chat function
+# Chat Function
 # ----------------------------
 def generate_response(message, history):
     yield "🤖 Thinking..."
-    time.sleep(0.3)
+    time.sleep(0.5)
+
     prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
     for h in history:
-        if isinstance(h, dict):
+        if isinstance(h, dict) and "role" in h and "message" in h:
             role = h["role"]
             msg = h["message"]
-            prompt += f"<|im_start|>{role}\n{msg}<|im_end|>\n"
-        elif isinstance(h, (list, tuple)) and len(h) == 2:
-            u, a = h
+            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|>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
 
-def copy_last(history):
-    if history:
-        last = history[-1]
-        if isinstance(last, dict):
-            return last["message"]
-        return last[1]
-    return ""
-
-def download_history(history):
-    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'
-    file_path = "chat_history.txt"
-    with open(file_path, "w", encoding="utf-8") as f:
-        f.write(text)
-    return file_path  # return path for gr.File output
-
-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
 # ----------------------------
 app = FastAPI()
+
 class ChatRequest(BaseModel):
     message: str
     history: list = []
@@ -102,6 +79,7 @@ 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}
@@ -109,47 +87,30 @@ def chat_endpoint(request: ChatRequest):
 # ----------------------------
 # Gradio UI
 # ----------------------------
-with gr.Blocks() as demo:
-    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
-    # ----------------------------
-    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_file.output = download_history  # assign output function
-    edit_btn.click(edit_message, inputs=[state, edit_index, edit_text], outputs=[chatbot])
+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)
+    )
+
+    # Rounded corners for main container
     demo.css = """
     .gradio-container {
-        width: 100vw !important;
-        height: 100vh !important;
-        border-radius: 0px !important;
-        overflow: hidden;
-        background-color: #0b0f19 !important;
-    }
-    .input-textbox textarea {
         border-radius: 25px !important;
+        max-width: 600px !important;
+        margin: auto !important;
+        overflow: hidden;
     }
-    .message.user { border-radius: 18px 18px 4px 18px !important; background:#2b6fff !important; color:white !important;}
-    .message.bot { border-radius: 18px 18px 18px 4px !important; background:#1c1f2a !important; color:white !important;}
+    .message.user { border-radius: 18px 18px 4px 18px !important; }
+    .message.bot { border-radius: 18px 18px 18px 4px !important; }
     """
 
+# ----------------------------
+# Run Gradio + FastAPI together
+# ----------------------------
 def run_gradio():
     demo.launch(server_name="0.0.0.0", server_port=7860)