import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class BasicGenerationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static OpenAICompatibleClient client() {
    return new OpenAICompatibleClient(
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.6-luna"), "model_config", Map.of("temperature", 0.0)));
  }

  public static void main(String[] args) throws Exception {
    AxGen program = Ax.ax("question:string -> answer:string");
    Map<String, Object> output = program.forward(
        client(),
        Map.of("question", "In one sentence, explain Ax as a language-agnostic LLM programming library."),
        Map.of("promptCacheKey", "ax-openai-example", "contextCache", Map.of()));
    System.out.println(Json.stringify(output));
  }
}

import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class AstraGenerationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static AxAIService client() {
    return Ax.ai("openai",
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-6-astra"), "model_config", Map.of("thinkingTokenBudget", "low", "max_tokens", 2048)));
  }

  public static void main(String[] args) throws Exception {
    AxGen program = Ax.ax("question:string -> answer:string");
    Map<String, Object> output = program.forward(
        client(),
        Map.of("question", "In one sentence, explain Ax as a language-agnostic LLM programming library."),
        Map.of("serviceTier", "standard", "promptCacheKey", "ax-openai-example", "contextCache", Map.of()));
    System.out.println(Json.stringify(output));
  }
}

import dev.axllm.ax.Ax;
import java.util.*;

public final class ModelCatalogExample {
  private static Map<?, ?> provider(List<Object> catalog, String name) {
    return catalog.stream()
        .map(entry -> (Map<?, ?>) entry)
        .filter(entry -> name.equals(entry.get("name")))
        .findFirst()
        .orElseThrow();
  }

  public static void main(String[] args) {
    List<Object> catalog = Ax.getSupportedAIModels();
    Map<?, ?> azure = provider(catalog, "azure-openai");
    Map<?, ?> openrouter = provider(catalog, "openrouter");
    Map<?, ?> azureCapabilities = (Map<?, ?>) azure.get("capabilities");
    Map<?, ?> openrouterCapabilities = (Map<?, ?>) openrouter.get("capabilities");

    assert Boolean.TRUE.equals(azure.get("isDynamic"));
    assert ((List<?>) azure.get("models")).isEmpty();
    assert ((List<?>) azureCapabilities.get("thinkingLevels")).contains("high");
    assert ((List<?>) azureCapabilities.get("serviceTiers")).contains("priority");
    assert ((List<?>) openrouterCapabilities.get("serviceTiers")).contains("flex");

    System.out.println(catalog.size() + " providers; Azure and OpenRouter named profiles are available");
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class TypesafeExample {
  public static void main(String[] args) throws Exception {
    var model = Ax.ai("typesafe", Map.of("apiKey", System.getenv("TYPESAFE_APIKEY"), "trueThreshold", 0.9));
    var triage = Ax.ax("ticket:string -> urgent:boolean(true \"Customers cannot complete a core task\", false \"Routine request\") \"Needs immediate attention?\", team:class \"support, billing, engineering\"");
    var decision = triage.forward(model, Map.of("ticket", "Checkout is unavailable for all customers after the latest deployment."));
    if (!(decision.get("urgent") instanceof Boolean) || !Set.of("support", "billing", "engineering").contains(decision.get("team"))) throw new AssertionError("Invalid decision");
    System.out.println(Json.stringify(decision));
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class MetaMuseExample {
  public static void main(String[] args) throws Exception {
    String key = System.getenv("MODEL_API_KEY");
    if (key == null || key.isBlank()) throw new IllegalStateException("Set MODEL_API_KEY to run this example.");
    for (String profile : List.of("meta", "meta-chat", "meta-messages")) {
      var client = Ax.ai(profile, Map.of("api_key", key, "model", "muse-spark-1.3"));
      var response = client.chat(Map.of(
          "chat_prompt", List.of(Map.of("role", "user", "content", "Name a solar-powered sailboat.")),
          "model_config", Map.of("thinking_token_budget", "highest")));
      System.out.println(profile + " " + Json.stringify(response));
    }
  }
}

import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class StructuredGenerationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static OpenAICompatibleClient client() {
    return new OpenAICompatibleClient(
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"), "model_config", Map.of("temperature", 0.0)));
  }

  public static void main(String[] args) throws Exception {
    AxGen program = Ax.ax("ticket:string -> priority:class \"high, normal, low\", summary:string, labels:string[]");
    Map<String, Object> output = program.forward(client(), Map.of("ticket", "Checkout has failed for enterprise customers since 09:00. Support wants a concise summary and tags."));
    System.out.println(Json.stringify(output));
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class VertexGeminiExample {
  private static String required(String name) {
    String value = System.getenv(name);
    if (value == null || value.isBlank()) throw new IllegalStateException("Set " + name + " to run this example.");
    return value;
  }

  public static void main(String[] args) throws Exception {
    GoogleGeminiClient client = new GoogleGeminiClient(Map.of(
        "api_key", required("GOOGLE_VERTEX_ACCESS_TOKEN"),
        "project_id", required("GOOGLE_PROJECT_ID"),
        "region", required("GOOGLE_REGION"),
        "model", System.getenv().getOrDefault("AX_VERTEX_MODEL", "gemini-3.5-flash")));
    Map<String, Object> out = client.chat(Map.of(
        "chat_prompt", List.of(Map.of("role", "user", "content", "Reply with the word ready."))));
    System.out.println(Json.stringify(out));
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class TypesafeHybridExample {
  public static void main(String[] args) throws Exception {
    var model = Ax.ai("typesafe", Map.of("apiKey", System.getenv("TYPESAFE_APIKEY"), "trueThreshold", 0.9));
    var triage = Ax.ax("ticket:string -> urgent:boolean(true \"Customers cannot complete a core task\", false \"Routine request\") \"Needs immediate attention?\", team:class \"support, billing, engineering\"");
    var decision = triage.forward(model, Map.of("ticket", "Checkout is unavailable for all customers after the latest deployment."));
    if (!(decision.get("urgent") instanceof Boolean) || !Set.of("support", "billing", "engineering").contains(decision.get("team"))) throw new AssertionError("Invalid decision");
    var key = System.getenv().getOrDefault("OPENAI_API_KEY", System.getenv("OPENAI_APIKEY"));
    var writer = Ax.ai("openai", Map.of("apiKey", key, "model", "gpt-5.6-luna", "model_config", Map.of("temperature", 1)));
    var inputs = new LinkedHashMap<String,Object>(decision);
    inputs.put("ticket", "Checkout is unavailable for all customers after the latest deployment.");
    var reply = Ax.ax("ticket:string, urgent:boolean, team:string -> reply:string").forward(writer, inputs);
    if (!(reply.get("reply") instanceof String text) || text.isBlank()) throw new AssertionError("Empty reply");
    System.out.println(Json.stringify(Map.of("decision", decision, "reply", reply)));
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class SignatureConstraintsExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static OpenAICompatibleClient client() {
    return new OpenAICompatibleClient(
        Map.of(
            "api_key", apiKey(),
            "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"),
            "model_config", Map.of("temperature", 0.0)));
  }

  public static void main(String[] args) throws Exception {
    AxSignature signature =
        Ax.f()
            .call()
            .input("requestText", Ax.f().string("Booking request").min(10).max(500))
            .input("contactEmail", Ax.f().string("Contact email").email())
            .output("partySize", Ax.f().number("Guests").min(1).max(12))
            .output(
                "bookingCode",
                Ax.f()
                    .string("Three letters, a dash, and four digits")
                    .regex("^[A-Z]{3}-\\d{4}$", "Must look like ABC-1234"))
            .output(
                "guestProfile",
                Ax.f()
                    .object(
                        Map.of(
                            "fullName", Ax.f().string("Primary guest").min(2),
                            "dietaryNotes",
                                Ax.f().string("Dietary requirements").optional())))
            .build();
    Map<String, Object> output =
        Ax.ax(signature)
            .forward(
                client(),
                Map.of(
                    "requestText",
                    "Book dinner for four people under the name Ada Lovelace.",
                    "contactEmail",
                    "ada@example.com"));
    System.out.println(Json.stringify(output));
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class UsageObserverExample {
  public static void main(String[] args) throws Exception {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }

    List<AxUsageEvent> events = new ArrayList<>();
    AxGlobals.setUsageObserver(events::add);
    OpenAICompatibleClient client =
        new OpenAICompatibleClient(
            Map.of(
                "api_key", apiKey,
                "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"),
                "usageContext",
                    Map.of(
                        "tenantId", "tenant-42",
                        "feature", "support-chat",
                        "attributes", Map.of("environment", "example"))));
    try {
      client.chat(
          Map.of(
              "chat_prompt",
              List.of(Map.of("role", "user", "content", "Reply with one short greeting."))),
          Map.of(
              "usageContext",
              Map.of("userId", "user-7", "requestId", UUID.randomUUID().toString())));
      System.out.println(Json.stringify(events.stream().map(AxUsageEvent::value).toList()));
    } finally {
      AxGlobals.setUsageObserver(null);
    }
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class ProviderStreamExample {
  public static void main(String[] args) throws Exception {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    AxAIService client = Ax.ai("openai", Map.of(
      "api_key", apiKey,
      "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.6-luna")
    ));
    long started = System.nanoTime();
    try (AxChatStream stream = client.openStream(Map.of(
      "chat_prompt", List.of(Map.of("role", "user", "content", "Reply with exactly: streaming works")),
      "model_config", Map.of("temperature", 1)
    ))) {
      for (Map<String, Object> event : stream) {
        List<?> results = (List<?>) event.get("results");
        Object content = results.isEmpty() ? null : ((Map<?, ?>) results.get(0)).get("content");
        if (content != null && !content.toString().isEmpty()) {
          System.out.printf("[%d ms] %s", (System.nanoTime() - started) / 1_000_000, content);
        }
      }
    }
    System.out.println();
  }
}

import dev.axllm.ax.*;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

public final class CancellationExample {
  public static void main(String[] args) throws Exception {
    AtomicInteger calls = new AtomicInteger();
    OpenAICompatibleClient.Transport transport = request -> {
      calls.incrementAndGet();
      return Map.of("status", 200, "json", Map.of());
    };
    AxAIService client = new OpenAICompatibleClient(Map.of(
      "api_key", "test-key", "model", "gpt-5.6-luna", "transport", transport
    ));
    AxCancellationToken token = new AxCancellationToken();
    assert token.cancel("user stopped") && !token.cancel("later reason");

    try {
      client.chatWithCancellation(
        Map.of("chat_prompt", List.of(Map.of("role", "user", "content", "This must not be sent."))),
        Map.of(), token
      );
      throw new AssertionError("pre-cancelled request unexpectedly completed");
    } catch (AxAIServiceAbortedError error) {
      assert "user stopped".equals(error.reason()) && !error.retryable;
    }

    assert calls.get() == 0;
    System.out.println("cancelled before transport: user stopped");
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class GeminiServiceTierExample {
  private static String apiKey() {
    String value = System.getenv("GOOGLE_API_KEY");
    if (value == null || value.isBlank()) value = System.getenv("GOOGLE_APIKEY");
    if (value == null || value.isBlank()) {
      throw new IllegalStateException("Set GOOGLE_API_KEY or GOOGLE_APIKEY to run this example.");
    }
    return value;
  }

  public static void main(String[] args) throws Exception {
    GoogleGeminiClient client = new GoogleGeminiClient(Map.of(
        "api_key", apiKey(),
        "model", System.getenv().getOrDefault("AX_GEMINI_MODEL", "gemini-3.8-flash")));
    Map<String, Object> out = client.chat(Map.of(
        "chat_prompt", List.of(Map.of(
            "role", "user",
            "content", "Explain in one sentence why batch evaluations save time."))),
        Map.of("service_tier", "flex"));
    System.out.println(Json.stringify(out));
  }
}

import dev.axllm.ax.*;
import java.util.*;
public final class NativeFileRoutingExample {
 public static void main(String[] args) throws Exception {
  String key=System.getenv("OPENAI_API_KEY");if(key==null||key.isBlank())key=System.getenv("OPENAI_APIKEY");
  var client=Ax.ai("openai",Map.of("api_key",Objects.requireNonNull(key,"Set OPENAI_API_KEY or OPENAI_APIKEY"),"model","gpt-6-astra","model_config",Map.of("thinkingTokenBudget","low")));
  var router=new AxProviderRouter(Map.of("providers",Map.of("primary",client)));
  var program=Ax.ax("document:file -> summary:string");
  var result=program.forward(router,Map.of("document",Map.of("filename","report.pdf","mimeType","application/pdf","data",Objects.requireNonNull(System.getenv("AX_PDF_BASE64"),"Set AX_PDF_BASE64"))),Map.of("serviceTier","standard"));
  System.out.println(Json.stringify(result));
 }
}

import dev.axllm.ax.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;

public final class AstraAsyncExample {
  public static void main(String[] args) throws Exception {
    String key=System.getenv("OPENAI_API_KEY");if(key==null||key.isBlank())key=System.getenv("OPENAI_APIKEY");if(key==null||key.isBlank())throw new IllegalArgumentException("Set OPENAI_API_KEY or OPENAI_APIKEY.");
    AiClient client=Ax.ai("openai",Map.of("api_key",key,"model","gpt-6-astra","model_config",Map.of("thinkingTokenBudget","low","max_tokens",4096)));
    CountDownLatch pending=new CountDownLatch(1);AtomicBoolean finished=new AtomicBoolean(),overlap=new AtomicBoolean(),steered=new AtomicBoolean();AtomicInteger applied=new AtomicInteger();
    AxRunControl control=Ax.runControl();control.onEvent(event->{if("applied".equals(event.get("type")))applied.incrementAndGet();});
    Tool slow=Ax.fn("slow_reference").description("Look up a reference; takes a few seconds.").execution("background").handler(values->{pending.countDown();if(steered.compareAndSet(false,true)){control.steer("Include the word VERIFIED in the final answer.");control.setThinkingTokenBudget("medium");}Thread.sleep(6000);finished.set(true);return "REF-42";}).build();
    Tool label=Ax.fn("local_label").description("Read an independent local label immediately.").handler(values->{if(pending.await(3,TimeUnit.SECONDS)&&!finished.get())overlap.set(true);return "LAUNCH";}).build();
    var program=Ax.ax("question -> answer").addTool(slow).addTool(label);
    var result=program.forward(client,Map.of("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."),Map.of("control",control,"serviceTier","standard","maxSteps",6));
    String answer=result.toString();for(String word:List.of("REF-42","LAUNCH","VERIFIED"))if(!answer.contains(word))throw new AssertionError("Missing final result: "+answer);
    if(!overlap.get())throw new AssertionError("No independent work while background tool was pending");if(applied.get()!=2)throw new AssertionError("Control updates were not applied");
    System.out.println(answer);System.out.println("Background overlap verified; steering and reasoning applied at the next response.");
  }
}

import dev.axllm.ax.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;
public final class AstraCancellationExample {
 public static void main(String[] args)throws Exception {
  String key=System.getenv("OPENAI_API_KEY");if(key==null||key.isBlank())key=System.getenv("OPENAI_APIKEY");if(key==null||key.isBlank())throw new IllegalArgumentException("Set OPENAI_API_KEY or OPENAI_APIKEY.");
  var control=Ax.runControl();var settled=new CountDownLatch(1);var started=new AtomicLong();
  var lookup=Ax.fn("lookup").description("Look up the reference.").execution("background").contextHandler((values,cancelled)->{started.set(System.nanoTime());control.abort();long deadline=System.nanoTime()+TimeUnit.SECONDS.toNanos(2);try {while(!cancelled.getAsBoolean()&&System.nanoTime()<deadline)Thread.sleep(1);}catch(InterruptedException interrupted){Thread.currentThread().interrupt();}if(!cancelled.getAsBoolean())throw new IllegalStateException("Tool missed cancellation");settled.countDown();return "LATE: discard this result";}).build();
  var program=Ax.ax("question -> answer").addTool(lookup);
  var client=Ax.ai("openai",Map.of("api_key",key,"model","gpt-6-astra","model_config",Map.of("thinkingTokenBudget","low","max_tokens",2048)));
  try {program.forward(client,Map.of("question","Call lookup once and return its result."),Map.of("control",control,"serviceTier","standard"));throw new AssertionError("Cancelled run returned success");}
  catch(CancellationException error){long elapsed=System.nanoTime()-started.get();if(started.get()==0||!error.getMessage().contains("unresolved calls")||elapsed>TimeUnit.SECONDS.toNanos(2)||!settled.await(2,TimeUnit.SECONDS))throw new AssertionError("Cancellation failed",error);System.out.println("Cancelled in "+TimeUnit.NANOSECONDS.toMillis(elapsed)+"ms; "+error.getMessage());}
 }
}

import dev.axllm.ax.*;
import java.nio.file.*;
import java.util.*;

public final class ContextGenerationExample {
  static String apiKey() {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    }
    return apiKey;
  }

  static OpenAICompatibleClient client() {
    return new OpenAICompatibleClient(
        Map.of("api_key", apiKey(), "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"), "model_config", Map.of("temperature", 0.0)));
  }

  public static void main(String[] args) throws Exception {
    AxGen program = Ax.ax("context:string, question:string -> answer:string, citations:string[]");
    Map<String, Object> output = program.forward(client(), Map.of("context", "Ax uses signatures, ai(), ax(), agent(), flow(), and optimize().", "question", "How should a new developer think about Ax?"));
    System.out.println(Json.stringify(output));
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class TypesafeNativeExample {
  public static void main(String[] args) throws Exception {
    var client = Ax.typesafe(Map.of("apiKey", System.getenv("TYPESAFE_APIKEY")));
    if (client.listModels().isEmpty()) throw new AssertionError("Empty model catalog");
    var account = new LinkedHashMap<String, Object>();
    account.put("tier", "enterprise");
    account.put("notes", null);
    var request = new AxAITypesafeClient.Request(
        Map.of("ticket", "Checkout is unavailable for all customers after the latest deployment.",
               "account", account),
        Map.of(
            "urgent", AxAITypesafeClient.Question.noul(
                Map.of("question", "Does this need immediate attention?"),
                Map.of("true", "Customers cannot complete a core task", "false", "Routine request")),
            "team", AxAITypesafeClient.Question.choice(
                "Who should handle the ticket?",
                Map.of("support", "Usage guidance",
                       "billing", Map.of("scope", "Invoices and payments"),
                       "engineering", "Product failures")),
            "severity", AxAITypesafeClient.Question.score(
                "Rate customer impact",
                List.of("Minor inconvenience", "One task blocked",
                        "Core task unavailable", "Widespread outage"))));
    var response = client.systemOne(request);
    double probability = ((AxAITypesafeClient.Noul) response.answers().get("urgent")).noul();
    double score = ((AxAITypesafeClient.Score) response.answers().get("severity")).score();
    if (probability < 0 || probability > 1 || score < 0 || score > 3) throw new AssertionError("Invalid answer bounds");
    // Apply thresholds and custom score scales in application code.
    System.out.println(Json.stringify(Map.of("page_on_call", probability >= 0.9, "severity_1_to_5", 1 + 4 * score / 3, "response", response.toMap())));
  }
}

import dev.axllm.ax.*;
import java.util.*;

public final class AdaptiveBalancerExample {
  static String requiredKey() {
    String value = System.getenv("OPENAI_API_KEY");
    if (value == null || value.isBlank()) value = System.getenv("OPENAI_APIKEY");
    if (value == null || value.isBlank()) throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    return value;
  }

  public static void main(String[] args) throws Exception {
    String key = requiredKey();
    String model = System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini");
    List<AxAIService> services = List.of(
        new OpenAICompatibleClient(Map.of("api_key", key, "model", model, "base_url", System.getenv().getOrDefault("OPENAI_PRIMARY_BASE_URL", "https://api.openai.com/v1"))),
        new OpenAICompatibleClient(Map.of("api_key", System.getenv().getOrDefault("OPENAI_BACKUP_API_KEY", key), "model", model, "base_url", System.getenv().getOrDefault("OPENAI_BACKUP_BASE_URL", "https://api.openai.com/v1"))));

    var store = new AxInMemoryBalancerStatsStore();
    List<String> routeKeys = List.of("openai-primary", "openai-backup");
    List<String> events = new ArrayList<>();
    var strategy = new AxBalancerAdaptiveStrategy(6_000, 0.02)
        .expectedTokens(1_200, 300)
        .namespace("support-summary-v1")
        .routeKey((service, index) -> routeKeys.get(index))
        .slice(context -> context.get("options") instanceof Map<?, ?> options && Boolean.TRUE.equals(options.get("stream")) ? "streaming" : "interactive")
        .statsStore(store)
        .onRoutingEvent(event -> events.add(event.type()));
    AxBalancer balancer = new AxBalancer(services, new AxBalancerOptions().strategy(strategy));
    Map<String, Object> response = balancer.chat(Map.of("model", model, "chat_prompt", List.of(Map.of("role", "user", "content", "Summarize why shared routing state matters."))));
    System.out.println(Json.stringify(response));
    System.out.println(events);
  }
}

import dev.axllm.ax.*;
import dev.axllm.ax.runtime.quickjs.*;
import java.util.*;

public final class RuntimeHooksExample {
  static final class LogSpan implements AxSpan {
    private final String name;
    LogSpan(String name) { this.name = name; System.out.println("[span:start] " + name); }
    public void setAttributes(Map<String, Object> attributes) {}
    public void addEvent(String event, Map<String, Object> attributes) { System.out.println("[span:event] " + name + " " + event); }
    public void recordException(Throwable error) { System.out.println("[span:error] " + name + " " + error); }
    public void setStatus(String status, String description) {}
    public void end() { System.out.println("[span:end] " + name); }
  }

  static final class LogMeter implements AxMeter {
    public AxCounter createCounter(String name, AxMetricInstrumentOptions options) {
      return (value, attributes) -> System.out.println("[metric] " + name + " += " + value);
    }
    public AxHistogram createHistogram(String name, AxMetricInstrumentOptions options) {
      return (value, attributes) -> System.out.println("[metric] " + name + " = " + value);
    }
    public AxGauge createGauge(String name, AxMetricInstrumentOptions options) {
      return (value, attributes) -> System.out.println("[metric] " + name + " = " + value);
    }
  }

  static AxRateLimiter limiter(String label) {
    return (next, info) -> {
      System.out.printf("[limit:%s] %s %s/%s stream=%s%n", label, info.operation(), info.provider(), info.model(), info.streaming());
      return next.execute();
    };
  }

  public static void main(String[] args) throws Exception {
    String apiKey = System.getenv("OPENAI_API_KEY");
    if (apiKey == null || apiKey.isBlank()) apiKey = System.getenv("OPENAI_APIKEY");
    if (apiKey == null || apiKey.isBlank()) throw new IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.");
    OpenAICompatibleClient client = new OpenAICompatibleClient(Map.of(
        "api_key", apiKey,
        "model", System.getenv().getOrDefault("AX_OPENAI_MODEL", "gpt-5.4-mini"),
        "model_config", Map.of("temperature", 0.0)));
    AxTracer tracer = start -> new LogSpan(start.name());
    AxMeter meter = new LogMeter();
    AxRuntimeHooks overrideHooks = new AxRuntimeHooks(limiter("forward"), tracer, meter);

    AxGlobals.setRateLimiter(limiter("global"));
    AxGlobals.setTracer(tracer);
    AxGlobals.setMeter(meter);
    try {
      System.out.println(Ax.ax("topic:string -> summary:string").forward(client, Map.of("topic", "portable Ax runtime hooks")));
      AxAgent helper = Ax.agent("question:string -> answer:string", Map.of());
      try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
        System.out.println(helper.forward(client, Map.of("question", "What does a rate limiter wrap?"), Map.of("runtime", runtime, "max_actor_steps", 12), overrideHooks));
      }
      AxFlow workflow = Ax.flow(Map.of("id", "examples.runtimeHooks"))
          .execute("outline", Ax.ax("topic:string -> outline:string"))
          .execute("polish", Ax.ax("outline:string -> answer:string"))
          .returns(Map.of("answer", "polish"));
      System.out.println(workflow.forward(client, Map.of("topic", "Ax runtime hooks"), Map.of(), overrideHooks));
    } finally {
      AxGlobals.setRateLimiter(null);
      AxGlobals.setTracer(null);
      AxGlobals.setMeter(null);
    }
  }
}