Tool calling and structured outputs: validate every response

"""contracts.py: JSON Schemas (draft 2020-12) for replies, arguments and results."""
from jsonschema import Draft202012Validator

UNIT_KIND = {"m": "length", "km": "length", "g": "mass", "kg": "mass"}
ENVELOPE = {
    "type": "object",
    "properties": {"tool": {"type": "string"}, "arguments": {"type": "object"}},
    "required": ["tool", "arguments"],
    "additionalProperties": False,
}
ARGUMENTS = {  # the allowlist: a tool not named here can never run
    "convert_units": {
        "type": "object",
        "properties": {
            "value": {"type": "number", "minimum": -1000000, "maximum": 1000000},
            "from_unit": {"enum": list(UNIT_KIND)},
            "to_unit": {"enum": list(UNIT_KIND)},
        },
        "required": ["value", "from_unit", "to_unit"],
        "additionalProperties": False,
    },
    "lookup_stock": {
        "type": "object",
        "properties": {"ticker": {"enum": ["ZEPHYR", "SNAIL"]}},
        "required": ["ticker"],
        "additionalProperties": False,
    },
}
RESULTS = {
    "lookup_stock": {
        "type": "object",
        "properties": {"ticker": {"type": "string"}, "price": {"type": "number"},
                       "summary": {"type": "string", "maxLength": 200}},
        "required": ["ticker", "price", "summary"],
        "additionalProperties": False,
    },
}


def problems(instance, schema):
    """Check the schema itself, then return every error message, sorted."""
    Draft202012Validator.check_schema(schema)
    return sorted(e.message for e in Draft202012Validator(schema).iter_errors(instance))


if __name__ == "__main__":
    print(problems({"value": "5", "from_unit": "miles"}, ARGUMENTS["convert_units"]))

python contracts.py
["'5' is not of type 'number'", "'miles' is not one of ['m', 'km', 'g', 'kg']", "'to_unit' is a required property"]

"""gate.py: parse, then validate, then check business rules."""
import json

from contracts import ARGUMENTS, ENVELOPE, UNIT_KIND, problems


def refuse(constant):
    raise ValueError(f"{constant} is not allowed in JSON")


def check_call(reply):
    """Return (call, None) if the reply may run, or (None, reason) if not."""
    if len(reply) > 2000:
        return None, "reply too long"
    try:
        call = json.loads(reply, parse_constant=refuse)
    except ValueError as error:  # json.JSONDecodeError is a ValueError
        return None, f"not JSON: {error}"
    found = problems(call, ENVELOPE)
    if found:
        return None, "envelope: " + "; ".join(found)
    tool, args = call["tool"], call["arguments"]
    if tool not in ARGUMENTS:
        return None, f"tool not allowed: {tool}"
    found = problems(args, ARGUMENTS[tool])
    if found:
        return None, "arguments: " + "; ".join(found)
    if tool == "convert_units" and UNIT_KIND[args["from_unit"]] != UNIT_KIND[args["to_unit"]]:
        return None, f"rule: cannot convert {args['from_unit']} to {args['to_unit']}"
    return call, None


if __name__ == "__main__":
    from fake_model import REPLIES
    for name, reply in REPLIES.items():
        call, reason = check_call(reply)
        print(f"{name:<13}", "RUN" if call else "REJECT " + reason)

python gate.py
good          RUN
malformed     REJECT not JSON: Expecting ',' delimiter: line 1 column 51 (char 50)
wrong_type    REJECT arguments: 'five' is not of type 'number'
extra_arg     REJECT arguments: Additional properties are not allowed ('round' was unexpected)
unknown_tool  REJECT tool not allowed: delete_all_notes
out_of_range  REJECT arguments: 1000000000.0 is greater than the maximum of 1000000
smuggled      REJECT envelope: Additional properties are not allowed ('system' was unexpected)

"""retry.py: send feedback and ask again, a limited number of times."""
from fake_model import fake_model
from gate import check_call


def get_valid_call(script, max_attempts=3):
    messages = [{"role": "user", "content": "Convert 5 km to metres."}]
    for attempt in range(max_attempts):
        reply = fake_model(script, attempt)
        call, reason = check_call(reply)
        print(f"attempt {attempt + 1}:", "accepted" if call else reason)
        if call:
            return call
        messages += [{"role": "assistant", "content": reply},
                     {"role": "user", "content": f"Rejected: {reason}. Reply with one "
                      "JSON object that matches the tool schema, and nothing else."}]
    return None  # give up and tell the person; never guess


print(get_valid_call(["unknown_tool", "wrong_type", "good"]))

python retry.py
attempt 1: tool not allowed: delete_all_notes
attempt 2: arguments: 'five' is not of type 'number'
attempt 3: accepted
{'tool': 'convert_units', 'arguments': {'value': 5, 'from_unit': 'km', 'to_unit': 'm'}}