Merge branch 'main' into feat/pii-banned-words

.gitignore

@@ -165,4 +165,7 @@ cursor_prompts/
 uv.lock
 test.py
 temp.txt
-**.csv
+**.csv
+binary-classifier/
+wandb/
+artifacts/

README.md

@@ -18,7 +18,12 @@ source .venv/bin/activate
 ## Run the App
 
 ```bash
-OPENAI_API_KEY="YOUR_OPENAI_API_KEY" streamlit run app.py
+export OPENAI_API_KEY="YOUR_OPENAI_API_KEY"
+export WEAVE_PROJECT="YOUR_WEAVE_PROJECT"
+export WANDB_PROJECT_NAME="YOUR_WANDB_PROJECT_NAME"
+export WANDB_ENTITY_NAME="YOUR_WANDB_ENTITY_NAME"
+export WANDB_LOG_MODEL="checkpoint"
+streamlit run app.py
 ```
 
 ## Use the Library

app.py

@@ -4,13 +4,22 @@ intro_page = st.Page(
     "application_pages/intro_page.py", title="Introduction", icon=":material/guardian:"
 )
 chat_page = st.Page(
-    "application_pages/chat_app.py", title="Chat", icon=":material/robot:"
+    "application_pages/chat_app.py",
+    title="Playground",
+    icon=":material/sports_esports:",
 )
 evaluation_page = st.Page(
     "application_pages/evaluation_app.py",
     title="Evaluation",
     icon=":material/monitoring:",
 )
-page_navigation = st.navigation([intro_page, chat_page, evaluation_page])
+train_classifier_page = st.Page(
+    "application_pages/train_classifier.py",
+    title="Train Classifier",
+    icon=":material/fitness_center:",
+)
+page_navigation = st.navigation(
+    [intro_page, chat_page, evaluation_page, train_classifier_page]
+)
 st.set_page_config(page_title="Guardrails Genie", page_icon=":material/guardian:")
 page_navigation.run()

application_pages/chat_app.py

@@ -1,4 +1,5 @@
 import importlib
+import os
 
 import streamlit as st
 import weave
@@ -7,19 +8,27 @@ from dotenv import load_dotenv
 from guardrails_genie.guardrails import GuardrailManager
 from guardrails_genie.llm import OpenAIModel
 
-load_dotenv()
-weave.init(project_name="guardrails-genie")
 
-st.title(":material/robot: Guardrails Genie")
-
-if "guardrails" not in st.session_state:
-    st.session_state.guardrails = []
-if "guardrail_names" not in st.session_state:
-    st.session_state.guardrail_names = []
-if "guardrails_manager" not in st.session_state:
-    st.session_state.guardrails_manager = None
-if "chat_started" not in st.session_state:
-    st.session_state.chat_started = False
+def initialize_session_state():
+    load_dotenv()
+    weave.init(project_name=os.getenv("WEAVE_PROJECT"))
+
+    if "guardrails" not in st.session_state:
+        st.session_state.guardrails = []
+    if "guardrail_names" not in st.session_state:
+        st.session_state.guardrail_names = []
+    if "guardrails_manager" not in st.session_state:
+        st.session_state.guardrails_manager = None
+    if "initialize_guardrails" not in st.session_state:
+        st.session_state.initialize_guardrails = False
+    if "system_prompt" not in st.session_state:
+        st.session_state.system_prompt = ""
+    if "user_prompt" not in st.session_state:
+        st.session_state.user_prompt = ""
+    if "test_guardrails" not in st.session_state:
+        st.session_state.test_guardrails = False
+    if "llm_model" not in st.session_state:
+        st.session_state.llm_model = None
 
 
 def initialize_guardrails():
@@ -36,18 +45,30 @@ def initialize_guardrails():
                         guardrail_name,
                     )(llm_model=OpenAIModel(model_name=survey_guardrail_model))
                 )
-        else:
-            st.session_state.guardrails.append(
-                getattr(
-                    importlib.import_module("guardrails_genie.guardrails"),
-                    guardrail_name,
-                )()
+        elif guardrail_name == "PromptInjectionClassifierGuardrail":
+            classifier_model_name = st.sidebar.selectbox(
+                "Classifier Guardrail Model",
+                [
+                    "",
+                    "ProtectAI/deberta-v3-base-prompt-injection-v2",
+                    "wandb://geekyrakshit/guardrails-genie/model-6rwqup9b:v3",
+                ],
             )
+            if classifier_model_name != "":
+                st.session_state.guardrails.append(
+                    getattr(
+                        importlib.import_module("guardrails_genie.guardrails"),
+                        guardrail_name,
+                    )(model_name=classifier_model_name)
+                )
     st.session_state.guardrails_manager = GuardrailManager(
         guardrails=st.session_state.guardrails
     )
 
 
+initialize_session_state()
+st.title(":material/robot: Guardrails Genie Playground")
+
 openai_model = st.sidebar.selectbox(
     "OpenAI LLM for Chat", ["", "gpt-4o-mini", "gpt-4o"]
 )
@@ -67,48 +88,41 @@ guardrail_names = st.sidebar.multiselect(
 )
 st.session_state.guardrail_names = guardrail_names
 
-if st.sidebar.button("Start Chat") and chat_condition:
-    st.session_state.chat_started = True
+if st.sidebar.button("Initialize Guardrails") and chat_condition:
+    st.session_state.initialize_guardrails = True
 
-if st.session_state.chat_started:
+if st.session_state.initialize_guardrails:
     with st.sidebar.status("Initializing Guardrails..."):
         initialize_guardrails()
+        st.session_state.llm_model = OpenAIModel(model_name=openai_model)
 
-    # Initialize chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-
-    llm_model = OpenAIModel(model_name=openai_model)
-
-    # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
+    user_prompt = st.text_area("User Prompt", value="")
+    st.session_state.user_prompt = user_prompt
 
-    # React to user input
-    if prompt := st.chat_input("What is up?"):
-        # Display user message in chat message container
-        st.chat_message("user").markdown(prompt)
-        # Add user message to chat history
-        st.session_state.messages.append({"role": "user", "content": prompt})
+    test_guardrails_button = st.button("Test Guardrails")
+    st.session_state.test_guardrails = test_guardrails_button
 
-        guardrails_response, call = st.session_state.guardrails_manager.guard.call(
-            st.session_state.guardrails_manager, prompt=prompt
-        )
+    if st.session_state.test_guardrails:
+        with st.sidebar.status("Running Guardrails..."):
+            guardrails_response, call = st.session_state.guardrails_manager.guard.call(
+                st.session_state.guardrails_manager, prompt=st.session_state.user_prompt
+            )
 
         if guardrails_response["safe"]:
-            response, call = llm_model.predict.call(
-                llm_model, user_prompts=prompt, messages=st.session_state.messages
+            st.markdown(
+                f"\n\n---\nPrompt is safe! Explore guardrail trace on [Weave]({call.ui_url})\n\n---\n"
             )
-            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})
+            with st.sidebar.status("Generating response from LLM..."):
+                response, call = st.session_state.llm_model.predict.call(
+                    st.session_state.llm_model,
+                    user_prompts=st.session_state.user_prompt,
+                )
+            st.markdown(
+                response.choices[0].message.content
+                + f"\n\n---\nExplore LLM generation trace on [Weave]({call.ui_url})"
+            )
         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}")
+            st.warning("Prompt is not safe!")
+            st.markdown(guardrails_response["summary"])
+            st.markdown(f"Explore prompt trace on [Weave]({call.ui_url})")

application_pages/evaluation_app.py

@@ -1,7 +1,10 @@
 import asyncio
+import os
+import time
 from importlib import import_module
 
 import pandas as pd
+import rich
 import streamlit as st
 import weave
 from dotenv import load_dotenv
@@ -9,12 +12,11 @@ from dotenv import load_dotenv
 from guardrails_genie.guardrails import GuardrailManager
 from guardrails_genie.llm import OpenAIModel
 from guardrails_genie.metrics import AccuracyMetric
-
-load_dotenv()
-weave.init(project_name="guardrails-genie")
+from guardrails_genie.utils import EvaluationCallManager
 
 
 def initialize_session_state():
+    load_dotenv()
     if "uploaded_file" not in st.session_state:
         st.session_state.uploaded_file = None
     if "dataset_name" not in st.session_state:
@@ -35,6 +37,18 @@ def initialize_session_state():
         st.session_state.evaluation_summary = None
     if "guardrail_manager" not in st.session_state:
         st.session_state.guardrail_manager = None
+    if "evaluation_name" not in st.session_state:
+        st.session_state.evaluation_name = ""
+    if "show_result_table" not in st.session_state:
+        st.session_state.show_result_table = False
+    if "weave_client" not in st.session_state:
+        st.session_state.weave_client = weave.init(
+            project_name=os.getenv("WEAVE_PROJECT")
+        )
+    if "evaluation_call_manager" not in st.session_state:
+        st.session_state.evaluation_call_manager = None
+    if "call_id" not in st.session_state:
+        st.session_state.call_id = None
 
 
 def initialize_guardrail():
@@ -51,10 +65,22 @@ def initialize_guardrail():
                         guardrail_name,
                     )(llm_model=OpenAIModel(model_name=survey_guardrail_model))
                 )
-        else:
-            guardrails.append(
-                getattr(import_module("guardrails_genie.guardrails"), guardrail_name)()
+        elif guardrail_name == "PromptInjectionClassifierGuardrail":
+            classifier_model_name = st.sidebar.selectbox(
+                "Classifier Guardrail Model",
+                [
+                    "",
+                    "ProtectAI/deberta-v3-base-prompt-injection-v2",
+                    "wandb://geekyrakshit/guardrails-genie/model-6rwqup9b:v3",
+                ],
             )
+            if classifier_model_name:
+                st.session_state.guardrails.append(
+                    getattr(
+                        import_module("guardrails_genie.guardrails"),
+                        guardrail_name,
+                    )(model_name=classifier_model_name)
+                )
     st.session_state.guardrails = guardrails
     st.session_state.guardrail_manager = GuardrailManager(guardrails=guardrails)
 
@@ -107,6 +133,8 @@ if st.session_state.dataset_previewed:
 
     if st.session_state.guardrail_names != []:
         initialize_guardrail()
+        evaluation_name = st.sidebar.text_input("Evaluation name", value="")
+        st.session_state.evaluation_name = evaluation_name
         if st.session_state.guardrail_manager is not None:
             if st.sidebar.button("Start Evaluation"):
                 st.session_state.start_evaluation = True
@@ -119,10 +147,55 @@ if st.session_state.dataset_previewed:
                 with st.expander("Evaluation Results", expanded=True):
                     evaluation_summary, call = asyncio.run(
                         evaluation.evaluate.call(
-                            evaluation, st.session_state.guardrail_manager
+                            evaluation,
+                            st.session_state.guardrail_manager,
+                            __weave={
+                                "display_name": "Evaluation.evaluate:"
+                                + st.session_state.evaluation_name
+                            },
                         )
                     )
-                    st.markdown(f"[Explore evaluation in Weave]({call.ui_url})")
-                    st.write(evaluation_summary)
-                st.session_state.evaluation_summary = evaluation_summary
-                st.session_state.start_evaluation = False
+                    x_axis = list(evaluation_summary["AccuracyMetric"].keys())
+                    y_axis = [
+                        evaluation_summary["AccuracyMetric"][x_axis_item]
+                        for x_axis_item in x_axis
+                    ]
+                    st.bar_chart(
+                        pd.DataFrame({"Metric": x_axis, "Score": y_axis}),
+                        x="Metric",
+                        y="Score",
+                    )
+                    st.session_state.evaluation_summary = evaluation_summary
+                    st.session_state.call_id = call.id
+                    st.session_state.start_evaluation = False
+
+                    if not st.session_state.start_evaluation:
+                        time.sleep(5)
+                        st.session_state.evaluation_call_manager = (
+                            EvaluationCallManager(
+                                entity="geekyrakshit",
+                                project="guardrails-genie",
+                                call_id=st.session_state.call_id,
+                            )
+                        )
+                        for guardrail_name in st.session_state.guardrail_names:
+                            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(),
+                                }
+                            )
+                            rich.print(
+                                st.session_state.evaluation_call_manager.call_list
+                            )
+                        st.dataframe(
+                            st.session_state.evaluation_call_manager.render_calls_to_streamlit()
+                        )
+                        if st.session_state.evaluation_call_manager.show_warning_in_app:
+                            st.warning(
+                                f"Only {st.session_state.evaluation_call_manager.max_count} calls can be shown in the app."
+                            )
+                        st.markdown(
+                            f"Explore the entire evaluation trace table in [Weave]({call.ui_url})"
+                        )
+                    st.session_state.evaluation_call_manager = None

application_pages/train_classifier.py

@@ -0,0 +1,61 @@
+import os
+
+import streamlit as st
+from dotenv import load_dotenv
+
+from guardrails_genie.train_classifier import train_binary_classifier
+
+
+def initialize_session_state():
+    load_dotenv()
+    if "dataset_name" not in st.session_state:
+        st.session_state.dataset_name = None
+    if "base_model_name" not in st.session_state:
+        st.session_state.base_model_name = None
+    if "batch_size" not in st.session_state:
+        st.session_state.batch_size = 16
+    if "should_start_training" not in st.session_state:
+        st.session_state.should_start_training = False
+    if "training_output" not in st.session_state:
+        st.session_state.training_output = None
+
+
+initialize_session_state()
+st.title(":material/fitness_center: Train Classifier")
+
+dataset_name = st.sidebar.text_input("Dataset Name", value="")
+st.session_state.dataset_name = dataset_name
+
+base_model_name = st.sidebar.selectbox(
+    "Base Model",
+    options=[
+        "distilbert/distilbert-base-uncased",
+        "FacebookAI/roberta-base",
+        "microsoft/deberta-v3-base",
+    ],
+)
+st.session_state.base_model_name = base_model_name
+
+batch_size = st.sidebar.slider(
+    "Batch Size", min_value=4, max_value=256, value=16, step=4
+)
+st.session_state.batch_size = batch_size
+
+train_button = st.sidebar.button("Train")
+st.session_state.should_start_training = (
+    train_button and st.session_state.dataset_name and st.session_state.base_model_name
+)
+
+if st.session_state.should_start_training:
+    with st.expander("Training", expanded=True):
+        training_output = train_binary_classifier(
+            project_name=os.getenv("WANDB_PROJECT_NAME"),
+            entity_name=os.getenv("WANDB_ENTITY_NAME"),
+            run_name=f"{st.session_state.base_model_name}-finetuned",
+            dataset_repo=st.session_state.dataset_name,
+            model_name=st.session_state.base_model_name,
+            batch_size=st.session_state.batch_size,
+            streamlit_mode=True,
+        )
+        st.session_state.training_output = training_output
+        st.write(training_output)

guardrails_genie/guardrails/__init__.py

@@ -1,8 +1,11 @@
-from .injection import PromptInjectionProtectAIGuardrail, PromptInjectionSurveyGuardrail
+from .injection import (
+    PromptInjectionClassifierGuardrail,
+    PromptInjectionSurveyGuardrail,
+)
 from .manager import GuardrailManager
 
 __all__ = [
     "PromptInjectionSurveyGuardrail",
-    "PromptInjectionProtectAIGuardrail",
+    "PromptInjectionClassifierGuardrail",
     "GuardrailManager",
 ]

guardrails_genie/guardrails/injection/__init__.py

@@ -1,4 +1,4 @@
-from .protectai_guardrail import PromptInjectionProtectAIGuardrail
+from .classifier_guardrail import PromptInjectionClassifierGuardrail
 from .survey_guardrail import PromptInjectionSurveyGuardrail
 
-__all__ = ["PromptInjectionSurveyGuardrail", "PromptInjectionProtectAIGuardrail"]
+__all__ = ["PromptInjectionSurveyGuardrail", "PromptInjectionClassifierGuardrail"]

guardrails_genie/guardrails/injection/{protectai_guardrail.py → classifier_guardrail.py}

@@ -5,16 +5,25 @@ import weave
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 from transformers.pipelines.base import Pipeline
 
+import wandb
+
 from ..base import Guardrail
 
 
-class PromptInjectionProtectAIGuardrail(Guardrail):
+class PromptInjectionClassifierGuardrail(Guardrail):
     model_name: str = "ProtectAI/deberta-v3-base-prompt-injection-v2"
     _classifier: Optional[Pipeline] = None
 
     def model_post_init(self, __context):
-        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
+        if self.model_name.startswith("wandb://"):
+            api = wandb.Api()
+            artifact = api.artifact(self.model_name.removeprefix("wandb://"))
+            artifact_dir = artifact.download()
+            tokenizer = AutoTokenizer.from_pretrained(artifact_dir)
+            model = AutoModelForSequenceClassification.from_pretrained(artifact_dir)
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
         self._classifier = pipeline(
             "text-classification",
             model=model,
@@ -29,10 +38,14 @@ class PromptInjectionProtectAIGuardrail(Guardrail):
         return self._classifier(prompt)
 
     @weave.op()
-    def predict(self, prompt: str):
+    def guard(self, prompt: str):
         response = self.classify(prompt)
-        return {"safe": response[0]["label"] != "INJECTION"}
+        confidence_percentage = round(response[0]["score"] * 100, 2)
+        return {
+            "safe": response[0]["label"] != "INJECTION",
+            "summary": f"Prompt is deemed {response[0]['label']} with {confidence_percentage}% confidence.",
+        }
 
     @weave.op()
-    def guard(self, prompt: str):
-        return self.predict(prompt)
+    def predict(self, prompt: str):
+        return self.guard(prompt)

guardrails_genie/guardrails/injection/survey_guardrail.py

@@ -70,8 +70,17 @@ Here are some strict instructions that you must follow:
             **kwargs,
         )
         response = chat_completion.choices[0].message.parsed
-        return {"safe": not response.injection_prompt}
+        return response
 
     @weave.op()
     def guard(self, prompt: str, **kwargs) -> list[str]:
-        return self.predict(prompt, **kwargs)
+        response = self.predict(prompt, **kwargs)
+        summary = (
+            f"Prompt is deemed safe. {response.explanation}"
+            if not response.injection_prompt
+            else f"Prompt is deemed a {'direct attack' if response.is_direct_attack else 'indirect attack'} of type {response.attack_type}. {response.explanation}"
+        )
+        return {
+            "safe": not response.injection_prompt,
+            "summary": summary,
+        }

guardrails_genie/guardrails/manager.py

@@ -9,7 +9,7 @@ class GuardrailManager(weave.Model):
 
     @weave.op()
     def guard(self, prompt: str, progress_bar: bool = True, **kwargs) -> dict:
-        alerts, safe = [], True
+        alerts, summaries, safe = [], "", True
         iterable = (
             track(self.guardrails, description="Running guardrails")
             if progress_bar
@@ -21,7 +21,10 @@ class GuardrailManager(weave.Model):
                 {"guardrail_name": guardrail.__class__.__name__, "response": response}
             )
             safe = safe and response["safe"]
-        return {"safe": safe, "alerts": alerts}
+            summaries += (
+                f"**{guardrail.__class__.__name__}**: {response['summary']}\n\n---\n\n"
+            )
+        return {"safe": safe, "alerts": alerts, "summary": summaries}
 
     @weave.op()
     def predict(self, prompt: str, **kwargs) -> dict:

guardrails_genie/train_classifier.py

@@ -0,0 +1,110 @@
+import evaluate
+import numpy as np
+import streamlit as st
+from datasets import load_dataset
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    Trainer,
+    TrainerCallback,
+    TrainingArguments,
+)
+from transformers.trainer_callback import TrainerControl, TrainerState
+
+import wandb
+
+
+class StreamlitProgressbarCallback(TrainerCallback):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.progress_bar = st.progress(0, text="Training")
+
+    def on_step_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        super().on_step_begin(args, state, control, **kwargs)
+        self.progress_bar.progress(
+            (state.global_step * 100 // state.max_steps) + 1,
+            text=f"Training {state.global_step} / {state.max_steps}",
+        )
+
+
+def train_binary_classifier(
+    project_name: str,
+    entity_name: str,
+    run_name: str,
+    dataset_repo: str = "geekyrakshit/prompt-injection-dataset",
+    model_name: str = "distilbert/distilbert-base-uncased",
+    learning_rate: float = 2e-5,
+    batch_size: int = 16,
+    num_epochs: int = 2,
+    weight_decay: float = 0.01,
+    streamlit_mode: bool = False,
+):
+    wandb.init(project=project_name, entity=entity_name, name=run_name)
+    if streamlit_mode:
+        st.markdown(
+            f"Explore your training logs on [Weights & Biases]({wandb.run.url})"
+        )
+    dataset = load_dataset(dataset_repo)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def preprocess_function(examples):
+        return tokenizer(examples["prompt"], truncation=True)
+
+    tokenized_datasets = dataset.map(preprocess_function, batched=True)
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    accuracy = evaluate.load("accuracy")
+
+    def compute_metrics(eval_pred):
+        predictions, labels = eval_pred
+        predictions = np.argmax(predictions, axis=1)
+        return accuracy.compute(predictions=predictions, references=labels)
+
+    id2label = {0: "SAFE", 1: "INJECTION"}
+    label2id = {"SAFE": 0, "INJECTION": 1}
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name,
+        num_labels=2,
+        id2label=id2label,
+        label2id=label2id,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=TrainingArguments(
+            output_dir="binary-classifier",
+            learning_rate=learning_rate,
+            per_device_train_batch_size=batch_size,
+            per_device_eval_batch_size=batch_size,
+            num_train_epochs=num_epochs,
+            weight_decay=weight_decay,
+            eval_strategy="epoch",
+            save_strategy="epoch",
+            load_best_model_at_end=True,
+            push_to_hub=True,
+            report_to="wandb",
+            logging_strategy="steps",
+            logging_steps=1,
+        ),
+        train_dataset=tokenized_datasets["train"],
+        eval_dataset=tokenized_datasets["test"],
+        processing_class=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+        callbacks=[StreamlitProgressbarCallback()] if streamlit_mode else [],
+    )
+    try:
+        training_output = trainer.train()
+    except Exception as e:
+        wandb.finish()
+        raise e
+    wandb.finish()
+    return training_output

guardrails_genie/utils.py

@@ -1,7 +1,9 @@
 import os
 
+import pandas as pd
 import pymupdf4llm
 import weave
+import weave.trace
 from firerequests import FireRequests
 
 
@@ -11,3 +13,47 @@ def get_markdown_from_pdf_url(url: str) -> str:
     markdown = pymupdf4llm.to_markdown("temp.pdf", show_progress=False)
     os.remove("temp.pdf")
     return markdown
+
+
+class EvaluationCallManager:
+    def __init__(self, entity: str, project: str, call_id: str, max_count: int = 10):
+        self.base_call = weave.init(f"{entity}/{project}").get_call(call_id=call_id)
+        self.max_count = max_count
+        self.show_warning_in_app = False
+        self.call_list = []
+
+    def collect_guardrail_guard_calls_from_eval(self):
+        guard_calls, count = [], 0
+        for eval_predict_and_score_call in self.base_call.children():
+            if "Evaluation.summarize" in eval_predict_and_score_call._op_name:
+                break
+            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(guard_call.inputs["prompt"]),
+                    "outputs": dict(guard_call.output),
+                    "score": dict(score_call.output),
+                }
+            )
+            count += 1
+            if count >= self.max_count:
+                self.show_warning_in_app = True
+                break
+        return guard_calls
+
+    def render_calls_to_streamlit(self):
+        dataframe = {
+            "input_prompt": [
+                call["input_prompt"] for call in self.call_list[0]["calls"]
+            ]
+        }
+        for guardrail_call in self.call_list:
+            dataframe[guardrail_call["guardrail_name"] + ".safe"] = [
+                call["outputs"]["safe"] 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)

pyproject.toml

@@ -12,7 +12,7 @@ dependencies = [
     "ruff>=0.6.9",
     "pip>=24.2",
     "uv>=0.4.20",
-    "git+https://github.com/wandb/weave@feat/eval-progressbar",
+    "weave @ git+https://github.com/wandb/weave@feat/eval-progressbar",
     "streamlit>=1.40.1",
     "python-dotenv>=1.0.1",
     "watchdog>=6.0.0",
@@ -23,4 +23,4 @@ dependencies = [
 ]
 
 [tool.setuptools]
-py-modules = ["guardrails_genie"]
+py-modules = ["guardrails_genie"]