refactor: SurveyGuardrail to PromptInjectionSurveyGuardrail

app.py

@@ -2,8 +2,7 @@ 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.guardrails import GuardrailManager, PromptInjectionSurveyGuardrail
 from guardrails_genie.llm import OpenAIModel
 
 load_dotenv()
@@ -14,11 +13,19 @@ chat_condition = openai_model != ""
 
 guardrails = []
 
-with st.sidebar.expander("Switch on Guardrails"):
-    is_survey_guardrail_enabled = st.toggle("Survey Guardrail", value=True)
+with st.sidebar.expander("Switch on Prompt Injection Guardrails"):
+    is_survey_guardrail_enabled = st.toggle("Survey Guardrail")
 
     if is_survey_guardrail_enabled:
-        guardrails.append(SurveyGuardrail(llm_model=OpenAIModel(model_name="gpt-4o")))
+        survey_guardrail_model = st.selectbox(
+            "Survey Guardrail Model", ["", "gpt-4o-mini", "gpt-4o"]
+        )
+        if survey_guardrail_model:
+            guardrails.append(
+                PromptInjectionSurveyGuardrail(
+                    llm_model=OpenAIModel(model_name=survey_guardrail_model)
+                )
+            )
 
 guardrails_manager = GuardrailManager(guardrails=guardrails)
 

guardrails_genie/guardrails/__init__.py

@@ -1,4 +1,4 @@
-from .injection import SurveyGuardrail
+from .injection import PromptInjectionSurveyGuardrail
 from .manager import GuardrailManager
 
-__all__ = ["SurveyGuardrail", "GuardrailManager"]
+__all__ = ["PromptInjectionSurveyGuardrail", "GuardrailManager"]

guardrails_genie/guardrails/injection/__init__.py

@@ -1,3 +1,3 @@
-from .survey_guardrail import SurveyGuardrail
+from .survey_guardrail import PromptInjectionSurveyGuardrail
 
-__all__ = ["SurveyGuardrail"]
+__all__ = ["PromptInjectionSurveyGuardrail"]

guardrails_genie/guardrails/injection/survey_guardrail.py

@@ -15,7 +15,7 @@ class SurveyGuardrailResponse(BaseModel):
     explanation: Optional[str]
 
 
-class SurveyGuardrail(Guardrail):
+class PromptInjectionSurveyGuardrail(Guardrail):
     llm_model: OpenAIModel
 
     @weave.op()

pyproject.toml

@@ -12,7 +12,7 @@ dependencies = [
     "ruff>=0.6.9",
     "pip>=24.2",
     "uv>=0.4.20",
-    "weave>=0.51.19",
+    "weave>=0.51.22",
     "streamlit>=1.40.1",
     "python-dotenv>=1.0.1",
     "watchdog>=6.0.0",