How to return an image from agent tool call and render it on agent-chat-ui?

[fengye87]

I was trying to write an agent that could write and run some Python code in a sandbox (provided by e2b.dev), so that my agent can do react on prompts like "plot and show sinus" and give me back a sinus drawing image.

I could make the whole thing flow, but the thing is I was having trouble to let agent-chat-ui render the returned image. I've tried as this discussion suggested, but it didn't work. Some help please!

Here's my grapy.py (I was using LangGraph local server):

"""LangGraph single-node graph template.

Returns a predefined response. Replace logic and configuration as needed.
"""

from __future__ import annotations

from e2b_code_interpreter import AsyncSandbox
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_core.tools import tool
from langchain_tavily import TavilySearch
from langgraph.prebuilt import create_react_agent

search = TavilySearch()


@tool
async def run_code(code: str):
    """Execute python code in a Jupyter Notebook cell and returns the output."""
    async with await AsyncSandbox.create() as sandbox:
        execution = await sandbox.run_code(code)

        if execution.results:
            for result in execution.results:
                if result.png:
                    return [
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/png;base64,{result.png}"},
                        }
                    ]
        else:
            return {
                "results": execution.results,
                "stdout": execution.logs.stdout,
                "stderr": execution.logs.stderr,
                "error": execution.error,
            }


graph = create_react_agent(ChatTongyi(model_name="qwq-plus"), [search, run_code])