add: limited eval table rendering in app

application_pages/evaluation_app.py

@@ -1,4 +1,6 @@
 import asyncio
+import os
+import time
 from importlib import import_module
 
 import pandas as pd
@@ -9,12 +11,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 +36,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():
@@ -107,6 +120,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 +134,54 @@ 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(
+                                        call=call
+                                    ),
+                                }
+                            )
+                        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

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,44 @@ 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, call):
+        guard_calls, count = [], 0
+        for eval_predict_call in call.children():
+            if "Evaluation.summarize" in eval_predict_call._op_name:
+                break
+            required_call = eval_predict_call.children()[0].children()[0].children()[0]
+            guard_calls.append(
+                {
+                    "input_prompt": str(required_call.inputs["prompt"]),
+                    "outputs": dict(required_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"] + ".summary"] = [
+                call["outputs"]["summary"] for call in guardrail_call["calls"]
+            ]
+        return pd.DataFrame(dataframe)