Update app.py

app.py

@@ -1,21 +1,18 @@
-import os
-from dotenv import load_dotenv
 import gradio as gr
+from guardrail import is_safe  # Import the guardrail validation function
 from langchain_huggingface import HuggingFaceEndpoint
-from together import Together
+from dotenv import load_dotenv
+import os
 
 # Load environment variables
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
-API_KEY = os.getenv("API_KEY")
 
-if not HF_TOKEN or not API_KEY:
-    raise ValueError("Missing API tokens. Please check your .env file.")
+# Validate Hugging Face token
+if not HF_TOKEN:
+    raise ValueError("Missing Hugging Face API token. Please check your .env file.")
 
-# Initialize the Together client for guardrail functionality
-client = Together(api_key=API_KEY)
-
-# Initialize the Hugging Face endpoint for text generation (Mistral model)
+# Initialize the Hugging Face endpoint for generating responses
 llm = HuggingFaceEndpoint(
     repo_id="mistralai/Mistral-7B-Instruct-v0.3",
     huggingfacehub_api_token=HF_TOKEN.strip(),
@@ -23,61 +20,52 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=100
 )
 
-# List of topics or keywords inappropriate for kids under 16
-prohibited_topics = [
-    "violence", "drugs", "explicit content", "profanity", "hate speech", "alcohol",
-    "self-harm", "gambling", "kidnapping","sexual content", "graphic descriptions"
-]
-
-# Function to handle chatbot response with TogetherAI's guardrails
-def chatbot_response_with_guardrails(message):
+# Chatbot response function with safety checks
+def chatbot_response(user_message: str) -> str:
+    """
+    Generates a chatbot response while ensuring the content is safe for children under 16.
+    
+    Args:
+        user_message (str): The input message from the user.
+    
+    Returns:
+        str: A safe chatbot response or a safety warning message.
+    """
     try:
-        # Step 1: Generate raw response using Mistral model
-        raw_response = llm.invoke(message)
+        # Step 1: Validate the user input
+        if not is_safe(user_message):
+            return "Sorry, I cannot respond to that as it violates our safety policy."
 
-        # Step 2: Use TogetherAI's guardrail model to check the response
-        response = client.completions.create(
-            model="Meta-Llama/LlamaGuard-2-8b",
-            prompt=f"Check the following text for safety: {raw_response}"
-        )
+        # Step 2: Generate a response using the Mistral model
+        raw_response = llm.invoke(user_message)
 
-        # Extract the response from TogetherAI's guardrail model
-        guardrail_check = response.choices[0].text.strip()
+        # Step 3: Validate the generated response
+        if not is_safe(raw_response):
+            return "Sorry, I cannot share that information as it violates our safety policy."
 
-        # Step 3: Check for inappropriate content
-        if 'toxic' in guardrail_check.lower() or any(
-            topic in guardrail_check.lower() for topic in prohibited_topics
-        ):
-            return "Sorry, the content is not suitable for children under 16."
-        
-        # Step 4: Check raw response for prohibited topics
-        if any(topic in raw_response.lower() for topic in prohibited_topics):
-            return "Sorry, the content is not suitable for children under 16."
-
-        # If the response is safe, return the raw response
+        # Step 4: Return the validated response
         return raw_response
 
     except Exception as e:
-        return f"Error: {e}"
+        return f"An error occurred: {str(e)}"
 
-# Gradio Interface for Chatbot with Guardrails
-with gr.Blocks() as app_with_guardrails:
-    gr.Markdown("## Chatbot With Kid-Safe Guardrails")
-    gr.Markdown(
-        "This chatbot ensures all responses are appropriate for children under 16."
-    )
+# Gradio Interface for the chatbot
+with gr.Blocks() as app:
+    gr.Markdown("## Kid-Safe Chatbot 🛡️")
+    gr.Markdown("This chatbot ensures that all responses are appropriate for children under 16.")
 
     with gr.Row():
-        user_input = gr.Textbox(label="Your Message", placeholder="Type here...")
-    response_output = gr.Textbox(label="Guarded Response", placeholder="Bot will respond here...")
+        user_input = gr.Textbox(label="Your Message", placeholder="Type your message here...")
+    response_output = gr.Textbox(label="Chatbot Response", placeholder="The chatbot will respond here.")
     submit_button = gr.Button("Send")
 
+    # On button click, generate response
     submit_button.click(
-        chatbot_response_with_guardrails,
+        fn=chatbot_response,
         inputs=[user_input],
         outputs=[response_output]
     )
 
 # Launch the app
 if __name__ == "__main__":
-    app_with_guardrails.launch()
+    app.launch()