Returning plots for llama index agent tools

[tituslhy]

Hello!

I'm designing a tool for an agent to use to generate plotly images. I'm able to make this work in Jupyter Notebook but I'm not sure about how to get the agent to return just the plotly graph figure to my frontend which I'm using chainlit to do.

Code snippet without chainlit:

## utils.py
from llama_index.core.agent import FunctionCallingAgentWorker, AgentRunner
from llama_index.core.tools.tool_spec.base import BaseToolSpec
import pandas as pd
pd.options.plotting.backend = "plotly"

class DBTools(BaseToolSpec):
    spec_functions = [
        "plot_data",
        ...
    ]
    
    def __init__(
        self,  
        db_path: Optional[str] = db_path 
        ):
        """Class constructor. Connects to sqlite database"""
        self.conn = sqlite3.connect(db_path, check_same_thread=False)

    def plot_data(self) ->str:
        df = self.get_df_function() ##dummy function to get dataframe
        fig = df.plot.line(
                title=title,
                labels=labels
            )
        fig.show()
        return "Plotted"
  
   ### Other methods ###

agent_worker = FunctionCallingAgentWorker.from_tools(
    tools = db_tool_list, llm = Settings.llm, verbose = True
)
agent = AgentRunner(agent_worker = agent_worker)
query = "Generate a plot of the general trend of Illumina, Apple and NVIDIA's closing prices from 2020 till date"
response = agent.chat(query)

This code snippet runs neatly in a notebook with the figure showing. When I tried leaving everything as is with chainlit, My plotly diagram opens in a new tab, and the llm replies "Here is the chart" with a blank image icon. When I returned fig instead of plotted there is an error because the llm doesn't actually know how to work with the graph object, and is still trying to reply in natural text.

Code snippet with chainlit

import chainlit as cl
from utils import agent

@cl.on_chat_start
async def on_chat_start():
      cl.user_session.set("agent", agent)
      await cl.Message("Hi!").send()

@cl.on_message
async def on_message(msg: cl.Message):
    agent = cl.user_session.get("agent")
    response = await cl.make_async(agent.chat)(message.content)
    await cl.Message(content=str(response)).send()

I was wondering is there a way to just send the figure as an output from the agent's tool to the frontend? Otherwise I'll have to create tmp folders to save images and render them as static images.

[dosubot[bot]]

To return Plotly figures from an agent tool to the frontend using Chainlit, you can extend the FunctionCallingAgentWorker or AgentRunner classes in LlamaIndex to handle and return Plotly figures. Here’s how you can do it:

1. Extend TaskStepOutput to include Plotly figures:

from typing import Union
import plotly.graph_objects as go
from llama_index.core.agent.runner.base import TaskStepOutput

class CustomTaskStepOutput(TaskStepOutput):
    output: Union[str, go.Figure]

2. Modify the run_step method to handle Plotly figures:

from llama_index.core.agent.runner.base import AgentRunner

class CustomAgentRunner(AgentRunner):
    def run_step(
        self,
        task_id: str,
        input: Optional[str] = None,
        step: Optional[TaskStep] = None,
        **kwargs: Any,
    ) -> CustomTaskStepOutput:
        """Run step."""
        step = validate_step_from_args(task_id, input, step, **kwargs)
        # Custom logic to handle Plotly figures
        if isinstance(step.output, go.Figure):
            return CustomTaskStepOutput(output=step.output)
        else:
            return super().run_step(task_id, input, step, **kwargs)

3. Modify the finalize_response method to handle Plotly figures:

class CustomAgentRunner(AgentRunner):
    def finalize_response(
        self,
        task_id: str,
        step_output: Optional[CustomTaskStepOutput] = None,
    ) -> AGENT_CHAT_RESPONSE_TYPE:
        """Finalize response."""
        if step_output is None:
            step_output = self.state.get_completed_steps(task_id)[-1]
        if not step_output.is_last:
            raise ValueError(
                "finalize_response can only be called on the last step output"
            )

        if isinstance(step_output.output, go.Figure):
            # Handle Plotly figure
            return step_output.output.to_json()
        else:
            return super().finalize_response(task_id, step_output)

4. Integrate with Chainlit:

Modify your Chainlit code to handle the JSON representation of the Plotly figure:

import chainlit as cl
from utils import agent
import plotly.graph_objects as go
import json

@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("agent", agent)
    await cl.Message("Hi!").send()

@cl.on_message
async def on_message(msg: cl.Message):
    agent = cl.user_session.get("agent")
    response = await cl.make_async(agent.chat)(msg.content)
    
    # Check if the response is a Plotly figure JSON
    try:
        fig = go.Figure(json.loads(response))
        await cl.Message(content=fig.to_html(), author="iFactory Assistant").send()
    except json.JSONDecodeError:
        await cl.Message(content=str(response), author="iFactory Assistant").send()

By following these steps, you can extend the LlamaIndex framework to handle and return Plotly figures, and integrate this functionality with Chainlit to display the figures in your frontend ^[1].

To continue talking to Dosu, mention @dosu.

[dosubot[bot]]

To address the pydantic error and ensure the @dispatcher.span decorator is applied correctly, you can follow these steps:

1. Allow Arbitrary Types in Pydantic Models:

Ensure that the Config class in your custom TaskStepOutput allows arbitrary types:

from pydantic import BaseModel, Field
from typing import Union, List
import plotly.graph_objects as go
from llama_index.core.agent.runner.base import TaskStep, TaskStepOutput

class CustomTaskStepOutput(BaseModel):
    output: Union[str, go.Figure] = Field(..., description="Task step output")
    task_step: TaskStep = Field(..., description="Task step input")
    next_steps: List[TaskStep] = Field(..., description="Next steps to be executed.")
    is_last: bool = Field(default=False, description="Is this the last step?")

    class Config:
        arbitrary_types_allowed = True

    def __str__(self) -> str:
        return str(self.output)

2. Extend AgentRunner with @dispatcher.span Decorator:

Ensure the run_step method in your custom AgentRunner class is decorated with @dispatcher.span:

from llama_index.core.agent.runner.base import AgentRunner, TaskStep
from llama_index.core.instrumentation import dispatcher
from typing import Optional, Any

class CustomAgentRunner(AgentRunner):
    @dispatcher.span
    def run_step(
        self,
        task_id: str,
        input: Optional[str] = None,
        step: Optional[TaskStep] = None,
        **kwargs: Any,
    ) -> CustomTaskStepOutput:
        step = self.validate_step_from_args(task_id, input, step, **kwargs)
        # Custom logic to handle Plotly figures
        if isinstance(step.output, go.Figure):
            return CustomTaskStepOutput(output=step.output, task_step=step, next_steps=[], is_last=True)
        else:
            return super().run_step(task_id, input, step, **kwargs)

    @dispatcher.span
    def finalize_response(
        self,
        task_id: str,
        step_output: Optional[CustomTaskStepOutput] = None,
    ) -> Any:
        if step_output is None:
            step_output = self.state.get_completed_steps(task_id)[-1]
        if not step_output.is_last:
            raise ValueError("finalize_response can only be called on the last step output")

        if isinstance(step_output.output, go.Figure):
            return step_output.output.to_json()
        else:
            return super().finalize_response(task_id, step_output)

3. Integrate with Chainlit:

Modify your Chainlit code to handle the JSON representation of the Plotly figure:

import chainlit as cl
from utils import agent
import plotly.graph_objects as go
import json

@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("agent", agent)
    await cl.Message("Hi!").send()

@cl.on_message
async def on_message(msg: cl.Message):
    agent = cl.user_session.get("agent")
    response = await cl.make_async(agent.chat)(msg.content)
    
    # Check if the response is a Plotly figure JSON
    try:
        fig = go.Figure(json.loads(response))
        await cl.Message(content=fig.to_html(), author="iFactory Assistant").send()
    except json.JSONDecodeError:
        await cl.Message(content=str(response), author="iFactory Assistant").send()

By following these steps, you should be able to handle and return Plotly figures without encountering the pydantic error and ensure the @dispatcher.span decorator is correctly applied to the run_step method.

To continue talking to Dosu, mention @dosu.

[tituslhy]

@dosu sorry this doesn't work. Json loads does not work on the AgentChatResponse class. And when I printed response.response, I still get a string response instead of a graph.

Just in case, I also edited my function tool to make it return a go.Figure instance:

from llama_index.core.agent import FunctionCallingAgentWorker, AgentRunner
from llama_index.core.tools.tool_spec.base import BaseToolSpec
import pandas as pd
pd.options.plotting.backend = "plotly"

class DBTools(BaseToolSpec):
    spec_functions = [
        "plot_data",
        ...
    ]
    
    def __init__(
        self,  
        db_path: Optional[str] = db_path 
        ):
        """Class constructor. Connects to sqlite database"""
        self.conn = sqlite3.connect(db_path, check_same_thread=False)

    def plot_data(self) ->str:
        df = self.get_df_function() ##dummy function to get dataframe
        fig = df.plot.line(
                title=title,
                labels=labels
            )
        return fig # <- This is the change
  
   ### Other methods ###

[dosubot[bot]]

To properly return Plotly figures from your agent tool to the frontend using Chainlit, you need to ensure that the agent can handle and return the Plotly figure object correctly. Here’s a revised approach:

1. Modify the plot_data method to return a JSON representation of the Plotly figure:

import pandas as pd
import plotly.graph_objects as go
import sqlite3
from typing import Optional
from llama_index.core.tools.tool_spec.base import BaseToolSpec

pd.options.plotting.backend = "plotly"

class DBTools(BaseToolSpec):
    spec_functions = [
        "plot_data",
        # ...
    ]
    
    def __init__(self, db_path: Optional[str] = None):
        """Class constructor. Connects to sqlite database"""
        self.conn = sqlite3.connect(db_path, check_same_thread=False)

    def plot_data(self) -> str:
        df = self.get_df_function()  # Dummy function to get dataframe
        fig = df.plot.line(
            title="Sample Title",
            labels={"index": "X-axis", "value": "Y-axis"}
        )
        return fig.to_json()  # Return the JSON representation of the figure

2. Update the Chainlit integration to handle the JSON representation of the Plotly figure:

import chainlit as cl
from utils import agent
import plotly.graph_objects as go
import json

@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("agent", agent)
    await cl.Message("Hi!").send()

@cl.on_message
async def on_message(msg: cl.Message):
    agent = cl.user_session.get("agent")
    response = await cl.make_async(agent.chat)(msg.content)
    
    # Check if the response is a Plotly figure JSON
    try:
        fig = go.Figure(json.loads(response))
        await cl.Message(content=fig.to_html(full_html=False), author="iFactory Assistant").send()
    except json.JSONDecodeError:
        await cl.Message(content=str(response), author="iFactory Assistant").send()

3. Ensure the agent's chat method correctly handles the JSON response:

Make sure that the chat method in your agent correctly returns the JSON string from the plot_data method.

By following these steps, you should be able to return Plotly figures from your agent tool to the frontend using Chainlit, with the figures being rendered correctly in the frontend.

To continue talking to Dosu, mention @dosu.

[tituslhy]

@dosu No this still won't work because the agent's output type will still be AgentChatResponse. So json.loads(response) will error.

Just in case, here are my implementations of agent tools, custom agent runner and task step output:

from llama_index.core.agent.runner.base import AgentRunner
import llama_index.core.instrumentation as instrument
from typing import Optional, Any, List
from llama_index.core.bridge.pydantic import BaseModel, Field
import plotly.graph_objects as go
from llama_index.core.base.agent.types import TaskStep
import pandas as pd
import plotly.graph_objects as go
import sqlite3
from llama_index.core.tools.tool_spec.base import BaseToolSpec

pd.options.plotting.backend = "plotly"

class DBTools(BaseToolSpec):
    spec_functions = [
        "plot_data",
        # ...
    ]
    
    def __init__(self, db_path: Optional[str] = None):
        """Class constructor. Connects to sqlite database"""
        self.conn = sqlite3.connect(db_path, check_same_thread=False)

    def plot_data(self) -> str:
        df = self.get_df_function()  # Dummy function to get dataframe
        fig = df.plot.line(
            title="Sample Title",
            labels={"index": "X-axis", "value": "Y-axis"}
        )
        return fig.to_json()  # Return the JSON representation of the figure

class CustomTaskStepOutput(BaseModel):
    output: Union[str, go.Figure] = Field(..., description="Task step output")
    task_step: TaskStep = Field(..., description = "Task step input")
    next_steps: List[TaskStep] = Field(..., description="Next steps to be executed")
    is_last: bool = Field(default=False, description="Is this the last step?")
    
    class Config:
        arbitrary_types_allowed = True
    
    def __str__(self) -> str:
        return str(self.output)

dispatcher = instrument.get_dispatcher(__name__)

class CustomAgentRunner(AgentRunner):
    
    @dispatcher.span
    def run_step(
        self,
        task_id: str,
        input: Optional[str] = None,
        step: Optional[TaskStep] = None,
        **kwargs
    ) -> CustomTaskStepOutput:
        step = self.validate_step_from_args(task_id, input, step, **kwargs)
        
        # Custom logic to handle plotly figures
        if isinstance(step.output, go.Figure):
            return CustomTaskStepOutput(output=step.output, 
                                        task_step=step,
                                        next_steps=[],
                                        is_last=True)
        return super().run_step(task_id, input, step, **kwargs)
    
    @dispatcher.span
    def finalize_response(
        self,
        task_id: str, 
        step_output: Optional[CustomTaskStepOutput] = None,
    ) -> Any:
        if step_output is None:
            step_output = self.state.get_completed_steps(task_id)[-1]
        if not step_output.is_last:
            raise ValueError("finalize_response can only be called on the last step output")
        if isinstance(step_output.output, go.Figure):
            return step_output.output.to_json()
        return super().finalize_response(task_id, step_output)


```. 

[dosubot[bot]]

To address the issue of AgentChatResponse not being able to handle Plotly figures directly, you can modify the finalize_response method to ensure it returns the JSON representation of the Plotly figure correctly. Here’s how you can adjust your implementation:

1. Ensure plot_data returns JSON:

Your plot_data method is already correctly returning the JSON representation of the Plotly figure.

2. Modify CustomAgentRunner to handle JSON responses:

Ensure that the finalize_response method in CustomAgentRunner correctly returns the JSON string from the plot_data method.

from llama_index.core.agent.runner.base import AgentRunner
import llama_index.core.instrumentation as instrument
from typing import Optional, Any, List, Union
from llama_index.core.bridge.pydantic import BaseModel, Field
import plotly.graph_objects as go
from llama_index.core.base.agent.types import TaskStep
import pandas as pd
import sqlite3
from llama_index.core.tools.tool_spec.base import BaseToolSpec

pd.options.plotting.backend = "plotly"

class DBTools(BaseToolSpec):
    spec_functions = [
        "plot_data",
        # ...
    ]
    
    def __init__(self, db_path: Optional[str] = None):
        """Class constructor. Connects to sqlite database"""
        self.conn = sqlite3.connect(db_path, check_same_thread=False)

    def plot_data(self) -> str:
        df = self.get_df_function()  # Dummy function to get dataframe
        fig = df.plot.line(
            title="Sample Title",
            labels={"index": "X-axis", "value": "Y-axis"}
        )
        return fig.to_json()  # Return the JSON representation of the figure

class CustomTaskStepOutput(BaseModel):
    output: Union[str, go.Figure] = Field(..., description="Task step output")
    task_step: TaskStep = Field(..., description="Task step input")
    next_steps: List[TaskStep] = Field(..., description="Next steps to be executed")
    is_last: bool = Field(default=False, description="Is this the last step?")
    
    class Config:
        arbitrary_types_allowed = True
    
    def __str__(self) -> str:
        return str(self.output)

dispatcher = instrument.get_dispatcher(__name__)

class CustomAgentRunner(AgentRunner):
    
    @dispatcher.span
    def run_step(
        self,
        task_id: str,
        input: Optional[str] = None,
        step: Optional[TaskStep] = None,
        **kwargs: Any
    ) -> CustomTaskStepOutput:
        step = self.validate_step_from_args(task_id, input, step, **kwargs)
        
        # Custom logic to handle plotly figures
        if isinstance(step.output, go.Figure):
            return CustomTaskStepOutput(output=step.output, 
                                        task_step=step,
                                        next_steps=[],
                                        is_last=True)
        return super().run_step(task_id, input, step, **kwargs)
    
    @dispatcher.span
    def finalize_response(
        self,
        task_id: str, 
        step_output: Optional[CustomTaskStepOutput] = None,
    ) -> Any:
        if step_output is None:
            step_output = self.state.get_completed_steps(task_id)[-1]
        if not step_output.is_last:
            raise ValueError("finalize_response can only be called on the last step output")
        if isinstance(step_output.output, go.Figure):
            return step_output.output.to_json()
        return super().finalize_response(task_id, step_output)

3. Update Chainlit integration:

Modify your Chainlit code to handle the JSON representation of the Plotly figure:

import chainlit as cl
from utils import agent
import plotly.graph_objects as go
import json

@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("agent", agent)
    await cl.Message("Hi!").send()

@cl.on_message
async def on_message(msg: cl.Message):
    agent = cl.user_session.get("agent")
    response = await cl.make_async(agent.chat)(msg.content)
    
    # Check if the response is a Plotly figure JSON
    try:
        fig = go.Figure(json.loads(response.response))
        await cl.Message(content=fig.to_html(full_html=False), author="iFactory Assistant").send()
    except json.JSONDecodeError:
        await cl.Message(content=str(response.response), author="iFactory Assistant").send()

By ensuring that the finalize_response method returns the JSON representation of the Plotly figure and updating the Chainlit integration to handle this JSON, you should be able to render the Plotly figures correctly in your frontend.

To continue talking to Dosu, mention @dosu.