# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Data structures for function/tool calling in LLM APIs.
These models use ``extra="allow"`` so unknown keys round-trip through
serialization unchanged. This matters for two reasons: (1) OpenAI may
add new fields to the tool wire format over time, and (2)
``transformers.utils.chat_template_utils.get_json_schema`` produces a
richer schema (e.g., a top-level ``return`` block on
``FunctionDefinition``) than the OpenAI Chat Completions API itself
defines, and chat templates may consume those extras. Dropping unknown
keys at validation time would silently lose information that
downstream code relies on.
"""
import json
from enum import Enum
from typing import Any, Literal
import pydantic
from pydantic import BaseModel, ConfigDict, JsonValue
# The seven primitive types defined by JSON Schema.
JSONSchemaType = Literal[
"object", "string", "number", "integer", "boolean", "array", "null"
]
[docs]
class JSONSchema(BaseModel):
"""A JSON Schema object describing the shape of a value.
Models the subset of JSON Schema commonly used in LLM tool
definitions. ``extra="allow"`` lets less-common keywords
(``$ref``, ``$defs``, ``additionalProperties``, ``anyOf``, numeric
constraints, etc.) round-trip unchanged, matching the rest of this
module — see the module docstring for why round-tripping matters.
"""
model_config = ConfigDict(frozen=True, extra="allow")
type: JSONSchemaType | list[JSONSchemaType] | None = None
"""JSON type(s) of this value. A list expresses a union
(e.g., ``["string", "null"]`` for a nullable string).
"""
description: str | None = None
"""Human-readable description, used by the model to choose values."""
title: str | None = None
"""Short human-readable label."""
properties: dict[str, "JSONSchema"] | None = None
"""For ``type="object"``: schema for each named property."""
required: list[str] | None = None
"""For ``type="object"``: names of properties that must be present."""
items: "JSONSchema | None" = None
"""For ``type="array"``: schema for array elements."""
enum: list[Any] | None = None
"""Restricts the value to a fixed set of allowed values."""
default: Any = None
"""Default value used when the field is omitted."""
format: str | None = None
"""Semantic format hint (e.g., ``"date-time"``, ``"email"``)."""
[docs]
class FunctionDefinition(BaseModel):
"""Definition of a function that can be called by the model."""
model_config = ConfigDict(frozen=True, extra="allow")
name: str
"""The name of the function to be called.
Must be a-z, A-Z, 0-9, underscores and dashes, with a maximum
length of 64.
"""
description: str | None = None
"""A description of what the function does.
Used by the model to choose when and how to call the function.
"""
parameters: JSONSchema | None = None
"""The parameters the function accepts, as a JSON Schema object.
See https://json-schema.org/understanding-json-schema/ for the
format. To describe a function that accepts no parameters,
provide ``{"type": "object", "properties": {}}``.
"""
strict: bool | None = None
"""Whether to enable strict schema adherence in the generated call.
If true, the model will follow the exact schema provided in
``parameters``. Only supported by OpenAI gpt-4o and later;
ignored by other providers.
"""
class _TemplateFieldAccess:
"""Dict-style field access, for chat templates that subscript and call ``.get()``.
Gemma 4's template reads its inputs as mappings (``tool_call.get('function')``,
``function['arguments']``) rather than as objects, so a Pydantic model reaches
it as something with no ``get`` and no ``__getitem__`` and rendering fails
outright. ``Message`` already carries its own ``get`` for this reason; these are
the same accommodation for the tool-call types. Only declared fields are exposed,
so a key that collides with a method name returns the field or the default rather
than the bound method.
"""
def __getitem__(self, key: str) -> Any:
if key in type(self).model_fields: # type: ignore[attr-defined]
return getattr(self, key)
raise KeyError(key)
def get(self, key: str, default: Any = None) -> Any:
"""Returns a field by name, dict-style, or ``default`` if absent."""
try:
return self[key]
except KeyError:
return default
[docs]
class FunctionCall(_TemplateFieldAccess, BaseModel):
"""A function call made by the model."""
model_config = ConfigDict(frozen=True, extra="allow")
name: str
"""The name of the function being called."""
arguments: str
"""The arguments to call the function with, as a JSON string.
OpenAI wire format keeps this as an unparsed JSON-encoded string
(NOT a dict). Some providers return malformed JSON; downstream
code is responsible for parsing and handling errors.
"""
@pydantic.field_validator("arguments", mode="before")
@classmethod
def _coerce_arguments_to_str(cls, raw: Any) -> str:
"""Normalize dict/list arguments to a JSON string.
Runs in ``mode="before"`` so we can intercept non-string inputs
(e.g., dicts from JSONL files) and serialize them before Pydantic's
core validator enforces the ``str`` type.
"""
if isinstance(raw, str):
return raw
if isinstance(raw, (dict, list)):
return json.dumps(raw)
raise ValueError(
f"arguments must be a str, dict, or list; got {type(raw).__name__}."
)
[docs]
def __getitem__(self, key: str) -> Any:
"""As `_TemplateFieldAccess`, except that ``arguments`` decodes to its mapping.
Gemma 4's template reads ``function['arguments']`` and raises unless it gets
a mapping, so the JSON string stored here fails to render at all. Decoding on
subscript hands the template the object it asks for while ``.arguments`` keeps
the OpenAI wire string as the storage contract, and a dataset written with
mappings still normalizes to that string on the way in.
A template that subscripts ``arguments`` expecting the raw JSON string sees a
dict instead. Piping it through ``tojson`` produces the same text, and printing
it bare would produce a Python repr. Attribute access is unchanged either way.
"""
if key == "arguments":
return self.get_arguments_dict()
return super().__getitem__(key)
[docs]
def get_arguments_dict(self) -> dict[str, Any]:
"""Returns ``arguments`` decoded from its JSON string form.
The string is the storage contract; every consumer that needs a dict
decodes it here rather than reimplementing the parse and its error
handling.
"""
if not self.arguments:
return {}
try:
parsed = json.loads(self.arguments)
except json.JSONDecodeError as e:
raise ValueError(
f"Function '{self.name}' has invalid JSON arguments: {e}"
) from e
if not isinstance(parsed, dict):
raise ValueError(
f"Function '{self.name}' arguments must be a JSON object, "
f"got {type(parsed).__name__}."
)
return parsed
