rai_core 2.0.0a0__tar.gz

rai_core-2.0.0a0/PKG-INFO

@@ -0,0 +1,129 @@
+Metadata-Version: 2.3
+Name: rai_core
+Version: 2.0.0a0
+Summary: Core functionality for RAI framework
+Author: Maciej Majek
+Author-email: maciej.majek@robotec.ai
+Requires-Python: >=3.10,<3.13
+Classifier: Development Status :: 4 - Beta
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: coloredlogs (>=15.0.1,<16.0.0)
+Requires-Dist: deprecated (>=1.2.14,<2.0.0)
+Requires-Dist: langchain
+Requires-Dist: langchain-aws
+Requires-Dist: langchain-community
+Requires-Dist: langchain-core (>=0.3,<0.4)
+Requires-Dist: langchain-ollama
+Requires-Dist: langchain-openai
+Requires-Dist: langfuse (>=2.60.2,<3.0.0)
+Requires-Dist: langgraph
+Requires-Dist: langgraph-prebuilt
+Requires-Dist: lark (>=1.1.9,<2.0.0)
+Requires-Dist: numpy (<2.0)
+Requires-Dist: opencv-python (>=4.9.0.80,<5.0.0.0)
+Requires-Dist: pillow (>=11.0.0,<12.0.0)
+Requires-Dist: pydub (>=0.25.1,<0.26.0)
+Requires-Dist: requests (>=2.32.2,<3.0.0)
+Requires-Dist: scipy (>=1.14.0,<2.0.0)
+Requires-Dist: sounddevice (>=0.4.7,<0.5.0)
+Requires-Dist: streamlit (>=1.44,<2.0)
+Requires-Dist: tomli (>=2.0.1,<3.0.0)
+Requires-Dist: tomli-w (>=1.1.0,<2.0.0)
+Requires-Dist: tqdm (>=4.66.4,<5.0.0)
+Requires-Dist: transforms3d (>=0.4.1,<0.5.0)
+Description-Content-Type: text/markdown
+
+# RAI
+
+RAI is a flexible AI agent framework to develop and deploy Embodied AI features for your robots.
+
+📚 Visit [robotecai.github.io/rai](https://robotecai.github.io/rai/) for the latest documentation, setup
+guide and tutorials. 📚
+
+---
+
+<div align="center">
+
+![rai-image](./docs/imgs/RAI_simple_diagram_medium.png)
+
+---
+
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+![GitHub Release](https://img.shields.io/github/v/release/RobotecAI/rai)
+![Contributors](https://img.shields.io/github/contributors/robotecai/rai)
+
+![Static Badge](https://img.shields.io/badge/Ubuntu-24.04-orange)
+![Static Badge](https://img.shields.io/badge/Ubuntu-22.04-orange)
+![Static Badge](https://img.shields.io/badge/Python-3.12-blue)
+![Static Badge](https://img.shields.io/badge/Python-3.10-blue)
+![Static Badge](https://img.shields.io/badge/ROS2-jazzy-blue)
+![Static Badge](https://img.shields.io/badge/ROS2-humble-blue)
+
+[![](https://dcbadge.limes.pink/api/server/https://discord.gg/3PGHgTaJSB)](https://discord.gg/3PGHgTaJSB)
+
+</div>
+
+---
+
+## 🎯 Overview
+
+| Category                       | Description                                       | Features                                                                                                                                                           |
+| ------------------------------ | ------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| 🤖 **Multi-Agent Systems**     | Empowering robotics with advanced AI capabilities | • Seamlessly integrate Gen AI capabilities into your robots<br>• Enable sophisticated agent-based architectures                                                    |
+| 🔄 **Robot Intelligence**      | Enhancing robotic systems with smart features     | • Add natural human-robot interaction capabilities<br>• Bring flexible problem-solving to your existing stack<br>• Provide ready-to-use AI features out of the box |
+| 🌟 **Multi-Modal Interaction** | Supporting diverse interaction capabilities       | • Handle diverse data types natively<br>• Enable rich sensory integration<br>• Process multiple input/output modalities simultaneously                             |
+
+## RAI framework
+
+-   [x] rai core: Core functionality for multi-agent system, human-robot interaction and
+        multi-modalities.
+-   [x] rai whoami: Tool to extract and synthesize robot embodiment information from a structured
+        directory of documentation, images, and URDFs.
+-   [x] rai_asr: Speech-to-text models and tools.
+-   [x] rai_tts: Text-to-speech models and tools.
+-   [x] rai_sim: Package for connecting RAI to simulation environments.
+-   [x] rai_bench: Benchmarking suite for RAI. Test agents, models, tools, simulators, etc.
+-   [x] rai_openset: Openset detection models and tools.
+-   [x] rai_nomad: Integration with NoMaD for navigation.
+-   [ ] rai_finetune: Finetune LLMs on your embodied data.
+
+### Simulation demos
+
+Try RAI yourself with these demos:
+| Application | Robot | Description | Docs Link |
+| ------------------------------------------ | ------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------- |
+| Mission and obstacle reasoning in orchards | Autonomous tractor | In a beautiful scene of a virtual orchard, RAI goes beyond obstacle detection to analyze best course of action for a given unexpected situation. | [link](https://robotecai.github.io/rai/demos/agriculture/) |
+| Manipulation tasks with natural language | Robot Arm (Franka Panda) | Complete flexible manipulation tasks thanks to RAI and Grounded SAM 2 | [link](https://robotecai.github.io/rai/demos/manipulation/) |
+| Autonomous mobile robot demo | Husarion ROSbot XL | Demonstrate RAI's interaction with an autonomous mobile robot platform for navigation and control | [link](https://robotecai.github.io/rai/demos/rosbot_xl/) |
+| Speech-to-speech interaction with autonomous taxi | Simulated car | Demonstrate RAI's speech-to-speech interaction capabilities for specifying destinations to an autonomous taxi in awsim with autoware environment | [link](https://robotecai.github.io/rai/demos/taxi/) |
+
+## Community
+
+### Embodied AI Community Group
+
+RAI is one of the main projects in focus of the
+[Embodied AI Community Group](https://github.com/ros-wg-embodied-ai). If you would like to join the
+next meeting, look for it in the
+[ROS Community Calendar](https://calendar.google.com/calendar/u/0/embed?src=c_3fc5c4d6ece9d80d49f136c1dcd54d7f44e1acefdbe87228c92ff268e85e2ea0@group.calendar.google.com&ctz=Etc/UTC).
+
+### Publicity
+
+-   A talk about [RAI at ROSCon 2024](https://vimeo.com/1026029511).
+
+### RAI Q&A
+
+Please take a look at [Q&A](https://github.com/RobotecAI/rai/discussions/categories/q-a).
+
+### Developer Resources
+
+See our [documentation](https://robotecai.github.io/rai/) for a deeper dive into RAI, including
+instructions on creating a configuration specifically for your robot.
+
+### Contributing
+
+You are welcome to contribute to RAI! Please see our [Contribution Guide](https://github.com/RobotecAI/rai/blob/development/CONTRIBUTING.md).
+

rai_core-2.0.0a0/README.md

@@ -0,0 +1,89 @@
+# RAI
+
+RAI is a flexible AI agent framework to develop and deploy Embodied AI features for your robots.
+
+📚 Visit [robotecai.github.io/rai](https://robotecai.github.io/rai/) for the latest documentation, setup
+guide and tutorials. 📚
+
+---
+
+<div align="center">
+
+![rai-image](./docs/imgs/RAI_simple_diagram_medium.png)
+
+---
+
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+![GitHub Release](https://img.shields.io/github/v/release/RobotecAI/rai)
+![Contributors](https://img.shields.io/github/contributors/robotecai/rai)
+
+![Static Badge](https://img.shields.io/badge/Ubuntu-24.04-orange)
+![Static Badge](https://img.shields.io/badge/Ubuntu-22.04-orange)
+![Static Badge](https://img.shields.io/badge/Python-3.12-blue)
+![Static Badge](https://img.shields.io/badge/Python-3.10-blue)
+![Static Badge](https://img.shields.io/badge/ROS2-jazzy-blue)
+![Static Badge](https://img.shields.io/badge/ROS2-humble-blue)
+
+[![](https://dcbadge.limes.pink/api/server/https://discord.gg/3PGHgTaJSB)](https://discord.gg/3PGHgTaJSB)
+
+</div>
+
+---
+
+## 🎯 Overview
+
+| Category                       | Description                                       | Features                                                                                                                                                           |
+| ------------------------------ | ------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| 🤖 **Multi-Agent Systems**     | Empowering robotics with advanced AI capabilities | • Seamlessly integrate Gen AI capabilities into your robots<br>• Enable sophisticated agent-based architectures                                                    |
+| 🔄 **Robot Intelligence**      | Enhancing robotic systems with smart features     | • Add natural human-robot interaction capabilities<br>• Bring flexible problem-solving to your existing stack<br>• Provide ready-to-use AI features out of the box |
+| 🌟 **Multi-Modal Interaction** | Supporting diverse interaction capabilities       | • Handle diverse data types natively<br>• Enable rich sensory integration<br>• Process multiple input/output modalities simultaneously                             |
+
+## RAI framework
+
+-   [x] rai core: Core functionality for multi-agent system, human-robot interaction and
+        multi-modalities.
+-   [x] rai whoami: Tool to extract and synthesize robot embodiment information from a structured
+        directory of documentation, images, and URDFs.
+-   [x] rai_asr: Speech-to-text models and tools.
+-   [x] rai_tts: Text-to-speech models and tools.
+-   [x] rai_sim: Package for connecting RAI to simulation environments.
+-   [x] rai_bench: Benchmarking suite for RAI. Test agents, models, tools, simulators, etc.
+-   [x] rai_openset: Openset detection models and tools.
+-   [x] rai_nomad: Integration with NoMaD for navigation.
+-   [ ] rai_finetune: Finetune LLMs on your embodied data.
+
+### Simulation demos
+
+Try RAI yourself with these demos:
+| Application | Robot | Description | Docs Link |
+| ------------------------------------------ | ------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------- |
+| Mission and obstacle reasoning in orchards | Autonomous tractor | In a beautiful scene of a virtual orchard, RAI goes beyond obstacle detection to analyze best course of action for a given unexpected situation. | [link](https://robotecai.github.io/rai/demos/agriculture/) |
+| Manipulation tasks with natural language | Robot Arm (Franka Panda) | Complete flexible manipulation tasks thanks to RAI and Grounded SAM 2 | [link](https://robotecai.github.io/rai/demos/manipulation/) |
+| Autonomous mobile robot demo | Husarion ROSbot XL | Demonstrate RAI's interaction with an autonomous mobile robot platform for navigation and control | [link](https://robotecai.github.io/rai/demos/rosbot_xl/) |
+| Speech-to-speech interaction with autonomous taxi | Simulated car | Demonstrate RAI's speech-to-speech interaction capabilities for specifying destinations to an autonomous taxi in awsim with autoware environment | [link](https://robotecai.github.io/rai/demos/taxi/) |
+
+## Community
+
+### Embodied AI Community Group
+
+RAI is one of the main projects in focus of the
+[Embodied AI Community Group](https://github.com/ros-wg-embodied-ai). If you would like to join the
+next meeting, look for it in the
+[ROS Community Calendar](https://calendar.google.com/calendar/u/0/embed?src=c_3fc5c4d6ece9d80d49f136c1dcd54d7f44e1acefdbe87228c92ff268e85e2ea0@group.calendar.google.com&ctz=Etc/UTC).
+
+### Publicity
+
+-   A talk about [RAI at ROSCon 2024](https://vimeo.com/1026029511).
+
+### RAI Q&A
+
+Please take a look at [Q&A](https://github.com/RobotecAI/rai/discussions/categories/q-a).
+
+### Developer Resources
+
+See our [documentation](https://robotecai.github.io/rai/) for a deeper dive into RAI, including
+instructions on creating a configuration specifically for your robot.
+
+### Contributing
+
+You are welcome to contribute to RAI! Please see our [Contribution Guide](https://github.com/RobotecAI/rai/blob/development/CONTRIBUTING.md).

rai_core-2.0.0a0/pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "rai_core"
+version = "2.0.0.a0"
+description = "Core functionality for RAI framework"
+authors = ["Maciej Majek <maciej.majek@robotec.ai>", "Bartłomiej Boczek <bartlomiej.boczek@robotec.ai>", "Kajetan Rachwał <kajetan.rachwal@robotec.ai>"]
+readme = "README.md"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Development Status :: 4 - Beta",
+    "License :: OSI Approved :: Apache Software License",
+]
+packages = [
+    { include = "rai", from = "." },
+]
+
+[tool.poetry.dependencies]
+python = "^3.10, <3.13"
+langchain-core = "^0.3"
+langgraph = "*"
+langgraph-prebuilt = "*"
+langchain = "*"
+langchain-aws = "*"
+langchain-openai = "*"
+langchain-ollama = "*"
+langchain-community = "*"
+
+requests = "^2.32.2"
+coloredlogs = "^15.0.1"
+tqdm = "^4.66.4"
+deprecated = "^1.2.14"
+tomli = "^2.0.1"
+tomli-w = "^1.1.0"
+opencv-python = "^4.9.0.80"
+lark = "^1.1.9"
+transforms3d = "^0.4.1"
+pillow = "^11.0.0"
+langfuse = "^2.60.2"
+pydub = "^0.25.1"
+scipy = "^1.14.0"
+sounddevice = "^0.4.7"
+streamlit = "^1.44"
+numpy = "<2.0"

rai_core-2.0.0a0/rai/__init__.py

@@ -0,0 +1,32 @@
+# Copyright (C) 2025 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .agents import AgentRunner, ReActAgent
+from .initialization import (
+    get_embeddings_model,
+    get_llm_model,
+    get_llm_model_config_and_vendor,
+    get_llm_model_direct,
+    get_tracing_callbacks,
+)
+
+__all__ = [
+    "AgentRunner",
+    "ReActAgent",
+    "get_embeddings_model",
+    "get_llm_model",
+    "get_llm_model_config_and_vendor",
+    "get_llm_model_direct",
+    "get_tracing_callbacks",
+]

rai_core-2.0.0a0/rai/agents/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (C) 2024 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from rai.agents.base import BaseAgent
+from rai.agents.conversational_agent import create_conversational_agent
+from rai.agents.langchain.react_agent import ReActAgent
+from rai.agents.runner import AgentRunner, wait_for_shutdown
+from rai.agents.tool_runner import ToolRunner
+
+__all__ = [
+    "AgentRunner",
+    "BaseAgent",
+    "ReActAgent",
+    "ToolRunner",
+    "create_conversational_agent",
+    "wait_for_shutdown",
+]

rai_core-2.0.0a0/rai/agents/base.py

@@ -0,0 +1,34 @@
+# Copyright (C) 2024 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from abc import ABC, abstractmethod
+
+
+class BaseAgent(ABC):
+    def __init__(self):
+        """Initializes a new agent instance and sets up logging with the class name."""
+        self.logger = logging.getLogger(self.__class__.__name__)
+
+    @abstractmethod
+    def run(self):
+        """Starts the agent's main execution loop.
+        In some cases, concrete run implementation may not be needed.
+        In that case use pass as a placeholder."""
+        pass
+
+    @abstractmethod
+    def stop(self):
+        """Gracefully terminates the agent's execution and cleans up resources."""
+        pass

rai_core-2.0.0a0/rai/agents/conversational_agent.py

@@ -0,0 +1,81 @@
+# Copyright (C) 2024 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+from functools import partial
+from typing import List, Optional, TypedDict
+
+from langchain.chat_models.base import BaseChatModel
+from langchain_core.messages import BaseMessage, SystemMessage
+from langchain_core.tools import BaseTool
+from langgraph.graph import START, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.prebuilt.tool_node import tools_condition
+
+from rai.agents.tool_runner import ToolRunner
+
+
+class State(TypedDict):
+    messages: List[BaseMessage]
+
+
+def agent(llm: BaseChatModel, logger: logging.Logger, system_prompt: str, state: State):
+    logger.info("Running thinker")
+
+    # If there are no messages, do nothing
+    if len(state["messages"]) == 0:
+        return state
+
+    # Insert system message if not already present
+    if not isinstance(state["messages"][0], SystemMessage):
+        state["messages"].insert(0, SystemMessage(content=system_prompt))
+    ai_msg = llm.invoke(state["messages"])
+    state["messages"].append(ai_msg)
+    return state
+
+
+def create_conversational_agent(
+    llm: BaseChatModel,
+    tools: List[BaseTool],
+    system_prompt: str,
+    logger: Optional[logging.Logger] = None,
+    debug=False,
+) -> CompiledStateGraph:
+    _logger = None
+    if logger:
+        _logger = logger
+    else:
+        _logger = logging.getLogger(__name__)
+
+    _logger.info("Creating state based agent")
+
+    llm_with_tools = llm.bind_tools(tools)
+    tool_node = ToolRunner(tools=tools, logger=_logger)
+
+    workflow = StateGraph(State)
+    workflow.add_node("tools", tool_node)
+    workflow.add_node("thinker", partial(agent, llm_with_tools, _logger, system_prompt))
+
+    workflow.add_edge(START, "thinker")
+    workflow.add_edge("tools", "thinker")
+
+    workflow.add_conditional_edges(
+        "thinker",
+        tools_condition,
+    )
+
+    app = workflow.compile(debug=debug)
+    _logger.info("State based agent created")
+    return app

rai_core-2.0.0a0/rai/agents/integrations/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (C) 2024 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

rai_core-2.0.0a0/rai/agents/integrations/streamlit.py

@@ -0,0 +1,183 @@
+# Copyright (C) 2024 Robotec.AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import base64
+import inspect
+from typing import Any, Callable, Dict, List, TypeVar
+
+import cv2
+import numpy as np
+import streamlit as st
+from langchain_core.callbacks.base import BaseCallbackHandler
+from langchain_core.messages import BaseMessage
+from langchain_core.runnables import Runnable, RunnableConfig
+from streamlit.delta_generator import DeltaGenerator
+from streamlit.runtime.scriptrunner import add_script_run_ctx, get_script_run_ctx
+
+# code inspired by (mostly copied, some changes were applied and might be updated in the future)
+# https://github.com/shiv248/Streamlit-x-LangGraph-Cookbooks/tree/b8e623bdc9821fc1cf581607454dae1afc054df2/tool_calling_via_callback
+
+
+# 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 tool: ", serialized["name"])
+                    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:
+                with self.tool_output_placeholder.container():
+                    st.code(output.content)  # Display the tool's output
+
+                    if output.artifact is not None:
+                        if "images" in output.artifact:  # Display available images
+                            for image in output.artifact["images"]:
+                                image_cv2 = cv2.imdecode(
+                                    np.frombuffer(base64.b64decode(image), np.uint8),
+                                    cv2.IMREAD_COLOR,
+                                )
+                                st.image(
+                                    image_cv2, channels="BGR", use_container_width=True
+                                )
+
+    # 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
+
+
+def streamlit_invoke(
+    graph: Runnable[Any, Any], messages: List[BaseMessage], callables: List[Callable]
+):
+    if not isinstance(callables, list):
+        raise TypeError("callables must be a list")
+    return graph.invoke(
+        {"messages": messages},
+        config=RunnableConfig({"callbacks": callables, "recursion_limit": 100}),
+    )