Python SDK reference

Every Python node handler receives the platform through a single parameter: ax, typed as AxiomContext. This page is the complete reference for that surface — structured logging, secrets, flow reflection, and flow mutation. There is no pip package to install: the Axiom CLI generates the AxiomContext interface into your package at gen/axiom_context.py (regenerated by the CLI; do not edit), and every capability reaches the platform through the sidecar — node code never calls platform services directly (see sandboxing and tenancy).

Prerequisites. Examples on this page assume a Python package scaffolded with axiom init demo/chat-tools --language python, two messages created with axiom create message ChatRequest --fields "session_id:string; text:string" and axiom create message ChatReply --fields "text:string", and a node created with axiom create node AnswerQuestion --input ChatRequest --output ChatReply (see create a node in Python).

Handler signatures#

A unary Python node is one function in nodes/. The function name is the snake_case of the node name, the handler file is nodes/<snake_name>.py, and the test file is nodes/<snake_name>_test.py. Local message types come from the generated gen.messages_pb2 module:

# nodes/answer_question.py — generated shape of a unary Python node
from gen.messages_pb2 import ChatRequest, ChatReply
from gen.axiom_context import AxiomContext


def answer_question(ax: AxiomContext, input: ChatRequest) -> ChatReply:
    """Answers a chat message."""
    return ChatReply()

The handler may also be declared async def — the platform detects an awaitable result and runs it to completion. Declare your handler async def whenever you await other async platform APIs.

A pipeline node (created with --type pipeline) is a generator: it takes an iterator of input frames and yields output frames. For the entry node of a flow running in pipeline mode, the iterator yields exactly one item:

# nodes/stream_replies.py — generated shape of a Python pipeline node
from typing import Iterator
from gen.messages_pb2 import ChatRequest, ChatReply
from gen.axiom_context import AxiomContext


def stream_replies(ax: AxiomContext, inputs: Iterator[ChatRequest]) -> Iterator[ChatReply]:
    """Streams replies for each incoming request frame."""
    for inp in inputs:
        yield ChatReply(text=inp.text)

The handler's full source — docstring included — is captured at publish time and shown in the Axiom registry's View source panel, and it feeds the AI-generated description displayed on the node's card. Replace the generated placeholder docstring before pushing, and set the node's description field in axiom.yaml (or pass --description to axiom create node) to state the description directly.

At-least-once invocation: design for idempotency#

Axiom guarantees at-least-once invocation of every node by default (see the axiom.yaml reference for the retry knobs). A transient invocation failure (NODE_ERROR_TRANSPORT / NODE_ERROR_TIMEOUT) can cause the platform to call your handler again for the same execution, with the same input. Write handlers so a repeat call is safe:

• Prefer naturally idempotent operations — an upsert on a stable key, a conditional write, a PUT instead of a bare POST.
• For a non-idempotent side effect (charging a customer, sending an email), derive a stable idempotency key from ax.execution_id and de-duplicate against it downstream, or use the target system's own idempotency-key API if it has one.
• Raising an exception from your handler produces NODE_ERROR_USER, which the platform never retries, at any layer — retrying a flaky call inside your own logic is your handler's job (a normal try/except retry loop around the flaky call), not a platform mechanism.

This is a design responsibility, not something the platform enforces.

AxiomContext at a glance#

New platform capabilities are added to AxiomContext — node function signatures never change to accommodate them.

Logging with ax.log#

ax.log has four methods, each taking a message string plus arbitrary keyword attributes:

# nodes/answer_question.py — inside the handler body
ax.log.debug("raw input", size=len(input.text))
ax.log.info("answering", session_id=input.session_id)
ax.log.warn("input truncated", limit=4096)
ax.log.error("model call failed", attempt=3)

In local development (axiom dev) the logger writes concise plain text to the terminal. In production it writes JSON records carrying level, msg, trace_id, span_id, execution_id, and the node name plus your keyword attributes — making every log line searchable by execution and linkable to its distributed trace. Use ax.log instead of print() for anything you want visible in production.

Each invocation receives its own logger instance pre-configured with that invocation's trace context — never share a logger across calls, and never construct one yourself.

Reading secrets with ax.secrets#

ax.secrets.get(name) returns a (value, ok) tuple: (value, True) when the named secret is present, ("", False) otherwise. Values are plaintext strings; the platform handles encryption and decryption.

# nodes/answer_question.py — inside the handler body
api_key, ok = ax.secrets.get("OPENAI_KEY")
if not ok:
    raise ValueError("secret OPENAI_KEY is not registered for this tenant")

Declare every secret name the node reads under the node's required_secrets list in axiom.yaml so users know what to register before invoking. axiom validate scans handler source for ax.secrets.get("NAME") calls and warns (without failing) when a referenced name is missing from required_secrets; the scan is best-effort and cannot see dynamically constructed names. Secrets are registered per tenant in the console — see manage secrets in a flow.

get() alone can't tell "revoked" from "never configured" — both return ("", False). Use ax.secrets.status(name) for that: it returns a SecretStatus (AVAILABLE, REVOKED, or UNSET). See manage secrets in a flow for the full accessor table across languages.

Agent memory removed for beta. There is no ax.agent.memory in the beta SDK — agentic memory is out of scope for beta and has been removed at compile time.

Inspecting the running flow with ax.reflection.flow#

ax.reflection.flow is a read-only view of the flow's compiled graph and the current invocation's position in it. Five properties:

• nodes — list[ReflectionNode], every node placement in the graph.
• edges — list[ReflectionEdge], the forward edges.
• loop_edges — list[ReflectionEdge], the loop-back edges.
• position — a FlowPosition for the current invocation.
• graph_id — the artifact ID of the graph this node runs in (the sub-flow's own ID when running inside a sub-flow).

# nodes/answer_question.py — inside the handler body
pos = ax.reflection.flow.position
downstream = [e for e in ax.reflection.flow.edges
              if e.src_instance == pos.current_instance]
iteration = pos.loop_iterations.get(pos.current_instance, 0)
ax.log.info("graph view",
            nodes=len(ax.reflection.flow.nodes),
            downstream=len(downstream),
            iteration=iteration,
            graph_id=ax.reflection.flow.graph_id)

If the platform does not supply reflection data for an invocation, the view degrades gracefully: the lists are empty, graph_id is "", and position is all zeros — reading reflection never raises.

Reflection data types#

ReflectionNode — one node placement: instance_id (int), node_ulid (str), name (str), package_name (str), package_version (str), node_type (node, subflow, or pipeline), input_message_name and output_message_name (fully-qualified Protocol Buffers message names), and canvas_node_id (str).

ReflectionEdge — one edge: src_instance (int), dst_instance (int), canvas_edge_id (str), has_condition (bool), has_adapter (bool), max_iterations (int — meaningful only on entries from loop_edges), and condition_summary (a ConditionSummary when the edge is conditional, else None). has_condition and has_adapter are structural flags only — the compiled adapter recipe is not exposed.

ConditionSummary — an agent-readable digest of a conditional edge's dispatch predicate: field (str), op (str), and operands (list[str]). For example field="tools", op="EQ", operands=["ToolX"] means the edge fires when "ToolX" is in the source output's tools field — enough for a node to make idempotent decisions (such as skipping a tool that is already wired) without parsing the compiled condition.

FlowPosition — where this invocation sits: current_instance (int), depth (int, 0 at the root flow), loop_iterations (dict[int, int] keyed by the loop's destination instance), and subflow_stack_graph_ids (list[str], ordered root → immediate parent).

Adding to the running flow with ax.mutation.flow#

Nodes declared with mutation_capable: true in axiom.yaml (default false) may append nodes and edges to the running flow during their handler:

• ax.mutation.flow.add_node(package, version, canvas_position=None) — buffer a new node placement; returns the int instance ID to use as the endpoint of subsequent add_edge calls in the same invocation.
• ax.mutation.flow.add_edge(src_instance, dst_instance, condition=None) — buffer an edge between two instance IDs (existing or just added). Pass condition={"op": "EQ", "field": "tools", "value": "ToolX"} to make the edge fire only when the predicate holds on the source node's output (on a repeated field, EQ has membership semantics); omit condition for an unconditional edge.

# nodes/plan_tools.py — inside a mutation-capable handler body
search_iid = ax.mutation.flow.add_node(
    package="axiom/search-tools", version="0.4.0",
)
ax.mutation.flow.add_edge(
    src_instance=ax.reflection.flow.position.current_instance,
    dst_instance=search_iid,
)

Calls buffer locally and are applied by the platform after the handler returns. If the platform rejects a buffered mutation, the SDK surfaces an AxiomMutationError whose .message attribute holds the human-readable reason.

Testing nodes that use AxiomContext#

axiom create node generates nodes/<snake_name>_test.py containing a _TestContext class — a minimal AxiomContext with a silent logger, secrets served from a dict, and execution_id set to a test value. Supply secrets your node reads via secrets_map:

# nodes/answer_question_test.py — replace the generated test function with:
def test_answer_question():
    ax = _TestContext(secrets_map={"OPENAI_KEY": "sk-test"})
    input_msg = ChatRequest(session_id="s1", text="hi")
    result = answer_question(ax, input_msg)
    assert isinstance(result, ChatReply)

A plain def handler is called directly, as above. For an async def handler, run it to completion with asyncio.run(...) inside the test.

Run tests with axiom test, which uses pytest for Python packages. Tests run again inside the publish build, and a failing test fails the push. axiom validate warns (without blocking) when a node has no test — a node with zero tests would otherwise pass silently. Assert output field values meaningfully, not just the return type.