general_function_calling.py

# MODEL_NAME="meta-llama/Meta-Llama-3.1-8B-Instruct" OPENAI_API_URL="http://localhost:1234/v1" OPENAI_API_TOKEN="..." python general_function_calling.py
# MODEL_NAME="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo" OPENAI_API_URL="https://api.together.xyz/v1" OPENAI_API_TOKEN="$TOGETHER_API_KEY" python general_function_calling.py
# MODEL_NAME="meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo" OPENAI_API_URL="https://api.together.xyz/v1" OPENAI_API_TOKEN="$TOGETHER_API_KEY" python general_function_calling.py

import copy
import os
import traceback
import json
import re
import inspect
from typing import Optional, Callable, List, Dict
import logging

import requests

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.WARNING)

FUNCTION_RE = re.compile(r"<function=([a-zA-Z0-9_]+)>({.*})")


def build_system_message(
    messages: list[dict], tools: Optional[List[Callable]]
) -> Optional[str]:
    system_messages = [m for m in messages if m["role"] == "system"]
    prev_system_message = system_messages[0]["content"] if system_messages else None
    if not tools:
        return prev_system_message

    tool_definitions = []
    for tool in tools:
        signature = inspect.signature(tool)
        parameters = {}
        for name, param in signature.parameters.items():
            param_type = (
                param.annotation
                if param.annotation != inspect.Parameter.empty
                else "Any"
            )
            param_type = str(param_type).replace("<class '", "").replace("'>", "")
            default = (
                param.default if param.default != inspect.Parameter.empty else None
            )
            parameters[name] = {"type": param_type, "default": default}

        tool_definitions.append(
            {
                "name": tool.__name__,
                "description": (
                    tool.__doc__.strip() if tool.__doc__ else "No description available"
                ),
                "parameters": parameters,
            }
        )

    resp = f"""You have access to the following tools:

{json.dumps(tool_definitions, indent=2)}

If you don't need to use a tool, just respond with your answer.

However, if you think a tool might help you answer the question,
you can call one function at a time by using a pseudo-XML tag called
`function`, which you pass a name, and whose content is a JSON object
containing the arguments to pass to the function.

For example:
<function=foobar>{{"a": "5"}}</function>
<function=sample>{{"city": "New York", "days": 3}}</function>
<function=baz>{{"amount": 100, "from_currency": "USD", "to_currency": "EUR"}}</function>
<function=blahblah>{{"category": "technology", "count": 5}}</function>
<function=check_something>{{"expression": "17 * 32"}}</function>
<function=get_some_value>{{}}</function>

You may only call one function at a time. If you need to use multiple functions,
please do so sequentially in 1-by-1 manner over conversational turns.

You should know: whenever you request a function call, your message will not
be shown to the user, but instead intercepted by a system to be executed and run.
This intercepting process will insert a message from the user, which contains the
function output. This message did not actually come from the user, so please use
it for informational purposes, but don't reference it in subsequent responses, because
your reference won't make sense to the user who never saw the function call/response
messages.

Example exchanges of hypothetical function calls and responses:

## Ex 1
User: What is the weather like today?
Assistant: <function=weather_forecast>{{"city": "New York", "days": 1}}</function>
User: {{"name": "weather_forecast", "arguments": {{"city": "New York", "days": 1}}, "output": {{"temperature": "72°F", "description": "Sunny with a chance of rain"}}}}
Assistant: The weather in New York today is sunny with a chance of rain and the temperature is 72°F.

## Ex 2
User: What's the current time?
Assistant: <function=get_current_time>{{}}</function>
User: {{"name": "get_current_time", "arguments": {{}}, "output": {{"time": "10:35 AM"}}}}
Assistant: The current time is 10:35 AM.

## Ex 3
User: Convert 100 USD to EUR.
Assistant: <function=currency_converter>{{"amount": 100, "from_currency": "USD", "to_currency": "EUR"}}</function>
User: {{"name": "currency_converter", "arguments": {{"amount": 100, "from_currency": "USD", "to_currency": "EUR"}}, "output": {{"converted_amount": 92.56}}}}
Assistant: 100 USD is equivalent to approximately 92.56 EUR.
"""
    if prev_system_message:
        resp += f"\n{prev_system_message}"
    return resp


def run_inference(
    messages: List[Dict[str, str]],
    base_url: str = os.environ.get("OPENAI_API_URL", "http://localhost:1234/v1"),
    model: str = os.environ.get("MODEL_NAME", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
    temperature: float = float(os.environ.get("MODEL_TEMPERATURE", 0.35)),
    tools: Optional[List[Callable]] = None,
    auth_token: Optional[str] = os.environ.get("OPENAI_API_TOKEN"),
    max_tries: int = 10,
) -> list[Dict[str, str]]:
    logger.info(f"Starting run_inference with {len(messages)} messages")

    msgs = []
    system_message = build_system_message(messages, tools)
    if system_message:
        msgs.append({"role": "system", "content": system_message})

    # Add user and assistant messages to the list of messages
    for message in messages:
        if message["role"] == "system":
            continue
        msgs.append({"role": message["role"], "content": message["content"]})

    headers = {"Content-Type": "application/json"}
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"

    prev_response = None
    prev_error = None
    to_return = []
    for _ in range(max_tries):
        call_msgs = copy.deepcopy(msgs)
        if prev_error and prev_response:
            call_msgs.append(prev_response)
            call_msgs.append(
                {
                    "role": "user",
                    "content": "\n\n".join(traceback.format_exception(prev_error)),
                }
            )
        try:
            response = requests.post(
                f"{base_url}/chat/completions",
                json={
                    "model": model,
                    "messages": call_msgs,
                    "temperature": temperature,
                },
                headers=headers,
            )
            response.raise_for_status()
            text = response.json()["choices"][0]["message"]["content"]
            prev_response = {"role": "assistant", "content": text}
        except (requests.RequestException, requests.exceptions.HTTPError) as e:
            logger.exception(e)
            prev_error = e
            continue

        text_prop = text.strip()

        match = FUNCTION_RE.search(text_prop)
        if not match:
            # No function call found, return the text
            to_return.append({"role": "assistant", "content": text})
            return to_return

        # Lookup tool by name
        tool = next((t for t in tools if t.__name__ == match.group(1)), None)

        # If there was no valid tool, we set a new error explaining that there's no
        # valid function of that name available, and then try again
        if not tool:
            prev_error = ValueError(f"No valid function named {match.group(1)}")
            continue

        # Parse the arguments
        try:
            args = json.loads(match.group(2))
        except json.decoder.JSONDecodeError as e:
            prev_error = e
            continue

        # Otherwise, call the tool
        try:
            tool_output = tool(**args)
        except Exception as e:
            prev_error = e
            continue

        tool_content = {
            "name": match.group(1),
            "arguments": args,
            "output": str(tool_output),
        }

        # Now we prepare for the next loop
        new_messages = [
            {
                "role": "assistant",
                "content": f"<function={match.group(1)}>{json.dumps(args)}</function>",  # Normalize
            },
            {"role": "user", "content": json.dumps(tool_content)},
        ]
        msgs += new_messages
        to_return += new_messages

        # Ensure that we clear prev_error and prev_response so that we don't
        # accidentally repeat the same error message or response
        prev_error = None
        prev_response = None

    # If we've exhausted all retries, raise the last exception
    if prev_error:
        raise prev_error

    # If there was no exception, raise
    raise ValueError(f"No valid function found in response {text}")


def main():
    import random
    from datetime import datetime, timedelta

    def calculate(expression: str) -> float:
        """Evaluates the given arithmetic expression."""
        return {"calculation": eval(expression)}

    def get_current_time() -> str:
        """Returns the current time in ISO format."""
        return {"time": datetime.now().isoformat()}

    def generate_random_number(min_value: int, max_value: int) -> int:
        """Generates a random integer between min_value and max_value (inclusive)."""
        return {"random_number": random.randint(min_value, max_value)}

    def forecast_weather(city: str, days: int) -> str:
        """Simulates a weather forecast for the given city and number of days."""
        weather_conditions = ["Sunny", "Cloudy", "Rainy", "Windy", "Snowy"]
        forecast = f"Weather forecast for {city} for the next {days} days:\n"

        for i in range(days):
            date = (datetime.now() + timedelta(days=i)).strftime("%Y-%m-%d")
            condition = random.choice(weather_conditions)
            temperature = random.randint(0, 30)
            forecast += f"{date}: {condition}, {temperature}°C\n"

        return forecast

    messages = [
        {
            "role": "system",
            "content": "You are a helpful assistant capable of performing various tasks using the provided tools. Use the appropriate tool for each task and engage in a multi-round conversation to demonstrate your capabilities.",
        },
        {
            "role": "user",
            "content": "Let's start with a simple math problem. What's the result of 15 * 7 + 22?",
        },
    ]

    tools = [calculate, get_current_time, generate_random_number, forecast_weather]

    for _ in range(6):  # Simulate rounds of conversation
        response = run_inference(messages, tools=tools)
        messages.extend(response)
        last_response = messages[-1]["content"]
        last_user_message = [
            m
            for m in messages
            if m["role"] == "user"
            and (m["content"][0] != "{" or m["content"][-1] != "}")
        ][-1]["content"]
        print(f"User: {last_user_message}")
        print(f"Assistant: {last_response}")

        # Generate follow-up questions based on the response
        if "calculate" in last_response.lower():
            messages.append(
                {
                    "role": "user",
                    "content": "Great! Now, can you tell me the current time and generate a random number between 1 and 100?",
                }
            )
        elif (
            "current time" in last_response.lower()
            and "random number" in last_response.lower()
        ):
            messages.append(
                {
                    "role": "user",
                    "content": "Interesting! Let's move on to something more complex. Can you provide a 3-day weather forecast for a city of your choice?",
                }
            )
        elif "weather forecast" in last_response.lower():
            messages.append(
                {
                    "role": "user",
                    "content": "That's helpful! Now, let's combine some of these tools. Can you generate a random number between 1 and 10, and then calculate its square root?",
                }
            )
        elif "square root" in last_response.lower():
            messages.append(
                {
                    "role": "user",
                    "content": "Excellent! For our final task, can you tell me the current time again and calculate how many hours have passed since midnight?",
                }
            )
        else:
            messages.append(
                {
                    "role": "user",
                    "content": "That's interesting. Can you demonstrate another capability using the tools available to you?",
                }
            )


if __name__ == "__main__":
    main()