ToolNode execution output streaming

[Soumadip-Saha]

Hello @JoshuaC215 and @peterkeppert,

First, I'd like to express my appreciation for your excellent project. As I delve into the internals of LangGraph for a personal project, I came across your repository and found it highly insightful.

In line with the ongoing discussion, I have a question regarding the design of a tool intended to execute Python code in a remote environment and stream its standard output via a /stream API endpoint. The goal is to transform the ToolNode into a generator that streams output progressively, ultimately returning the complete result upon code execution completion.

For remote code execution, I'm utilizing the Jupyter Kernel Gateway's WebSocket method. The rationale behind streaming the tool's output is to handle time-consuming processes—such as model training or computationally intensive tasks—without requiring users to wait for the entire execution to finish before receiving updates. This approach aims to provide real-time progress feedback.

Currently, I have a synchronous function that displays the output in a Streamlit UI:

def python_executor_sync(code: str) -> str:
    """Execute python code synchronously and return the output
    
    Useful for when you need to execute any python code,
    and return the output. Works similar to Jupyter Cell execution way,
    where each code is getting executed in new cells.
    
    Args:
        code (str): The python code to execute

    Returns:
        str: The output of the code
    """
    async def async_python_executor(code: str) -> str:
        uri = f"ws://jupyter-kernel-to-execute-code"
        answer = ""
        try:
            async with websockets.connect(uri, additional_headers={"Authorization": "Basic ZmFrZXVzZXI6ZmFrZXBhc3M="}) as ws:
                msg_id = str(ULID())
                execute_msg = {
                    "header": {
                        "username": "",
                        "version": "5.0",
                        "session": "",
                        "msg_id": msg_id,
                        "msg_type": "execute_request",
                    },
                    "parent_header": {},
                    "channel": "shell",
                    "content": {
                        "code": code,
                        "silent": False,
                        "store_history": False,
                        "user_expressions": {},
                        "allow_stdin": False,
                    },
                    "metadata": {},
                    "buffers": {},
                }
                await ws.send(json.dumps(execute_msg))

                while True:
                    try:
                        message = await ws.recv()
                        data = json.loads(message)

                        msg_type = data.get("header", {}).get("msg_type")

                        if msg_type == "stream":
                            answer += data["content"]["text"]
                        elif msg_type == "error":
                            error_content = {
                                "ename": data["content"]["ename"],
                                "evalue": data["content"]["evalue"],
                                "traceback": data["content"]["traceback"],
                            }
                            answer += f"Error: {json.dumps(error_content)}"
                            break
                        elif msg_type == "execute_reply":
                            if data["content"]["status"] == "error":
                                continue
                            break
                        elif msg_type == "status":
                            state = data["content"]["execution_state"]
                            continue
                        elif msg_type == "display_data":
                            answer += str(data["content"])
                        elif msg_type == "execute_result":
                            answer += data['content']['data']['text/plain']

                    except Exception as e:
                        answer += f"Error: {str(e)}"
                        break
            return answer
        except Exception as e:
            answer += f"Error: {str(e)}"
            return answer

    # Use asyncio.run to execute the async function in a synchronous manner
    return asyncio.run(async_python_executor(code))

While this synchronous method works, my objective is to transition to an asynchronous model to facilitate streaming. This transition is essential for handling long-running tasks efficiently.

During my exploration, I observed that the ToolNode triggers on_tool_start and on_tool_end events:

2024-11-28 20:15:45,726 - Parameter value: {'event': 'on_tool_start', 'data': {'input': {'code': 'def fibonacci(n):\n    if n <= 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        a, b = 0, 1\n        for _ in range(2, n + 1):\n            a, b = b, a + b\n        return b\n\nfibonacci(10)'}}, 'name': 'Python_Executor', 'tags': ['seq:step:1'], 'run_id': '965cf078-1366-4d6e-bb05-b7e334fa1688', 'metadata': {'thread_id': '6948cac7-e9e6-4407-b058-f5282c7e94ca', 'model': 'gpt-4o-mini', 'langgraph_step': 2, 'langgraph_node': 'tools', 'langgraph_triggers': ['branch:model:pending_tool_calls:tools'], 'langgraph_path': ('__pregel_pull', 'tools'), 'langgraph_checkpoint_ns': 'tools:2e415f9d-9abd-067f-8301-48fdceb372dd', 'checkpoint_ns': 'tools:2e415f9d-9abd-067f-8301-48fdceb372dd'}, 'parent_ids': ['617ffa28-7dff-48e3-9961-f18a0f82bddb', 'c6454658-3abe-4fca-8c55-c0565bc1c16f']}
2024-11-28 20:15:46,918 - Parameter value: {'event': 'on_tool_end', 'data': {'output': ToolMessage(content='55', name='Python_Executor', tool_call_id='call_2kivIOwrw5u2mxc0O8ho8L7y'), 'input': {'code': 'def fibonacci(n):\n    if n <= 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        a, b = 0, 1\n        for _ in range(2, n + 1):\n            a, b = b, a + b\n        return b\n\nfibonacci(10)'}}, 'run_id': '965cf078-1366-4d6e-bb05-b7e334fa1688', 'name': 'Python_Executor', 'tags': ['seq:step:1'], 'metadata': {'thread_id': '6948cac7-e9e6-4407-b058-f5282c7e94ca', 'model': 'gpt-4o-mini', 'langgraph_step': 2, 'langgraph_node': 'tools', 'langgraph_triggers': ['branch:model:pending_tool_calls:tools'], 'langgraph_path': ('__pregel_pull', 'tools'), 'langgraph_checkpoint_ns': 'tools:2e415f9d-9abd-067f-8301-48fdceb372dd', 'checkpoint_ns': 'tools:2e415f9d-9abd-067f-8301-48fdceb372dd'}, 'parent_ids': ['617ffa28-7dff-48e3-9961-f18a0f82bddb', 'c6454658-3abe-4fca-8c55-c0565bc1c16f']}

To achieve the desired streaming behavior, I am considering introducing an on_tool_stream event. This event would allow the tool to emit intermediate outputs during execution, enabling real-time progress updates.

Additionally, in a previous message within the discussion thread #97, it was mentioned that

Background Tasks don't work well along the default Tool Calls status info solution provided in the streamlit_app.py.

Since I do not plan to use Streamlit as a frontend and intend to develop a custom solution, I am exploring alternative approaches that might be more suitable for my use case.

I would appreciate any insights or recommendations you can provide on implementing this streaming functionality within the ToolNode framework.

Thank you for your time and assistance.

[JoshuaC215]

Hi @Soumadip-Saha

You may want to implement your own version of ToolNode. Within agent-service-toolkit, I think the BG Task might actually be pretty well suited to your use case, or doing something similar with CustomData.adispatch() which enables you to send an arbitrary message blob as a stream chunk in the middle of a node execution.

That's about the extent of insight I can offer although Peter might have other thoughts. Cheers and good luck!

[Soumadip-Saha]

Ok yeah I was also leaning towards that. Thanks for the help, let me try out and post the updates. Might be able to contribute to the repo something if succeed.

[leonardozilli]

@JoshuaC215 thank you for the insight, i'm also trying to tackle a similar problem where I want to show the user the input and output of each intermediate agent in the workflow. The TaskData solution looks great, but how would I go about hiding the streaming data of the intermediate nodes and only show the input and output in the task dropdown?

E.g. the following node prints in the streamlit app both as an AI message and a Task message

async def filter_definitions(self, state: AgentState, config: RunnableConfig) -> Dict[str, Dict]:
    """
    Filters the retrieved definitions based on their relevance to the question.
    """
    await Task("Filter definitions").start(config=config)
    parser = JsonOutputParser(pydantic_object=schema.DefinitionsList)
    prompt = PromptTemplate(
        template="""
        ...
        """,
        input_variables=["context", "definendum"],
        partial_variables={"format_instructions": parser.get_format_instructions()}
    )

    chain = prompt | self.model | parser
    definendum = state["definendum"]
    retrieved_definitions = state["retrieved_definitions"]
    response = chain.invoke({"context": retrieved_definitions, "definendum": definendum})

    await Task("Filter definitions").finish(result="success", data={"output": response['relevant_definitions']}, config=config)
    return {"messages": [utils.json_to_aimessage(response)], "relevant_defs": response['relevant_definitions']}

[JoshuaC215]

Hey @leonardozilli I'm not sure if the Task as provided will meet your use case of sending back the input and output of each agent node to the client.

Two ideas:

1. You could add your logical block in the service for getting and sending node start events, similar to the current handler for node end events. This would probably need some experimentation and debugging through events as they are received to find the right event name and keys to set (that's how I wrote it in the first place), or check the LangGraph documentation in more detail.
2. You could dispatch a custom data message manually during the node start, similar to what Task implementation does. For the end event, it's already streamed by the return value as you observed. Something roughly like this:

async def filter_definitions(self, state: AgentState, config: RunnableConfig) -> Dict[str, Dict]:
    # ...
    start_event = CustomData(
        type="filter_start",
        data={"context": retrieved_definitions, "definendum": definendum}, # or whatever you want to send
    )
    await start_event.adispatch(config)
    # ...
    return [some_output]

In either case you'd need to handle it in the client / app appropriately. A third possible solution would be to emit the intermediate state NOT in the "messages" key, or handle messages differently in the service (handled just below the section of code I linked in (1) ) to get the desired effect.

Hope it helps!