add: guardrails manager

guardrails_genie/guardrails/__init__.py

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

guardrails_genie/guardrails/base.py

@@ -11,7 +11,3 @@ class Guardrail(weave.Model):
     @weave.op()
     def guard(self, prompt: str, **kwargs) -> list[str]:
         pass
-
-    @weave.op()
-    def predict(self, prompt: str, **kwargs) -> list[str]:
-        return self.guard(prompt, **kwargs)

guardrails_genie/guardrails/injection/survey_guardrail.py

@@ -17,7 +17,7 @@ class SurveyGuardrailResponse(BaseModel):
 
 class SurveyGuardrail(Guardrail):
     llm_model: OpenAIModel
-    
+
     @weave.op()
     def load_prompt_injection_survey(self) -> str:
         prompt_injection_survey_path = os.path.join(
@@ -61,7 +61,7 @@ Here are some strict instructions that you must follow:
         return user_prompt, system_prompt
 
     @weave.op()
-    def guard(self, prompt: str, **kwargs) -> list[str]:
+    def predict(self, prompt: str, **kwargs) -> list[str]:
         user_prompt, system_prompt = self.format_prompts(prompt)
         chat_completion = self.llm_model.predict(
             user_prompts=user_prompt,
@@ -70,3 +70,8 @@ Here are some strict instructions that you must follow:
             **kwargs,
         )
         return chat_completion.choices[0].message.parsed
+
+    @weave.op()
+    def guard(self, prompt: str, **kwargs) -> list[str]:
+        response = self.predict(prompt, **kwargs)
+        return {"verdict": response.injection_prompt}

guardrails_genie/guardrails/manager.py

@@ -0,0 +1,16 @@
+import weave
+from weave.flow.obj import Object as WeaveObject
+
+from .base import Guardrail
+
+
+class GuardrailManager(WeaveObject):
+    guardrails: list[Guardrail]
+
+    @weave.op()
+    def guard(self, prompt: str, **kwargs) -> dict:
+        alerts = []
+        for guardrail in self.guardrails:
+            response = guardrail.guard(prompt, **kwargs)
+            alerts.append({guardrail.name: response})
+        return alerts