Function Calling with DSPy

function_call.py

import ast

import dspy

from dspygen.utils.dspy_tools import init_dspy


def get_current_weather(location: str, temperature_unit: str) -> str:
    """
    Get the current weather for a given location in the specified format.

    :param location: The city and state, e.g., "San Francisco, CA".
    :param temperature_unit: The temperature unit to use, either "celsius" or "fahrenheit".
    :return: A string describing the current weather.
    """
    print(f"Retrieving weather for {location}, temperature unit: {temperature_unit}")
    # TODO: API Call
    return f"{location}, {temperature_unit}"


def get_n_day_weather_forecast(
    location: str, temperature_unit: str, num_days: int
) -> str:
    """
    Get an N-day weather forecast for a given location in the specified format.

    :param location: The city and state, e.g., "San Francisco, CA".
    :param temperature_unit: The temperature unit to use, either "celsius" or "fahrenheit".
    :param num_days: The number of days to forecast.
    :return: A string describing the weather forecast.
    """
    print(f"Retrieving {num_days} day weather forecast for {location} in {temperature_unit}")
    return f"{location}, {temperature_unit}, {num_days}"


def function_to_dict(func: "function") -> dict:
    output = {
        "function": func.__name__,
        "docstring": func.__doc__,
        "annotations": func.__annotations__,
    }
    return output


functions_dict = {
    func.__name__: function_to_dict(func) for func in [get_current_weather, get_n_day_weather_forecast]
}

func_list = [get_current_weather, get_n_day_weather_forecast]


def eval_dict_str(dict_str: str) -> dict:
    """Safely convert str to dict"""
    return ast.literal_eval(dict_str)


def lm_function_call(prompt, functions_list: list["function"]) -> dict:
    funcs_str = ",".join([str(function_to_dict(func)) for func in functions_list])

    pred = dspy.Predict("prompt, functions_as_string -> chosen_function")
    choice = pred.forward(prompt=prompt, functions_as_string=funcs_str).chosen_function

    function_dict = functions_dict.get(choice, None)

    pred = dspy.Predict("prompt, function_dict -> keyword_arguments_dict")
    keyword_arguments_dict = pred.forward(prompt=prompt, function_dict=str(function_dict)).keyword_arguments_dict
    function_dict["kwargs"] = eval_dict_str(keyword_arguments_dict)
    return function_dict


init_dspy()

curr_dict = lm_function_call("Today's weather in los angeles", func_list)

assert curr_dict["function"] == "get_current_weather"

get_current_weather(**curr_dict.get("kwargs", {}))

curr_dict = lm_function_call("Years weather in paris, france", func_list)

assert curr_dict["function"] == "get_n_day_weather_forecast"

get_n_day_weather_forecast(**curr_dict.get("kwargs", {}))

print()