AI Agent Service

from cmdop import AsyncCMDOPClient
from pydantic import BaseModel
 
# Define a Pydantic model for the expected structured response
class TaskResult(BaseModel):
    success: bool
    output: str
    error: str | None
 
# Connect to a remote machine using your API key
async with AsyncCMDOPClient.remote(api_key="cmd_xxx") as client:
    # Select the target machine for command execution
    await client.terminal.set_machine("my-server")
 
    # Run a natural-language prompt and get a typed result back
    result = await client.agent.run(
        prompt="Check if nginx is running",
        output_schema=TaskResult  # AI response is validated against this schema
    )
 
    print(f"Success: {result.output.success}")

from pydantic import BaseModel, Field
 
# Define a detailed schema with Field descriptions to guide the AI's output
class ServerHealth(BaseModel):
    hostname: str
    cpu_percent: float = Field(description="CPU usage percentage")
    memory_percent: float = Field(description="Memory usage percentage")
    disk_percent: float = Field(description="Disk usage percentage")
    services_running: list[str]
    issues: list[str]
 
# The AI runs commands, collects data, and structures the response
result = await client.agent.run(
    prompt="Check server health and identify issues",
    output_schema=ServerHealth
)
 
# result.output is a fully typed ServerHealth instance
health: ServerHealth = result.output
 
# Use typed fields directly for conditional logic
if health.cpu_percent > 90:
    alert(f"High CPU on {health.hostname}")
 
# Iterate over structured issue data from the AI's analysis
for issue in health.issues:
    create_ticket(issue)

# Configure the agent run with fine-grained options
result = await client.agent.run(
    prompt="Deploy version 2.0",
    output_schema=DeployResult,
    timeout=300,              # Timeout in seconds (5 minutes for long tasks)
    max_tokens=4096,          # Max response tokens from the AI model
    temperature=0.1,          # Lower = more deterministic and consistent results
    context={                 # Additional context passed to the AI for decision-making
        "version": "2.0",
        "environment": "production"
    }
)

# Stream events as the AI thinks, executes commands, and produces results
async for event in client.agent.run_stream(
    prompt="Analyze and fix the issue",
    output_schema=FixResult
):
    if event.type == "thinking":
        print(f"AI: {event.content}")       # AI's reasoning process
    elif event.type == "tool_call":
        print(f"Running: {event.tool}")      # Command the AI is about to execute
    elif event.type == "tool_result":
        print(f"Output: {event.result[:100]}...")  # Truncated command output
    elif event.type == "result":
        fix_result = event.output            # Final structured result (FixResult)

# Terminal agent (default) — executes shell commands on the machine
result = await client.agent.run(
    prompt="Check disk usage",
    output_schema=DiskInfo,
    agent_type="terminal"
)
 
# Planner agent — creates step-by-step plans without executing them
result = await client.agent.run(
    prompt="Plan a migration from MySQL to PostgreSQL",
    output_schema=MigrationPlan,
    agent_type="planner"
)
 
# Router agent — decides which service or handler should process the request
result = await client.agent.run(
    prompt="Handle this support request",
    output_schema=RoutingDecision,
    agent_type="router"
)

# Apply safety restrictions to limit what the AI agent can do
result = await client.agent.run(
    prompt="Clean up old files",
    output_schema=CleanupResult,
    restrictions={
        "no_delete": True,     # Prevent the agent from deleting any files
        "no_sudo": True,       # Block all root/sudo commands
        "dry_run": True,       # Plan only — no commands are actually executed
        "allowed_dirs": ["/tmp", "/var/cache"],  # Restrict to these directories
        "forbidden_commands": ["rm -rf", "shutdown"]  # Explicitly block dangerous commands
    }
)

# Pass structured context so the AI has relevant information before running commands
result = await client.agent.run(
    prompt="Fix the memory issue",
    output_schema=FixResult,
    context={
        "service": "api-server",        # Which service is affected
        "memory_limit": "4GB",          # Configured memory ceiling
        "current_usage": "3.8GB",       # Current memory consumption
        "recent_events": [              # Timeline of recent activity for diagnosis
            "Deploy v2.0.5 (2 hours ago)",
            "Memory spike detected (1 hour ago)"
        ]
    }
)

# Create a conversation to maintain context across multiple turns
conversation = client.agent.conversation()
 
# First turn — the AI discovers what services exist
result1 = await conversation.run(
    prompt="What services are running?",
    output_schema=ServiceList
)
 
# Follow-up — the AI remembers previous results (context preserved automatically)
result2 = await conversation.run(
    prompt="Which one is using the most memory?",
    output_schema=ServiceInfo
)
 
# Another follow-up — "it" refers to the service from the previous turn
result3 = await conversation.run(
    prompt="Restart it",
    output_schema=RestartResult
)

# Register a custom tool that the AI agent can invoke during execution
@client.agent.tool
async def check_database(query: str) -> dict:
    """Run database query and return results."""
    # Your implementation — the AI calls this function when it needs DB data
    return {"rows": [...]}
 
# Pass the tool name so the AI knows it can use check_database
result = await client.agent.run(
    prompt="Check if user exists in database",
    output_schema=UserCheckResult,
    tools=["check_database"]  # Make this custom tool available to the AI
)

# Import specific exception types for granular error handling
from cmdop.exceptions import (
    AgentError,
    SchemaValidationError,
    TimeoutError
)
 
try:
    result = await client.agent.run(
        prompt="Task",
        output_schema=Result
    )
except SchemaValidationError as e:
    print(f"AI returned invalid data: {e}")  # Output didn't match the Pydantic schema
except TimeoutError:
    print("Operation timed out")              # Agent exceeded the timeout limit
except AgentError as e:
    print(f"Agent error: {e}")                # Catch-all for other agent failures

# Step 1: Run in dry_run mode to get a plan without executing anything
plan = await client.agent.run(
    prompt="Plan database migration",
    output_schema=MigrationPlan,
    restrictions={"dry_run": True}  # No commands are executed
)
 
# Step 2: Present the plan to the user for review
print("Planned actions:")
for step in plan.output.steps:
    print(f"  - {step}")
 
# Step 3: Only execute after explicit user approval
if user_approves():
    # Execute the approved plan — dry_run is now omitted
    result = await client.agent.run(
        prompt=f"Execute this plan: {plan.output.steps}",
        output_schema=MigrationResult
    )

# Good — specific instructions with clear output expectations
prompt = """
Check CPU usage for the last minute.
If above 90%, identify the top process.
Report hostname, CPU percent, and top process name.
"""
 
# Bad — vague prompts lead to unpredictable output structure
prompt = "Check CPU"

# Field descriptions guide the AI on what data to collect and how to format it
class Result(BaseModel):
    latency_ms: float = Field(description="P99 latency in milliseconds")
    success_rate: float = Field(description="Percentage of successful requests")

try:
    result = await client.agent.run(...)
except SchemaValidationError:
    # Fall back to unstructured text if the AI can't match the schema
    result = await client.agent.run(prompt, output_schema=None)
    print(result.text)  # Raw text response without Pydantic validation

# First pass: generate the plan without executing any commands
plan = await client.agent.run(..., restrictions={"dry_run": True})
 
# Validate the plan before committing to execution
if looks_good(plan.output):
    # Second pass: execute with dry_run disabled
    result = await client.agent.run(..., restrictions={"dry_run": False})