geekyrakshit commited on
Commit
a62b646
·
1 Parent(s): 351e5f6

fix: tracing for PromptInjectionProtectAIGuardrail

Browse files
guardrails_genie/guardrails/injection/protectai_guardrail.py CHANGED
@@ -23,10 +23,14 @@ class PromptInjectionProtectAIGuardrail(Guardrail):
23
  max_length=512,
24
  device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
25
  )
 
 
 
 
26
 
27
  @weave.op()
28
  def predict(self, prompt: str):
29
- response = weave.op()(self._classifier)(prompt)
30
  return {"safe": response[0]["label"] != "INJECTION"}
31
 
32
  @weave.op()
 
23
  max_length=512,
24
  device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
25
  )
26
+
27
+ @weave.op()
28
+ def classify(self, prompt: str):
29
+ return self._classifier(prompt)
30
 
31
  @weave.op()
32
  def predict(self, prompt: str):
33
+ response = self.classify(prompt)
34
  return {"safe": response[0]["label"] != "INJECTION"}
35
 
36
  @weave.op()