Building a Custom Agent Image

📓 Try it yourself! This example is available as an executable Jupyter notebook.

This example walks through creating a custom Pydantic AI agent with custom tools, packaging it as a Docker image, and deploying it to KAOS.

Prerequisites

• KAOS operator installed (Installation Guide)
• kaos-cli installed (pip install kaos-cli)
• Docker available for building images
• kubectl configured to your cluster

Setup

import os
os.environ['NAMESPACE'] = 'custom-agent-example'
REPO_ROOT = os.path.abspath("../../")
os.environ['REPO_ROOT'] = REPO_ROOT

# Base agent image — set KAOS_AGENT_IMAGE for local/CI builds
AGENT_IMAGE = os.environ.get("KAOS_AGENT_IMAGE", "axsauze/kaos-agent:latest")
os.environ['KAOS_AGENT_IMAGE'] = AGENT_IMAGE

kubectl create namespace $NAMESPACE 2>/dev/null || true
kubectl config set-context --current --namespace=$NAMESPACE

Step 1: Initialize the Agent Project

Use the KAOS CLI to scaffold a new custom agent project:

!kaos agent init custom-math-agent

This creates agent.py, pyproject.toml, and README.md with a template Pydantic AI agent.

Step 2: Customize the Agent

Replace the template agent with a math agent that has custom tools:

%%writefile custom-math-agent/agent.py
"""Custom Agent — Pydantic AI agent with math tools."""

import random
from pydantic_ai import Agent

agent = Agent(
    model="test",  # Overridden by KAOS env vars at runtime
    instructions="You are a helpful math and utility assistant.",
    name="custom-agent",
    defer_model_check=True,
)


@agent.tool_plain
def add(a: float, b: float) -> str:
    """Add two numbers together.

    Args:
        a: First number
        b: Second number

    Returns:
        The sum as a string
    """
    return str(a + b)


@agent.tool_plain
def multiply(a: float, b: float) -> str:
    """Multiply two numbers.

    Args:
        a: First number
        b: Second number

    Returns:
        The product as a string
    """
    return str(a * b)


@agent.tool_plain
def random_number(min_val: int = 1, max_val: int = 100) -> str:
    """Generate a random number in a range.

    Args:
        min_val: Minimum value (inclusive)
        max_val: Maximum value (inclusive)

    Returns:
        A random integer as a string
    """
    return str(random.randint(min_val, max_val))

Step 3: Build the Docker Image

Build the container image using the CLI:

!cd custom-math-agent && kaos agent build --image custom-agent:test --create-dockerfile --base-image $KAOS_AGENT_IMAGE

For KIND clusters, load the image directly:

!cd custom-math-agent && kaos agent build --image custom-agent:test --create-dockerfile --base-image $KAOS_AGENT_IMAGE --kind-load

Step 4: Create a ModelAPI

Create a ModelAPI in Proxy mode (we'll use mock responses so no real LLM needed):

kaos modelapi deploy custom-api --mode Proxy --wait

Step 5: Deploy the Custom Agent

Deploy the agent with custom image and mock responses that exercise the add tool:

import json, subprocess, os

mock1 = json.dumps({"tool_calls": [{"id": "call_1", "name": "add", "arguments": {"a": 5, "b": 3}}]})
mock2 = "The result of 5 + 3 is 8."
mock_responses = json.dumps([mock1, mock2])

namespace = os.environ["NAMESPACE"]
result = subprocess.run([
    "kaos", "agent", "deploy", "custom-math-agent",
    "--modelapi", "custom-api",
    "--model", "mock-model",
    "--image", "custom-agent:test",
    "--description", "Custom math agent with add, multiply, and random tools",
    "--instructions", "You are a helpful math and utility assistant.",
    "--env", "AGENT_LOG_LEVEL=DEBUG",
    "--env", f"DEBUG_MOCK_RESPONSES={mock_responses}",
    "-n", namespace,
], capture_output=True, text=True)
print(result.stdout or result.stderr)
assert result.returncode == 0, f"Deploy failed: {result.stderr}"

Wait for the agent to be ready:

import subprocess, time

for i in range(60):
    result = subprocess.run(
        ["kubectl", "get", "agent/custom-math-agent", "-o", "jsonpath={.status.phase}"],
        capture_output=True, text=True
    )
    if result.stdout.strip() == "Ready":
        print(f"Agent ready after ~{i*2}s")
        break
    time.sleep(2)
else:
    raise RuntimeError("Agent did not become Ready within 120s")

Step 6: Test the Agent

Verify the agent card shows custom tools:

import httpx
import subprocess
import json

# Get the Gateway URL
gateway_url = os.environ.get("GATEWAY_URL", "http://localhost:8888")
namespace = os.environ["NAMESPACE"]
agent_url = f"{gateway_url}/{namespace}/agent/custom-math-agent"

# Wait for agent to be accessible
import time
for _ in range(30):
    try:
        r = httpx.get(f"{agent_url}/health", timeout=2.0)
        if r.status_code == 200:
            break
    except Exception:
        pass
    time.sleep(1)

# Check agent card
response = httpx.get(f"{agent_url}/.well-known/agent.json", timeout=10.0)
assert response.status_code == 200, f"Agent card failed: {response.status_code}"
card = response.json()
skill_names = [s.get("name") for s in card.get("skills", [])]
print(f"Agent skills: {skill_names}")
assert "add" in skill_names, f"add not in skills: {skill_names}"
assert "multiply" in skill_names, f"multiply not in skills: {skill_names}"
assert "random_number" in skill_names, f"random_number not in skills: {skill_names}"
print("SUCCESS: Custom tools discovered!")

Now invoke the agent:

response = httpx.post(
    f"{agent_url}/v1/chat/completions",
    json={
        "model": "custom-math-agent",
        "messages": [{"role": "user", "content": "Add 5 and 3"}],
    },
    timeout=30.0,
)
assert response.status_code == 200, f"Chat failed: {response.text}"
data = response.json()
content = data["choices"][0]["message"]["content"]
print(f"Agent response: {content}")
assert len(content) > 0, "Empty response"
print("SUCCESS: Custom agent responded!")

Verify memory has tool events:

response = httpx.get(f"{agent_url}/memory/events", timeout=10.0)
memory = response.json()
event_types = [e["event_type"] for e in memory["events"]]
print(f"Memory event types: {event_types}")
assert "tool_call" in event_types, f"Missing tool_call in {event_types}"
assert "tool_result" in event_types, f"Missing tool_result in {event_types}"
print("SUCCESS: Tool events recorded in memory!")

What You Get

Custom agent images automatically include:

• Health/Ready probes — GET /health, GET /ready
• A2A agent card — GET /.well-known/agent.json with custom tool discovery
• Memory endpoints — GET /memory/events, GET /memory/sessions
• OpenAI-compatible API — POST /v1/chat/completions
• Session management — X-Session-ID header support
• OpenTelemetry — set OTEL_ENABLED=true in the CRD env

Cleanup

kubectl delete namespace $NAMESPACE --wait=false

# Clean up local files
import shutil, os
if os.path.exists("custom-math-agent"):
    shutil.rmtree("custom-math-agent")

Next Steps

• Custom MCP Server — Build custom MCP tool servers
• Multi-Agent Telemetry — Add observability
• Agent CRD Reference — Full CRD documentation