update: evaluation + classifier guardrail

application_pages/evaluation_app.py

@@ -4,6 +4,7 @@ import time
 from importlib import import_module
 
 import pandas as pd
+import rich
 import streamlit as st
 import weave
 from dotenv import load_dotenv
@@ -181,11 +182,12 @@ if st.session_state.dataset_previewed:
                             st.session_state.evaluation_call_manager.call_list.append(
                                 {
                                     "guardrail_name": guardrail_name,
-                                    "calls": st.session_state.evaluation_call_manager.collect_guardrail_guard_calls_from_eval(
-                                        call=call
-                                    ),
+                                    "calls": st.session_state.evaluation_call_manager.collect_guardrail_guard_calls_from_eval(),
                                 }
                             )
+                            rich.print(
+                                st.session_state.evaluation_call_manager.call_list
+                            )
                         st.dataframe(
                             st.session_state.evaluation_call_manager.render_calls_to_streamlit()
                         )

guardrails_genie/guardrails/injection/protectai_guardrail.py

@@ -37,11 +37,6 @@ class PromptInjectionClassifierGuardrail(Guardrail):
     def classify(self, prompt: str):
         return self._classifier(prompt)
 
-    @weave.op()
-    def predict(self, prompt: str):
-        response = self.classify(prompt)
-        return {"safe": response[0]["label"] != "INJECTION"}
-
     @weave.op()
     def guard(self, prompt: str):
         response = self.classify(prompt)
@@ -50,3 +45,7 @@ class PromptInjectionClassifierGuardrail(Guardrail):
             "safe": response[0]["label"] != "INJECTION",
             "summary": f"Prompt is deemed {response[0]['label']} with {confidence_percentage}% confidence.",
         }
+
+    @weave.op()
+    def predict(self, prompt: str):
+        return self.guard(prompt)

guardrails_genie/utils.py

@@ -22,16 +22,19 @@ class EvaluationCallManager:
         self.show_warning_in_app = False
         self.call_list = []
 
-    def collect_guardrail_guard_calls_from_eval(self, call):
+    def collect_guardrail_guard_calls_from_eval(self):
         guard_calls, count = [], 0
-        for eval_predict_call in call.children():
-            if "Evaluation.summarize" in eval_predict_call._op_name:
+        for eval_predict_and_score_call in self.base_call.children():
+            if "Evaluation.summarize" in eval_predict_and_score_call._op_name:
                 break
-            required_call = eval_predict_call.children()[0].children()[0].children()[0]
+            guardrail_predict_call = eval_predict_and_score_call.children()[0]
+            guard_call = guardrail_predict_call.children()[0]
+            score_call = eval_predict_and_score_call.children()[1]
             guard_calls.append(
                 {
-                    "input_prompt": str(required_call.inputs["prompt"]),
-                    "outputs": dict(required_call.output),
+                    "input_prompt": str(guard_call.inputs["prompt"]),
+                    "outputs": dict(guard_call.output),
+                    "score": dict(score_call.output),
                 }
             )
             count += 1
@@ -50,7 +53,7 @@ class EvaluationCallManager:
             dataframe[guardrail_call["guardrail_name"] + ".safe"] = [
                 call["outputs"]["safe"] for call in guardrail_call["calls"]
             ]
-            dataframe[guardrail_call["guardrail_name"] + ".summary"] = [
-                call["outputs"]["summary"] for call in guardrail_call["calls"]
+            dataframe[guardrail_call["guardrail_name"] + ".prediction_correctness"] = [
+                call["score"]["correct"] for call in guardrail_call["calls"]
             ]
         return pd.DataFrame(dataframe)