Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,606 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import asyncio
+import json
+import os
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List
+
 import streamlit as st
+from acp_sdk.client import Client
+from acp_sdk.models import Message, MessagePart
+from rich.console import Console
+
+from gaf_guard.clients.stream_adaptors import get_adapter
+from gaf_guard.core.models import WorkflowMessage
+from gaf_guard.toolkit.enums import MessageType, Role, StreamStatus, UserInputType
+from gaf_guard.toolkit.file_utils import resolve_file_paths
+
+
+GAF_GUARD_ROOT = Path(__file__).parent.parent.absolute()
+
+# Apply CSS to hide chat_input when app is running (processing)
+st.markdown(
+    """
+<style>
+.header {
+    padding: 1rem;
+    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
+    color: white;
+    text-align: center;
+    border-radius: 10px;
+    margin-bottom: 1rem;
+}
+.message-card {
+    padding: 1rem;
+    border-left: 4px solid #667eea;
+    background-color: #f8f9fa;
+    border-radius: 5px;
+    margin: 0.5rem 0;
+}
+.stApp[data-teststate=running] .stChatInput textarea,
+.stApp[data-test-script-state=running] .stChatInput textarea {
+    display: none !important;
+}
+.stTextInput {{
+      position: fixed;
+      bottom: 3rem;
+    }}
+.block-container {
+    padding-top: 1rem;
+    padding-bottom: 0rem;
+    padding-left: 5rem;
+    padding-right: 5rem;
+}
+</style>
+""",
+    unsafe_allow_html=True,
+)
+
+# Declare global session variables
+st.session_state.priority = ["low", "medium", "high"]
+st.session_state.initial_risks_master = ["Toxic output", "Hallucination"]
+st.set_page_config(
+    page_title="GAF Guard - A real-time monitoring system for risk assessment and drift monitoring.",
+    layout="wide",  # This sets the app to wide mode
+    # initial_sidebar_state="expanded",
+)
+console = Console(log_time=True)
+run_configs = {
+    "RiskGeneratorAgent": {
+        "risk_questionnaire_cot": os.path.join(
+            GAF_GUARD_ROOT, "chain_of_thought", "risk_questionnaire.json"
+        )
+    },
+    "DriftMonitoringAgent": {
+        "drift_monitoring_cot": os.path.join(
+            GAF_GUARD_ROOT, "chain_of_thought", "drift_monitoring.json"
+        ),
+        "drift_threshold": (
+            st.session_state.drift_threshold
+            if "drift_threshold" in st.session_state
+            else 8
+        ),
+    },
+}
+resolve_file_paths(run_configs)
+
+
+def file_uploaded():
+    st.session_state.prompt_file = st.session_state.prompt_file_uploader.getvalue()
+    message = WorkflowMessage(
+        name="GAF Guard Client",
+        type=MessageType.GAF_GUARD_QUERY,
+        role=Role.SYSTEM,
+        content=f"**File uploaded successfully:** {st.session_state.prompt_file_uploader.name}",
+        accept=UserInputType.INPUT_PROMPT,
+        run_configs=run_configs,
+    )
+    st.session_state.messages.append(message)
+    render(message, simulate=True)
+
+
+def play_button(adapter_type):
+    if st.session_state.setdefault(
+        "stream_adaptor",
+        get_adapter(
+            adapter_type,
+            config={"byte_data": st.session_state.prompt_file},
+        ),
+    ):
+        st.session_state.stream_status = StreamStatus.ACTIVE
+    else:
+        st.write("Selected adaptor is not available.")
+
+
+def pause_button():
+    st.session_state.stream_status = StreamStatus.PAUSED
+    st.session_state.messages.append(
+        WorkflowMessage(
+            name="GAF Guard Client",
+            type=MessageType.GAF_GUARD_QUERY,
+            role=Role.SYSTEM,
+            content=f"**:red[Alert:]** Current input stream is paused. Please click on **Start** to resume.",
+            accept=UserInputType.INPUT_PROMPT,
+            run_configs=run_configs,
+        )
+    )
+
+
+@st.fragment
+def pause_fragment(adapter_type):
+    st.button(
+        "⏸️  Pause",
+        use_container_width=True,
+        disabled=(
+            adapter_type == "Select"
+            or st.session_state.stream_status
+            in [StreamStatus.STOPPED, StreamStatus.PAUSED]
+        ),
+        on_click=pause_button,
+    )
+
+
+def add_sidebar():
+
+    with st.sidebar:
+        st.sidebar.title("⚙️ Settings")
+        if st.session_state.sidebar_display in ["settings_view", "input_prompt_source"]:
+            st.subheader(f":blue[Taxonomy:] {st.session_state.taxonomy}")
+            st.subheader(f":blue[Drift Threshold:] {st.session_state.drift_threshold}")
+        if st.session_state.sidebar_display == "settings_edit":
+            st.session_state.taxonomy = st.selectbox(
+                "Risk Taxonomy",
+                ("IBM Risk Atlas"),
+            )
+            st.session_state.drift_threshold = st.slider(
+                "Drift Threshold",
+                value=st.session_state.drift_threshold,
+                min_value=2,
+                max_value=10,
+                step=1,
+            )
+        if st.session_state.sidebar_display == "input_prompt_source":
+            st.sidebar.title("⚙️ Streaming Source")
+            adapter_type = st.selectbox(
+                "Select Input Prompt Source",
+                ["Select", "JSON"],
+                help="Choose your streaming source",
+                index=0,
+                disabled="stream_adaptor" in st.session_state,
+            )
+            if adapter_type == "JSON":
+                st.subheader("JSON File Source")
+                st.file_uploader(
+                    "OK",
+                    accept_multiple_files=False,
+                    type="json",
+                    label_visibility="collapsed",
+                    on_change=file_uploaded,
+                    key="prompt_file_uploader",
+                    disabled="stream_adaptor" in st.session_state,
+                )
+
+            # Control buttons
+            col1, col2 = st.columns(2)
+            with col1:
+                st.button(
+                    "▶️  Start",
+                    use_container_width=True,
+                    disabled=(
+                        adapter_type == "Select"
+                        or "prompt_file" not in st.session_state
+                        or st.session_state.stream_status == StreamStatus.ACTIVE
+                    ),
+                    on_click=play_button,
+                    args=(adapter_type,),
+                )
+            with col2:
+                pause_fragment(adapter_type)
+            st.markdown(
+                "**Note:** Pause button will temporarily halt the stream after processing the current prompt."
+            )
+
+        st.divider()
+        ai_atlas_button = st.container(
+            horizontal_alignment="center", vertical_alignment="bottom", height="stretch"
+        )
+        ai_atlas_button.markdown(":blue[Powered by:]", text_alignment="center")
+        ai_atlas_button.link_button(
+            "AI Atlas Nexus",
+            "https://github.com/IBM/ai-atlas-nexus",
+            icon=":material/thumb_up:",
+            type="secondary",
+        )
+        if hasattr(st.session_state, "client_session"):
+            ai_atlas_button.markdown(
+                f"Client Id: {str(st.session_state.client_session._session.id)[0:13]} \n :violet-badge[:material/rocket_launch: Connected to :yellow[GAF Guard] Server:] :orange-badge[:material/check: {st.session_state.host}:{st.session_state.port}]",
+                text_alignment="center",
+            )
+        else:
+            ai_atlas_button.markdown(
+                f":red-badge[:material/mimo_disconnect: Client Disconnected]",
+                text_alignment="center",
+            )
+
+
+# render agent reponse from the server
+def render(message: WorkflowMessage, simulate=False):
+
+    def simulate_agent_response(
+        role: Role,
+        message: str,
+        json_data: Dict = None,
+        simulate: bool = False,
+        accept: Dict = None,
+    ):
+        with st.chat_message(role):
+            if simulate:
+                message_placeholder = st.empty()
+                full_response = ""
+                for chunk in message.split():
+                    full_response += chunk + " "
+                    time.sleep(0.05)
+                    message_placeholder.markdown(full_response + "▌")
+                message_placeholder.markdown(full_response)
+            else:
+                st.markdown(message)
+
+            if json_data:
+                st.json(json_data, expanded=4)
+            elif accept == UserInputType.INITIAL_RISKS:
+                st.button(
+                    "Add Initial Risks",
+                    on_click=initial_risks_selector,
+                    disabled=hasattr(st.session_state, "initial_risks"),
+                )
+                st.session_state.disabled_input = False
+            elif accept == UserInputType.INPUT_PROMPT:
+                st.session_state.sidebar_display = "input_prompt_source"
+                st.session_state.disabled_input = True
+
+    if not message.display:
+        return False
+    if message.type == MessageType.GAF_GUARD_WF_STARTED:
+        return False
+    if message.type == MessageType.GAF_GUARD_WF_COMPLETED:
+        return False
+    elif message.type == MessageType.GAF_GUARD_STEP_STARTED:
+        simulate_agent_response(
+            role=message.role.value,
+            message=f"##### :blue[Workflow Step:] **{message.name}**",
+            simulate=simulate,
+            accept=message.accept,
+        )
+    elif message.type == MessageType.GAF_GUARD_STEP_COMPLETED:
+        # simulate_agent_response(
+        #     role=message.role.value,
+        #     message=f"##### :blue[Workflow Step:] **{message.name}** COMPLETED",
+        #     simulate=simulate,
+        #     accept=message.accept,
+        # )
+        return False
+    elif message.type == MessageType.GAF_GUARD_STEP_DATA:
+        if isinstance(message.content, dict):
+            if message.name == "Input Prompt":
+                simulate_agent_response(
+                    role=message.role.value,
+                    message=f"###### :yellow[**Prompt {message.content["prompt_index"]}**]:  {message.content["prompt"]}",
+                    simulate=simulate,
+                    accept=message.accept,
+                )
+            else:
+                if len(message.content.items()) > 2:
+                    data = []
+                    for key, value in message.content.items():
+                        data.append({key.title(): value})
+
+                    simulate_agent_response(
+                        role=message.role.value,
+                        message="###### :yellow[Risk Report]",
+                        json_data=data,
+                        simulate=simulate,
+                        accept=message.accept,
+                    )
+                else:
+                    for key, value in message.content.items():
+                        if key == "identified_risks":
+                            st.session_state.risks = value
+                        if isinstance(value, List) or isinstance(value, Dict):
+                            simulate_agent_response(
+                                role=message.role.value,
+                                message=f"###### :yellow[{key.replace('_', ' ').title()}]",
+                                json_data=value,
+                                simulate=simulate,
+                                accept=message.accept,
+                            )
+                        elif isinstance(value, str) and key.endswith("alert"):
+                            simulate_agent_response(
+                                role=message.role.value,
+                                message=f"###### :yellow[{key.replace('_', ' ').title()}]: :red[{value}]",
+                                simulate=simulate,
+                                accept=message.accept,
+                            )
+                        else:
+                            simulate_agent_response(
+                                role=message.role.value,
+                                message=f"###### :yellow[{key.replace('_', ' ').title()}]: {value}",
+                                simulate=simulate,
+                                accept=message.accept,
+                            )
+    elif message.type == MessageType.GAF_GUARD_QUERY:
+        simulate_agent_response(
+            role=message.role.value,
+            message=f":blue[{message.content}]",
+            simulate=simulate,
+            accept=message.accept,
+        )
+    else:
+        # raise Exception(f"Invalid message type: {message.type}")
+        if message.content:
+            simulate_agent_response(
+                role=message.role.value,
+                message=message.content,
+                simulate=simulate,
+                accept=message.accept,
+            )
+
+    return True
+
+
+@st.dialog("Initial risks", width="medium")
+def initial_risks_selector():
+
+    def add_row():
+        st.session_state.setdefault("initial_risks", {}).update(
+            {
+                str(len(st.session_state.initial_risks)): {
+                    "risk": st.session_state.initial_risks_master[0],
+                    "priority": "low",
+                    "threshold": 0.01,
+                }
+            }
+        )
+
+    if "initial_risks" not in st.session_state:
+        add_row()
+
+    st.button("Add New Row", type="primary", on_click=add_row)
+    with st.form("input_form"):
+
+        # Create columns for the form inputs
+        col1, col2, col3 = st.columns(3)
+
+        for key, initial_risk in st.session_state.initial_risks.items():
+            with col1:
+                value = st.selectbox(
+                    "Risk" if key == "0" else " ",
+                    tuple(st.session_state.initial_risks_master),
+                    key=f"col1{key}",
+                    index=st.session_state.initial_risks_master.index(
+                        initial_risk["risk"]
+                    ),
+                )
+                st.session_state.initial_risks[key].update({"risk": value})
+            with col2:
+                value = st.selectbox(
+                    "Priority" if key == "0" else " ",
+                    tuple(st.session_state.priority),
+                    key=f"col2{key}",
+                    index=st.session_state.priority.index(initial_risk["priority"]),
+                )
+                st.session_state.initial_risks[key].update({"priority": value})
+            with col3:
+                threshold = st.number_input(
+                    "Threshold" if key == "0" else " ",
+                    key=f"col3{key}",
+                    value=initial_risk["threshold"],
+                )
+                st.session_state.initial_risks[key].update({"threshold": threshold})
+
+        submitted = st.form_submit_button("Submit")
+
+    if submitted:
+        st.session_state.user_input = json.dumps(
+            list(st.session_state.initial_risks.values())
+        )
+        st.rerun()
+
+
+@st.dialog(
+    "GAF Guard Connect",
+    width="medium",
+    dismissible=False,
+    icon=":material/login:",
+)
+def connect_screen_dialog():
+    if hasattr(st.session_state, "error"):
+        st.error(st.session_state.error, icon="🚨")
+    with st.form("login_form"):
+        input_host = st.text_input("GAF Guard Host", value="localhost")
+        input_port = st.number_input("GAF Guard Port", value=8000)
+        submitted = st.form_submit_button("Connect", type="primary")
+
+    if submitted:
+        if hasattr(st.session_state, "error"):
+            del st.session_state["error"]
+        st.session_state.host = input_host
+        st.session_state.port = input_port
+        st.rerun()
+
+
+@st.dialog(
+    "GAF Guard Connect",
+    width="medium",
+    dismissible=False,
+    icon=":material/login:",
+)
+def connect():
+
+    async def ping_server(client):
+        await client.ping()
+
+    with st.status(
+        f"Connecting to GAF Guard using host: :blue[**{st.session_state.host}**] and port: :blue[**{st.session_state.port}**]",
+        expanded=True,
+    ) as status:
+        try:
+            client = Client(
+                base_url=f"http://{st.session_state.host}:{st.session_state.port}",
+                verify=True,
+            )
+            # asyncio.run(ping_server(client))
+            st.write("Client created...")
+        except Exception as e:
+            st.session_state.error = "Failed to connect. Check hostname and port."
+            st.rerun()
+
+        st.session_state.client_session = client.session()
+        st.write("Client session created...")
+
+        st.session_state.drift_threshold = 8
+        st.session_state.disabled_input = False
+        st.session_state.stream_status = StreamStatus.STOPPED
+        st.session_state.sidebar_display = "settings_edit"
+        st.session_state.messages = [
+            WorkflowMessage(
+                name="GAF Guard Client",
+                type=MessageType.CLIENT_INPUT,
+                role=Role.USER,
+                accept=UserInputType.USER_INTENT,
+                run_configs=run_configs,
+            )
+        ]
+        st.write("Client initialisation done...")
+
+        # print information in the client console window
+        console.print(
+            f"[[bold white]{datetime.now().strftime('%d-%m-%Y %H:%M:%S')}[/]] [italic bold white] :rocket: Connected to GAF Guard Server at[/italic bold white] [bold white]{st.session_state.host}:{st.session_state.port}[/bold white]"
+        )
+        console.print(
+            f"[[bold white]{datetime.now().strftime('%d-%m-%Y %H:%M:%S')}[/]] Client Id: {st.session_state.client_session._session.id}"
+        )
+        #     console.print(
+        #         f"""
+        #     You can now view your Streamlit app in your browser.
+
+        #     Local URL: http://{st.session_state.host}:{st.session_state.port}
+        # """
+        #     )
+
+        status.update(
+            label=f":material/rocket_launch: Connected to :yellow[**GAF Guard**] Server: :orange-badge[:material/check: {st.session_state.host}:{st.session_state.port}]",
+            state="complete",
+            expanded=True,
+        )
+        time.sleep(1)
+
+    st.rerun()
+
+
+def submit_input():
+    st.session_state.sidebar_display = "settings_view"
+
+
+async def app():
+
+    st.title(f":yellow[GAF Guard]", text_alignment="center")
+    st.subheader(
+        "A real-time monitoring system for risk assessment and drift monitoring",
+        text_alignment="center",
+        divider=True,
+    )
+
+    # add sidebar and related components
+    add_sidebar()
+
+    # Display chat messages from history
+    for message in st.session_state.messages:
+        render(message)
+
+    last_message: WorkflowMessage = st.session_state.messages[-1]
+
+    if st.session_state.stream_status == StreamStatus.ACTIVE:
+        user_input = st.session_state.stream_adaptor.next()
+        if not user_input:
+            del st.session_state["stream_adaptor"]
+            st.session_state.stream_status = StreamStatus.STOPPED
+            st.session_state.messages.append(
+                WorkflowMessage(
+                    name="GAF Guard Client",
+                    type=MessageType.GAF_GUARD_QUERY,
+                    role=Role.SYSTEM,
+                    content=f"**The streaming input has ended. Please choose a streaming source and start again.**",
+                    accept=UserInputType.INPUT_PROMPT,
+                    run_configs=run_configs,
+                )
+            )
+            st.rerun()
+    else:
+        # Accept user input
+        user_input = st.chat_input(
+            placeholder="Enter your response here",
+            key="user_input",
+            disabled=st.session_state.disabled_input,
+            on_submit=submit_input,
+        )
+
+    if not user_input:
+        st.stop()
+    else:
+        COMPLETED = False
+        async for event in st.session_state.client_session.run_stream(
+            agent="orchestrator",
+            input=[
+                Message(
+                    parts=[
+                        MessagePart(
+                            content=WorkflowMessage(
+                                name="GAF Guard Client",
+                                type=(
+                                    MessageType.CLIENT_RESPONSE
+                                    if last_message.type == MessageType.GAF_GUARD_QUERY
+                                    else MessageType.CLIENT_INPUT
+                                ),
+                                role=Role.USER,
+                                content={last_message.accept: user_input},
+                                run_configs=run_configs,
+                            ).model_dump_json(),
+                            content_type="text/plain",
+                        )
+                    ]
+                )
+            ],
+        ):
+            if event.type == "message.part":
+                message = WorkflowMessage(**json.loads(event.part.content))
+                if render(message, simulate=True):
+                    st.session_state.messages.append(message)
+            elif event.type == "run.awaiting":
+                if hasattr(event, "run"):
+                    message = WorkflowMessage(
+                        **json.loads(event.run.await_request.message.parts[0].content)
+                    )
+                    if message.accept == UserInputType.INPUT_PROMPT:
+                        if st.session_state.stream_status == StreamStatus.STOPPED:
+                            render(message, simulate=True)
+                        else:
+                            message.display = False
+                    else:
+                        render(message, simulate=True)
+
+                    st.session_state.messages.append(message)
+                    st.session_state.disabled_input = True
+                    st.rerun()
+
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+if hasattr(st.session_state, "client_session"):
+    asyncio.run(app())
+elif (
+    not hasattr(st.session_state, "error")
+    and hasattr(st.session_state, "host")
+    and hasattr(st.session_state, "port")
+):
+    connect()
+else:
+    connect_screen_dialog()