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

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

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