from pydantic_ai import Agent

from ..config import RegulatorConfig
from .agent import RegulatedAgent
from .mcp import RegulatedMCPAdapter
from .toolset import RegulatedToolset


def regulate(
    agent: Agent,
    *,
    agent_id: str,
    config: RegulatorConfig | None = None,
    wrap_agent_toolset: bool = True,
    govern_builtin_tools: bool = False,
) -> RegulatedAgent:
    return RegulatedAgent(
        agent,
        agent_id=agent_id,
        config=config,
        wrap_agent_toolset=wrap_agent_toolset,
        govern_builtin_tools=govern_builtin_tools,
    )


__all__ = ["RegulatedAgent", "RegulatedMCPAdapter", "RegulatedToolset", "regulate"]
