src/agent_server/chat/tools.py - master - man4j/agent-server

src/agent_server/chat/tools.py
139 строк5 KB
Vladimir
fixes
21 май 2026, 11:09
21 май 2026, 11:09ae5cb53
import re

import chainlit as cl

from agent_server.chat.session import (
    ToolSessionState,
    apply_tool_session_state,
    get_tool_session_state,
)
from agent_server.chat.types import McpState, OpenAIToolSpec, ToolResultPayload
from agent_server.chat.profile_selection import (
    get_current_bot_name,
    get_selected_agent,
    get_selected_profile,
)
from agent_server.local_tools import (
    dispatch_local_tool,
    get_local_tools_by_names,
    is_local_tool_name,
    validate_local_tool_names,
)
from agent_server.mcp_client.client import (
    all_openai_tools_from_state,
    call_mcp_tool,
    connect_mcp_servers,
    disconnect_mcp_servers,
)
from agent_server.config.agents import AgentDefinition
from agent_server.config.profiles import LlmProfile


def render_plain_text_block(text: str) -> str:
    longest_backtick_run = max((len(match.group(0)) for match in re.finditer(r"`+", text)), default=0)
    fence = "`" * max(4, longest_backtick_run + 1)
    return f"{fence}markdown\n{text}\n{fence}"


def build_tool_registry(
    selected_profile: LlmProfile,
    selected_agent: AgentDefinition,
    mcp_state: McpState | None,
) -> list[OpenAIToolSpec]:
    validate_local_tool_names(selected_profile.local_tool_names)
    return all_openai_tools_from_state(mcp_state) + get_local_tools_by_names(
        selected_profile.local_tool_names
    )


def get_selected_agent_local_tools() -> list[OpenAIToolSpec]:
    return build_tool_registry(get_selected_profile(), get_selected_agent(), None)


def build_allowed_local_tool_names(selected_profile: LlmProfile) -> set[str]:
    return set(selected_profile.local_tool_names or ())


async def load_tool_session_state_for_profile(
    selected_profile: LlmProfile,
    selected_agent: AgentDefinition,
) -> ToolSessionState:
    mcp_state = await connect_mcp_servers(list(selected_profile.mcp_targets))
    return ToolSessionState(
        agent_id=selected_agent.id,
        profile_id=selected_profile.id,
        mcp_state=mcp_state,
        tools=build_tool_registry(selected_profile, selected_agent, mcp_state),
    )


async def send_mcp_errors(mcp_state: McpState) -> None:
    errors = mcp_state.get("errors", {})
    if not errors:
        return

    msg = cl.Message(
        author=get_current_bot_name(),
        content=(
            "Не удалось подключить часть MCP серверов:\n\n"
            + "\n".join(f"- `{name}`: {err}" for name, err in errors.items())
        ),
    )
    await msg.send()


async def send_loaded_context_step(system_prompt: str) -> None:
    async with cl.Step(name="Загруженный контекст", type="tool") as step:
        step.output = render_plain_text_block(system_prompt)


async def connect_selected_agent_tools() -> McpState:
    selected_profile = get_selected_profile()
    selected_agent = get_selected_agent()
    previous_state = get_tool_session_state().mcp_state
    tool_state = await load_tool_session_state_for_profile(selected_profile, selected_agent)
    apply_tool_session_state(tool_state)
    await disconnect_mcp_servers(previous_state)
    return tool_state.mcp_state or {}


async def ensure_mcp_ready() -> None:
    tool_state = get_tool_session_state()
    selected_profile = get_selected_profile()
    selected_agent = get_selected_agent()

    if (
        tool_state.mcp_state is None
        or tool_state.agent_id != selected_agent.id
        or tool_state.profile_id != selected_profile.id
    ):
        mcp_state = await connect_selected_agent_tools()
        await send_mcp_errors(mcp_state)
        return

    tool_state.tools = build_tool_registry(selected_profile, selected_agent, tool_state.mcp_state)
    apply_tool_session_state(tool_state)


async def shutdown_selected_agent_tools() -> None:
    tool_state = get_tool_session_state()
    await disconnect_mcp_servers(tool_state.mcp_state)


async def dispatch_tool_call(
    tool_name: str,
    tool_input: dict[str, object],
) -> ToolResultPayload:
    selected_profile = get_selected_profile()
    selected_agent = get_selected_agent()
    allowed_local_tool_names = build_allowed_local_tool_names(selected_profile)

    if tool_name in allowed_local_tool_names:
        local_result = await dispatch_local_tool(tool_name, tool_input)
        if local_result is not None:
            return local_result

    if is_local_tool_name(tool_name) and tool_name not in allowed_local_tool_names:
        return {"error": f"Local tool is not enabled for current agent: {tool_name}"}

    return await call_mcp_tool(get_tool_session_state().mcp_state, tool_name, tool_input)