main4-function-calling.py

import json
import os

from openai import AzureOpenAI
from rich import print as rprint

from tools import book_flight, search_flights

background = (
    "You are an expert travel planner with 25+ years of experience. Leverage your expertise to provide recommendations,"
    + "while considering the user's personal interests and preferences."
)

tools = [
    {
        "type": "function",
        "function": {
            "name": "search_flights",
            "description": "Search for flights between the provided cities on the given date.",
            "parameters": {
                "type": "object",
                "properties": {
                    "departure_city": {"type": "string"},
                    "destination_city": {"type": "string"},
                    "travel_date": {"type": "string"},
                },
                "required": ["departure_city", "destination_city", "travel_date"],
                "additionalProperties": False,
            },
            "strict": True,
        },
    },
    {
        "type": "function",
        "function": {
            "name": "book_flight",
            "description": "Book a flight based on the provided flight number and travel date.",
            "parameters": {
                "type": "object",
                "properties": {
                    "flight_number": {"type": "string"},
                    "travel_date": {"type": "string"},
                },
                "required": ["flight_number", "travel_date"],
                "additionalProperties": False,
            },
            "strict": True,
        },
    },
]


def format_msg(msg: dict):
    if msg.get("tool_calls"):
        return msg["tool_calls"]

    return msg["content"]


def main():
    user_request = input("How can I help you today? ")

    azure_endpoint = os.environ["AZURE_OPENAI_ENDPOINT"]
    azure_deployment = os.environ["AZURE_OPENAI_DEPLOYMENT"]

    client = AzureOpenAI(
        azure_endpoint=azure_endpoint,
        api_version="2024-08-01-preview",
        azure_deployment=azure_deployment,
    )

    messages = [
        {
            "role": "system",
            "content": background,
        },
        {
            "role": "user",
            "content": user_request,
        },
    ]

    while True:
        chat_completion = client.chat.completions.create(
            messages=messages,
            model="gpt-4o",
            temperature=0.1,
            tools=tools,
        )
        messages.append(chat_completion.choices[0].message.to_dict())

        rprint(
            "\n".join(
                [
                    f"[red]{format_msg(m)}[/red]"
                    for m in messages
                    if m["role"] != "system"
                ]
            )
        )
        print()

        if chat_completion.choices[0].message.tool_calls:
            for tool_call in chat_completion.choices[0].message.tool_calls:
                func_name = tool_call.function.name
                func_args = json.loads(tool_call.function.arguments)

                match func_name:
                    case "search_flights":
                        result = search_flights(*func_args)
                        messages.append(
                            {
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": f"{result}",
                            }
                        )
                    case "book_flight":
                        booking_result = book_flight(*func_args)
                        messages.append(
                            {
                                "role": "tool",
                                "tool_call_id": tool_call.id,
                                "content": f"{booking_result}",
                            }
                        )
                    case _:
                        raise ValueError(f"Unsupported function: {func_name}")

        else:
            print(chat_completion.choices[0].message.content)
            human_response = input("> ")
            messages.append({"role": "user", "content": human_response})


if __name__ == "__main__":
    main()