from pydantic_ai import Agent, FunctionToolset, RunContext
from pydantic_ai_backends import ConsoleDeps, create_console_toolset

from app.agent.deps import Deps


def configure_console_toolset(
    agent: Agent[Deps, str],
    enable_execute: bool = False,
) -> Agent[Deps, str]:
    console_toolset = create_console_toolset(
        id="deep-console",
        include_execute=enable_execute,
        require_write_approval=False,
        require_execute_approval=False,
        max_retries=2,
        image_support=True,
    )

    @agent.toolset
    def _console_toolset(_ctx: RunContext[Deps]) -> FunctionToolset[ConsoleDeps]:
        return console_toolset

    return agent
