Create a node in Python

A node is a plain Python function with a typed input message and a typed output message. This guide takes one node from empty directory to a pushed package: scaffold, messages, implementation, tests, push.

The five commands#

The whole path, end to end. Each step is explained in the sections below.

# Run from the directory that will contain your package
axiom init your-handle/order-tools --language python --description "Order processing utilities"
cd order-tools
touch requirements.txt   # axiom validate requires this file for Python packages
axiom create message OrderRequest --fields "order_id:string; quantity:int32; unit_price_cents:int64"
axiom create message OrderConfirmation --fields "order_id:string; total_cents:int64; accepted:bool"
axiom create node ProcessOrder --input OrderRequest --output OrderConfirmation
# ...implement nodes/process_order.py, edit nodes/process_order_test.py...
axiom test
axiom push

Before you start#

• Install the Axiom CLI and authenticate with axiom login — see Installation. axiom login uses the OAuth Device Flow and stores the resulting API key in ~/.axiom/credentials; in CI, set AXIOM_API_KEY instead.
• Python 3.10 or newer (the generated code uses dict | None syntax).
• pip install grpcio-tools pytest — axiom generate compiles your .proto files with python -m grpc_tools.protoc (falling back to a standalone protoc if installed), and axiom test runs your tests with pytest.
• For axiom push: the package must live in a git repository with an origin remote, because the platform builds from the remote at your HEAD commit — not from your working tree.

Create the package and its messages#

axiom init <handle>/<name> --language python creates a subdirectory named after the package (the part after the last /) containing axiom.yaml, messages/messages.proto, nodes/conftest.py, and a .gitignore. Package names must be scoped (your-handle/package-name); pushing rejects unscoped names. Use --description to set the package description and --version to override the default 0.1.0.

Every node input and output is a message defined in messages/messages.proto. All messages live in that single file so they can reference each other without imports. Scaffold them with axiom create message:

# Run inside the package directory (where axiom.yaml lives)
axiom create message OrderRequest --fields "order_id:string; quantity:int32; unit_price_cents:int64"
axiom create message OrderConfirmation --fields "order_id:string; total_cents:int64; accepted:bool"

--fields takes semicolon-separated definitions in colon shorthand (name:string), proto3 without field numbers (string name), or canonical proto3 (string name = 1); field numbers are auto-assigned when omitted. Omit --fields to get a placeholder block with syntax hints to edit by hand. After appending the message, the command runs axiom generate automatically to compile bindings into gen/ (skip with --no-generate when creating several messages at once).

The result in messages/messages.proto:

// messages/messages.proto
message OrderRequest {
  string order_id = 1;
  int32 quantity = 2;
  int64 unit_price_cents = 3;
}

message OrderConfirmation {
  string order_id = 1;
  int64 total_cents = 2;
  bool accepted = 3;
}

Comments written directly above a message or field declaration (with no blank line in between) are extracted at publish time and shown in the Axiom registry as documentation; a comment separated by a blank line is ignored. See The type system for how messages type-check flow edges, and Import message types from another package to reuse messages published by others.

Scaffold the node#

axiom create node writes the implementation file and a test file into nodes/, then records the node in axiom.yaml:

# Run inside the package directory
axiom create node ProcessOrder --input OrderRequest --output OrderConfirmation

The node name must be PascalCase. Input and output messages must be defined in messages/messages.proto or available from an imported package. If you omit --input or --output in an interactive terminal, the command prompts with a numbered list of available messages. The default node type is unary (one input message in, one output message out per invocation); pass --type pipeline for a streaming node instead.

This creates nodes/process_order.py and nodes/process_order_test.py, generates gen/messages_pb2.py (your message bindings) and gen/axiom_context.py (the AxiomContext type), and updates axiom.yaml:

# axiom.yaml
name: your-handle/order-tools
version: 0.1.0
language: python
description: Order processing utilities
nodes:
    - name: ProcessOrder
      input: OrderRequest
      output: OrderConfirmation

See the axiom.yaml reference for every manifest field.

Implement the node function#

A Python node is a function whose first parameter is the AxiomContext (ax) and whose second parameter is the input message; it returns the output message. Declare it as async def if you need await — for example with the ax.agent.memory APIs.

# nodes/process_order.py
from gen.messages_pb2 import OrderRequest, OrderConfirmation
from gen.axiom_context import AxiomContext


def process_order(ax: AxiomContext, input: OrderRequest) -> OrderConfirmation:
    """Validates an order request and returns a confirmation with the
    calculated total. Orders with zero quantity are rejected.
    """
    if input.quantity <= 0:
        ax.log.warn("rejected order", order_id=input.order_id)
        return OrderConfirmation(order_id=input.order_id, accepted=False)

    total = input.quantity * input.unit_price_cents
    ax.log.info("order accepted", order_id=input.order_id, total_cents=total)
    return OrderConfirmation(
        order_id=input.order_id,
        total_cents=total,
        accepted=True,
    )

The docstring directly inside the function is extracted at publish time and shown in the Axiom registry as the node's documentation — write a plain description of what the node does.

ax is the single injection point for every platform capability:

• ax.log — structured logger with debug / info / warn / error, taking keyword attributes. Use it instead of print(): in local development it writes plain text; in production it writes JSON with trace_id, span_id, and execution_id baked in.
• ax.secrets.get(name) — returns a (value, found) tuple. Use it instead of environment variables for API keys; list the secret names a node reads under required_secrets in axiom.yaml so they are validated during publish. See Manage secrets in a flow.
• ax.agent.memory — durable agent memory; see Agent memory.
• ax.execution_id, ax.flow_id, ax.tenant_id — identifiers for the current execution, the flow, and the tenant.

The full surface is documented in the Python SDK reference.

Test the node#

axiom create node generated nodes/process_order_test.py containing a _TestContext class — a minimal in-memory AxiomContext with a no-op logger, a secrets map you can populate (_TestContext(secrets_map={"OPENAI_KEY": "sk-test"})), and stub agent memory. Replace the generated placeholder test with assertions on real output fields:

# nodes/process_order_test.py — replaces the generated placeholder test;
# keep the generated _TestContext class and imports above it.
def test_process_order_accepts_and_totals():
    ax = _TestContext()
    input_msg = OrderRequest(order_id="ord-1", quantity=3, unit_price_cents=250)
    result = process_order(ax, input_msg)
    assert result.order_id == "ord-1"
    assert result.total_cents == 750
    assert result.accepted is True


def test_process_order_rejects_zero_quantity():
    ax = _TestContext()
    result = process_order(ax, OrderRequest(order_id="ord-2", quantity=0))
    assert result.accepted is False

Create a requirements.txt listing your Python dependencies — it may be empty, but axiom validate fails without it. Then run:

# Run inside the package directory
axiom test

axiom test compiles proto bindings (axiom generate), validates the package (axiom validate — axiom.yaml schema, proto definitions, node signatures, node tests, and language files such as requirements.txt), then runs pytest nodes/ with output streamed live. Pass arguments to pytest after --, e.g. axiom test -- -k test_process_order.

Tests run again inside the push build: the platform's image build runs pytest, and a failing test fails the push. axiom validate warns — without blocking — when a node has no test. Assert output fields meaningfully — not just type-checked — so the gate actually catches regressions.

Push the package#

axiom push deploys the package to the Axiom platform, visible only to your own tenant — it does not appear in the public marketplace. The platform clones your repository at your current HEAD commit and builds from that, so commit and push first:

# Run inside the package directory
git add -A && git commit -m "ProcessOrder node" && git push
axiom push

Push validates the package locally, then streams the platform pipeline's progress: server-side validation, code generation, Docker image build, and deployment. It requires a prior axiom login, a git remote named origin, and the current HEAD pushed to that remote. You can push the same version repeatedly — each push overwrites the previous one — which lets you iterate on deployed code before publishing. When you are satisfied, publish the package through the Axiom UI to make it available to others as an immutable versioned release. Use --json for a single machine-readable result object instead of progress output.

After pushing, axiom info your-handle/order-tools shows the package's nodes, messages, and live endpoint.

Next steps#

• Place the node in a flow on the canvas: Build your first flow.
• Invoke the flow over HTTP: Invoke a flow via the API.
• Iterate locally with hot reload: axiom dev starts a local development server that watches nodes/, messages/, and axiom.yaml, exposes each node at POST http://localhost:8083/nodes/<NodeName>, and converts JSON payloads to and from protobuf automatically.
• Understand how nodes, packages, and flows relate: Nodes, packages, and flows.