Advanced Start

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.4-mini"), "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."));
    System.out.println(Json.stringify(output));
  }
}

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

public final class BasicAgentExample {
  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)));
  }

  // The handbook can be arbitrarily large. Listing it in `contextFields` keeps it
  // in the agent's runtime so it never inflates the model prompt -- the agent reads
  // it through code, not through tokens. That is the whole point of an Ax agent
  // over a plain gen() call: the source material stays out of the context window.
  static final String HANDBOOK = String.join("\n",
      "# Acme Cloud -- Support Handbook",
      "",
      "## Billing",
      "- Invoices are issued on the 1st of each month and are due net-15.",
      "- Plan downgrades take effect at the END of the current billing cycle, not immediately.",
      "- Refunds are issued to the original payment method within 5 business days.",
      "",
      "## Access",
      "- Seats can be added by any workspace Owner under Settings -> Members.",
      "- SSO (SAML) is available on Enterprise; SCIM provisioning is Owner-only.",
      "",
      "## Incidents",
      "- Status and uptime are published at status.acme.example.",
      "- Sev-1 incidents page the on-call within 5 minutes; updates post every 30 minutes.",
      "",
      "## Data",
      "- Exports are available in CSV and JSON from Settings -> Data.",
      "- Deleted workspaces are recoverable for 30 days, then permanently purged.");

  public static void main(String[] args) throws Exception {
    AxAgent assistant = Ax.agent(
        "question:string, handbook:string -> answer:string, citations:string[] \"Handbook sections the answer relies on\"",
        // Keep the handbook in the runtime, out of the prompt.
        Map.of("contextFields", List.of("handbook"), "runtime", Map.of("language", "JavaScript")));

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      Map<String, Object> result = assistant.forward(
          client(),
          Map.of(
              "question", "A customer downgraded their plan today. When does it take effect, and can they get a refund for the current cycle?",
              "handbook", HANDBOOK),
          Map.of("runtime", runtime, "max_actor_steps", 12));

      System.out.println(Json.pretty(result));
    }
  }
}

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

public final class SequentialFlowExample {
  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 step = Ax.ax("documentText:string -> summaryText:string");
    AxFlow program =
        Ax.flow(Map.of("id", "examples.sequentialFlow"))
            .execute("step", step)
            .map("note", state -> Map.of("note", "Mapped flow state after the provider-backed step."))
            .returns(Map.of("step", "step", "note", "note"));
    Map<String, Object> output = program.forward(client(), Map.of("documentText", "Ax gives developers signatures, provider clients, agents, flows, tracing, and optimization."));
    System.out.println(Json.stringify(output));
  }
}

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

public final class SpeechAudioExample {
  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 OpenAIResponsesClient client() {
    return new OpenAIResponsesClient(
        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 {
    OpenAIResponsesClient audio = client();
    Map<String, Object> speech = audio.speak(Map.of("text", "Ax turns LLM prompts into typed programs.", "voice", "alloy", "format", "mp3"));
    System.out.println(Json.stringify(Map.of("format", speech.get("format"), "audioBytesBase64", String.valueOf(speech.get("audio")).length())));
  }
}

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 java.nio.file.*;
import java.util.*;

public final class AxgenOptimizationExample {
  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)));
  }

  static final class ExampleOptimizer implements OptimizerEngine {
    public String name() { return "example"; }
    public String version() { return "1"; }
    public Map<String, Object> optimize(Map<String, Object> request) {
      return Map.of("componentMap", Map.of("priority::instruction", "Classify operational risk. Use high for production-impacting urgency."), "metadata", Map.of("source", "axgen"));
    }
  }

  public static void main(String[] args) throws Exception {
    AxGen program = new AxGen(Ax.s("emailText:string -> priority:class \"high, normal, low\", rationale:string"), Map.of("id", "priority", "instruction", "Classify the email priority."));
    Map<String, Object> baseline = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    Map<String, Object> artifact = program.optimizeWith(new ExampleOptimizer(), List.of(Map.of("emailText", "URGENT: checkout is down", "priority", "high")), Map.of("apply", false));
    program.applyOptimization(Json.stringify(artifact));
    Map<String, Object> after = program.forward(client(), Map.of("emailText", "Production checkout is failing for enterprise customers."));
    System.out.println(Json.stringify(Map.of("baseline", baseline, "after", after)));
  }
}

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

public final class NativeMCPToolsExample {
  public static void main(String[] args) {
    String key = Optional.ofNullable(System.getenv("OPENAI_API_KEY")).orElse(System.getenv("OPENAI_APIKEY"));
    String endpoint = System.getenv("MCP_URL");
    if (key == null || endpoint == null) throw new IllegalStateException("Set OPENAI_API_KEY and MCP_URL.");
    AxMCPClient mcp = new AxMCPClient(new AxMCPStreamableHTTPTransport(endpoint), Map.of("namespace", "inventory"));
    AxGen program = new AxGen(Ax.s("request:string -> answer:string"), Map.of("mcp", mcp));
    OpenAICompatibleClient llm = new OpenAICompatibleClient(Map.of("api_key", key, "model", "gpt-5.4-mini"));
    try {
      AxMCPClient.CatalogSnapshot catalog = mcp.inspectCatalog();
      System.out.println(Json.stringify(Map.of(
          "tools", catalog.tools().stream().map(tool -> tool.get("name")).toList(),
          "resources", catalog.resources(),
          "resourceTemplates", catalog.resourceTemplates())));
      System.out.println(Json.stringify(program.forward(llm, Map.of("request", "Reindex inventory."))));
    } finally {
      mcp.close();
    }
  }
}

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

public final class ResourceWakeAgentExample {
  public static void main(String[] args) throws Exception {
    String key=Optional.ofNullable(System.getenv("OPENAI_API_KEY")).orElse(System.getenv("OPENAI_APIKEY"));String endpoint=System.getenv("AX_MCP_ENDPOINT");if(key==null||endpoint==null)throw new IllegalStateException("Set OPENAI_API_KEY and AX_MCP_ENDPOINT.");boolean local=endpoint.startsWith("http://127.0.0.1");
    AxMCPStreamableHTTPTransport transport=new AxMCPStreamableHTTPTransport(endpoint,Map.of("ssrfProtection",Map.of("requireHttps",!local,"allowLocalhost",local,"allowPrivateNetworks",local)));AxMCPClient client=new AxMCPClient(transport,Map.of("namespace","inventory"));AxMCPEventSource source=new AxMCPEventSource(client,"inventory","tenant:demo","authenticated",AxMCPEventSource.all());AxAgent agent=Ax.agent("uri:string -> summary:string",Map.of("runtime",Map.of("language","JavaScript")));OpenAICompatibleClient llm=new OpenAICompatibleClient(Map.of("api_key",key,"model","gpt-5.4-mini"));CountDownLatch completed=new CountDownLatch(1);
    AxEventRuntime runtime=new AxEventRuntime(List.of(new AxEventRoute("resource-wake","wake",Map.of("types",List.of("mcp.resource.updated")),"inventory-agent",true,"strict",0))).registerTarget(new AxEventRuntime.Target("inventory-agent",(input,context)->{try(AxQuickJsCodeRuntime js=new AxQuickJsCodeRuntime()){Object output=agent.forward(llm,castMap(input),Map.of("runtime",js));System.out.println(Json.stringify(output));completed.countDown();return output;}}).mapInput((event,continuation)->Map.of("uri",castMap(event.data()).get("uri"))).retrySafety("idempotent")).addSource(source);runtime.start();if(!completed.await(60,TimeUnit.SECONDS))throw new IllegalStateException("Timed out waiting for an MCP resource notification");runtime.close();client.close();
  }
  @SuppressWarnings("unchecked")private static Map<String,Object> castMap(Object value){return(Map<String,Object>)value;}
}

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

public final class TaskResumeFlowExample {
  public static void main(String[] args)throws Exception{
    String key=Optional.ofNullable(System.getenv("OPENAI_API_KEY")).orElse(System.getenv("OPENAI_APIKEY"));String endpoint=System.getenv("AX_MCP_ENDPOINT");if(key==null||endpoint==null)throw new IllegalStateException("Set OPENAI_API_KEY and AX_MCP_ENDPOINT.");boolean local=endpoint.startsWith("http://127.0.0.1");AxMCPStreamableHTTPTransport transport=new AxMCPStreamableHTTPTransport(endpoint,Map.of("ssrfProtection",Map.of("requireHttps",!local,"allowLocalhost",local,"allowPrivateNetworks",local)));AxMCPClient client=new AxMCPClient(transport,Map.of("namespace","inventory"));client.addNotificationListener(message->{if("notifications/progress".equals(message.get("method")))System.out.println("MCP task progress");});client.init();String taskId=String.valueOf(castMap(client.callTool("start_reindex",Map.of("scope","all")).get("task")).get("taskId"));
    AxFlow flow=Ax.flow(Map.of("id","reindex-flow")).execute("status",Ax.ax("taskId:string -> status:string")).returns(Map.of("status","status"));OpenAICompatibleClient llm=new OpenAICompatibleClient(Map.of("api_key",key,"model","gpt-5.4-mini"));AtomicInteger calls=new AtomicInteger();CountDownLatch completed=new CountDownLatch(1);AxEventRuntime.Target target=new AxEventRuntime.Target("reindex-flow",(input,context)->{Object output=flow.forward(llm,castMap(input));System.out.println(Json.stringify(output));if(calls.incrementAndGet()>=2)completed.countDown();return output;}).mapInput((event,continuation)->Map.of("taskId",continuation==null?castMap(event.data()).get("taskId"):continuation.metadata.get("taskId"))).waitFor("mcp.task","taskKey",Map.of("taskId",taskId)).retrySafety("idempotent");
    AxEventRuntime runtime=new AxEventRuntime(List.of(new AxEventRoute("task-start","wake",Map.of("types",List.of("app.task.started")),"reindex-flow",false,"strict",0),new AxEventRoute("task-progress","observe",Map.of("types",List.of("mcp.progress")),null,false,"strict",0),new AxEventRoute("task-resume","resume",Map.of("types",List.of("mcp.task.status")),"reindex-flow",false,"strict",0))).registerTarget(target);runtime.start();runtime.publish(new AxEventEnvelope("task-start","app://tasks","app.task.started",Map.of("taskId",taskId,"taskKey","inventory:"+taskId)),"tenant:demo","authenticated");AxMCPEventSource source=new AxMCPEventSource(client,"inventory","tenant:demo","authenticated",List.of());source.start(event->runtime.publish(event,source.identityScope(),source.trust()));System.out.println("Waiting for terminal MCP task notification "+taskId);if(!completed.await(60,TimeUnit.SECONDS))throw new IllegalStateException("Timed out waiting for the MCP task continuation");source.close();runtime.close();client.close();
  }
  @SuppressWarnings("unchecked")private static Map<String,Object> castMap(Object value){return(Map<String,Object>)value;}
}