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

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

import json
import os
import threading
import time
from axllm import ai, ax, fn, flow, run_control

key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY.")
client = ai("openai", api_key=key, model="gpt-6-astra",
            model_config={"thinkingTokenBudget": "low", "max_tokens": 4096})
pending, finished, overlap = threading.Event(), threading.Event(), threading.Event()
control = run_control()
applied = []

def slow_reference(_):
    pending.set()
    time.sleep(6)
    finished.set()
    return "REF-42"

def local_label(_):
    if pending.wait(3) and not finished.is_set():
        overlap.set()
    return "LAUNCH"

def on_event(event):
    if event["type"] == "tool.started" and not applied:
        applied.append("queued")
        control.steer("Include the word VERIFIED in the final answer.")
        control.set_thinking_token_budget("medium")
    if event["type"] == "applied":
        applied.append(event["timing"])

control.on_event(on_event)
program = ax("question -> answer", {"functions": [
    fn("slow_reference").description("Look up a reference; takes a few seconds.")
        .execution("background").handler(slow_reference).build(),
    fn("local_label").description("Read an independent local label immediately.")
        .handler(local_label).build(),
]})
workflow = flow().execute("lookup", program, {"writes":["answer"]}).execute("verify", ax('answer -> report "Repeat the exact reference, label, and verification word from the answer."'), {"reads":["answer"]}).returns({"answer":"report"})
result = workflow.forward(client, {"question": "First call slow_reference. While it is pending, call local_label. Call each tool only once; do not call a tool again while its result is pending. If a required tool result is still pending, end this response with a brief progress message. The application will continue with the result when it arrives; do not spend reasoning tokens waiting for it. Once both results arrive, return them in one sentence."},
                         {"control": control, "serviceTier": "standard", "maxSteps": 6})
assert overlap.is_set(), "The model did not perform independent work while the background tool was pending"
assert all(word in result["answer"] for word in ("REF-42", "LAUNCH", "VERIFIED")), result
assert applied.count("next-response") == 4, applied
print(json.dumps({"result": result, "background_overlap": True, "control_timing": applied[1:]}, indent=2))

import json
import os
import threading
from axllm import ai, ax, fn, flow, run_control

key = os.getenv("OPENAI_API_KEY") or os.getenv("OPENAI_APIKEY")
if not key:
    raise SystemExit("Set OPENAI_API_KEY or OPENAI_APIKEY.")
client = ai("openai", api_key=key, model="gpt-6-astra",
            model_config={"thinkingTokenBudget": "low", "max_tokens": 4096})
barrier = threading.Barrier(2, timeout=45)
control = run_control()
updates_ready = threading.Event()
started_paths, applied = set(), []

def lookup(_):
    # Both nodes must start their own tool before either can finish.
    barrier.wait()
    if not updates_ready.wait(5):
        raise RuntimeError("Controller did not observe both active nodes")
    return "REF-42"

def observe(event):
    if event["type"] == "tool.started":
        started_paths.add(event["path"])
        if len(started_paths) == 2:
            control.steer("Include VERIFIED with the exact reference in your final answer.")
            control.set_thinking_token_budget("medium", target="root/left")
            updates_ready.set()
    if event["type"] == "applied":
        applied.append({"path": event["path"], "timing": event["timing"]})

control.on_event(observe)
program = ax("question -> answer", {"functions": [
    fn("lookup").description("Look up the exact reference once.")
      .execution("background").handler(lookup).build()
]})
workflow = flow().execute("left", program).execute("right", program).returns({
    "left": "leftResult", "right": "rightResult"
})
result = workflow.forward(client, {
    "question": "Call lookup exactly once. If its result is pending, return a brief progress message without calling it again. Return the exact reference when its result arrives."
}, {"control": control, "serviceTier": "standard", "maxSteps": 6})
assert started_paths == {"root/left", "root/right"}, started_paths
assert all("REF-42" in result[node]["answer"] and "VERIFIED" in result[node]["answer"] for node in ("left", "right")), result
assert len(applied) == 3, applied
print(json.dumps({"result": result, "parallel_overlap": True, "applied_controls": applied}, indent=2))

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

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