Advanced Start

import json
import os

from axllm import OpenAICompatibleClient, ax


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},
)
program = ax('question:string -> answer:string')
out = program.forward(client, {"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."})
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os

from axllm import OpenAICompatibleClient, agent
from axllm.runtime_quickjs import AxQuickJsCodeRuntime

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},
)

# The handbook can be arbitrarily large. Listing it in `contextFields` keeps it
# in the agent's runtime so it never inflates the model prompt -- the agent reads
# it through code, not through tokens. That is the whole point of an Ax agent
# over a plain gen() call: the source material stays out of the context window.
handbook = """
# Acme Cloud -- Support Handbook

## Billing
- Invoices are issued on the 1st of each month and are due net-15.
- Plan downgrades take effect at the END of the current billing cycle, not immediately.
- Refunds are issued to the original payment method within 5 business days.

## Access
- Seats can be added by any workspace Owner under Settings -> Members.
- SSO (SAML) is available on Enterprise; SCIM provisioning is Owner-only.

## Incidents
- Status and uptime are published at status.acme.example.
- Sev-1 incidents page the on-call within 5 minutes; updates post every 30 minutes.

## Data
- Exports are available in CSV and JSON from Settings -> Data.
- Deleted workspaces are recoverable for 30 days, then permanently purged.
""".strip()

assistant = agent(
    'question:string, handbook:string -> answer:string, citations:string[] "Handbook sections the answer relies on"',
    # Keep the handbook in the runtime, out of the prompt.
    {"contextFields": ["handbook"], "runtime": {"language": "JavaScript"}},
)

result = assistant.forward(
    client,
    {
        "question": "A customer downgraded their plan today. When does it take effect, and can they get a refund for the current cycle?",
        "handbook": handbook,
    },
    {"runtime": AxQuickJsCodeRuntime(), "max_actor_steps": 12},
)

print(json.dumps(result, indent=2, sort_keys=True))

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))

import base64
import json
import os
from pathlib import Path

from axllm import OpenAIResponsesClient


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 = OpenAIResponsesClient(
    api_key=api_key,
    model=os.getenv("AX_OPENAI_AUDIO_MODEL", "gpt-4o-mini-tts"),
    model_config={"temperature": 0},
)
speech = client.speak({"text": "Ax turns LLM prompts into typed programs.", "voice": "alloy", "format": "mp3"})
print(json.dumps({"format": speech.get("format"), "transcript": speech.get("transcript"), "audioBytesBase64": len(speech.get("audio") or speech.get("data") or "")}, indent=2, sort_keys=True))

import os

from axllm import (
    AxBalancer,
    AxBalancerAdaptiveStrategy,
    AxBalancerOptions,
    AxInMemoryBalancerStatsStore,
    OpenAICompatibleClient,
)


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.")

model = os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini")
services = [
    OpenAICompatibleClient(
        api_key=api_key,
        model=model,
        base_url=os.getenv("OPENAI_PRIMARY_BASE_URL", "https://api.openai.com/v1"),
    ),
    OpenAICompatibleClient(
        api_key=os.getenv("OPENAI_BACKUP_API_KEY", api_key),
        model=model,
        base_url=os.getenv("OPENAI_BACKUP_BASE_URL", "https://api.openai.com/v1"),
    ),
]

# Reuse this store across balancers in one process. A Redis/database adapter can
# implement the same atomic get/observe contract for multi-process state.
stats_store = AxInMemoryBalancerStatsStore()
route_keys = ["openai-primary", "openai-backup"]
events = []
strategy = AxBalancerAdaptiveStrategy(
    deadline_ms=6_000,
    bad_outcome_cost=0.02,
    expected_tokens={"promptTokens": 1_200, "completionTokens": 300},
    namespace="support-summary-v1",
    route_key=lambda _service, index: route_keys[index],
    slice=lambda context: "streaming" if context["options"].get("stream") else "interactive",
    stats_store=stats_store,
    on_routing_event=lambda event: events.append(event),
)
balancer = AxBalancer(services, AxBalancerOptions(strategy=strategy))
response = balancer.chat(
    {"model": model, "chat_prompt": [{"role": "user", "content": "Summarize why shared routing state matters."}]}
)
print(response["results"][0]["content"])
print([event["type"] for event in events])

import json
import os

from axllm import OpenAICompatibleClient, ax, OptimizerEngine


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},
)
program = ax('emailText:string -> priority:class "high, normal, low", rationale:string', {"id": "priority", "instruction": "Classify the email priority."})
baseline = program.forward(client, {"emailText": "Production checkout is failing for enterprise customers."})


class ExampleOptimizer(OptimizerEngine):
    name = "example"
    version = "1"

    def optimize(self, request, evaluator=None):
        return {"componentMap": {"priority::instruction": "Classify operational risk. Use high for production-impacting urgency."}, "metadata": {"source": "axgen"}}


artifact = program.optimize_with(ExampleOptimizer(), [{"emailText": "URGENT: checkout is down", "priority": "high"}], {"apply": False})
program.apply_optimization(json.dumps(artifact))
after = program.forward(client, {"emailText": "Production checkout is failing for enterprise customers."})
print(json.dumps({"baseline": baseline, "after": after}, indent=2, sort_keys=True))

import os

from axllm import AxMCPClient, AxMCPStreamableHTTPTransport, OpenAICompatibleClient, ax

api_key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
endpoint = os.getenv("MCP_URL")
if not api_key or not endpoint:
    raise SystemExit("Set OPENAI_API_KEY and MCP_URL.")

mcp = AxMCPClient(AxMCPStreamableHTTPTransport(endpoint), {"namespace": "inventory"})
llm = OpenAICompatibleClient(api_key=api_key, model="gpt-5.4-mini")
program = ax(
    'request:string -> answer:string "Use the inventory MCP tool."',
    {"mcp": mcp},
)
try:
    catalog = mcp.inspect_catalog()
    print({
        "tools": [tool["name"] for tool in catalog["tools"]],
        "resources": [
            {"name": resource["name"], "uri": resource["uri"]}
            for resource in catalog["resources"]
        ],
        "resourceTemplates": catalog["resourceTemplates"],
    })
    print(program.forward(llm, {"request": "Reindex inventory."}))
finally:
    mcp.close()

import os
import threading
import urllib.request

from axllm import (
    AxEventRoute,
    AxEventRuntime,
    AxEventTarget,
    AxMCPClient,
    AxMCPEventSource,
    AxMCPStreamableHTTPTransport,
    OpenAICompatibleClient,
    agent,
)
from axllm.runtime_quickjs import AxQuickJsCodeRuntime

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

endpoint = os.environ.get("AX_MCP_ENDPOINT")
if not endpoint:
    raise SystemExit("Set AX_MCP_ENDPOINT to a Streamable HTTP MCP server.")
transport = AxMCPStreamableHTTPTransport(
    endpoint,
    {
        "ssrfProtection": {
            "requireHttps": not endpoint.startswith("http://127.0.0.1"),
            "allowLocalhost": endpoint.startswith("http://127.0.0.1"),
            "allowPrivateNetworks": endpoint.startswith("http://127.0.0.1"),
        }
    },
)
client = AxMCPClient(transport, {"namespace": "inventory"})
source = AxMCPEventSource(
    client,
    "inventory",
    identity_scope="tenant:demo",
    trust="authenticated",
    resource_subscriptions="all",
)
llm = OpenAICompatibleClient(api_key=api_key, model="gpt-5.4-mini")
program = agent("uri:string -> summary:string", {"runtime": {"language": "JavaScript"}})
completed = threading.Event()


def invoke(input, _context):
    output = program.forward(llm, input, {"runtime": AxQuickJsCodeRuntime()})
    print(output)
    completed.set()
    return output


target = AxEventTarget(
    "inventory-agent",
    invoke,
    mapInput=lambda event, _continuation: {"uri": event.data["uri"]},
    retrySafety="idempotent",
)
runtime = AxEventRuntime(
    [
        AxEventRoute(
            "resource-wake",
            "wake",
            {"types": ["mcp.resource.updated"]},
            "inventory-agent",
            True,
        )
    ],
    {"targets": [target], "sources": [source]},
)
runtime.start()
if os.getenv("AX_MCP_DEMO_AUTO") == "1":
    urllib.request.urlopen(
        urllib.request.Request(
            endpoint.replace("/mcp", "/control/resource"), data=b"", method="POST"
        )
    ).close()
if not completed.wait(60):
    raise RuntimeError("Timed out waiting for an MCP resource notification")
runtime.close()
client.close()

import os
import threading
import urllib.request

from axllm import (
    AxEventEnvelope,
    AxEventRoute,
    AxEventRuntime,
    AxEventTarget,
    AxMCPClient,
    AxMCPEventSource,
    AxMCPStreamableHTTPTransport,
    AxPushEventSource,
    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.")

endpoint = os.environ.get("AX_MCP_ENDPOINT")
if not endpoint:
    raise SystemExit("Set AX_MCP_ENDPOINT to a Streamable HTTP MCP server.")
transport = AxMCPStreamableHTTPTransport(
    endpoint,
    {
        "ssrfProtection": {
            "requireHttps": not endpoint.startswith("http://127.0.0.1"),
            "allowLocalhost": endpoint.startswith("http://127.0.0.1"),
            "allowPrivateNetworks": endpoint.startswith("http://127.0.0.1"),
        }
    },
)
client = AxMCPClient(transport, {"namespace": "inventory"})
mcp = AxMCPEventSource(
    client, "inventory", identity_scope="tenant:demo", trust="authenticated"
)
started = AxPushEventSource("task-started")
llm = OpenAICompatibleClient(api_key=api_key, model="gpt-5.4-mini")
step = ax("taskId:string -> status:string")
program = (
    flow({"id": "reindex-flow"}).execute("status", step).returns({"status": "status"})
)
completed = threading.Event()
calls = 0


def invoke(input, _context):
    global calls
    output = program.forward(llm, input)
    calls += 1
    print(output)
    if calls >= 2:
        completed.set()
    return output


target = AxEventTarget(
    "reindex-flow",
    invoke,
    mapInput=lambda event, continuation: {
        "taskId": (
            continuation.metadata["taskId"] if continuation else event.data["taskId"]
        )
    },
    retrySafety="idempotent",
    waitFor=[{"kind": "mcp.task", "value": "taskKey", "metadata": {"taskId": "42"}}],
)
runtime = AxEventRuntime(
    [
        AxEventRoute(
            "task-start", "wake", {"types": ["app.task.started"]}, "reindex-flow"
        ),
        AxEventRoute("task-progress", "observe", {"types": ["mcp.progress"]}),
        AxEventRoute(
            "task-resume", "resume", {"types": ["mcp.task.status"]}, "reindex-flow"
        ),
    ],
    {"targets": [target], "sources": [started, mcp]},
)
runtime.start()
task_id = client.call_tool("start_reindex", {"scope": "all"})["task"]["taskId"]
target.waitFor[0]["metadata"] = {"taskId": task_id}
started.publish(
    AxEventEnvelope(
        "task-start",
        "app://tasks",
        "app.task.started",
        {"taskId": task_id, "taskKey": f"inventory:{task_id}"},
    ),
    identity_scope="tenant:demo",
    trust="authenticated",
)
print(f"Task {task_id} is waiting for a terminal MCP notification.")
if os.getenv("AX_MCP_DEMO_AUTO") == "1":
    urllib.request.urlopen(
        urllib.request.Request(
            endpoint.replace("/mcp", "/control/task/complete"), data=b"", method="POST"
        )
    ).close()
if not completed.wait(60):
    raise RuntimeError("Timed out waiting for the MCP task continuation")
runtime.close()
client.close()