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

src/agent_server/local_tools/plotly_tool.py
72 строки2 KB
Vladimir
initial
18 апр 2026, 19:36
18 апр 2026, 19:36a95f341
from typing import Any

import chainlit as cl
import plotly.graph_objects as go

from agent_server.chat.types import PlotlyToolResult, ToolResultPayload
from .base import LocalToolModule

LOCAL_PLOTLY_TOOL = {
    "type": "function",
    "function": {
        "name": "render_plotly_chart",
        "description": (
            "Показывает интерактивный график Plotly в интерфейсе. "
            "Передавай готовую Plotly figure spec в формате {data, layout[, frames]}."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "figure": {"type": "object"},
                "title": {"type": "string"},
            },
            "required": ["figure"],
        },
    },
}


async def dispatch_plotly_tool(
    tool_name: str,
    tool_input: dict[str, Any],
) -> PlotlyToolResult | dict[str, Any] | None:
    return await render_plotly_chart(tool_input)


async def render_plotly_chart(tool_input: dict[str, Any]) -> PlotlyToolResult | dict[str, Any]:
    figure_spec = tool_input.get("figure")
    title = tool_input.get("title") or "График"

    if not isinstance(figure_spec, dict):
        return {"error": "field 'figure' must be an object with Plotly figure spec"}

    try:
        go.Figure(figure_spec)
    except Exception as e:
        return {"error": f"Invalid Plotly figure spec: {type(e).__name__}: {e}"}

    return {
        "ok": True,
        "ui_type": "plotly",
        "title": title,
        "figure": figure_spec,
    }


def build_plotly_element(result: ToolResultPayload, index: int):
    if not result.get("ok") or result.get("ui_type") != "plotly":
        return None

    fig = go.Figure(result["figure"])
    return cl.Plotly(
        name=f"plotly_chart_{index}",
        figure=fig,
        display="inline",
    )


PLOTLY_TOOL_MODULE = LocalToolModule(
    tools=[LOCAL_PLOTLY_TOOL],
    dispatcher=dispatch_plotly_tool,
    ui_builder=build_plotly_element,
)