Update: Streamlit modification

app.py

@@ -1,59 +1,67 @@
+import os
+
 import streamlit as st
-from modules.data_class import DataState
-from modules.tools import data_node
-from modules.nodes import chatbot_with_tools, human_node, maybe_exit_human_node, maybe_route_to_tools
+from langchain_core.messages import AIMessage, HumanMessage
 
-from langgraph.graph import StateGraph, START, END
+from modules.graph import invoke_our_graph
+from modules.st_callable_util import get_streamlit_cb  # Utility function to get a Streamlit callback handler with context
 
-from IPython.display import Image, display
-from pprint import pprint
-from typing import Literal
 
-from langgraph.prebuilt import ToolNode
+# Streamlit UI
+st.title("Paintrek Medical Assistant")
+st.markdown("Chat with an AI-powered health assistant.")
 
-from collections.abc import Iterable
-from IPython.display import display, clear_output
-import sys
 
+# Initialize the expander state
+if "expander_open" not in st.session_state:
+    st.session_state.expander_open = True
 
-# Define the LangGraph chatbot
-graph_builder = StateGraph(DataState)
+# Check if the OpenAI API key is set
+if not os.getenv('GOOGLE_API_KEY'):
+    # If not, display a sidebar input for the user to provide the API key
+    st.sidebar.header("GOOGLE_API_KEY Setup")
+    api_key = st.sidebar.text_input(label="API Key", type="password", label_visibility="collapsed")
+    os.environ["GOOGLE_API_KEY"] = api_key
+    # If no key is provided, show an info message and stop further execution and wait till key is entered
+    if not api_key:
+        st.info("Please enter your GOOGLE_API_KEY in the sidebar.")
+        st.stop()
 
-# Add nodes
-graph_builder.add_node("chatbot_healthassistant", chatbot_with_tools)
-graph_builder.add_node("patient", human_node)
-graph_builder.add_node("documenting", data_node)
 
-# Define edges
-graph_builder.add_conditional_edges("chatbot_healthassistant", maybe_route_to_tools)
-graph_builder.add_conditional_edges("patient", maybe_exit_human_node)
-graph_builder.add_edge("documenting", "chatbot_healthassistant")
-graph_builder.add_edge(START, "chatbot_healthassistant")
+# Capture user input from chat input
+prompt = st.chat_input()
 
-# Compile the graph
-graph_with_order_tools = graph_builder.compile()
+# Toggle expander state based on user input
+if prompt is not None:
+    st.session_state.expander_open = False  # Close the expander when the user starts typing
 
-# Streamlit UI
-st.title("LangGraph Chatbot")
-st.markdown("Chat with an AI-powered health assistant.")
+# st write magic
+with st.expander(label="Paintrek Bot", expanded=st.session_state.expander_open):
+    """
+    At any time you can type 'q' or 'quit' to quit.
+    """
 
-# Initialize session state
+# Initialize chat messages in session state
 if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-user_input = st.text_input("You:", key="input")
+    st.session_state["messages"] = [AIMessage(content="Welcome to the Paintrek world. I am a health assistant, an interactive clinical recording system. I will ask you questions about your pain and related symptoms and record your responses.  I will then store this information securely. At any time, you can type `q` to quit.")]
 
-if st.button("Send"):
-    if user_input:
-        # Add user input to history
-        st.session_state.messages.append(("user", user_input))
+# Loop through all messages in the session state and render them as a chat on every st.refresh mech
+for msg in st.session_state.messages:
+    # https://docs.streamlit.io/develop/api-reference/chat/st.chat_message
+    # we store them as AIMessage and HumanMessage as its easier to send to LangGraph
+    if isinstance(msg, AIMessage):
+        st.chat_message("assistant").write(msg.content)
+    elif isinstance(msg, HumanMessage):
+        st.chat_message("user").write(msg.content)
 
-        # Run LangGraph chatbot
-        state = DataState(messages=st.session_state.messages, data={}, finished=False)
-        for output in graph_with_order_tools.stream(state):
-            response = output["messages"][-1]  # Get the last chatbot response
-            st.session_state.messages.append(("ai", response))
+# Handle user input if provided
+if prompt:
+    st.session_state.messages.append(HumanMessage(content=prompt))
+    st.chat_message("user").write(prompt)
 
-# Display chat history
-for sender, message in st.session_state.messages:
-    st.write(f"**{sender}:** {message}")
+    with st.chat_message("assistant"):
+        # create a new placeholder for streaming messages and other events, and give it context
+        st_callback = get_streamlit_cb(st.container())
+        response = invoke_our_graph(st.session_state.messages, [st_callback])
+        st.session_state.messages.append(AIMessage(content=response["messages"][-1].content))   # Add that last message to the st_message_state
+        st.write(response["messages"][-1].content)  # Write the message inside the chat_message context

modules/graph.py

@@ -0,0 +1,26 @@
+from modules.data_class import DataState
+from modules.tools import data_node
+from modules.nodes import chatbot_with_tools, human_node, maybe_exit_human_node, maybe_route_to_tools
+from langgraph.graph import StateGraph, START, END
+
+# Define the LangGraph chatbot
+graph_builder = StateGraph(DataState)
+
+# Add nodes
+graph_builder.add_node("chatbot_healthassistant", chatbot_with_tools)
+graph_builder.add_node("patient", human_node)
+graph_builder.add_node("documenting", data_node)
+
+# Define edges
+graph_builder.add_conditional_edges("chatbot_healthassistant", maybe_route_to_tools)
+graph_builder.add_conditional_edges("patient", maybe_exit_human_node)
+graph_builder.add_edge("documenting", "chatbot_healthassistant")
+graph_builder.add_edge(START, "chatbot_healthassistant")
+
+# Compile the state graph into a runnable object
+graph_with_order_tools = graph_builder.compile()
+
+def invoke_our_graph(st_messages, callables):
+    if not isinstance(callables, list):
+        raise TypeError("callables must be a list")
+    return graph_with_order_tools.invoke({"messages": st_messages}, config={"callbacks": callables})

modules/st_callable_util.py

@@ -0,0 +1,119 @@
+from typing import Callable, TypeVar, Any, Dict, Optional
+import inspect
+
+from streamlit.runtime.scriptrunner import add_script_run_ctx, get_script_run_ctx
+from streamlit.delta_generator import DeltaGenerator
+
+from langchain_core.callbacks.base import BaseCallbackHandler
+import streamlit as st
+
+
+# Define a function to create a callback handler for Streamlit that updates the UI dynamically
+def get_streamlit_cb(parent_container: DeltaGenerator) -> BaseCallbackHandler:
+    """
+    Creates a Streamlit callback handler that updates the provided Streamlit container with new tokens.
+    Args:
+        parent_container (DeltaGenerator): The Streamlit container where the text will be rendered.
+    Returns:
+        BaseCallbackHandler: An instance of a callback handler configured for Streamlit.
+    """
+
+    # Define a custom callback handler class for managing and displaying stream events in Streamlit
+    class StreamHandler(BaseCallbackHandler):
+        """
+        Custom callback handler for Streamlit that updates a Streamlit container with new tokens.
+        """
+
+        def __init__(self, container: st.delta_generator.DeltaGenerator, initial_text: str = ""):
+            """
+            Initializes the StreamHandler with a Streamlit container and optional initial text.
+            Args:
+                container (st.delta_generator.DeltaGenerator): The Streamlit container where text will be rendered.
+                initial_text (str): Optional initial text to start with in the container.
+            """
+            self.container = container  # The Streamlit container to update
+            self.thoughts_placeholder = self.container.container()  # container to hold tool_call renders
+            self.tool_output_placeholder = None # placeholder for the output of the tool call to be in the expander
+            self.token_placeholder = self.container.empty()  # for token streaming
+            self.text = initial_text  # The text content to display, starting with initial text
+
+        def on_llm_new_token(self, token: str, **kwargs) -> None:
+            """
+            Callback method triggered when a new token is received (e.g., from a language model).
+            Args:
+                token (str): The new token received.
+                **kwargs: Additional keyword arguments.
+            """
+            self.text += token  # Append the new token to the existing text
+            self.token_placeholder.write(self.text)
+
+        def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> None:
+            """
+            Run when the tool starts running.
+            Args:
+                serialized (Dict[str, Any]): The serialized tool.
+                input_str (str): The input string.
+                kwargs (Any): Additional keyword arguments.
+            """
+            with self.thoughts_placeholder:
+                status_placeholder = st.empty()   # Placeholder to show the tool's status
+                with status_placeholder.status("Calling Tool...", expanded=True) as s:
+                    st.write("called ", serialized["name"])  # Show which tool is being called
+                    st.write("tool description: ", serialized["description"])
+                    st.write("tool input: ")
+                    st.code(input_str)   # Display the input data sent to the tool
+                    st.write("tool output: ")
+                    # Placeholder for tool output that will be updated later below
+                    self.tool_output_placeholder = st.empty()
+                    s.update(label="Completed Calling Tool!", expanded=False)   # Update the status once done
+
+        def on_tool_end(self, output: Any, **kwargs: Any) -> Any:
+            """
+            Run when the tool ends.
+            Args:
+                output (Any): The output from the tool.
+                kwargs (Any): Additional keyword arguments.
+            """
+            # We assume that `on_tool_end` comes after `on_tool_start`, meaning output_placeholder exists
+            if self.tool_output_placeholder:
+                self.tool_output_placeholder.code(output.content)   # Display the tool's output
+
+    # Define a type variable for generic type hinting in the decorator, to maintain
+    # input function and wrapped function return type
+    fn_return_type = TypeVar('fn_return_type')
+
+    # Decorator function to add the Streamlit execution context to a function
+    def add_streamlit_context(fn: Callable[..., fn_return_type]) -> Callable[..., fn_return_type]:
+        """
+        Decorator to ensure that the decorated function runs within the Streamlit execution context.
+        Args:
+            fn (Callable[..., fn_return_type]): The function to be decorated.
+        Returns:
+            Callable[..., fn_return_type]: The decorated function that includes the Streamlit context setup.
+        """
+        ctx = get_script_run_ctx()  # Retrieve the current Streamlit script execution context
+
+        def wrapper(*args, **kwargs) -> fn_return_type:
+            """
+            Wrapper function that adds the Streamlit context and then calls the original function.
+            Args:
+                *args: Positional arguments to pass to the original function.
+                **kwargs: Keyword arguments to pass to the original function.
+            Returns:
+                fn_return_type: The result from the original function.
+            """
+            add_script_run_ctx(ctx=ctx)  # Add the Streamlit context to the current execution
+            return fn(*args, **kwargs)  # Call the original function with its arguments
+
+        return wrapper
+
+    # Create an instance of the custom StreamHandler with the provided Streamlit container
+    st_cb = StreamHandler(parent_container)
+
+    # Iterate over all methods of the StreamHandler instance
+    for method_name, method_func in inspect.getmembers(st_cb, predicate=inspect.ismethod):
+        if method_name.startswith('on_'):  # Identify callback methods
+            setattr(st_cb, method_name, add_streamlit_context(method_func))  # Wrap and replace the method
+
+    # Return the fully configured StreamHandler instance with the context-aware callback methods
+    return st_cb