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.*;

public final class BranchFlowExample {
  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 classifier = Ax.ax("request:string -> route:class \"support, sales, engineering\"");
    AxGen responder = Ax.ax("request:string, route:string -> response:string");
    AxFlow program =
        Ax.flow(Map.of("id", "examples.branchFlow"))
            .execute(
                "classifier",
                classifier,
                Map.of(
                    "reads", List.of("request"),
                    "writes", List.of("classifierResult", "route")))
            .execute(
                "responder",
                responder,
                Map.of(
                    "reads", List.of("request", "route"),
                    "writes", List.of("responderResult", "response")))
            .returns(Map.of("route", "route", "response", "response"));
    Map<String, Object> output = program.forward(client(), Map.of("request", "A customer says checkout is down for their enterprise account."));
    System.out.println(Json.stringify(output));
  }
}

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

public final class ParallelFlowExample {
  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 research = Ax.ax("topicText:string -> factList:string[]");
    AxGen audience = Ax.ax("topicText:string -> audienceAngle:string");
    AxGen join = Ax.ax("factList:string[], audienceAngle:string -> briefText:string");
    AxFlow program =
        Ax.flow(Map.of("id", "examples.parallelFlow"))
            .execute(
                "research",
                research,
                Map.of(
                    "reads", List.of("topicText"),
                    "writes", List.of("researchResult", "factList")))
            .execute(
                "audience",
                audience,
                Map.of(
                    "reads", List.of("topicText"),
                    "writes", List.of("audienceResult", "audienceAngle")))
            .execute(
                "join",
                join,
                Map.of(
                    "reads", List.of("factList", "audienceAngle"),
                    "writes", List.of("joinResult", "briefText")))
            .returns(Map.of("briefText", "briefText"));
    Map<String, Object> output =
        program.forward(
            client(),
            Map.of(
                "topicText",
                "Why typed contracts make multi-step LLM systems easier to maintain"));
    System.out.println(Json.stringify(output));
  }
}

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

public final class AstraAsyncFlowExample {
  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 workflow=Ax.flow().execute("lookup",program,Map.of("writes",List.of("answer"))).execute("verify",Ax.ax("answer -> report \"Repeat the exact reference, label, and verification word from the answer.\""),Map.of("reads",List.of("answer"))).returns(Map.of("answer","report"));
    var result=workflow.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()!=4)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.*;

public class AstraParallelFlowExample {
  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 both=new CyclicBarrier(2);var updatesReady=new CountDownLatch(1);var paths=new HashSet<String>();var applied=new ArrayList<Map<String,Object>>();
    control.onEvent(event->{
      if("tool.started".equals(event.get("type"))){paths.add(String.valueOf(event.get("path")));if(paths.size()==2){control.steer("Include VERIFIED with the exact reference in your final answer.");control.setThinkingTokenBudget("medium","root/left");updatesReady.countDown();}}
      if("applied".equals(event.get("type")))applied.add(event);
    });
    var lookup=Ax.fn("lookup").description("Look up the exact reference once.").execution("background").handler(values->{
      both.await(45,TimeUnit.SECONDS);if(!updatesReady.await(5,TimeUnit.SECONDS))throw new IllegalStateException("Controller did not observe both nodes");return "REF-42";
    }).build();
    var program=Ax.ax("question -> answer").addTool(lookup);
    var workflow=Ax.flow().execute("left",program).execute("right",program).returns(Map.of("left","leftResult","right","rightResult"));
    var result=workflow.forward(client,Map.of("question","Call lookup exactly once. If its result is pending, return a brief progress message without calling it again. Return the exact reference when its result arrives."),Map.of("control",control,"serviceTier","standard","maxSteps",6));
    if(!paths.equals(Set.of("root/left","root/right"))||applied.size()!=3)throw new AssertionError("Missing scoped controls: "+applied);
    for(String node:List.of("left","right")){String answer=String.valueOf(result.get(node));if(!answer.contains("REF-42")||!answer.contains("VERIFIED"))throw new AssertionError("Missing final result: "+answer);}
    System.out.println(Json.stringify(Map.of("result",result,"parallel_overlap",true,"applied_controls",applied)));
  }
}

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

public final class ComposedFlowExample {
  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("topic:string -> outline:string[]");
    AxFlow program =
        Ax.flow(Map.of("id", "examples.composedFlow"))
            .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("topic", "How Ax moves from typed generation to agents, flows, and optimization"));
    System.out.println(Json.stringify(output));
  }
}

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

public final class RefineFlowExample {
  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 draft = Ax.ax("topicText:string -> draftText:string");
    AxGen critique = Ax.ax("draftText:string -> critiqueText:string");
    AxGen revise = Ax.ax("draftText:string, critiqueText:string -> revisedText:string");
    AxFlow program =
        Ax.flow(Map.of("id", "examples.refineFlow"))
            .execute(
                "draft",
                draft,
                Map.of(
                    "reads", List.of("topicText"),
                    "writes", List.of("draftResult", "draftText")))
            .execute(
                "critique",
                critique,
                Map.of(
                    "reads", List.of("draftText"),
                    "writes", List.of("critiqueResult", "critiqueText")))
            .execute(
                "revise",
                revise,
                Map.of(
                    "reads", List.of("draftText", "critiqueText"),
                    "writes", List.of("reviseResult", "revisedText")))
            .returns(Map.of("revisedText", "revisedText"));
    Map<String, Object> output =
        program.forward(
            client(),
            Map.of("topicText", "Explain automatic flow parallelism to a backend engineer."));
    System.out.println(Json.stringify(output));
  }
}