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.6-luna"),
    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."},
    {"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))

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 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 OpenAICompatibleClient, 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 = OpenAICompatibleClient(
    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 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('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", model_config={"temperature": 1})
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 OpenAICompatibleClient, 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 = OpenAICompatibleClient(
    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 OpenAICompatibleClient, 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 = OpenAICompatibleClient(
    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"}],
        "model_config": {"temperature": 1},
    }
)
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()

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 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 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('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 os

from axllm import (
    AxRuntimeHooks,
    OpenAICompatibleClient,
    agent,
    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 = OpenAICompatibleClient(
    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": "polish"})
    print(workflow.forward(client, {"topic": "Ax runtime hooks"}, hooks=override_hooks))
finally:
    set_rate_limiter(None)
    set_tracer(None)
    set_meter(None)