import json
import os

from axllm import ai, 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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6-luna"),
)
program = ax('question:string -> answer:string')
out = program.forward(
    client,
    {"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."},
    {"timeout": 30000, "promptCacheKey": "ax-openai-example", "contextCache": {}},
)
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os

from axllm import ai, 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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6-astra"),
    model_config={"thinkingTokenBudget": "low", "max_tokens": 2048},
)
program = ax('question:string -> answer:string')
out = program.forward(
    client,
    {"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."},
    {"serviceTier": "standard", "promptCacheKey": "ax-openai-example", "contextCache": {}},
)
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os

from axllm import ai, 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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6.1-sol"),
    model_config={"thinkingTokenBudget": "low", "max_tokens": 2048},
)
program = ax('question:string -> answer:string')
out = program.forward(
    client,
    {"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."},
    {"serviceTier": "standard", "promptCacheKey": "ax-gpt61-sol-example", "contextCache": {}},
)
print(json.dumps(out, indent=2, sort_keys=True))

from axllm import get_supported_ai_models


catalog = get_supported_ai_models()
providers = {entry["name"]: entry for entry in catalog}
azure = providers["azure-openai"]
openrouter = providers["openrouter"]

assert azure["isDynamic"] is True and azure["models"] == []
assert "high" in azure["capabilities"]["thinkingLevels"]
assert "priority" in azure["capabilities"]["serviceTiers"]
assert "flex" in openrouter["capabilities"]["serviceTiers"]

print(f"{len(catalog)} providers; Azure and OpenRouter named profiles are available")

import json
import os
from axllm import ai, ax, typesafe

model = ai("typesafe", api_key=os.environ["TYPESAFE_APIKEY"], trueThreshold=0.9)
triage = ax("ticket:string -> urgent:boolean(true \"Customers cannot complete a core task\", false \"Routine request\") \"Needs immediate attention?\", team:class \"support, billing, engineering\"")
decision = triage.forward(model, {"ticket": "Checkout is unavailable for all customers after the latest deployment."})
assert isinstance(decision["urgent"], bool)
assert decision["team"] in ("support", "billing", "engineering")
print(json.dumps(decision, indent=2))

import json
import os
from axllm import ai, ax

model = ai("openai-decisions", api_key=(os.getenv("OPENAI_API_KEY") or os.environ["OPENAI_APIKEY"]), trueThreshold=0.9)
triage = ax("ticket:string -> urgent:boolean(true \"Customers cannot complete a core task\", false \"Routine request\") \"Needs immediate attention?\", team:class \"support, billing, engineering\"")
decision = triage.forward(model, {"ticket": "Checkout is unavailable for all customers after the latest deployment."})
assert isinstance(decision["urgent"], bool)
assert decision["team"] in ("support", "billing", "engineering")
print(json.dumps(decision, indent=2))

import os

from axllm import ai

key = os.getenv("MODEL_API_KEY")
if not key:
    raise SystemExit("Set MODEL_API_KEY to run this example.")

for profile in ("meta", "meta-chat", "meta-messages"):
    client = ai(profile, api_key=key, model="muse-spark-1.3")
    response = client.chat({
        "chat_prompt": [{"role": "user", "content": "Name a solar-powered sailboat."}],
        "model_config": {"thinking_token_budget": "highest"},
    })
    print(profile, response["results"][0].get("content"))

import json
import os

from axllm import ai, ax


def pick_highest_score(samples):
    return max(range(len(samples)), key=lambda index: samples[index]["sample"]["score"])


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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0.8},
)
program = ax(
    "topic:string -> answer:string, score:number",
    sample_count=3,
    result_picker=pick_highest_score,
)
out = program.forward(
    client,
    {"topic": "Explain why typed signatures make LLM programs easier to maintain."},
)
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os

from axllm import ai, 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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
program = ax('ticket:string -> priority:class "high, normal, low", summary:string, labels:string[]')
out = program.forward(client, {"ticket": "Checkout has failed for enterprise customers since 09:00. Support wants a concise summary and tags."})
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os

from axllm import GoogleGeminiClient


def required(name):
    value = os.getenv(name)
    if not value:
        raise SystemExit(f"Set {name} to run this example.")
    return value


client = GoogleGeminiClient(
    api_key=required("GOOGLE_VERTEX_ACCESS_TOKEN"),
    project_id=required("GOOGLE_PROJECT_ID"),
    region=required("GOOGLE_REGION"),
    model=os.getenv("AX_VERTEX_MODEL", "gemini-3.5-flash"),
)
out = client.chat({"chat_prompt": [{"role": "user", "content": "Reply with the word ready."}]})
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os
from axllm import ai, ax, typesafe

model = ai("typesafe", api_key=os.environ["TYPESAFE_APIKEY"], trueThreshold=0.9)
triage = ax("ticket:string -> urgent:boolean(true \"Customers cannot complete a core task\", false \"Routine request\") \"Needs immediate attention?\", team:class \"support, billing, engineering\"")
decision = triage.forward(model, {"ticket": "Checkout is unavailable for all customers after the latest deployment."})
assert isinstance(decision["urgent"], bool)
assert decision["team"] in ("support", "billing", "engineering")
writer = ai("openai", api_key=os.environ.get("OPENAI_API_KEY") or os.environ["OPENAI_APIKEY"], model="gpt-5.6-luna")
reply = ax("ticket:string, urgent:boolean, team:string -> reply:string").forward(writer, {"ticket": "Checkout is unavailable for all customers after the latest deployment.", **decision})
assert isinstance(reply["reply"], str) and reply["reply"].strip()
print(json.dumps({"decision": decision, **reply}, indent=2))

import json
import os

from axllm import ai, ax, f


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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
signature = (
    f()
    .input("requestText", f.string("Booking request").min(10).max(500))
    .input("contactEmail", f.string("Contact email").email())
    .output("partySize", f.number("Guests").min(1).max(12))
    .output(
        "bookingCode",
        f.string("Three letters, a dash, and four digits").regex(
            r"^[A-Z]{3}-\d{4}$", "Must look like ABC-1234"
        ),
    )
    .output(
        "guestProfile",
        f.object(
            {
                "fullName": f.string("Primary guest").min(2),
                "dietaryNotes": f.string("Dietary requirements").optional(),
            }
        ),
    )
    .build()
)
output = ax(signature).forward(
    client,
    {
        "requestText": "Book dinner for four people under the name Ada Lovelace.",
        "contactEmail": "ada@example.com",
    },
)
print(json.dumps(output, indent=2, sort_keys=True))

import json
import os
import uuid

from axllm import ai, set_usage_observer


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

events = []
set_usage_observer(events.append)
client = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    usage_context={
        "tenantId": "tenant-42",
        "feature": "support-chat",
        "attributes": {"environment": "example"},
    },
)

try:
    client.chat(
        {"chat_prompt": [{"role": "user", "content": "Reply with one short greeting."}]},
        {"usageContext": {"userId": "user-7", "requestId": str(uuid.uuid4())}},
    )
    print(json.dumps(events, indent=2, sort_keys=True))
finally:
    set_usage_observer(None)

import os
import time

from axllm import ai


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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.6-luna"),
)
started = time.perf_counter()
stream = client.stream(
    {
        "chat_prompt": [{"role": "user", "content": "Reply with exactly: streaming works"}],
    }
)
try:
    for event in stream:
        results = event.get("results") or []
        content = (results[0].get("content") or "") if results else ""
        if content:
            print(f"[{(time.perf_counter() - started) * 1000:.0f} ms] {content}", end="", flush=True)
finally:
    stream.close()
print()

import json
import os
from axllm import ai, ax

model = ai("openai-decisions", api_key=os.getenv("OPENAI_API_KEY") or os.environ["OPENAI_APIKEY"])
judge = ax('query:string, document:string -> relevant:boolean(evidence, threshold 0.7), relevance:score "Irrelevant, Related, Partial answer, Direct answer"(evidence)')
query = "How do I request a refund?"
candidates = [{'id': 'refund', 'text': 'Request a refund within 30 days through your account billing page.'}, {'id': 'password', 'text': 'Reset your password using the sign-in page.'}]
ranked = [{"id": item["id"], **judge.forward(model, {"query": query, "document": item["text"]})} for item in candidates]
ranked.sort(key=lambda item: item["relevance"]["value"], reverse=True)
print(json.dumps(ranked, indent=2))

from axllm import AxAIServiceAbortedError, AxCancellationToken, OpenAICompatibleClient


class CountingTransport:
    def __init__(self):
        self.calls = 0

    def __call__(self, _request):
        self.calls += 1
        return {"status": 200, "json": {}}


transport = CountingTransport()
client = OpenAICompatibleClient(api_key="test-key", model="gpt-5.6-luna", transport=transport)
token = AxCancellationToken()
assert token.cancel("user stopped")
assert not token.cancel("later reason")

try:
    client.chat(
        {"chat_prompt": [{"role": "user", "content": "This must not be sent."}]},
        {"cancellation": token},
    )
except AxAIServiceAbortedError as error:
    assert error.reason == "user stopped" and not error.retryable
else:
    raise AssertionError("pre-cancelled request unexpectedly completed")

assert transport.calls == 0
print("cancelled before transport: user stopped")

import os
from axllm import ai

api_key = os.getenv("ANTHROPIC_API_KEY") or os.getenv("ANTHROPIC_APIKEY")
if not api_key:
    raise SystemExit("Set ANTHROPIC_API_KEY or ANTHROPIC_APIKEY to run this example.")
client = ai("anthropic", api_key=api_key, model="claude-sonnet-5-5")
response = client.chat({
    "chat_prompt": [{"role": "user", "content": "Reply with exactly: Sonnet 5.5 works"}],
    "model_config": {"thinkingTokenBudget": "none", "effort": "high", "maxTokens": 128},
})
print(response["results"][0]["content"])

import os

from axllm import ai, 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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6-luna"),
)
story = ax('topic:string -> title:string, story:string "Three short sentences"')

# Each delta holds the new text of one field. Merge a sample's deltas (strings
# and lists append, other values replace) and start over when the version
# changes: a retry or a replaced step starts a new version.
merged, version = {}, 0
for delta in story.streaming_forward(client, {"topic": "a lighthouse keeper's cat"}):
    if delta["version"] != version:
        merged, version = {}, delta["version"]
        print("\n[retry: starting over]")
    for field, value in delta["delta"].items():
        merged[field] = merged.get(field, "") + value if isinstance(value, str) else value
        if field == "story":
            print(value, end="", flush=True)
print()
print(f"Title: {merged.get('title')}")

import datetime
import os

from axllm import ai, 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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6-luna"),
)

# Date outputs parse by default; set parseDates (or parse_dates) to false to keep text.
planner = ax(
    'emailText:string, sentAt:datetime -> meetingStartsAt:datetime "Start time with its time zone", '
    'meetingDay:date, travelWindow:dateRange "First and last day away"',
)
out = planner.forward(
    client,
    {
        "emailText": "Can we meet next Tuesday at 3pm New York time? I'm travelling from the 8th to the 12th.",
        # A native datetime is rendered as TypeScript renders a Date.
        "sentAt": datetime.datetime(2024, 5, 2, 16, 30, tzinfo=datetime.timezone.utc),
    },
)
print(out["meetingStartsAt"])  # e.g. 2024-05-07T19:00:00.000Z
print(out["meetingDay"])  # e.g. 2024-05-07T00:00:00.000Z
print(out["travelWindow"])  # {'start': '2024-05-08T00:00:00.000Z', 'end': ...}

import os

from axllm import ai, 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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-6-luna"),
)
summarize = ax("text:string -> summary:string")


def keep_it_short(summary, context):
    # A non-empty result goes back to the model as a user message, and the
    # next step's answer replaces this one.
    words = len(summary.split())
    return f"That summary has {words} words; answer again in at most 12 words." if words > 12 else None


summarize.add_field_processor("summary", keep_it_short)
summarize.add_field_transform("summary", "trim")

text = (
    "The committee met on Tuesday to review the budget. After a long debate about "
    "the new library wing, they approved the plan and asked staff to find a builder "
    "who can start in spring, while keeping the reading room open during the work."
)
print(summarize.forward(client, {"text": text})["summary"])

import os
from axllm import ai

api_key = os.getenv("ANTHROPIC_API_KEY") or os.getenv("ANTHROPIC_APIKEY")
if not api_key:
    raise SystemExit("Set ANTHROPIC_API_KEY or ANTHROPIC_APIKEY to run this example.")
client = ai("anthropic", api_key=api_key, model="claude-haiku-5-5")
response = client.chat({
    "chat_prompt": [{"role": "user", "content": "Reply with exactly: Haiku 5.5 works"}],
    "model_config": {"thinkingTokenBudget": "low", "showThoughts": False, "maxTokens": 2048},
})
print(response["results"][0]["content"])

import os
from axllm import ai, ax, fn

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=os.getenv("AX_OPENAI_MODEL", "gpt-6-luna"))
inventory = fn("inventory").description("Read the current stock count.").handler(
    lambda _: {"available": 12, "warehouse": "A"}
).build()
program = ax("question:string -> answer:string", {"functions": [inventory]})
# The model receives this text; tool traces retain the original object.
program.set_function_result_formatter(lambda result: f"{result['available']} units available")
result = program.forward(client, {"question": "Call inventory and report how many units are available."})
print(result["answer"])

import json
import os

from axllm import GoogleGeminiClient


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

client = GoogleGeminiClient(
    api_key=api_key,
    model=os.getenv("AX_GEMINI_MODEL", "gemini-3.8-flash"),
)
out = client.chat(
    {
        "chat_prompt": [
            {
                "role": "user",
                "content": "Explain in one sentence why batch evaluations save time.",
            }
        ]
    },
    {"service_tier": "flex"},
)
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os
from axllm import ai, ax, ProviderRouter

client = ai("openai", api_key=os.getenv("OPENAI_API_KEY") or os.environ["OPENAI_APIKEY"], model="gpt-6-astra", model_config={"thinkingTokenBudget": "low"})
router = ProviderRouter({"providers": {"primary": client}})
program = ax("document:file -> summary:string")
result = program.forward(router, {"document": {"filename": "report.pdf", "mimeType": "application/pdf", "data": os.environ["AX_PDF_BASE64"]}}, {"serviceTier": "standard"})
print(json.dumps(result, indent=2))

import json
import os
import threading
import time
from axllm import ai, ax, fn, 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(),
]})
result = program.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 "next-response" in applied, applied
print(json.dumps({"result": result, "background_overlap": True, "control_timing": applied[1:]}, indent=2))

import json
import os
import threading
import time
from axllm import ai, ax, fn, 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.")
control = run_control()
settled = threading.Event()
started = {}
def lookup(args, context):
    started.update(call_id=context["call_id"], at=time.monotonic())
    control.abort()
    if not context["signal"].wait(2):
        raise RuntimeError("Tool did not receive cancellation")
    settled.set()
    return "LATE: discard this result"
program = ax("question -> answer", {"functions": [
    fn("lookup").description("Look up the reference.").execution("background")
      .context_handler(lookup).build()
]})
client = ai("openai", api_key=key, model="gpt-6-astra",
            model_config={"thinkingTokenBudget":"low", "max_tokens":2048})
try:
    program.forward(client, {"question":"Call lookup once and return its result."},
                    {"control":control, "serviceTier":"standard"})
except RuntimeError as error:
    assert started and started["call_id"] in str(error), error
    elapsed = time.monotonic() - started["at"]
    assert elapsed < 2 and settled.wait(2), "Cancellation did not settle promptly"
    print(json.dumps({"cancelled":True, "unresolved_call_id":started["call_id"], "seconds":elapsed}))
else:
    raise AssertionError("Cancelled run returned a successful answer")

import json
import os

from axllm import ai, 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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
program = ax('context:string, question:string -> answer:string, citations:string[]')
out = program.forward(client, {"context": "Ax uses signatures, ai(), ax(), agent(), flow(), and optimize() for production LLM programs.", "question": "How should a new developer think about Ax?"})
print(json.dumps(out, indent=2, sort_keys=True))

import json
import os
from axllm import typesafe

client = typesafe(api_key=os.environ["TYPESAFE_APIKEY"])
models = client.list_models()
assert models and all(model["name"] for model in models)
response = client.system_one(
    {
        "state": {
            "ticket": "Checkout is unavailable for all customers after the latest deployment.",
            "account": {
                "tier": "enterprise",
                "notes": None
            }
        },
        "questions": {
            "urgent": {
                "type": "noul",
                "instructions": {
                    "question": "Does this need immediate attention?"
                },
                "criteria": {
                    "true": "Customers cannot complete a core task",
                    "false": "Routine request"
                }
            },
            "team": {
                "type": "choice",
                "instructions": "Who should handle the ticket?",
                "criteria": {
                    "support": "Usage guidance",
                    "billing": {
                        "scope": "Invoices and payments"
                    },
                    "engineering": "Product failures"
                }
            },
            "severity": {
                "type": "score",
                "instructions": "Rate customer impact",
                "criteria": [
                    "Minor inconvenience",
                    "One task blocked",
                    "Core task unavailable",
                    "Widespread outage"
                ]
            }
        }
    }
)
answers = response["answers"]
urgent, severity, team = answers["urgent"], answers["severity"], answers["team"]
assert urgent["type"] == "noul"
assert severity["type"] == "score"
assert team["type"] == "choice"
assert 0 <= urgent["noul"] <= 1
assert 0 <= severity["score"] <= 3
assert team["choice"] in ("support", "billing", "engineering")
# Probability policies and score scales are application decisions.
print(json.dumps({"page_on_call": urgent["noul"] >= 0.9, "severity_1_to_5": 1 + 4 * severity["score"] / 3, "response": response}, indent=2))

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 openai_decisions

client = openai_decisions(api_key=(os.getenv("OPENAI_API_KEY") or os.environ["OPENAI_APIKEY"]))
response = client.create(json.loads('{"input": "Checkout is unavailable for all customers after the latest deployment.", "questions": [{"type": "predicate", "name": "urgent", "instructions": "Are customers unable to complete a core task?"}, {"type": "choice", "name": "team", "instructions": "Who should handle the ticket?", "choices": [{"value": "support", "description": "Usage guidance"}, {"value": "billing", "description": "Invoices and payments"}, {"value": "engineering", "description": "Product failures"}]}, {"type": "score", "name": "severity", "instructions": "Rate customer impact", "levels": [{"label": "Minor inconvenience"}, {"label": "One task blocked"}, {"label": "Core task unavailable"}, {"label": "Widespread outage"}]}]}'))
# Handle per-question refusals before using probability, choice, or score.
assert [answer["name"] for answer in response["answers"]] == ["urgent", "team", "severity"]
print(json.dumps(response, indent=2))

import os

from axllm import AxRuntimeHooks, agent, ai, ax, flow, set_meter, set_rate_limiter, set_tracer
from axllm.runtime_quickjs import AxQuickJsCodeRuntime


class Span:
    def __init__(self, name):
        self.name = name
        print(f"[span:start] {name}")

    def set_attributes(self, attributes): pass
    def add_event(self, name, attributes=None): print(f"[span:event] {self.name} {name}")
    def record_exception(self, error): print(f"[span:error] {self.name} {error}")
    def set_status(self, status, description=None): pass
    def end(self): print(f"[span:end] {self.name}")


class Tracer:
    def start_span(self, start):
        return Span(start.name)


class Instrument:
    def __init__(self, name): self.name = name
    def add(self, value, attributes=None): print(f"[metric] {self.name} += {value}")
    def record(self, value, attributes=None): print(f"[metric] {self.name} = {value}")


class Meter:
    def create_counter(self, name, options=None): return Instrument(name)
    def create_histogram(self, name, options=None): return Instrument(name)
    def create_gauge(self, name, options=None): return Instrument(name)


def limiter(label):
    def run(next_request, info):
        print(f"[limit:{label}] {info.operation} {info.provider}/{info.model} stream={info.streaming}")
        return next_request()
    return run


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 = ai(
    "openai",
    api_key=api_key,
    model=os.getenv("AX_OPENAI_MODEL", "gpt-5.4-mini"),
    model_config={"temperature": 0},
)
tracer = Tracer()
meter = Meter()
override_hooks = AxRuntimeHooks(limiter("forward"), tracer, meter)

set_rate_limiter(limiter("global"))
set_tracer(tracer)
set_meter(meter)
try:
    print(ax("topic:string -> summary:string").forward(client, {"topic": "portable Ax runtime hooks"}))
    print(agent("question:string -> answer:string").forward(
        client,
        {"question": "What does a rate limiter wrap?"},
        {"runtime": AxQuickJsCodeRuntime(), "max_actor_steps": 12},
        override_hooks,
    ))
    workflow = flow({"id": "examples.runtimeHooks"}).execute(
        "outline", ax("topic:string -> outline:string")
    ).execute("polish", ax("outline:string -> answer:string")).returns({"answer": "answer"})
    result = workflow.forward(client, {"topic": "Ax runtime hooks"}, hooks=override_hooks)
    assert result.get("answer"), result
    print(result)
finally:
    set_rate_limiter(None)
    set_tracer(None)
    set_meter(None)