Flows | Ax

These Python examples are real runnable files. Edit the source file first; this page is rebuilt from the checked-in example and its metadata header.

Python Sequential Flow

Runs a two-step Ax flow against OpenAI.

Provider: openai
Env: OPENAI_API_KEY, OPENAI_APIKEY
Level: beginner
Run: npm run example -- python src/examples/python/flows/flow-openai.py
Source: src/examples/python/flows/flow-openai.py

Python

import json
import os

from axllm import OpenAICompatibleClient, ax, flow


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
step = ax('documentText:string -> summaryText:string')
program = (
    flow({"id": "examples.sequentialFlow"})
    .execute("step", step)
    .map("note", lambda state: {"note": "Mapped flow state after the provider-backed step."})
    .returns({"summary": "step", "note": "note"})
)
output = program.forward(client, {"documentText": "Ax gives developers signatures, provider clients, agents, flows, tracing, and optimization."})
print(json.dumps(output, indent=2, sort_keys=True))

Python Branching Flow

Routes a classification through follow-up flow logic backed by OpenAI.

Provider: openai
Env: OPENAI_API_KEY, OPENAI_APIKEY
Level: intermediate
Run: npm run example -- python src/examples/python/flows/branch-flow.py
Source: src/examples/python/flows/branch-flow.py

Python

import json
import os

from axllm import OpenAICompatibleClient, ax, flow


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
classifier = ax('request:string -> route:class "support, sales, engineering"')
responder = ax("request:string, route:string -> response:string")
program = (
    flow({"id": "examples.branchFlow"})
    .execute(
        "classifier",
        classifier,
        {"reads": ["request"], "writes": ["classifierResult", "route"]},
    )
    .execute(
        "responder",
        responder,
        {
            "reads": ["request", "route"],
            "writes": ["responderResult", "response"],
        },
    )
    .returns({"route": "route", "response": "response"})
)
output = program.forward(
    client,
    {"request": "A customer says checkout is down for their enterprise account."},
)
print(json.dumps(output, indent=2, sort_keys=True))

Python Parallel Flow

Runs two independent OpenAI-backed steps in parallel before joining their results.

Provider: openai
Env: OPENAI_API_KEY, OPENAI_APIKEY
Level: intermediate
Run: npm run example -- python src/examples/python/flows/parallel-flow.py
Source: src/examples/python/flows/parallel-flow.py

Python

import json
import os

from axllm import OpenAICompatibleClient, ax, flow


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
research = ax("topicText:string -> factList:string[]")
audience = ax("topicText:string -> audienceAngle:string")
join = ax("factList:string[], audienceAngle:string -> briefText:string")
program = (
    flow({"id": "examples.parallelFlow"})
    .execute(
        "research",
        research,
        {"reads": ["topicText"], "writes": ["researchResult", "factList"]},
    )
    .execute(
        "audience",
        audience,
        {"reads": ["topicText"], "writes": ["audienceResult", "audienceAngle"]},
    )
    .execute(
        "join",
        join,
        {
            "reads": ["factList", "audienceAngle"],
            "writes": ["joinResult", "briefText"],
        },
    )
    .returns({"briefText": "briefText"})
)
output = program.forward(
    client,
    {"topicText": "Why typed contracts make multi-step LLM systems easier to maintain"},
)
print(json.dumps(output, indent=2, sort_keys=True))

Python Composed Flow

Composes multiple typed programs into one OpenAI-backed flow.

Provider: openai
Env: OPENAI_API_KEY, OPENAI_APIKEY
Level: advanced
Run: npm run example -- python src/examples/python/flows/composed-flow.py
Source: src/examples/python/flows/composed-flow.py

Python

import json
import os

from axllm import OpenAICompatibleClient, ax, flow


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
step = ax('topic:string -> outline:string[]')
program = (
    flow({"id": "examples.composedFlow"})
    .execute("step", step)
    .map("note", lambda state: {"note": "Mapped flow state after the provider-backed step."})
    .returns({"outline": "step", "brief": "note"})
)
output = program.forward(client, {"topic": "How Ax moves from typed generation to agents, flows, and optimization"})
print(json.dumps(output, indent=2, sort_keys=True))

Drafts, critiques, and revises an answer through three OpenAI-backed steps.

Provider: openai
Env: OPENAI_API_KEY, OPENAI_APIKEY
Level: advanced
Run: npm run example -- python src/examples/python/flows/refine-flow.py
Source: src/examples/python/flows/refine-flow.py

Python

import json
import os

from axllm import OpenAICompatibleClient, ax, flow


api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not api_key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")

client = OpenAICompatibleClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
draft = ax("topicText:string -> draftText:string")
critique = ax("draftText:string -> critiqueText:string")
revise = ax("draftText:string, critiqueText:string -> revisedText:string")
program = (
    flow({"id": "examples.refineFlow"})
    .execute(
        "draft",
        draft,
        {"reads": ["topicText"], "writes": ["draftResult", "draftText"]},
    )
    .execute(
        "critique",
        critique,
        {"reads": ["draftText"], "writes": ["critiqueResult", "critiqueText"]},
    )
    .execute(
        "revise",
        revise,
        {
            "reads": ["draftText", "critiqueText"],
            "writes": ["reviseResult", "revisedText"],
        },
    )
    .returns({"revisedText": "revisedText"})
)
output = program.forward(
    client,
    {"topicText": "Explain automatic flow parallelism to a backend engineer."},
)
print(json.dumps(output, indent=2, sort_keys=True))

Python Sequential Flow

Python Branching Flow

Python Parallel Flow

Python Composed Flow

Python Refinement Flow