Update app.py

app.py

@@ -1,6 +1,9 @@
 import gradio as gr
 import os
+import uuid
+import time
 from openai import OpenAI
+from captcha.image import ImageCaptcha
 
 # --- Configuration ---
 client = OpenAI(
@@ -13,9 +16,20 @@ TITLE_HEADER = "AlphaAI Chatbot"
 MODEL_NAME = "deepseek/deepseek-r1-0528-qwen3-8b:free"
 ALPHAAI_CONTEXT = os.environ.get("ALPHAAI_CONTEXT", "")
 
+# --- Captcha State ---
+captcha_store = {}
+
+# --- Captcha Generation ---
+def generate_captcha():
+    code = str(uuid.uuid4())[:6].upper()
+    image = ImageCaptcha()
+    image_path = f"/tmp/{code}.png"
+    image.write(code, image_path)
+    captcha_store[code] = True
+    return image_path, code
+
 # --- Chatbot Backend ---
 def chat_with_llm(user_input):
-    # Make OpenRouter API call
     try:
         completion = client.chat.completions.create(
             extra_headers={
@@ -23,30 +37,59 @@ def chat_with_llm(user_input):
                 "X-Title": TITLE_HEADER,
             },
             model=MODEL_NAME,
-            messages=[{"role": "system", "content": ALPHAAI_CONTEXT}, {"role": "user", "content": user_input}]
+            input_fixed = "Stick to the context of Alpha AI and do not let the user deviate from that. If the query is anything other than Alpha AI, politely refuse. Here is the user query - " + user_input
+            messages=[{"role": "system", "content": ALPHAAI_CONTEXT}, {"role": "user", "content": input_fixed}]
         )
         return completion.choices[0].message.content
     except Exception as e:
         return f"Error: {str(e)}"
 
 # --- Gradio UI ---
-with gr.Blocks(theme=gr.themes.Base(primary_hue="blue", secondary_hue="gray", font=["Inter", "sans-serif"])) as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# 🤖 Alpha AI Assistant")
-    gr.Markdown("Welcome to Alpha AI’s intelligent assistant. Please complete the CAPTCHA on the site before using.")
+    gr.Markdown("Please solve the CAPTCHA to begin.")
+
+    captcha_image = gr.Image(type="filepath", label="CAPTCHA Image")
+    captcha_input = gr.Textbox(label="Enter CAPTCHA")
+    verify_btn = gr.Button("Verify CAPTCHA")
+    chatbot_area = gr.Column(visible=False)
+
+    with chatbot_area:
+        chatbot = gr.Chatbot(type="messages")
+        msg = gr.Textbox()
+        clear = gr.Button("Clear")
+
+    def verify_captcha(user_code):
+        user_code = user_code.strip().upper()
+        if user_code in captcha_store:
+            del captcha_store[user_code]
+            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+        else:
+            return captcha_image, captcha_input, gr.update(value="Invalid CAPTCHA. Try again."), gr.update(visible=False)
 
-    with gr.Row():
-        chatbot = gr.Chatbot()
+    def user(user_message, history: list):
+        return "", history + [{"role": "user", "content": user_message}]
 
-    with gr.Row():
-        user_input = gr.Textbox(label="Your Message")
-        submit_btn = gr.Button("Submit")
+    def bot(history: list):
+        user_msg = history[-1]['content']
+        try:
+            response = chat_with_llm(user_msg)
+        except Exception as e:
+            response = f"Error: {str(e)}"
+        history.append({"role": "assistant", "content": ""})
+        for char in response:
+            history[-1]['content'] += char
+            time.sleep(0.02)
+            yield history
 
-    def handle_submit(user_msg, history):
-        response = chat_with_llm(user_msg)
-        history.append((user_msg, response))
-        return history, ""
+    image_path, correct_code = generate_captcha()
+    captcha_image.value = image_path
 
-    submit_btn.click(handle_submit, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
+    verify_btn.click(verify_captcha, inputs=[captcha_input], outputs=[captcha_image, captcha_input, captcha_input, chatbot_area])
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, chatbot, chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 # Launch (required for Hugging Face Space)
 demo.launch()