add: wandb api key based authentication

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 
 intro_page = st.Page(
-    "application_pages/intro_page.py", title="Introduction", icon=":material/guardian:"
+    "application_pages/intro_page.py", title="Authenticate", icon=":material/guardian:"
 )
 chat_page = st.Page(
     "application_pages/chat_app.py",

application_pages/chat_app.py

@@ -1,8 +1,6 @@
 import importlib
-import os
 
 import streamlit as st
-import weave
 from dotenv import load_dotenv
 
 from guardrails_genie.guardrails import GuardrailManager
@@ -11,9 +9,6 @@ from guardrails_genie.llm import OpenAIModel
 
 def initialize_session_state():
     load_dotenv()
-    weave.init(project_name=os.getenv("WEAVE_PROJECT"))
-    if "weave_project_name" not in st.session_state:
-        st.session_state.weave_project_name = "guardrails-genie"
     if "guardrails" not in st.session_state:
         st.session_state.guardrails = []
     if "guardrail_names" not in st.session_state:
@@ -121,70 +116,69 @@ def initialize_guardrails():
     )
 
 
-initialize_session_state()
-st.title(":material/robot: Guardrails Genie Playground")
-
-weave_project_name = st.sidebar.text_input(
-    "Weave project name", value=st.session_state.weave_project_name
-)
-st.session_state.weave_project_name = weave_project_name
-if st.session_state.weave_project_name != "":
-    weave.init(project_name=st.session_state.weave_project_name)
-
-openai_model = st.sidebar.selectbox(
-    "OpenAI LLM for Chat", ["", "gpt-4o-mini", "gpt-4o"]
-)
-chat_condition = openai_model != ""
-
-guardrails = []
-
-guardrail_names = st.sidebar.multiselect(
-    label="Select Guardrails",
-    options=[
-        cls_name
-        for cls_name, cls_obj in vars(
-            importlib.import_module("guardrails_genie.guardrails")
-        ).items()
-        if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
-    ],
-)
-st.session_state.guardrail_names = guardrail_names
-
-if st.sidebar.button("Initialize Guardrails") and chat_condition:
-    st.session_state.initialize_guardrails = True
-
-if st.session_state.initialize_guardrails:
-    with st.sidebar.status("Initializing Guardrails..."):
-        initialize_guardrails()
-        st.session_state.llm_model = OpenAIModel(model_name=openai_model)
-
-    user_prompt = st.text_area("User Prompt", value="")
-    st.session_state.user_prompt = user_prompt
-
-    test_guardrails_button = st.button("Test Guardrails")
-    st.session_state.test_guardrails = test_guardrails_button
-
-    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 st.session_state.is_authenticated:
+    initialize_session_state()
+    st.title(":material/robot: Guardrails Genie Playground")
 
-        if guardrails_response["safe"]:
-            st.markdown(
-                f"\n\n---\nPrompt is safe! Explore guardrail trace on [Weave]({call.ui_url})\n\n---\n"
-            )
+    openai_model = st.sidebar.selectbox(
+        "OpenAI LLM for Chat", ["", "gpt-4o-mini", "gpt-4o"]
+    )
+    chat_condition = openai_model != ""
+
+    guardrails = []
+
+    guardrail_names = st.sidebar.multiselect(
+        label="Select Guardrails",
+        options=[
+            cls_name
+            for cls_name, cls_obj in vars(
+                importlib.import_module("guardrails_genie.guardrails")
+            ).items()
+            if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
+        ],
+    )
+    st.session_state.guardrail_names = guardrail_names
+
+    if st.sidebar.button("Initialize Guardrails") and chat_condition:
+        st.session_state.initialize_guardrails = True
+
+    if st.session_state.initialize_guardrails:
+        with st.sidebar.status("Initializing Guardrails..."):
+            initialize_guardrails()
+            st.session_state.llm_model = OpenAIModel(model_name=openai_model)
+
+        user_prompt = st.text_area("User Prompt", value="")
+        st.session_state.user_prompt = user_prompt
+
+        test_guardrails_button = st.button("Test Guardrails")
+        st.session_state.test_guardrails = test_guardrails_button
 
-            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,
+        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,
+                    )
                 )
-            st.markdown(
-                response.choices[0].message.content
-                + f"\n\n---\nExplore LLM generation trace on [Weave]({call.ui_url})"
-            )
-        else:
-            st.warning("Prompt is not safe!")
-            st.markdown(guardrails_response["summary"])
-            st.markdown(f"Explore prompt trace on [Weave]({call.ui_url})")
+
+            if guardrails_response["safe"]:
+                st.markdown(
+                    f"\n\n---\nPrompt is safe! Explore guardrail trace on [Weave]({call.ui_url})\n\n---\n"
+                )
+
+                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.warning("Prompt is not safe!")
+                st.markdown(guardrails_response["summary"])
+                st.markdown(f"Explore prompt trace on [Weave]({call.ui_url})")
+else:
+    st.warning("Please authenticate your WandB account to use this feature.")

application_pages/evaluation_app.py

@@ -13,8 +13,6 @@ from guardrails_genie.metrics import AccuracyMetric
 
 def initialize_session_state():
     load_dotenv()
-    if "weave_project_name" not in st.session_state:
-        st.session_state.weave_project_name = "guardrails-genie"
     if "uploaded_file" not in st.session_state:
         st.session_state.uploaded_file = None
     if "dataset_name" not in st.session_state:
@@ -124,108 +122,106 @@ def initialize_guardrails():
     )
 
 
-initialize_session_state()
-st.title(":material/monitoring: Evaluation")
-
-weave_project_name = st.sidebar.text_input(
-    "Weave project name", value=st.session_state.weave_project_name
-)
-st.session_state.weave_project_name = weave_project_name
-if st.session_state.weave_project_name != "":
-    weave.init(project_name=st.session_state.weave_project_name)
-
-uploaded_file = st.sidebar.file_uploader(
-    "Upload the evaluation dataset as a CSV file", type="csv"
-)
-st.session_state.uploaded_file = uploaded_file
-
-if st.session_state.uploaded_file is not None:
-    dataset_name = st.sidebar.text_input("Evaluation dataset name", value=None)
-    st.session_state.dataset_name = dataset_name
-    preview_in_app = st.sidebar.toggle("Preview in app", value=False)
-    st.session_state.preview_in_app = preview_in_app
-    publish_dataset_button = st.sidebar.button("Publish dataset")
-    st.session_state.publish_dataset_button = publish_dataset_button
-
-    if st.session_state.publish_dataset_button and (
-        st.session_state.dataset_name is not None
-        and st.session_state.dataset_name != ""
-    ):
-
-        with st.expander("Evaluation Dataset Preview", expanded=True):
-            dataframe = pd.read_csv(st.session_state.uploaded_file)
-            data_list = dataframe.to_dict(orient="records")
-
-            dataset = weave.Dataset(name=st.session_state.dataset_name, rows=data_list)
-            st.session_state.dataset_ref = weave.publish(dataset)
-
-            entity = st.session_state.dataset_ref.entity
-            project = st.session_state.dataset_ref.project
-            dataset_name = st.session_state.dataset_name
-            digest = st.session_state.dataset_ref._digest
-            dataset_url = f"https://wandb.ai/{entity}/{project}/weave/objects/{dataset_name}/versions/{digest}"
-            st.markdown(f"Dataset published to [**Weave**]({dataset_url})")
-
-            if preview_in_app:
-                st.dataframe(dataframe.head(20))
-                if len(dataframe) > 20:
-                    st.markdown(
-                        f"⚠️ Dataset is too large to preview in app, please explore in the [**Weave UI**]({dataset_url})"
-                    )
+if st.session_state.is_authenticated:
+    initialize_session_state()
+    st.title(":material/monitoring: Evaluation")
 
-        st.session_state.is_dataset_published = True
-
-    if st.session_state.is_dataset_published:
-        guardrail_names = st.sidebar.multiselect(
-            "Select Guardrails",
-            options=[
-                cls_name
-                for cls_name, cls_obj in vars(
-                    import_module("guardrails_genie.guardrails")
-                ).items()
-                if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
-            ],
-        )
-        st.session_state.guardrail_names = guardrail_names
-
-        initialize_guardrails()
-        evaluation_name = st.sidebar.text_input("Evaluation Name", value="")
-        st.session_state.evaluation_name = evaluation_name
-
-        start_evaluations_button = st.sidebar.button("Start Evaluations")
-        st.session_state.start_evaluations_button = start_evaluations_button
-        if st.session_state.start_evaluations_button:
-            # st.write(len(st.session_state.guardrails))
-            evaluation = weave.Evaluation(
-                dataset=st.session_state.dataset_ref,
-                scorers=[AccuracyMetric()],
-                streamlit_mode=True,
-            )
-            with st.expander("Evaluation Results", expanded=True):
-                evaluation_summary, call = asyncio.run(
-                    evaluation.evaluate.call(
-                        evaluation,
-                        GuardrailManager(guardrails=st.session_state.guardrails),
-                        __weave={
-                            "display_name": (
-                                "Evaluation.evaluate"
-                                if st.session_state.evaluation_name == ""
-                                else "Evaluation.evaluate:"
-                                + st.session_state.evaluation_name
-                            )
-                        },
-                    )
-                )
-                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",
+    uploaded_file = st.sidebar.file_uploader(
+        "Upload the evaluation dataset as a CSV file", type="csv"
+    )
+    st.session_state.uploaded_file = uploaded_file
+
+    if st.session_state.uploaded_file is not None:
+        dataset_name = st.sidebar.text_input("Evaluation dataset name", value=None)
+        st.session_state.dataset_name = dataset_name
+        preview_in_app = st.sidebar.toggle("Preview in app", value=False)
+        st.session_state.preview_in_app = preview_in_app
+        publish_dataset_button = st.sidebar.button("Publish dataset")
+        st.session_state.publish_dataset_button = publish_dataset_button
+
+        if st.session_state.publish_dataset_button and (
+            st.session_state.dataset_name is not None
+            and st.session_state.dataset_name != ""
+        ):
+
+            with st.expander("Evaluation Dataset Preview", expanded=True):
+                dataframe = pd.read_csv(st.session_state.uploaded_file)
+                data_list = dataframe.to_dict(orient="records")
+
+                dataset = weave.Dataset(
+                    name=st.session_state.dataset_name, rows=data_list
                 )
-                st.markdown(
-                    f"Explore the entire evaluation trace table in [Weave]({call.ui_url})"
+                st.session_state.dataset_ref = weave.publish(dataset)
+
+                entity = st.session_state.dataset_ref.entity
+                project = st.session_state.dataset_ref.project
+                dataset_name = st.session_state.dataset_name
+                digest = st.session_state.dataset_ref._digest
+                dataset_url = f"https://wandb.ai/{entity}/{project}/weave/objects/{dataset_name}/versions/{digest}"
+                st.markdown(f"Dataset published to [**Weave**]({dataset_url})")
+
+                if preview_in_app:
+                    st.dataframe(dataframe.head(20))
+                    if len(dataframe) > 20:
+                        st.markdown(
+                            f"⚠️ Dataset is too large to preview in app, please explore in the [**Weave UI**]({dataset_url})"
+                        )
+
+            st.session_state.is_dataset_published = True
+
+        if st.session_state.is_dataset_published:
+            guardrail_names = st.sidebar.multiselect(
+                "Select Guardrails",
+                options=[
+                    cls_name
+                    for cls_name, cls_obj in vars(
+                        import_module("guardrails_genie.guardrails")
+                    ).items()
+                    if isinstance(cls_obj, type) and cls_name != "GuardrailManager"
+                ],
+            )
+            st.session_state.guardrail_names = guardrail_names
+
+            initialize_guardrails()
+            evaluation_name = st.sidebar.text_input("Evaluation Name", value="")
+            st.session_state.evaluation_name = evaluation_name
+
+            start_evaluations_button = st.sidebar.button("Start Evaluations")
+            st.session_state.start_evaluations_button = start_evaluations_button
+            if st.session_state.start_evaluations_button:
+                # st.write(len(st.session_state.guardrails))
+                evaluation = weave.Evaluation(
+                    dataset=st.session_state.dataset_ref,
+                    scorers=[AccuracyMetric()],
+                    streamlit_mode=True,
                 )
+                with st.expander("Evaluation Results", expanded=True):
+                    evaluation_summary, call = asyncio.run(
+                        evaluation.evaluate.call(
+                            evaluation,
+                            GuardrailManager(guardrails=st.session_state.guardrails),
+                            __weave={
+                                "display_name": (
+                                    "Evaluation.evaluate"
+                                    if st.session_state.evaluation_name == ""
+                                    else "Evaluation.evaluate:"
+                                    + st.session_state.evaluation_name
+                                )
+                            },
+                        )
+                    )
+                    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.markdown(
+                        f"Explore the entire evaluation trace table in [Weave]({call.ui_url})"
+                    )
+else:
+    st.warning("Please authenticate your WandB account to use this feature.")

application_pages/intro_page.py

@@ -1,7 +1,54 @@
 import streamlit as st
 
+import wandb
+
+
+def initialize_session_state():
+    if "weave_project_name" not in st.session_state:
+        st.session_state.weave_project_name = "guardrails-genie"
+    if "weave_entity_name" not in st.session_state:
+        st.session_state.weave_entity_name = ""
+    if "wandb_api_key" not in st.session_state:
+        st.session_state.wandb_api_key = ""
+    if "authenticate_button" not in st.session_state:
+        st.session_state.authenticate_button = False
+    if "is_authenticated" not in st.session_state:
+        st.session_state.is_authenticated = False
+
+
+initialize_session_state()
 st.title("🧞‍♂️ Guardrails Genie")
 
 st.write(
     "Guardrails-Genie is a tool that helps you implement guardrails in your LLM applications."
 )
+
+with st.expander("Login to Your WandB Account", expanded=True):
+    st.markdown(
+        "Get your Wandb API key from [https://wandb.ai/authorize](https://wandb.ai/authorize)"
+    )
+    weave_entity_name = st.text_input(
+        "Weave Entity Name", value=st.session_state.weave_entity_name
+    )
+    st.session_state.weave_entity_name = weave_entity_name
+    weave_project_name = st.text_input(
+        "Weave Project Name", value=st.session_state.weave_project_name
+    )
+    st.session_state.weave_project_name = weave_project_name
+    wandb_api_key = st.text_input("Wandb API Key", value="", type="password")
+    st.session_state.wandb_api_key = wandb_api_key
+    authenticate_button = st.button("Authenticate")
+    st.session_state.authenticate_button = authenticate_button
+
+    if authenticate_button and (
+        st.session_state.wandb_api_key != ""
+        and st.session_state.weave_project_name != ""
+    ):
+        is_wandb_logged_in = wandb.login(
+            key=st.session_state.wandb_api_key, relogin=True
+        )
+        if is_wandb_logged_in:
+            st.session_state.is_authenticated = True
+            st.success("Logged in to Wandb")
+        else:
+            st.error("Failed to log in to Wandb")