lfm-audio-demo/tools.py

"""Test tools for the "Voice + tools" mode, plus the parser for the model's tool-call format."""

import ast
import hashlib
import json
import operator
import re
from datetime import datetime
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

SYSTEM_PROMPT = "Respond with function calls, and interleaved text and audio.\n\nList of tools:"

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "Get the current local date and time in a time zone.",
            "parameters": {
                "type": "object",
                "properties": {
                    "timezone": {"type": "string", "description": "IANA time zone, e.g. Europe/Stockholm. Defaults to UTC."}
                },
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get the current weather for a city.",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string", "description": "City name, e.g. Paris."}},
                "required": ["city"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "Evaluate an arithmetic expression with + - * / ** % and parentheses.",
            "parameters": {
                "type": "object",
                "properties": {"expression": {"type": "string", "description": "The expression, e.g. (17 * 23) / 4."}},
                "required": ["expression"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "set_theme_color",
            "description": "Change the accent color of the user's screen.",
            "parameters": {
                "type": "object",
                "properties": {"color": {"type": "string", "description": "A CSS color name or hex code, e.g. green or #ff8800."}},
                "required": ["color"],
            },
        },
    },
]

CALL_RE = re.compile(r"<tool_call>\s*<function=([^>\n]+)>(.*?)</function>\s*</tool_call>", re.DOTALL)
PARAM_RE = re.compile(r"<parameter=([^>\n]+)>\n?(.*?)\n?</parameter>", re.DOTALL)
OPEN_TAG = "<tool_call>"


def parse_tool_calls(text: str) -> list[dict]:
    calls = []
    for name, body in CALL_RE.findall(text):
        arguments = {}
        for key, value in PARAM_RE.findall(body):
            try:
                arguments[key] = json.loads(value)
            except json.JSONDecodeError:
                arguments[key] = value  # plain string
        calls.append({"type": "function", "function": {"name": name.strip(), "arguments": arguments}})
    return calls


def visible_text(text: str) -> str:
    """The text with tool-call blocks removed, cut before an unfinished one, so the chat bubble never shows markup."""
    text = CALL_RE.sub("", text)
    cut = text.find("<tool_call")
    if cut >= 0:
        return text[:cut]
    for k in range(len(OPEN_TAG) - 1, 0, -1):  # a tag still arriving token by token
        if text.endswith(OPEN_TAG[:k]):
            return text[:-k]
    return text


_OPS = {
    ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv,
    ast.Pow: operator.pow, ast.Mod: operator.mod, ast.FloorDiv: operator.floordiv,
    ast.USub: operator.neg, ast.UAdd: operator.pos,
}


def _eval(node):
    if isinstance(node, ast.Expression):
        return _eval(node.body)
    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
        return node.value
    if isinstance(node, ast.BinOp) and type(node.op) in _OPS:
        if isinstance(node.op, ast.Pow) and abs(_eval(node.right)) > 100:
            raise ValueError("exponent too large")
        return _OPS[type(node.op)](_eval(node.left), _eval(node.right))
    if isinstance(node, ast.UnaryOp) and type(node.op) in _OPS:
        return _OPS[type(node.op)](_eval(node.operand))
    raise ValueError("unsupported expression")


def get_current_time(timezone: str = "UTC", **_):
    try:
        now = datetime.now(ZoneInfo(timezone))
    except (ZoneInfoNotFoundError, ValueError):
        return {"error": f"unknown time zone {timezone!r}"}
    return {"timezone": timezone, "datetime": now.strftime("%Y-%m-%d %H:%M"), "weekday": now.strftime("%A")}


def get_weather(city: str, **_):
    # simulated but stable per city, so repeated tests give the same answer
    h = int(hashlib.sha1(city.strip().lower().encode()).hexdigest(), 16)
    conditions = ["sunny", "partly cloudy", "overcast", "light rain", "windy", "snowing"]
    return {"city": city, "temperature_c": h % 35 - 5, "condition": conditions[h % len(conditions)], "simulated": True}


def calculate(expression: str, **_):
    try:
        return {"expression": expression, "result": _eval(ast.parse(str(expression), mode="eval"))}
    except Exception as e:
        return {"expression": expression, "error": str(e)}


def run_tool(name: str, arguments: dict, ui) -> dict:
    """Run one call; `ui(action, value)` pushes a visible effect to the browser."""
    if name == "set_theme_color":
        ui("accent", str(arguments.get("color", "")))
        return {"ok": True, "color": arguments.get("color")}
    funcs = {"get_current_time": get_current_time, "get_weather": get_weather, "calculate": calculate}
    if name not in funcs:
        return {"error": f"unknown tool {name!r}"}
    try:
        return funcs[name](**arguments)
    except TypeError as e:
        return {"error": f"bad arguments: {e}"}