feat: add llm streamlit app

app.py

@@ -1,4 +1,50 @@
+import os
+
 import streamlit as st
+from langchain_core.messages import AIMessage, HumanMessage
+
+from graph import invoke_our_graph
+from st_callable_util import get_streamlit_cb  # Utility function to get a Streamlit callback handler with context
+
+st.title("StreamLit 🤝 LangGraph")
+st.markdown("#### Simple Chat Streaming")
+
+# st write magic
+"""
+In this example, we're going to be creating our own [`BaseCallbackHandler`](https://api.python.langchain.com/en/latest/callbacks/langchain_core.callbacks.base.BaseCallbackHandler.html) called StreamHandler 
+to stream our [_LangGraph_](https://langchain-ai.github.io/langgraph/) invocations and leveraging callbacks in our 
+graph's [`RunnableConfig`](https://api.python.langchain.com/en/latest/runnables/langchain_core.runnables.config.RunnableConfig.html).
+
+The BaseCallBackHandler is a [Mixin](https://www.wikiwand.com/en/articles/Mixin) overloader function which we will use
+to implement only `on_llm_new_token`, a method that run on every new generation of a token from the ChatLLM model.
+
+--- 
+"""
+
+# Check if the API key is available as an environment variable
+if "messages" not in st.session_state:
+    # default initial message to render in message state
+    st.session_state["messages"] = [AIMessage(content="How can I help you?")]
+
+# 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 type(msg) == AIMessage:
+        st.chat_message("assistant").write(msg.content)
+    if type(msg) == HumanMessage:
+        st.chat_message("user").write(msg.content)
+
+# takes new input in chat box from user and invokes the graph
+if prompt := st.chat_input():
+    st.session_state.messages.append(HumanMessage(content=prompt))
+    st.chat_message("user").write(prompt)
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+    # Process the AI's response and handles graph events using the callback mechanism
+    with st.chat_message("assistant"):
+        # create a new container for streaming messages only, and give it context
+        st_callback = get_streamlit_cb(st.container())
+        response = invoke_our_graph(st.session_state.messages, [st_callback])
+        # Add that last message to the st_message_state
+        # Streamlit's refresh the message will automatically be visually rendered bc of the msg render for loop above
+        st.session_state.messages.append(AIMessage(content=response["messages"][-1].content))

graph.py

@@ -0,0 +1,40 @@
+import os
+from typing import Annotated, TypedDict
+
+from langgraph.graph import START, END, StateGraph
+from langgraph.graph.message import AnyMessage, add_messages
+from langchain_openai import ChatOpenAI
+
+# This is the default state same as "MessageState" TypedDict but allows us accessibility to custom keys
+class GraphsState(TypedDict):
+    messages: Annotated[list[AnyMessage], add_messages]
+    # Custom keys for additional data can be added here such as - conversation_id: str
+
+graph = StateGraph(GraphsState)
+
+# Core invocation of the model
+def _call_model(state: GraphsState):
+    messages = state["messages"]
+    llm = ChatOpenAI(
+                model=os.environ["LLM_MODEL_ID"],
+                max_retries=2,
+                api_key="None",
+                base_url=os.environ["LLM_API_BASE"],
+            )
+    response = llm.invoke(messages)
+    return {"messages": [response]}# add the response to the messages using LangGraph reducer paradigm
+
+# Define the structure (nodes and directional edges between nodes) of the graph
+graph.add_edge(START, "modelNode")
+graph.add_node("modelNode", _call_model)
+graph.add_edge("modelNode", END)
+
+# Compile the state graph into a runnable object
+graph_runnable = graph.compile()
+
+def invoke_our_graph(st_messages, callables):
+    # Ensure the callables parameter is a list as you can have multiple callbacks
+    if not isinstance(callables, list):
+        raise TypeError("callables must be a list")
+    # Invoke the graph with the current messages and callback configuration
+    return graph_runnable.invoke({"messages": st_messages}, config={"callbacks": callables})

requirements.txt

@@ -1 +1,4 @@
-langchain_core
+langgraph
+streamlit
+langchain-openai
+python-dotenv

st_callable_util.py

@@ -0,0 +1,95 @@
+from typing import Callable, TypeVar
+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
+
+
+# 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 from LangGraph in Streamlit
+    class StreamHandler(BaseCallbackHandler):
+        """
+        Custom callback handler for Streamlit that updates a Streamlit container with new tokens.
+        """
+
+        def __init__(self, container: DeltaGenerator, initial_text: str = ""):
+            """
+            Initializes the StreamHandler with a Streamlit container and optional initial text.
+
+            Args:
+                container (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.token_placeholder = self.container.empty()  # Placeholder for dynamic token updates
+            self.text = initial_text  # Initialize the text content, starting with any 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)  # Update the Streamlit container with the full text
+
+    # Define a type variable for generic type hinting in the decorator, to maintain
+    # the return type of the input function and the wrapped function
+    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.
+        """
+        # Retrieve the current Streamlit script execution context.
+        # This context holds session information necessary for Streamlit operations.
+        ctx = get_script_run_ctx()
+
+        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)  # Set the correct Streamlit context for 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 that respond to LLM events
+            setattr(st_cb, method_name,
+                    add_streamlit_context(method_func))  # Wrap and replace the method with the context-aware version
+
+    # Return the fully configured StreamlitCallbackHandler instance, now context-aware and integrated with any ChatLLM
+    return st_cb