test_tool_calling.py

#!/usr/bin/env python

from llama_index.core.agent import AgentRunner
from llama_index.core.base.llms.types import ChatMessage, MessageRole
from llama_index.core.tools import FunctionTool
from llama_index.llms.openai import OpenAI


def weather_tool(city: str):
    """
    :param city: city name
    :return: the weather in the city tomorrow
    """
    print(f"Getting weather for {city}...")
    return "sunny"


SYSTEM_PROMPT = """
You are a helpful weather assistant. The users can ask you about the weather in any city in the world.
You can use the `weather_tool` tool to get the weather in a city, given the city name as the argument.

When the user asks you about the weather in a city, respond with the following MANDATORY_RESPONSE:

It sure is a great day today! Let me check the weather in {city} tomorrow for you. Please hold on for a moment.

and then use the `weather_tool` tool to get the weather in the city. Once you obtain the weather, respond with the following SECONDARY_RESPONSE:

The weather in {city} tomorrow will be {weather}.

You *MUST* first respond with the MANDATORY_RESPONSE, then use `weather_tool`, then respond with the SECONDARY_RESPONSE.
"""


def gather_response(chat_engine, message):
    system_message = ChatMessage(content=SYSTEM_PROMPT, role=MessageRole.SYSTEM)
    response = chat_engine.stream_chat(message, [system_message])

    result = ""
    for token in response.response_gen:
        result += token

    return result


def main():
    llm = OpenAI(model="gpt-4o")
    weather = FunctionTool.from_defaults(
        fn=weather_tool,
        name="weather_tool",
        description="Returns the weather in the given city tomorrow.",
    )
    chat_engine = AgentRunner.from_llm(llm=llm, tools=[weather])

    response = gather_response(chat_engine, "What is the weather in New York tomorrow?")

    assert ("sunny" in response.lower()), f"'sunny' not found in:\n{response}"


if __name__ == "__main__":
    main()