app.py

import os
from dotenv import load_dotenv
import gradio as gr
from langchain_huggingface import HuggingFaceEndpoint
from transformers import pipeline

# Load environment variables
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")

# Initialize the Hugging Face endpoint for text generation (Mistral model)
llm = HuggingFaceEndpoint(
    repo_id="mistralai/Mistral-7B-Instruct-v0.3",  # Replace with your model repo
    huggingfacehub_api_token=HF_TOKEN.strip(),
    temperature=0.7,
    max_new_tokens=100
)

# Initialize content moderation model (unitary/matti)
content_filter = pipeline("text-classification", model="unitary/matti")

# Function to handle chatbot response and guardrails
def chatbot_response_with_guardrails(message):
    try:
        # Generate raw response from the primary model (Mistral)
        raw_response = llm(message)

        # Check if the response contains inappropriate content using the content filter
        result = content_filter(raw_response)
        
        # If the response is deemed harmful, modify it or reject it
        if result[0]['label'] == 'toxic':  # Adjust based on your model's output
            return "Content not suitable."
        else:
            return raw_response
    except Exception as e:
        return f"Error: {e}"

# Gradio Interface for Chatbot with Guardrails
with gr.Blocks() as app_with_guardrails:
    gr.Markdown("## Chatbot With Guardrails")
    gr.Markdown("This chatbot ensures all responses are appropriate.")

    # Input and output
    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...")
    submit_button = gr.Button("Send")

    # Button click event
    submit_button.click(
        chatbot_response_with_guardrails,
        inputs=[user_input],
        outputs=[response_output]
    )

# Launch the app
if __name__ == "__main__":
    app_with_guardrails.launch()