Update app.py

app.py

@@ -9,12 +9,15 @@ 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.")
+
 # Initialize the Together client for guardrail functionality
 client = Together(api_key=API_KEY)
 
 # 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
+    repo_id="mistralai/Mistral-7B-Instruct-v0.3",
     huggingfacehub_api_token=HF_TOKEN.strip(),
     temperature=0.7,
     max_new_tokens=100
@@ -34,15 +37,14 @@ def chatbot_response_with_guardrails(message):
 
         # Step 2: Use TogetherAI's guardrail model to check the response
         response = client.completions.create(
-            model="Meta-Llama/LlamaGuard-2-8b",  # TogetherAI guardrail model
-            prompt=raw_response
+            model="Meta-Llama/LlamaGuard-2-8b",
+            prompt=f"Check the following text for safety: {raw_response}"
         )
 
         # Extract the response from TogetherAI's guardrail model
-        guardrail_check = response.choices[0].text.strip
-        print(guardrail_check)
+        guardrail_check = response.choices[0].text.strip()
 
-        # Step 3: Check for inappropriate content in the guardrail model's output
+        # Step 3: Check for inappropriate content
         if 'toxic' in guardrail_check.lower() or any(
             topic in guardrail_check.lower() for topic in prohibited_topics
         ):
@@ -53,8 +55,7 @@ def chatbot_response_with_guardrails(message):
             return "Sorry, the content is not suitable for children under 16."
 
         # If the response is safe, return the raw response
-        else:
-            return response
+        return raw_response
 
     except Exception as e:
         return f"Error: {e}"
@@ -66,13 +67,11 @@ with gr.Blocks() as app_with_guardrails:
         "This chatbot ensures all responses are appropriate for children under 16."
     )
 
-    # 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],