How-To: Streaming events from tools and using them in the client

In this document, we will build a tool that is able to send custom commands to the UI and display a Markdown table to the user.

Define a streaming tool

To achieve this, define a tool as a Generator or AsyncGenerator that yields an event that the client is able to handle. In our case, we will use a TextReponse event, which is supported by Ragbits UI.

Note, that the custom events yielded from the tool won't be automatically passed to the LLM. To indicate the output of the tool that needs to be passed to the LLM, use ToolReturn.

from collections.abc import Generator

from ragbits.agents.tool import ToolReturn
from ragbits.chat.interface.types import TextContent, TextResponse


def display_revenue_table(year: int) -> Generator[TextResponse | ToolReturn]:
    """Fetches revenue data (in thousands of dollars) for the given year and displays it as a Markdown table.

    Args:
        year: The year to fetch revenue data for.
    """
    data = fetch_revenue_data(year)
    if data is None:
        yield ToolReturn(value=f"No data for year {year}")
    else:
        rows = "\n".join(f"| {q} | {v} |" for q, v in zip(data.quarters, data.revenue, strict=False))
        table = f"\n\n| Quarter | Value |\n|---------|-------|\n{rows}\n\n"
        yield TextResponse(content=TextContent(text=table))
        yield ToolReturn(value={"year": year, "total": sum(data.revenue)})

You can also define custom events by inheriting ChatResponse class with a custom ResponseContent. The content can be arbitrary, as long as your client understands it and is able to handle it. To do that, you will need to extend Ragbits UI with your own event handlers.

Use it in a ChatInterface

Now, we will implement a ChatInterface that uses the agent with the streaming tool. You can serve and test it via RagbitsAPI. Since we used a TextResponse that is understood by the Ragbits UI, we can deploy the app and test the agent right away.

from collections.abc import AsyncGenerator

from ragbits.agents import Agent
from ragbits.chat.api import RagbitsAPI
from ragbits.chat.interface import ChatInterface
from ragbits.chat.interface.types import ChatContext, ChatResponse, TextResponse
from ragbits.core.llms import LiteLLM
from ragbits.core.prompt.base import ChatFormat


class RevenueChatInterface(ChatInterface):
    """Chat interface with a tool that streams revenue tables."""

    def __init__(self) -> None:
        self.llm = LiteLLM(model_name="gpt-4o")
        self.agent = Agent(
            llm=self.llm,
            prompt="""
            You are a helpful revenue analytics assistant. Use `display_revenue_table` to show revenue data to the user.
            Invoking that tool means that the user sees the table in the web interface. When invoking that tool,
            do not write the data yourself.

            You have access to data from 2023 and 2024.""",
            tools=[display_revenue_table],
        )

    async def chat(
        self,
        message: str,
        history: ChatFormat,
        context: ChatContext,
    ) -> AsyncGenerator[ChatResponse, None]:
        """Handles interaction with the user."""
        async for chunk in self.agent.run_streaming(message):
            if isinstance(chunk, str):
                yield self.create_text_response(chunk)
            elif isinstance(chunk, TextResponse):
                yield chunk

Note, that the system receiving the events may be completely different from Ragbits frontend.