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 dev.axllm.ax.runtime.quickjs.*;
import java.util.*;

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

  // A raw, noisy incident report. It lives in `contextFields`, so the agent works
  // it inside the runtime; `contextPolicy: lean` keeps the prompt compact by
  // preferring live runtime state and summaries over replaying the raw text.
  static final String REPORT = String.join("\n",
      "[2026-03-02 14:01:22Z] INFO  gateway       deploy svc-checkout-edge v812 -> prod (channel: canary 10%)",
      "[2026-03-02 14:03:10Z] WARN  checkout-api  p95 latency 1180ms (baseline 240ms) region=eu-west-1",
      "[2026-03-02 14:04:55Z] ERROR checkout-api  502 from svc-payments-gw: upstream timeout (10s) tenant_tier=enterprise",
      "[2026-03-02 14:05:01Z] ERROR checkout-api  502 from svc-payments-gw: upstream timeout (10s) tenant_tier=enterprise",
      "[2026-03-02 14:05:40Z] WARN  payments-gw   circuit half-open, 3 retries exhausted for order=ord_99214",
      "[2026-03-02 14:06:12Z] INFO  gateway       canary widened 10% -> 50% for svc-checkout-edge v812",
      "[2026-03-02 14:07:33Z] ERROR checkout-api  502 from svc-payments-gw: upstream timeout (10s) tenant_tier=enterprise",
      "[2026-03-02 14:08:02Z] ERROR checkout-api  user-visible: \"Payment could not be processed\" shown to 1,284 sessions",
      "[2026-03-02 14:09:48Z] WARN  payments-gw   connection pool exhausted (max=64) waiting=210",
      "[2026-03-02 14:11:20Z] INFO  on-call       paged: SEV-2 opened (eu-west-1 checkout error rate 38%)",
      "[2026-03-02 14:14:05Z] INFO  gateway       rollback svc-checkout-edge v812 -> v811 (channel: prod 100%)",
      "[2026-03-02 14:17:41Z] INFO  checkout-api  p95 latency 260ms, error rate 0.4% region=eu-west-1",
      "[2026-03-02 14:19:10Z] INFO  on-call       SEV-2 mitigated, monitoring for 30m");

  public static void main(String[] args) throws Exception {
    AxAgent triage = Ax.agent(
        "report:string, question:string -> severity:class \"low, medium, high, critical\", rootCause:string, nextSteps:string[], evidence:string[] \"Quoted log lines that support the assessment\"",
        Map.of(
            "contextFields", List.of("report"),
            "contextPolicy", Map.of("preset", "lean", "budget", "balanced"),
            "runtime", Map.of("language", "JavaScript")));

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      Map<String, Object> result = triage.forward(
          client(),
          Map.of(
              "report", REPORT,
              "question", "What happened, how bad was it, and what should the on-call do next? Cite the lines you relied on."),
          Map.of("runtime", runtime, "max_actor_steps", 12));

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

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

public final class AstraAsyncAgentExample {
  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.").execution("background").handler(values->{if(pending.await(3,TimeUnit.SECONDS)&&!finished.get())overlap.set(true);return "LAUNCH";}).build();
    var program=Ax.agent("question -> answer",Map.of("runtime",Map.of("language","JavaScript"),"directResponse","off","functions",List.of(slow,label)));
    try (AxQuickJsCodeRuntime runtime=new AxQuickJsCodeRuntime()) {
    var result=program.forward(client,Map.of("question","Use the native tools tools_slow_reference and tools_local_label. First call tools_slow_reference. While it is pending, call tools_local_label. Call each tool only once; do not call a tool again while its result is pending. In the executor, call the native tools directly rather than invoking them from actor code; then use final(...) in the code runtime to pass their results to the responder. Return both exact results in one sentence."),Map.of("runtime",runtime,"max_actor_steps",12,"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 dev.axllm.ax.runtime.quickjs.AxQuickJsCodeRuntime;
import java.util.*;

public final class AstraChildAgentExample {
 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 IllegalStateException("Set OPENAI_API_KEY or OPENAI_APIKEY.");
  var client=Ax.ai("openai",Map.of("api_key",key,"model","gpt-6-astra","model_config",Map.of("thinkingTokenBudget","low","max_tokens",4096)));
  var control=Ax.runControl();var events=new ArrayList<Map<String,Object>>();control.onEvent(events::add);
  control.steer("Include VERIFIED in the final answer.");control.steer("Include CHILD-CHECK in the final answer.","root/team.researcher");control.setThinkingTokenBudget("medium","root/team.researcher/executor");
  try(var parentRuntime=new AxQuickJsCodeRuntime(Map.of("timeoutMs",180000));var childRuntime=new AxQuickJsCodeRuntime(Map.of("timeoutMs",180000))){
   var child=Ax.agent("question -> answer",Map.of("directResponse","off","runtime",childRuntime));
   var parent=Ax.agent("question -> answer",Map.of("directResponse","off","runtime",parentRuntime)).addChildAgent("team","researcher",child);
   var result=parent.forward(client,Map.of("question","Delegate to team.researcher exactly once by calling await team.researcher({question: \"Compute 37 + 5 and return the exact sum.\"}) in actor code. Pass the complete child answer as evidence to final(...), then report it in your final answer."),Map.of("control",control,"serviceTier","standard","max_actor_steps",8));
   for(String word:List.of("42","VERIFIED","CHILD-CHECK"))if(!result.toString().contains(word))throw new AssertionError(result);
   if(events.stream().noneMatch(event->"applied".equals(event.get("type"))&&"root/team.researcher/executor".equals(event.get("path"))))throw new AssertionError("Child control did not apply");
   System.out.println(Json.stringify(Map.of("result",result,"usage",parent.getUsage())));
  }
 }
}

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

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

  // A long, messy brief -- exactly the kind of input you do not want replayed into
  // the prompt on every turn. `contextFields` holds it in the runtime, the
  // `checkpointed` policy compacts older turns once the prompt grows, and
  // `maxRuntimeChars` caps how much runtime output is echoed back.
  static final String BRIEF = String.join("\n",
      "# Migration brief: monolith -> services (draft, unordered notes)",
      "",
      "Current: single Rails monolith, Postgres primary + 1 replica, Sidekiq for jobs.",
      "Pain: deploys take 40m, one bad migration locks the orders table, on-call burnout.",
      "Constraints: no downtime windows > 5m, PCI scope must shrink, team of 6, 2 quarters.",
      "Hot paths: checkout (writes orders, payments), search (read-heavy), notifications (async).",
      "Known landmines: payments code has no tests; search shares the orders DB; a nightly",
      "cron rebuilds the catalog and pins CPU for ~20m; the replica lags up to 90s under load.",
      "Org wants: independent deploys for checkout, smaller blast radius, an audit trail.",
      "Nice to have: event log for orders, read-model for search, feature flags.",
      "Hard no: a big-bang rewrite; introducing Kubernetes this year.");

  public static void main(String[] args) throws Exception {
    AxAgent specialist = Ax.agent(
        "brief:string, goal:string -> plan:string[] \"Ordered, concrete steps\", answer:string, risks:string[]",
        Map.of(
            "contextFields", List.of("brief"),
            "contextPolicy", Map.of("preset", "checkpointed", "budget", "balanced"),
            "maxRuntimeChars", 3000,
            "runtime", Map.of("language", "JavaScript")));

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      Map<String, Object> result = specialist.forward(
          client(),
          Map.of(
              "brief", BRIEF,
              "goal", "Propose a safe, incremental 2-quarter plan to split checkout out first, respecting the hard constraints."),
          Map.of("runtime", runtime, "max_actor_steps", 12));

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

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

public final class ModelPanelExample {
  static String firstNonBlank(String... values) {
    for (String value : values) {
      if (value != null && !value.isBlank()) return value;
    }
    return null;
  }

  public static void main(String[] args) throws Exception {
    String openaiKey = firstNonBlank(System.getenv("OPENAI_API_KEY"), System.getenv("OPENAI_APIKEY"));
    String googleKey = firstNonBlank(System.getenv("GOOGLE_APIKEY"), System.getenv("GOOGLE_API_KEY"));
    String anthropicKey = firstNonBlank(System.getenv("ANTHROPIC_APIKEY"), System.getenv("ANTHROPIC_API_KEY"));
    if (openaiKey == null || googleKey == null || anthropicKey == null) {
      throw new IllegalStateException(
          "Set OPENAI_APIKEY, GOOGLE_APIKEY, and ANTHROPIC_APIKEY to run this multi-provider panel.");
    }

    // A panel of three different providers, each answering the same question
    // independently. Plain ax() composition (no agent runtime): fan out to the
    // panel, judge the candidates, then synthesize one grounded answer.
    List<String> panelModels = List.of(
        "openai/gpt-5.4-mini",
        "google/gemini-3.5-flash",
        "anthropic/claude-haiku-4.5");
    List<AiClient> panelClients = List.of(
        new OpenAICompatibleClient(Map.of(
            "api_key", openaiKey, "model", "gpt-5.4-mini", "model_config", Map.of("temperature", 0.0))),
        new GoogleGeminiClient(Map.of(
            "api_key", googleKey, "model", "gemini-3.5-flash")),
        new AnthropicClient(Map.of(
            "api_key", anthropicKey, "model", "claude-haiku-4-5")));

    AxGen researcher = Ax.ax(
        "question:string -> answer:string, keyFindings:string[], citations:string[], confidence:number");
    researcher.setInstruction(
        "Answer independently. Use evidence. Call out uncertainty. Do not optimize for consensus.");

    AxGen judge = Ax.ax(
        "question:string, candidates:json -> consensus:string[], contradictions:string[], uniqueInsights:string[], blindSpots:string[]");
    judge.setInstruction(
        "Compare the candidates. Find agreement, conflicts, missing coverage, and unique useful points.");

    AxGen synthesizer = Ax.ax(
        "question:string, candidates:json, review:json -> answer:string, citations:string[], caveats:string[]");
    synthesizer.setInstruction(
        "Write one final answer grounded in the candidates and review. Resolve conflicts explicitly.");

    String question = "What are the strongest arguments for and against a national carbon tax?";

    List<Object> candidates = new ArrayList<>();
    for (int i = 0; i < panelClients.size(); i++) {
      Map<String, Object> response = researcher.forward(panelClients.get(i), Map.of("question", question));
      Map<String, Object> candidate = new LinkedHashMap<>();
      candidate.put("model", panelModels.get(i));
      candidate.putAll(response);
      candidates.add(candidate);
      System.out.println("[panel] " + panelModels.get(i) + " responded.");
    }

    // The judge + synthesizer run on one of the panel clients (OpenAI here).
    AiClient orchestrator = panelClients.get(0);
    Map<String, Object> review = judge.forward(orchestrator, Map.of("question", question, "candidates", candidates));
    Map<String, Object> finalAnswer = synthesizer.forward(
        orchestrator, Map.of("question", question, "candidates", candidates, "review", review));

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