src/agent_server/local_tools/__init__.py - master - man4j/agent-server

src/agent_server/local_tools/__init__.py
135 строк4 KB
Vladimir
initial
18 апр 2026, 19:36
18 апр 2026, 19:36a95f341
from __future__ import annotations

from typing import Any

from agent_server.chat.types import ImageListToolResult, OpenAIToolSpec, ToolResultPayload
from .analytics_recipes import RECIPES_TOOL_MODULE
from .base import LocalToolModule
from .diff_tool import DIFF_TOOL_MODULE
from .plotly_tool import PLOTLY_TOOL_MODULE

import chainlit as cl

LOCAL_TOOL_MODULES: tuple[LocalToolModule, ...] = (
    PLOTLY_TOOL_MODULE,
    DIFF_TOOL_MODULE,
    RECIPES_TOOL_MODULE,
)

LOCAL_TOOL_REGISTRY: dict[str, OpenAIToolSpec] = {}
LOCAL_TOOL_TO_MODULE: dict[str, LocalToolModule] = {}

for module in LOCAL_TOOL_MODULES:
    for tool in module.tools:
        name = str(tool.get("function", {}).get("name") or "").strip()
        if not name:
            continue

        if name in LOCAL_TOOL_REGISTRY:
            raise ValueError(f"Duplicate local tool name registered: {name}")

        LOCAL_TOOL_REGISTRY[name] = tool
        LOCAL_TOOL_TO_MODULE[name] = module

LOCAL_TOOLS: list[OpenAIToolSpec] = list(LOCAL_TOOL_REGISTRY.values())


def get_local_tools_by_names(tool_names: list[str] | tuple[str, ...]) -> list[OpenAIToolSpec]:
    result: list[OpenAIToolSpec] = []

    for raw_name in tool_names:
        name = str(raw_name or "").strip()
        if not name:
            continue

        tool = LOCAL_TOOL_REGISTRY.get(name)
        if tool is not None:
            result.append(tool)

    return result


def get_local_tool_names() -> set[str]:
    return set(LOCAL_TOOL_REGISTRY.keys())


def is_local_tool_name(tool_name: str) -> bool:
    return str(tool_name or "").strip() in LOCAL_TOOL_REGISTRY


def validate_local_tool_names(tool_names: list[str] | tuple[str, ...]) -> None:
    unknown = [
        str(raw_name).strip()
        for raw_name in tool_names
        if str(raw_name or "").strip() and str(raw_name).strip() not in LOCAL_TOOL_REGISTRY
    ]
    if unknown:
        raise ValueError(f"Unknown local tools: {', '.join(sorted(set(unknown)))}")


async def dispatch_local_tool(
    tool_name: str,
    tool_input: dict[str, Any],
) -> ToolResultPayload | None:
    name = str(tool_name or "").strip()
    if not name:
        return None

    module = LOCAL_TOOL_TO_MODULE.get(name)
    if module is None:
        return None

    return await module.dispatcher(name, tool_input)


def build_local_ui_element(result: ToolResultPayload, index: int):
    for module in LOCAL_TOOL_MODULES:
        if module.ui_builder is None:
            continue
        element = module.ui_builder(result, index)
        if element is not None:
            return element
    return None


def build_local_ui_elements(result: ToolResultPayload, start_index: int) -> list:
    if result.get("ui_type") == "image_list":
        image_result = result
        built = []

        elements = image_result.get("elements")
        if isinstance(elements, list):
            for offset, item in enumerate(elements):
                url = item.get("url")
                if not isinstance(url, str) or not url.strip():
                    continue

                built.append(
                    cl.Image(
                        url=url,
                        name=item.get("name") or f"image-{start_index + offset}",
                        display=item.get("display") or "inline",
                    )
                )

        if built:
            return built

        images = image_result.get("images")
        if isinstance(images, list):
            for offset, url in enumerate(images):
                if not isinstance(url, str) or not url.strip():
                    continue

                built.append(
                    cl.Image(
                        url=url,
                        name=f"image-{start_index + offset}",
                        display="inline",
                    )
                )

        return built

    single = build_local_ui_element(result, start_index)
    return [single] if single is not None else []