Workflow for a Function Calling Agent¶

This notebook walks through setting up a Workflow to construct a function calling agent from scratch.

Function calling agents work by using an LLM that supports tools/functions in its API (OpenAI, Ollama, Anthropic, etc.) to call functions an use tools.

Our workflow will be stateful with memory, and will be able to call the LLM to select tools and process incoming user messages.

In [ ]:

!pip install -U llama-index

!pip install -U llama-index

In [ ]:

import os

os.environ["OPENAI_API_KEY"] = "sk-proj-..."

import os os.environ["OPENAI_API_KEY"] = "sk-proj-..."

[Optional] Set up observability with Llamatrace¶

Set up tracing to visualize each step in the workflow.

In [ ]:

!pip install "llama-index-core>=0.10.43" "openinference-instrumentation-llama-index>=2.2.2" "opentelemetry-proto>=1.12.0" opentelemetry-exporter-otlp opentelemetry-sdk

!pip install "llama-index-core>=0.10.43" "openinference-instrumentation-llama-index>=2.2.2" "opentelemetry-proto>=1.12.0" opentelemetry-exporter-otlp opentelemetry-sdk

In [ ]:

from opentelemetry.sdk import trace as trace_sdk
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
    OTLPSpanExporter as HTTPSpanExporter,
)
from openinference.instrumentation.llama_index import LlamaIndexInstrumentor


# Add Phoenix API Key for tracing
PHOENIX_API_KEY = "<YOUR-PHOENIX-API-KEY>"
os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"api_key={PHOENIX_API_KEY}"

# Add Phoenix
span_phoenix_processor = SimpleSpanProcessor(
    HTTPSpanExporter(endpoint="https://app.phoenix.arize.com/v1/traces")
)

# Add them to the tracer
tracer_provider = trace_sdk.TracerProvider()
tracer_provider.add_span_processor(span_processor=span_phoenix_processor)

# Instrument the application
LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)

from opentelemetry.sdk import trace as trace_sdk from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( OTLPSpanExporter as HTTPSpanExporter, ) from openinference.instrumentation.llama_index import LlamaIndexInstrumentor # Add Phoenix API Key for tracing PHOENIX_API_KEY = "" os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"api_key={PHOENIX_API_KEY}" # Add Phoenix span_phoenix_processor = SimpleSpanProcessor( HTTPSpanExporter(endpoint="https://app.phoenix.arize.com/v1/traces") ) # Add them to the tracer tracer_provider = trace_sdk.TracerProvider() tracer_provider.add_span_processor(span_processor=span_phoenix_processor) # Instrument the application LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)

Since workflows are async first, this all runs fine in a notebook. If you were running in your own code, you would want to use asyncio.run() to start an async event loop if one isn't already running.

async def main():
    <async code>

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

Designing the Workflow¶

An agent consists of several steps

1. Handling the latest incoming user message, including adding to memory and getting the latest chat history
2. Calling the LLM with tools + chat history
3. Parsing out tool calls (if any)
4. If there are tool calls, call them, and loop until there are none
5. When there is no tool calls, return the LLM response

The Workflow Events¶

To handle these steps, we need to define a few events:

1. An event to handle new messages and prepare the chat history
2. An event to trigger tool calls
3. An event to handle the results of tool calls

The other steps will use the built-in StartEvent and StopEvent events.

In [ ]:

from llama_index.core.llms import ChatMessage
from llama_index.core.tools import ToolSelection, ToolOutput
from llama_index.core.workflow import Event


class InputEvent(Event):
    input: list[ChatMessage]


class ToolCallEvent(Event):
    tool_calls: list[ToolSelection]


class FunctionOutputEvent(Event):
    output: ToolOutput

from llama_index.core.llms import ChatMessage from llama_index.core.tools import ToolSelection, ToolOutput from llama_index.core.workflow import Event class InputEvent(Event): input: list[ChatMessage] class ToolCallEvent(Event): tool_calls: list[ToolSelection] class FunctionOutputEvent(Event): output: ToolOutput

The Workflow Itself¶

With our events defined, we can construct our workflow and steps.

Note that the workflow automatically validates itself using type annotations, so the type annotations on our steps are very helpful!

In [ ]:

from typing import Any, List

from llama_index.core.llms.function_calling import FunctionCallingLLM
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.tools.types import BaseTool
from llama_index.core.workflow import Workflow, StartEvent, StopEvent, step


class FuncationCallingAgent(Workflow):
    def __init__(
        self,
        *args: Any,
        llm: FunctionCallingLLM | None = None,
        tools: List[BaseTool] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.tools = tools or []

        self.llm = llm or OpenAI()
        assert self.llm.metadata.is_function_calling_model

        self.memory = ChatMemoryBuffer.from_defaults(llm=llm)
        self.sources = []

    @step
    async def prepare_chat_history(self, ev: StartEvent) -> InputEvent:
        # clear sources
        self.sources = []

        # get user input
        user_input = ev.input
        user_msg = ChatMessage(role="user", content=user_input)
        self.memory.put(user_msg)

        # get chat history
        chat_history = self.memory.get()
        return InputEvent(input=chat_history)

    @step
    async def handle_llm_input(
        self, ev: InputEvent
    ) -> ToolCallEvent | StopEvent:
        chat_history = ev.input

        response = await self.llm.achat_with_tools(
            self.tools, chat_history=chat_history
        )
        self.memory.put(response.message)

        tool_calls = self.llm.get_tool_calls_from_response(
            response, error_on_no_tool_call=False
        )

        if not tool_calls:
            return StopEvent(
                result={"response": response, "sources": [*self.sources]}
            )
        else:
            return ToolCallEvent(tool_calls=tool_calls)

    @step
    async def handle_tool_calls(self, ev: ToolCallEvent) -> InputEvent:
        tool_calls = ev.tool_calls
        tools_by_name = {tool.metadata.get_name(): tool for tool in self.tools}

        tool_msgs = []

        # call tools -- safely!
        for tool_call in tool_calls:
            tool = tools_by_name.get(tool_call.tool_name)
            additional_kwargs = {
                "tool_call_id": tool_call.tool_id,
                "name": tool.metadata.get_name(),
            }
            if not tool:
                tool_msgs.append(
                    ChatMessage(
                        role="tool",
                        content=f"Tool {tool_call.tool_name} does not exist",
                        additional_kwargs=additional_kwargs,
                    )
                )
                continue

            try:
                tool_output = tool(**tool_call.tool_kwargs)
                self.sources.append(tool_output)
                tool_msgs.append(
                    ChatMessage(
                        role="tool",
                        content=tool_output.content,
                        additional_kwargs=additional_kwargs,
                    )
                )
            except Exception as e:
                tool_msgs.append(
                    ChatMessage(
                        role="tool",
                        content=f"Encountered error in tool call: {e}",
                        additional_kwargs=additional_kwargs,
                    )
                )

        for msg in tool_msgs:
            self.memory.put(msg)

        chat_history = self.memory.get()
        return InputEvent(input=chat_history)

from typing import Any, List from llama_index.core.llms.function_calling import FunctionCallingLLM from llama_index.core.memory import ChatMemoryBuffer from llama_index.core.tools.types import BaseTool from llama_index.core.workflow import Workflow, StartEvent, StopEvent, step class FuncationCallingAgent(Workflow): def __init__( self, *args: Any, llm: FunctionCallingLLM | None = None, tools: List[BaseTool] | None = None, **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) self.tools = tools or [] self.llm = llm or OpenAI() assert self.llm.metadata.is_function_calling_model self.memory = ChatMemoryBuffer.from_defaults(llm=llm) self.sources = [] @step async def prepare_chat_history(self, ev: StartEvent) -> InputEvent: # clear sources self.sources = [] # get user input user_input = ev.input user_msg = ChatMessage(role="user", content=user_input) self.memory.put(user_msg) # get chat history chat_history = self.memory.get() return InputEvent(input=chat_history) @step async def handle_llm_input( self, ev: InputEvent ) -> ToolCallEvent | StopEvent: chat_history = ev.input response = await self.llm.achat_with_tools( self.tools, chat_history=chat_history ) self.memory.put(response.message) tool_calls = self.llm.get_tool_calls_from_response( response, error_on_no_tool_call=False ) if not tool_calls: return StopEvent( result={"response": response, "sources": [*self.sources]} ) else: return ToolCallEvent(tool_calls=tool_calls) @step async def handle_tool_calls(self, ev: ToolCallEvent) -> InputEvent: tool_calls = ev.tool_calls tools_by_name = {tool.metadata.get_name(): tool for tool in self.tools} tool_msgs = [] # call tools -- safely! for tool_call in tool_calls: tool = tools_by_name.get(tool_call.tool_name) additional_kwargs = { "tool_call_id": tool_call.tool_id, "name": tool.metadata.get_name(), } if not tool: tool_msgs.append( ChatMessage( role="tool", content=f"Tool {tool_call.tool_name} does not exist", additional_kwargs=additional_kwargs, ) ) continue try: tool_output = tool(**tool_call.tool_kwargs) self.sources.append(tool_output) tool_msgs.append( ChatMessage( role="tool", content=tool_output.content, additional_kwargs=additional_kwargs, ) ) except Exception as e: tool_msgs.append( ChatMessage( role="tool", content=f"Encountered error in tool call: {e}", additional_kwargs=additional_kwargs, ) ) for msg in tool_msgs: self.memory.put(msg) chat_history = self.memory.get() return InputEvent(input=chat_history)

And thats it! Let's explore the workflow we wrote a bit.

prepare_chat_history(): This is our main entry point. It handles adding the user message to memory, and uses the memory to get the latest chat history. It returns an InputEvent.

handle_llm_input(): Triggered by an InputEvent, it uses the chat history and tools to prompt the llm. If tool calls are found, a ToolCallEvent is emitted. Otherwise, we say the workflow is done an emit a StopEvent

handle_tool_calls(): Triggered by ToolCallEvent, it calls tools with error handling and returns tool outputs. This event triggers a loop since it emits an InputEvent, which takes us back to handle_llm_input()

Run the Workflow!¶

NOTE: With loops, we need to be mindful of runtime. Here, we set a timeout of 120s.

In [ ]:

from llama_index.core.tools import FunctionTool
from llama_index.llms.openai import OpenAI


def add(x: int, y: int) -> int:
    """Useful function to add two numbers."""
    return x + y


def multiply(x: int, y: int) -> int:
    """Useful function to multiply two numbers."""
    return x * y


tools = [
    FunctionTool.from_defaults(add),
    FunctionTool.from_defaults(multiply),
]

agent = FuncationCallingAgent(
    llm=OpenAI(model="gpt-4o-mini"), tools=tools, timeout=120, verbose=True
)

ret = await agent.run(input="Hello!")

from llama_index.core.tools import FunctionTool from llama_index.llms.openai import OpenAI def add(x: int, y: int) -> int: """Useful function to add two numbers.""" return x + y def multiply(x: int, y: int) -> int: """Useful function to multiply two numbers.""" return x * y tools = [ FunctionTool.from_defaults(add), FunctionTool.from_defaults(multiply), ] agent = FuncationCallingAgent( llm=OpenAI(model="gpt-4o-mini"), tools=tools, timeout=120, verbose=True ) ret = await agent.run(input="Hello!")

Running step prepare_chat_history
Step prepare_chat_history produced event InputEvent
Running step handle_llm_input
Step handle_llm_input produced event StopEvent

In [ ]:

print(ret["response"])

print(ret["response"])

assistant: Hello! How can I assist you today?

In [ ]:

ret = await agent.run(input="What is (2123 + 2321) * 312?")

ret = await agent.run(input="What is (2123 + 2321) * 312?")

Running step prepare_chat_history
Step prepare_chat_history produced event InputEvent
Running step handle_llm_input
Step handle_llm_input produced event ToolCallEvent
Running step handle_tool_calls
Step handle_tool_calls produced event InputEvent
Running step handle_llm_input
Step handle_llm_input produced event ToolCallEvent
Running step handle_tool_calls
Step handle_tool_calls produced event InputEvent
Running step handle_llm_input
Step handle_llm_input produced event StopEvent

In [ ]:

print(ret["response"])

print(ret["response"])

assistant: The result of \((2123 + 2321) \times 312\) is \(1,386,528\).