The amc_sdk Python package

pip install "./amc_sdk[serve]"

import os

from fastapi import HTTPException, Request

from amc_sdk import Agent, AgentType, HealthReport, InputParam, RunContext, StepDef, build_app

API_KEY = os.environ.get("AGENT_API_KEY", "")


async def verify_caller(request: Request) -> None:
    if API_KEY and request.headers.get("api-key") != API_KEY:
        raise HTTPException(status_code=401, detail="unauthorized")


def health() -> HealthReport:
    if not os.environ.get("MODEL_API_KEY"):
        return HealthReport(ready=False, detail="MODEL_API_KEY is not set")
    return HealthReport(ready=True)


async def summarize(invoice_id: str) -> tuple[str, dict]:
    # Call your model here. Return the text and the usage it reported.
    return f"Summary of {invoice_id}", {"input": 4200, "output": 610, "cached": 3000}


agent = Agent(
    workflow_id="invoice-summarizer",
    name="Invoice Summarizer",
    description="Summarizes one invoice.",
    type="workflow",
    steps=[
        StepDef("fetch", "Fetch invoice"),
        StepDef("summarize", "Summarize"),
    ],
    agent_types=[AgentType("invoice-summarizer", "Invoice Summarizer", ["text"])],
    input_parameters=[
        InputParam("invoiceId", type="string", required=True, description="Invoice number"),
        InputParam("ownerEmail", type="string", role="assignee", description="Who approves it"),
    ],
    health=health,
)


@agent.run
async def run(ctx: RunContext) -> None:
    invoice_id = ctx.inputs["invoiceId"]
    ctx.run_renamed(f"invoice-{invoice_id}")
    if ctx.inputs.get("ownerEmail"):
        ctx.set_assignee(ctx.inputs["ownerEmail"])

    ctx.step_started("fetch")
    ctx.tool_activity("fetch", tool="get_invoice", text=f'{{"id": "{invoice_id}"}}', result='{"total": 1200}')
    ctx.step_completed("fetch", token_count=0, delivered=f"Fetched {invoice_id}")

    ctx.step_started("summarize", agent_name="Summarizer")
    ctx.thinking("summarize", "Grouping line items by cost center")
    summary, usage = await summarize(invoice_id)
    ctx.step_completed(
        "summarize",
        token_count=usage["input"] + usage["output"],
        delivered=summary,
        input_tokens=usage["input"],
        output_tokens=usage["output"],
        cached_input_tokens=usage["cached"],
        model="<model-id>",
    )
    ctx.set_output({"invoiceId": invoice_id, "summary": summary})


app = build_app(agent, verify_caller=verify_caller, max_concurrency=4)

AGENT_API_KEY=dev-key MODEL_API_KEY=<your-model-key> uvicorn app:app --host 0.0.0.0 --port 8080

    try:
        summary, usage = await summarize(invoice_id)
    except Exception:
        ctx.set_step_tokens("summarize", token_count=spent_in + spent_out,
                            input_tokens=spent_in, output_tokens=spent_out,
                            model="<model-id>", cached_input_tokens=spent_cached)
        ctx.set_cost(spent_usd)
        raise  # the run fails, with the error as its reason

import time

from fastapi.testclient import TestClient

from app import app


def wait_for(client, run_id, status):
    for _ in range(100):
        body = client.get(f"/runs/{run_id}/status").json()
        if body["status"] == status:
            return body
        time.sleep(0.02)
    raise AssertionError(body)


def test_run_completes():
    with TestClient(app) as client:
        run = client.post("/runs", json={"inputs": {"invoiceId": "INV-1042"}}).json()
        status = wait_for(client, run["orchestratorRunId"], "completed")
        assert status["runName"] == "invoice-INV-1042"
        tokens = client.get(f"/runs/{run['orchestratorRunId']}/tokens").json()
        assert tokens["inputTokens"] == 4200


def test_missing_input_fails_the_run():
    with TestClient(app) as client:
        run = client.post("/runs", json={"inputs": {}}).json()
        status = wait_for(client, run["orchestratorRunId"], "failed")
        assert "KeyError" in status["steps"][0]["failureReason"]