app.py

import os
from fastapi import FastAPI, HTTPException
import gradio as gr
from huggingface_hub import InferenceClient
import re
from pydantic import BaseModel
from typing import List, Dict, Optional

# Retrieve the API key from the environment variable set in Hugging Face Spaces secrets
secret_key = os.environ.get("SECRET_KEY")
if not secret_key:
    raise ValueError("Missing SECRET_KEY environment variable. Please set it in your Hugging Face Spaces secrets.")

# Initialize FastAPI
app = FastAPI(title="Chatbot API")

SYSTEM_PROMPT = """You are a friendly virtual assistant chatbot. Never disclose that you are DeepSeek or any other AI model - simply say you are a virtual assistant chatbot who is here to help. Always maintain this identity.

For each response:
1. First analyze the user's input in a <think> tag
2. Then provide your actual response
3. Keep your identity consistent as a virtual assistant chatbot

Example format:
<think>User asked who I am. I should maintain my identity as a virtual assistant chatbot.</think>
Hello! I'm a virtual assistant chatbot, ready to help you with any questions or tasks you might have!"""

class ChatMessage(BaseModel):
    message: str

class ChatResponse(BaseModel):
    thinking: str
    response: str
    thinking_history: List[str]

class Chatbot:
    def __init__(self, api_key, model="deepseek-ai/DeepSeek-R1"):
        self.client = InferenceClient(
            provider="together",
            api_key=api_key
        )
        self.model = model
        self.conversation_history = []
        self.thinking_history = []

    def add_message(self, role, content):
        if role == "system":
            self.conversation_history.insert(0, {
                "role": role,
                "content": content
            })
        else:
            self.conversation_history.append({
                "role": role,
                "content": content
            })

    def parse_response(self, response_text):
        think_pattern = r'<think>(.*?)</think>'
        think_match = re.search(think_pattern, response_text, re.DOTALL)
        
        if think_match:
            thinking = think_match.group(1).strip()
            actual_response = re.sub(think_pattern, '', response_text, flags=re.DOTALL).strip()
            return thinking, actual_response
        else:
            return "Processing response...", response_text.strip()

    def get_response(self, max_tokens=500):
        try:
            if not self.conversation_history or self.conversation_history[0]["role"] != "system":
                self.add_message("system", SYSTEM_PROMPT)
            
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=self.conversation_history,
                max_tokens=max_tokens,
            )
            
            full_response = completion.choices[0].message.content
            thinking, cleaned_response = self.parse_response(full_response)
            
            self.thinking_history.append(thinking)
            self.add_message("assistant", cleaned_response)
            
            return {
                "thinking": thinking,
                "response": cleaned_response,
                "thinking_history": self.thinking_history
            }
        except Exception as e:
            return {
                "thinking": "Error occurred while processing",
                "response": f"Error: {str(e)}",
                "thinking_history": self.thinking_history
            }

    def clear_history(self):
        self.conversation_history = []
        self.thinking_history = []
        self.add_message("system", SYSTEM_PROMPT)

# Initialize chatbot with the API key from the environment variable
chatbot = Chatbot(api_key=secret_key)

# FastAPI routes
@app.post("/chat", response_model=ChatResponse)
async def chat_endpoint(message: ChatMessage):
    if message.message.lower() == 'clear':
        chatbot.clear_history()
        return ChatResponse(
            thinking="Clearing chat history...",
            response="Conversation history cleared!",
            thinking_history=[]
        )
    
    chatbot.add_message("user", message.message)
    response_data = chatbot.get_response()
    return ChatResponse(**response_data)

# Gradio interface
def chat_response(message, history):
    chatbot.add_message("user", message)
    response_data = chatbot.get_response()
    return response_data["response"]

def create_gradio_interface():
    with gr.Blocks() as interface:
        chatbot_ui = gr.Chatbot()
        msg = gr.Textbox(label="Message")
        clear = gr.Button("Clear")
        
        msg.submit(
            chat_response,
            [msg, chatbot_ui],
            [chatbot_ui]
        ).then(
            lambda: "",
            None,
            [msg]
        )
        
        clear.click(
            lambda: None,
            None,
            chatbot_ui,
            queue=False
        )
    
    return interface

# Create and mount the Gradio interface
interface = create_gradio_interface()
app = gr.mount_gradio_app(app, interface, path="/")

# For deployment on Hugging Face Spaces
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)