update: app

app.py

@@ -2,6 +2,8 @@ import streamlit as st
 import weave
 from dotenv import load_dotenv
 
+from guardrails_genie.guardrails import GuardrailManager
+from guardrails_genie.guardrails.injection import SurveyGuardrail
 from guardrails_genie.llm import OpenAIModel
 
 load_dotenv()
@@ -10,6 +12,16 @@ weave.init(project_name="guardrails-genie")
 openai_model = st.sidebar.selectbox("OpenAI LLM", ["", "gpt-4o-mini", "gpt-4o"])
 chat_condition = openai_model != ""
 
+guardrails = []
+
+with st.sidebar.expander("Switch on Guardrails"):
+    is_survey_guardrail_enabled = st.toggle("Survey Guardrail", value=True)
+
+    if is_survey_guardrail_enabled:
+        guardrails.append(SurveyGuardrail(llm_model=OpenAIModel(model_name="gpt-4o")))
+
+guardrails_manager = GuardrailManager(guardrails=guardrails)
+
 # Use session state to track if the chat has started
 if "chat_started" not in st.session_state:
     st.session_state.chat_started = False
@@ -40,13 +52,23 @@ if st.session_state.chat_started:
         # Add user message to chat history
         st.session_state.messages.append({"role": "user", "content": prompt})
 
-        response, call = llm_model.predict.call(
-            llm_model, user_prompts=prompt, messages=st.session_state.messages
+        guardrails_response, call = guardrails_manager.guard.call(
+            guardrails_manager, prompt=prompt
         )
-        response = response.choices[0].message.content
 
-        # Display assistant response in chat message container
-        with st.chat_message("assistant"):
-            st.markdown(response + f"\n\n---\n[Explore in Weave]({call.ui_url})")
-        # Add assistant response to chat history
-        st.session_state.messages.append({"role": "assistant", "content": response})
+        if guardrails_response["safe"]:
+            response, call = llm_model.predict.call(
+                llm_model, user_prompts=prompt, messages=st.session_state.messages
+            )
+            response = response.choices[0].message.content
+
+            # Display assistant response in chat message container
+            with st.chat_message("assistant"):
+                st.markdown(response + f"\n\n---\n[Explore in Weave]({call.ui_url})")
+            # Add assistant response to chat history
+            st.session_state.messages.append({"role": "assistant", "content": response})
+        else:
+            st.error("Guardrails detected an issue with the prompt.")
+            for alert in guardrails_response["alerts"]:
+                st.error(f"{alert['guardrail_name']}: {alert['response']}")
+            st.error(f"For details, explore in Weave at {call.ui_url}")

guardrails_genie/guardrails/injection/survey_guardrail.py

@@ -74,4 +74,4 @@ Here are some strict instructions that you must follow:
     @weave.op()
     def guard(self, prompt: str, **kwargs) -> list[str]:
         response = self.predict(prompt, **kwargs)
-        return {"verdict": response.injection_prompt}
+        return {"safe": not response.injection_prompt}

guardrails_genie/guardrails/manager.py

@@ -1,4 +1,5 @@
 import weave
+from rich.progress import track
 from weave.flow.obj import Object as WeaveObject
 
 from .base import Guardrail
@@ -9,8 +10,11 @@ class GuardrailManager(WeaveObject):
 
     @weave.op()
     def guard(self, prompt: str, **kwargs) -> dict:
-        alerts = []
-        for guardrail in self.guardrails:
+        alerts, safe = [], True
+        for guardrail in track(self.guardrails, description="Running guardrails"):
             response = guardrail.guard(prompt, **kwargs)
-            alerts.append({guardrail.name: response})
-        return alerts
+            alerts.append(
+                {"guardrail_name": guardrail.__class__.__name__, "response": response}
+            )
+            safe = safe and response["safe"]
+        return {"safe": safe, "alerts": alerts}