Update app.py

app.py

@@ -1,22 +1,5 @@
-from fastapi import FastAPI, Request
-from fastapi.middleware.cors import CORSMiddleware
+import gradio as gr
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-from pydantic import BaseModel
-import uuid
-
-# Initialize FastAPI app
-app = FastAPI(title="Mental Health Chatbot API", 
-              description="API for mental health chatbot with emotion detection",
-              version="1.0.0")
-
-# Add CORS middleware
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # In production, restrict this to your app's domain
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
 
 # Load chatbot model
 print("Loading DialoGPT model...")
@@ -28,73 +11,137 @@ model = AutoModelForCausalLM.from_pretrained(chatbot_model)
 print("Loading emotion detection model...")
 emotion_pipeline = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")
 
-# Store chat histories by session ID
+# Store chat histories
 chat_histories = {}
 
-# Request models
-class ChatRequest(BaseModel):
-    message: str
-    session_id: str = None
-
-class ChatResponse(BaseModel):
-    response: str
-    emotion: str
-    emotion_score: float
-    session_id: str
-
-@app.get("/")
-async def root():
-    return {"message": "Mental Health Chatbot API is running. Use /api/chat to interact with the chatbot."}
-
-@app.post("/api/chat", response_model=ChatResponse)
-async def chat(request: ChatRequest):
-    # Create a new session ID if not provided
-    session_id = request.session_id if request.session_id else str(uuid.uuid4())
-    
-    # Initialize chat history for new sessions
+def chatbot_response(message, history=None, session_id="default"):
+    """Generate a chatbot response and detect emotion from user message"""
+    # Initialize session if it doesn't exist
     if session_id not in chat_histories:
         chat_histories[session_id] = []
     
-    user_input = request.message
-    
     # Generate chatbot response
-    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
+    input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")
     output = model.generate(input_ids, max_length=200, pad_token_id=tokenizer.eos_token_id)
     response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
     
     # Detect emotion
-    emotion_result = emotion_pipeline(user_input)[0]
-    emotion = emotion_result["label"]
-    score = emotion_result["score"]
+    emotion_result = emotion_pipeline(message)
+    emotion = emotion_result[0]["label"]
+    score = float(emotion_result[0]["score"])
+    
+    # Store in chat history
+    chat_histories[session_id].append((message, response))
     
-    # Store chat history
-    chat_histories[session_id].append({"role": "user", "content": user_input})
-    chat_histories[session_id].append({"role": "bot", "content": response})
+    return response, emotion, score, chat_histories[session_id]
+
+def api_chatbot_response(message, session_id="default"):
+    """API endpoint version that returns a structured response"""
+    response, emotion, score, _ = chatbot_response(message, None, session_id)
     
-    # Return the response
     return {
-        "response": response,
-        "emotion": emotion,
-        "emotion_score": float(score),
+        "bot_response": response, 
+        "emotion": emotion, 
+        "emotion_score": score,
         "session_id": session_id
     }
 
-@app.get("/api/history/{session_id}")
-async def get_chat_history(session_id: str):
-    if session_id not in chat_histories:
-        return {"error": "Session not found"}
-    
-    return {"session_id": session_id, "history": chat_histories[session_id]}
+def get_chat_history(session_id="default"):
+    """Get chat history for a specific session"""
+    if session_id in chat_histories:
+        return chat_histories[session_id]
+    return []
 
-@app.delete("/api/history/{session_id}")
-async def clear_chat_history(session_id: str):
+def clear_history(session_id="default"):
+    """Clear chat history for a specific session"""
     if session_id in chat_histories:
-        chat_histories.pop(session_id)
-        return {"message": f"Chat history for session {session_id} cleared"}
+        chat_histories[session_id] = []
+        return f"History cleared for session {session_id}"
+    return f"Session {session_id} not found"
+
+# Define UI interface
+with gr.Blocks(title="Mental Health Chatbot") as ui_interface:
+    gr.Markdown("# 🧠 Mental Health Chatbot")
+    
+    with gr.Row():
+        with gr.Column(scale=3):
+            chatbot = gr.Chatbot(height=400, label="Conversation")
+            
+            with gr.Row():
+                message = gr.Textbox(placeholder="Type your message here...", label="You", show_label=False)
+                submit_btn = gr.Button("Send")
+            
+            with gr.Row():
+                session_id = gr.Textbox(value="default", label="Session ID")
+                clear_btn = gr.Button("Clear Chat")
+        
+        with gr.Column(scale=1):
+            emotion_label = gr.Textbox(label="Emotion Detected")
+            emotion_score = gr.Number(label="Confidence Score")
+    
+    # Set up event handlers
+    def respond(message, chat_history, session_id):
+        response, emotion, score, _ = chatbot_response(message, chat_history, session_id)
+        chat_history.append((message, response))
+        return "", chat_history, emotion, score
+    
+    submit_btn.click(
+        respond, 
+        [message, chatbot, session_id], 
+        [message, chatbot, emotion_label, emotion_score]
+    )
     
-    return {"error": "Session not found"}
+    message.submit(
+        respond, 
+        [message, chatbot, session_id], 
+        [message, chatbot, emotion_label, emotion_score]
+    )
+    
+    clear_btn.click(
+        lambda s: ([], clear_history(s), "", 0), 
+        [session_id], 
+        [chatbot, emotion_label, emotion_score]
+    )
+
+# Define API interface
+api_interface = gr.Interface(
+    fn=api_chatbot_response,
+    inputs=[
+        gr.Textbox(label="Message"),
+        gr.Textbox(label="Session ID", value="default")
+    ],
+    outputs=gr.JSON(label="Response"),
+    title="Mental Health Chatbot API",
+    description="Send a message to get chatbot response with emotion analysis",
+    examples=[
+        ["I'm feeling sad today", "user1"],
+        ["I'm so excited about my new job!", "user2"],
+        ["I'm worried about my exam tomorrow", "user3"]
+    ]
+)
+
+history_api = gr.Interface(
+    fn=get_chat_history,
+    inputs=gr.Textbox(label="Session ID", value="default"),
+    outputs=gr.JSON(label="Chat History"),
+    title="Chat History API",
+    description="Get chat history for a specific session"
+)
+
+clear_api = gr.Interface(
+    fn=clear_history,
+    inputs=gr.Textbox(label="Session ID", value="default"),
+    outputs=gr.Textbox(label="Result"),
+    title="Clear History API",
+    description="Clear chat history for a specific session"
+)
+
+# Combine all interfaces
+demo = gr.TabbedInterface(
+    [ui_interface, api_interface, history_api, clear_api], 
+    ["Chat UI", "Chat API", "History API", "Clear API"]
+)
 
-# Launch the API server
+# Launch the app
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.launch()