agr-assistent/src/agr_assistent/llm/tools.py

"""Инструменты, которые модель может вызывать: описание, проверка аргументов, выполнение."""

from __future__ import annotations

import json
import logging
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from typing import Any

log = logging.getLogger(__name__)


class ToolError(Exception):
    """Ожидаемая ошибка инструмента: текст уходит модели и в журнал."""


@dataclass(frozen=True)
class ToolResult:
    ok: bool
    content: str  # что увидит модель
    display: str  # что увидит пользователь в журнале


@dataclass(frozen=True)
class Tool:
    name: str
    description: str
    parameters: dict[str, Any]  # JSON Schema объекта аргументов
    handler: Callable[[dict[str, Any]], ToolResult]
    # Выполнять только после явного «да» пользователя
    confirm: bool = False
    # Человекочитаемое описание вызова для вопроса-подтверждения и журнала
    describe: Callable[[dict[str, Any]], str] | None = None

    def summary(self, arguments: dict[str, Any]) -> str:
        return self.describe(arguments) if self.describe else self.description

    def schema(self) -> dict[str, Any]:
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


class ToolRegistry:
    def __init__(self, tools: Iterable[Tool] = ()) -> None:
        self._tools: dict[str, Tool] = {}
        for tool in tools:
            self.register(tool)

    def __len__(self) -> int:
        return len(self._tools)

    def register(self, tool: Tool) -> None:
        self._tools[tool.name] = tool

    def schemas(self) -> list[dict[str, Any]]:
        return [tool.schema() for tool in self._tools.values()]

    def get(self, name: str) -> Tool | None:
        return self._tools.get(name)

    def execute(self, name: str, arguments_json: str) -> ToolResult:
        prepared = self.prepare(name, arguments_json)
        if isinstance(prepared, ToolResult):
            return prepared
        return self.run(*prepared)

    def prepare(self, name: str, arguments_json: str) -> tuple[Tool, dict[str, Any]] | ToolResult:
        """Находит инструмент и проверяет аргументы; ToolResult — описание ошибки."""
        tool = self._tools.get(name)
        if tool is None:
            return ToolResult(False, f"Инструмента {name} не существует", f"Неизвестный инструмент {name}")
        try:
            raw_arguments = json.loads(arguments_json or "{}")
        except json.JSONDecodeError as exc:
            return ToolResult(False, f"Аргументы не являются JSON: {exc}", f"{name}: некорректные аргументы")
        if not isinstance(raw_arguments, dict):
            return ToolResult(False, "Аргументы должны быть JSON-объектом", f"{name}: некорректные аргументы")

        arguments, errors = validate_arguments(tool.parameters, raw_arguments)
        if errors:
            message = "; ".join(errors)
            return ToolResult(False, f"Некорректные аргументы: {message}", f"{name}: {message}")
        return tool, arguments

    @staticmethod
    def run(tool: Tool, arguments: dict[str, Any]) -> ToolResult:
        name = tool.name
        try:
            return tool.handler(arguments)
        except ToolError as exc:
            return ToolResult(False, str(exc), str(exc))
        except Exception as exc:
            log.exception("Сбой инструмента %s", name)
            return ToolResult(False, f"Ошибка выполнения: {exc}", f"{name}: ошибка — {exc}")


def validate_arguments(
    schema: dict[str, Any], arguments: dict[str, Any]
) -> tuple[dict[str, Any], list[str]]:
    """Проверяет аргументы по упрощённой JSON Schema.

    Небольшие локальные модели часто присылают числа и логические значения строками,
    поэтому такие значения приводятся к нужному типу.
    """
    properties: dict[str, Any] = schema.get("properties", {})
    errors = [
        f"не указан параметр {name}" for name in schema.get("required", []) if name not in arguments
    ]
    result: dict[str, Any] = {}
    for name, value in arguments.items():
        spec = properties.get(name)
        if spec is None:
            continue  # лишние параметры молча игнорируем
        converted, error = _convert(value, spec.get("type"))
        if error:
            errors.append(f"{name}: {error}")
            continue
        if "enum" in spec and converted not in spec["enum"]:
            allowed = ", ".join(str(option) for option in spec["enum"])
            errors.append(f"{name}: недопустимое значение {converted!r} (допустимы: {allowed})")
            continue
        if "minimum" in spec and converted < spec["minimum"]:
            errors.append(f"{name}: значение меньше {spec['minimum']}")
            continue
        if "maximum" in spec and converted > spec["maximum"]:
            errors.append(f"{name}: значение больше {spec['maximum']}")
            continue
        result[name] = converted
    return result, errors


def _convert(value: Any, expected: str | None) -> tuple[Any, str | None]:
    if expected == "string":
        if isinstance(value, (dict, list)):
            return None, "ожидалась строка"
        return str(value), None
    if expected == "integer":
        if isinstance(value, bool):
            return None, "ожидалось целое число"
        if isinstance(value, int):
            return value, None
        if isinstance(value, float) and value.is_integer():
            return int(value), None
        if isinstance(value, str):
            try:
                return int(value.strip()), None
            except ValueError:
                pass
        return None, "ожидалось целое число"
    if expected == "number":
        if isinstance(value, bool):
            return None, "ожидалось число"
        if isinstance(value, (int, float)):
            return value, None
        if isinstance(value, str):
            try:
                return float(value.strip().replace(",", ".")), None
            except ValueError:
                pass
        return None, "ожидалось число"
    if expected == "boolean":
        if isinstance(value, bool):
            return value, None
        if isinstance(value, str) and value.strip().lower() in ("true", "false"):
            return value.strip().lower() == "true", None
        return None, "ожидалось true или false"
    return value, None