Update app.py

app.py

@@ -40,6 +40,11 @@ except Exception:
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login, HfApi
 
+# -------------------
+# NEW: Safety imports (from your snippet / safety.py)
+# -------------------
+from safety import safety_filter, refusal_reply
+
 # -------------------
 # Config
 # -------------------
@@ -184,17 +189,39 @@ def local_generate(model, tokenizer, input_ids, max_new_tokens=350):
     return tokenizer.decode(gen_only, skip_special_tokens=True).strip()
 
 # -------------------
-# Chat Function
+# Chat Function (now with Safety layer)
 # -------------------
 def chat_fn(message, history, user_tz):
     try:
-        if is_identity_query(message, history):
+        # ---- INPUT SAFETY (your snippet semantics) ----
+        safe_in, blocked_in, reason_in = safety_filter(message, mode="input")
+        if blocked_in:
+            return refusal_reply(reason_in)
+
+        # Identity short-circuit (use sanitized input)
+        if is_identity_query(safe_in, history):
             return "I am ClarityOps, your strategic decision making AI partner."
+
+        # ---- GENERATION using sanitized input ----
         if USE_HOSTED_COHERE:
-            return cohere_chat(message, history)
-        model, tokenizer = load_local_model()
-        inputs = build_inputs(tokenizer, message, history)
-        return local_generate(model, tokenizer, inputs, max_new_tokens=350)
+            out = cohere_chat(safe_in, history)
+        else:
+            model, tokenizer = load_local_model()
+            inputs = build_inputs(tokenizer, safe_in, history)
+            out = local_generate(model, tokenizer, inputs, max_new_tokens=350)
+
+        # Tidy echoes (from your snippet)
+        if isinstance(out, str):
+            for tag in ("Assistant:", "System:", "User:"):
+                if out.startswith(tag):
+                    out = out[len(tag):].strip()
+
+        # ---- OUTPUT SAFETY ----
+        safe_out, blocked_out, reason_out = safety_filter(out, mode="output")
+        if blocked_out:
+            return refusal_reply(reason_out)
+
+        return safe_out
     except Exception as e:
         return f"Error: {e}"
 
@@ -214,19 +241,16 @@ custom_css = """
   --brand-text: #0f172a;
   --brand-text-light: #ffffff;
 }
-
 /* Page background */
 .gradio-container {
   background: var(--brand-bg);
 }
-
 /* Title */
 h1 {
   color: var(--brand-text);
   font-weight: 700;
   font-size: 28px !important;
 }
-
 /* Try to hide the default Chatbot label via CSS for multiple Gradio builds */
 .chatbot header,
 .chatbot .label,
@@ -236,7 +260,6 @@ h1 {
 .chatbot > .wrap > header {
   display: none !important;
 }
-
 /* Both bot and user bubbles teal with white text */
 .message.user, .message.bot {
   background: var(--brand-accent) !important;
@@ -244,7 +267,6 @@ h1 {
   border-radius: 12px !important;
   padding: 8px 12px !important;
 }
-
 /* Inputs a bit softer */
 textarea, input, .gr-input {
   border-radius: 12px !important;
@@ -308,9 +330,3 @@ if __name__ == "__main__":
         max_threads=8,       # optional: avoid thread-starvation on tiny CPUs
     )
 
-
-
-
-
-
-