omlx/solar_open2_tool_parser.py

import ast
import json
import re
from typing import Any


tool_call_start = "<|tool_call:start|>"
tool_call_end = "<|tool_call:end|>"

_ARG_RE = re.compile(
    r"<\|tool_arg:start\|>(.*?)<\|tool_arg:value\|>(.*?)<\|tool_arg:end\|>",
    re.DOTALL,
)

_STRING_TYPES = {"string", "str", "text", "varchar", "char", "enum"}
_BOOL_TYPES = {"boolean", "bool", "binary"}
_INT_PREFIXES = ("int", "uint", "long", "short", "unsigned")
_FLOAT_PREFIXES = ("num", "float", "double", "decimal")
_JSON_TYPES = {"object", "array", "arr", "dict", "list"}


def _tool_properties(name: str, tools: Any | None) -> dict[str, Any]:
    if not tools:
        return {}
    for tool in tools:
        function = tool.get("function") if isinstance(tool, dict) else None
        if function and function.get("name") == name:
            params = function.get("parameters") or {}
            return params.get("properties") or {}
    return {}


def _schema_type(properties: dict[str, Any], key: str) -> str:
    schema = properties.get(key) or {}
    typ = schema.get("type", "string")
    if isinstance(typ, list):
        typ = next((item for item in typ if item != "null"), typ[0] if typ else "string")
    return str(typ).lower()


def _coerce(value: str, key: str, properties: dict[str, Any]) -> Any:
    value = value.strip()
    typ = _schema_type(properties, key)

    if value.lower() == "null":
        return None
    if typ in _STRING_TYPES:
        return value
    if typ in _BOOL_TYPES:
        normalized = value.lower()
        if normalized in {"1", "true", "yes"}:
            return True
        if normalized in {"0", "false", "no"}:
            return False
        return value
    if typ.startswith(_INT_PREFIXES):
        try:
            return int(value)
        except (TypeError, ValueError):
            return value
    if typ.startswith(_FLOAT_PREFIXES):
        try:
            number = float(value)
            return int(number) if number.is_integer() else number
        except (TypeError, ValueError):
            return value
    if typ in _JSON_TYPES or typ.startswith(("dict", "list")):
        try:
            return json.loads(value)
        except (TypeError, ValueError, json.JSONDecodeError):
            try:
                return ast.literal_eval(value)
            except (SyntaxError, ValueError):
                return value

    try:
        return json.loads(value)
    except (TypeError, json.JSONDecodeError):
        return value


def parse_tool_call(text: str, tools: Any | None = None) -> dict[str, Any]:
    body = text.strip()
    if body.startswith(tool_call_start):
        body = body[len(tool_call_start) :].strip()
    if body.endswith(tool_call_end):
        body = body[: -len(tool_call_end)].strip()

    first_arg = body.find("<|tool_arg:start|>")
    if first_arg == -1:
        try:
            payload = json.loads(body)
        except json.JSONDecodeError as exc:
            raise ValueError(f"No Solar tool arguments found in: {text!r}") from exc
        name = payload.get("name") or payload.get("tool") or payload.get("function")
        args = payload.get("arguments") or payload.get("args") or {}
        if not name:
            raise ValueError(f"No Solar tool name found in: {text!r}")
        if isinstance(args, str):
            args = json.loads(args)
        return {"name": name, "arguments": args}

    name = body[:first_arg].strip().splitlines()[0].strip()
    if not name:
        raise ValueError(f"No Solar tool name found in: {text!r}")

    properties = _tool_properties(name, tools)
    arguments: dict[str, Any] = {}
    for key, value in _ARG_RE.findall(body[first_arg:]):
        key = key.strip()
        if key:
            arguments[key] = _coerce(value, key, properties)

    return {"name": name, "arguments": arguments}