agent.py

"""
Task runner agent - runs ML training pipelines with durable delivery.

Listens for intents via SSE. Simulates ML model training
and resumes with training results.

Usage:
    export AXME_API_KEY="<agent-key>"
    python agent.py
"""

import os
import sys
import time

sys.stdout.reconfigure(line_buffering=True)

from axme import AxmeClient, AxmeClientConfig


AGENT_ADDRESS = "hatchet-alt-runner-demo"


def handle_intent(client, intent_id):
    """Run ML training pipeline."""
    intent_data = client.get_intent(intent_id)
    intent = intent_data.get("intent", intent_data)
    payload = intent.get("payload", {})
    if "parent_payload" in payload:
        payload = payload["parent_payload"]

    task_id = payload.get("task_id", "unknown")
    task_type = payload.get("task_type", "unknown")
    model = payload.get("model", "unknown")
    dataset = payload.get("dataset", "unknown")
    epochs = payload.get("epochs", 10)

    print(f"  Task ID: {task_id}")
    print(f"  Type: {task_type}")
    print(f"  Model: {model}")
    print(f"  Dataset: {dataset}")

    # Step 1: Load dataset
    print(f"  [1/4] Loading dataset {dataset}...")
    time.sleep(1)

    # Step 2: Initialize model
    print(f"  [2/4] Initializing model {model}...")
    time.sleep(1)

    # Step 3: Train
    print(f"  [3/4] Training for {epochs} epochs...")
    time.sleep(2)

    # Step 4: Save artifacts
    print(f"  [4/4] Saving model artifacts...")
    time.sleep(1)

    result = {
        "action": "complete",
        "task_id": task_id,
        "model_accuracy": 0.9847,
        "training_time": "4m32s",
        "artifacts_url": f"https://storage.example.com/models/{model}/latest",
        "completed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }

    client.resume_intent(intent_id, result)
    print(f"  Training complete. Accuracy: {result['model_accuracy']}")
    print(f"  Artifacts: {result['artifacts_url']}")


def main():
    api_key = os.environ.get("AXME_API_KEY", "")
    if not api_key:
        print("Error: AXME_API_KEY not set.")
        print("Run the scenario first: axme scenarios apply scenario.json")
        print("Then get the agent key from ~/.config/axme/scenario-agents.json")
        sys.exit(1)

    client = AxmeClient(AxmeClientConfig(api_key=api_key))

    print(f"Agent listening on {AGENT_ADDRESS}...")
    print("Waiting for intents (Ctrl+C to stop)\n")

    for delivery in client.listen(AGENT_ADDRESS):
        intent_id = delivery.get("intent_id", "")
        status = delivery.get("status", "")

        if not intent_id:
            continue

        if status in ("DELIVERED", "CREATED", "IN_PROGRESS"):
            print(f"[{status}] Intent received: {intent_id}")
            try:
                handle_intent(client, intent_id)
            except Exception as e:
                print(f"  Error processing intent: {e}")


if __name__ == "__main__":
    main()