Add a Tool to AI-OPS¶
You can extend the AI-OPS agent by subclassing the generic Tool, registering it in the application, and adding it to the tools list in AgentConfig when you construct AgentRunner.
Extension model: this only works when using
ai_ops.coredirectly in your own code. The API server doesn't currently expose a way to register custom tools at runtime (the only way to add a tool there is by contributing it to the project). MCP-based tool extension is not currently planned.
Basic Tool¶
Register the tool at module load time, before you build AgentConfig (register_tool populates a global registry, so if it runs after AgentRunner is constructed, the tool won't be found).
from typing import Annotated
from pydantic import BaseModel, Field
from ai_ops.core.tools import Tool, register_tool
from ai_ops.core.runner import AgentConfig
class MyInput(BaseModel):
var: Annotated[int, Field(description="The value to double")]
class MyOutput(BaseModel):
result: int
class MyTool(Tool[MyInput, MyOutput]):
name = "my_tool"
description = "Doubles the given integer."
# execution logic
def __call__(self, tool_args: MyInput) -> MyOutput:
return MyOutput(result=tool_args.var * 2)
# converting to text for LLMs
@staticmethod
def format_result(tool_result: MyOutput) -> str:
return f"Result: {tool_result.result}"
register_tool(MyTool, lambda _: MyTool())
AgentConfig(tools=[MyTool, ...])
Tool State¶
register_tool(tool_class, factory) takes the class plus a factory lambda, not an instance, so AgentRunner can build the tool without knowing its constructor signature. The input/output models are read from the Tool[In, Out] base, so you don't pass them to register_tool separately. If your tool needs configuration or shared state, receive it through ToolContext:
ai_ops.core.tools.__init__
@dataclass
class ToolContext:
session_id: str
model_id: str | None = None
is_new_conversation: bool = True
command_policies: tuple[CommandAdmissionPolicy] = field(default_factory=list)
# anything that doesn't deserve a first-class field
extra: dict[str, Any] | None = None
from typing import Annotated
from pydantic import BaseModel, Field
from ai_ops.core.tools import Tool, ToolContext, register_tool
from ai_ops.core.runner import AgentConfig
class MyInput(BaseModel):
var: Annotated[int, Field(description="The value to multiply")]
class MyOutput(BaseModel):
result: int
class MyTool(Tool[MyInput, MyOutput]):
name = "my_tool"
description = "Multiplies the given integer by a configured factor."
def __init__(self, session_id: str, multiplier: int):
self.session_id = session_id
self.multiplier = multiplier
# execution logic
def __call__(self, tool_args: MyInput) -> MyOutput:
return MyOutput(result=tool_args.var * self.multiplier)
# converting to text for LLMs
@staticmethod
def format_result(tool_result: MyOutput) -> str:
return f"Result: {tool_result.result}"
register_tool(
MyTool,
lambda ctx: MyTool(session_id=ctx.session_id, multiplier=ctx.extra["multiplier"]),
)
AgentConfig(tools=[MyTool, ...])
Human-in-the-Loop Confirmation¶
Tools that perform sensitive operations should require user confirmation when the agent runs in supervised mode. A tool that requires confirmation:
- sets
requires_confirmation = True - implements
evaluate(tool_args) -> bool, returningTruewhen this specific call needs confirmation (or can't run at all in unsupervised mode) - implements
not_admitted_result(tool_args), building the result reported when the call is denied or times out
from typing import Annotated
from pydantic import BaseModel, Field
from ai_ops.core.tools import Tool, register_tool
from ai_ops.core.runner import AgentConfig
class MyInput(BaseModel):
var: Annotated[int, Field(description="The value to double")]
class MyOutput(BaseModel):
result: int
error: str | None = None
class MyTool(Tool[MyInput, MyOutput]):
name = "my_tool"
description = "Doubles the given integer; requires confirmation above 10."
requires_confirmation = True
def evaluate(self, tool_args: MyInput) -> bool:
return tool_args.var > 10
def not_admitted_result(self, tool_args: MyInput) -> MyOutput:
return MyOutput(result=-1, error="var can't be more than 10 without confirmation")
# execution logic (always assumes the call was already admitted)
def __call__(self, tool_args: MyInput) -> MyOutput:
return MyOutput(result=tool_args.var * 2)
# converting to text for LLMs
@staticmethod
def format_result(tool_result: MyOutput) -> str:
return f"Result: {tool_result.result}"
register_tool(MyTool, lambda _: MyTool())
AgentConfig(tools=[MyTool, ...])
evaluate and not_admitted_result are only ever consulted when requires_confirmation = True; a denied or timed-out call still flows back to the model as a normal tool result via not_admitted_result, so the conversation stays well-formed.