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

public final class IncidentLogForensicsExample {
  static GoogleGeminiClient client() {
    String apiKey = System.getenv("GOOGLE_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set GOOGLE_APIKEY to run this example.");
    }
    return new GoogleGeminiClient(Map.of(
        "api_key", apiKey,
        "model", System.getenv().getOrDefault("AX_GEMINI_MODEL", "gemini-3.5-flash")));
  }

  // ---------------------------------------------------------------------------
  // Synthetic CloudWatch-style export -- generated large on purpose. Dumping these
  // raw events into a prompt would blow the context window. The agent keeps them
  // in its runtime (contextFields) and only the *evidence it extracts* ever
  // reaches the model. Deterministic so the example is reproducible.
  // ---------------------------------------------------------------------------
  static List<Object> buildLogDump() {
    OffsetDateTime start = OffsetDateTime.of(2026, 3, 2, 13, 0, 0, 0, ZoneOffset.UTC);
    List<Object> events = new ArrayList<>();

    class Push {
      void at(int i, Map<String, Object> event) {
        Map<String, Object> e = new LinkedHashMap<>(event);
        e.put("timestamp", start.plusSeconds((long) i * 2).format(DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss'Z'")));
        e.put("requestId", "req-" + (100000 + i));
        events.add(e);
      }
    }
    Push push = new Push();

    for (int i = 0; i < 1600; i++) {
      // Routine, healthy traffic across the fleet.
      push.at(i, Map.of("level", "INFO", "service", "gateway", "statusCode", 200, "latencyMs", 40 + (i % 30), "message", "route ok GET /checkout"));
      push.at(i, Map.of("level", "INFO", "service", "search-api", "statusCode", 200, "latencyMs", 70 + (i % 50), "message", "query ok q=shoes"));

      // Window A: payments-gw upstream timeouts spill into checkout-api 502s for
      // enterprise tenants, with retry storms + pool exhaustion.
      if (i >= 300 && i < 520) {
        push.at(i, Map.of("level", "ERROR", "service", "payments-gw", "statusCode", 504, "latencyMs", 10000, "tenantTier", "enterprise", "message", "upstream timeout calling acquirer (10s)"));
        push.at(i, Map.of("level", "ERROR", "service", "checkout-api", "statusCode", 502, "tenantTier", "enterprise", "message", "bad gateway from svc-payments-gw"));
        if (i % 3 == 0) {
          push.at(i, Map.of("level", "WARN", "service", "payments-gw", "message", "connection pool exhausted (max=64) waiting=200+"));
          push.at(i, Map.of("level", "WARN", "service", "checkout-api", "tenantTier", "enterprise", "message", "user-visible: \"Payment could not be processed\""));
        }
      }

      // Window B: the nightly catalog-cron pins CPU and search-api returns 429s.
      if (i >= 1000 && i < 1120) {
        push.at(i, Map.of("level", "WARN", "service", "catalog-cron", "latencyMs", 0, "message", "rebuild step pinning CPU at 95% on shared node"));
        push.at(i, Map.of("level", "ERROR", "service", "search-api", "statusCode", 429, "message", "rate limited: downstream catalog unavailable"));
      }
    }

    return events;
  }

  public static void main(String[] args) throws Exception {
    GoogleGeminiClient client = client();

    List<Object> logs = buildLogDump();
    System.out.println("Generated " + logs.size() + " log events (kept out of the prompt).");

    AxAgent logRLM = Ax.agent(
        "task:string, logs:json \"Raw CloudWatch export; keep this out of the prompt\" -> architecture:string[] \"Services and how they call each other\", findings:json[] \"Each: issue, count, window, evidence, impact\", overallHealth:string, nextActions:string[]",
        Map.of(
            // The export stays in the runtime; only extracted evidence reaches the model.
            "contextFields", List.of("logs"),
            "contextPolicy", Map.of("preset", "lean", "budget", "balanced"),
            "maxRuntimeChars", 12000,
            "runtime", Map.of("language", "JavaScript")));

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      Map<String, Object> report = logRLM.forward(
          client,
          Map.of(
              "logs", logs,
              "task", "Infer the service architecture from the logs alone. Then find repeated errors, throttles, retries, and bad user states -- with the affected time window, an occurrence count, and concrete log evidence for each."),
          Map.of("runtime", runtime, "max_actor_steps", 40));

      System.out.println("\n=== Report ===");
      System.out.println(Json.pretty(report));
      System.out.println("\n=== Usage ===");
      System.out.println(Json.pretty(logRLM.getUsage()));
    }
  }
}

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

public final class CodebasePeekMapExample {
  static GoogleGeminiClient client() {
    String apiKey = System.getenv("GOOGLE_APIKEY");
    if (apiKey == null || apiKey.isBlank()) {
      throw new IllegalStateException("Set GOOGLE_APIKEY to run this example.");
    }
    return new GoogleGeminiClient(Map.of(
        "api_key", apiKey,
        "model", System.getenv().getOrDefault("AX_GEMINI_MODEL", "gemini-3.5-flash")));
  }

  // ---------------------------------------------------------------------------
  // A large module-dependency index for a monorepo. Each block is a record the
  // agent must *search* to answer -- the answers cannot be guessed, only computed
  // by filtering the index. Generated large so it would not fit comfortably in a
  // prompt; it lives in contextFields and is queried from the runtime.
  // ---------------------------------------------------------------------------
  static List<Map<String, Object>> buildModuleIndex() {
    List<Map<String, Object>> core = new ArrayList<>(List.of(
        module("packages/api/middleware/auth.ts", List.of("packages/shared"), "-"),
        module("packages/api/middleware/rateLimit.ts", List.of("packages/db"), "-"),
        module("packages/api/routes/checkout.ts", List.of("packages/api/middleware/auth.ts", "packages/services/orders/createOrder.ts", "packages/services/payments/charge.ts"), "-"),
        module("packages/api/routes/search.ts", List.of("packages/api/middleware/auth.ts", "packages/services/catalog/searchCatalog.ts"), "-"),
        module("packages/services/orders/createOrder.ts", List.of("packages/db", "packages/clients/bus"), "orders"),
        module("packages/services/orders/orderRepo.ts", List.of("packages/db"), "orders"),
        module("packages/services/payments/charge.ts", List.of("packages/clients/acquirer", "packages/db"), "payments"),
        module("packages/services/payments/refund.ts", List.of("packages/clients/acquirer", "packages/db"), "refunds"),
        module("packages/services/catalog/searchCatalog.ts", List.of("packages/db"), "-"),
        module("packages/clients/acquirer/index.ts", List.of("packages/shared"), "-"),
        module("packages/clients/bus/index.ts", List.of("packages/shared"), "-")));

    // Filler modules so the index is genuinely large; some also depend on the acquirer.
    for (int i = 0; i < 110; i++) {
      core.add(module(
          "packages/services/feature" + i + "/handler.ts",
          List.of(i % 4 == 0 ? "packages/clients/acquirer" : "packages/db", "packages/shared"),
          i % 6 == 0 ? "audit" : "-"));
    }
    return core;
  }

  static Map<String, Object> module(String path, List<String> imports, String writes) {
    Map<String, Object> m = new LinkedHashMap<>();
    m.put("path", path);
    m.put("imports", imports);
    m.put("writes", writes);
    return m;
  }

  @SuppressWarnings("unchecked")
  public static void main(String[] args) throws Exception {
    GoogleGeminiClient client = client();

    List<Map<String, Object>> modules = buildModuleIndex();
    StringBuilder sb = new StringBuilder();
    for (Map<String, Object> m : modules) {
      if (sb.length() > 0) sb.append("\n\n");
      sb.append("PATH: ").append(m.get("path")).append("\n")
        .append("IMPORTS: ").append(String.join(", ", (List<String>) m.get("imports"))).append("\n")
        .append("WRITES: ").append(m.get("writes"));
    }
    String codebaseIndex = sb.toString();
    System.out.println("Module index: " + modules.size() + " records (kept out of the prompt).");

    AxAgent analyst = Ax.agent(
        "context:string, question:string -> answer:string, paths:string[] \"Exact PATH values from the index that answer the question\"",
        Map.of(
            "contextFields", List.of("context"),
            "contextPolicy", Map.of("preset", "adaptive", "budget", "balanced"),
            "contextOptions", Map.of(
                "description", "The context is a module index of \"PATH / IMPORTS / WRITES\" records. Answer by filtering those records in code -- never guess. Return exact PATH values verbatim."),
            // The Peek context map: small, persistent orientation reused across queries.
            "contextMap", Map.of("maxChars", 1800, "infiniteEvolve", false, "evolveSteps", 1),
            "runtime", Map.of("language", "JavaScript")));

    List<String> questions = List.of(
        "Which modules import 'packages/clients/acquirer'? Give the exact PATH values.",
        "Which modules write to the 'orders' table?",
        "What are the direct IMPORTS of packages/api/routes/checkout.ts?");

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      for (String question : questions) {
        Map<String, Object> result = analyst.forward(
            client,
            Map.of("context", codebaseIndex, "question", question),
            Map.of("runtime", runtime, "max_actor_steps", 24));
        System.out.println("\nQ: " + question);
        System.out.println("A: " + result.get("answer"));
        Object paths = result.get("paths");
        List<Object> pathList = paths instanceof List ? (List<Object>) paths : List.of();
        List<String> pathStrings = new ArrayList<>();
        for (Object p : pathList) pathStrings.add(String.valueOf(p));
        System.out.println("Paths: " + String.join(", ", pathStrings));
      }
    }

    System.out.println("\nThe context map evolved on the first query and was reused for the rest.");
  }
}

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

public final class DataAnalystWithToolsExample {
  static final String[] MONTHS = {
      "Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"};

  @SuppressWarnings("unchecked")
  static Map<String, Object> asMap(Object value) {
    return value instanceof Map<?, ?> ? (Map<String, Object>) value : new LinkedHashMap<>();
  }

  // ---------------------------------------------------------------------------
  // The "warehouse": a few hundred rows that live in the host process and are
  // reachable only through tools. The model never sees the rows -- it queries
  // them. Deterministic so the example is reproducible.
  // ---------------------------------------------------------------------------
  static List<Map<String, Object>> buildWarehouse() {
    String[] regions = {"North", "South", "East", "West", "Central", "NW", "NE", "SE"};
    String[] products = {"Widget-A", "Widget-B", "Gadget-X", "Gadget-Y"};
    List<Map<String, Object>> rows = new ArrayList<>();
    long[] seed = {7};

    java.util.function.DoubleSupplier rand = () -> {
      seed[0] = (seed[0] * 1103515245L + 12345L) & 0x7FFFFFFFL;
      return (double) seed[0] / 0x7FFFFFFFL;
    };

    for (String region : regions) {
      for (String product : products) {
        int trend = (product.equals("Gadget-X") && region.equals("East")) ? 90 : 25; // a planted winner
        for (int m = 0; m < MONTHS.length; m++) {
          long units = Math.round(400 + rand.getAsDouble() * 1200 + m * trend);
          int price = product.startsWith("Gadget") ? 60 : 38;
          double returnRate = Math.round(
              (0.01 + rand.getAsDouble() * 0.05 + (product.equals("Widget-B") ? 0.03 : 0)) * 1000.0) / 1000.0;
          Map<String, Object> row = new LinkedHashMap<>();
          row.put("region", region);
          row.put("product", product);
          row.put("monthIndex", m);
          row.put("month", MONTHS[m]);
          row.put("units", units);
          row.put("revenue", units * price);
          row.put("returnRate", returnRate);
          rows.add(row);
        }
      }
    }
    return rows;
  }

  // The schema/data dictionary is large-ish and goes into contextFields so the
  // agent orients on column meaning + business rules without the doc entering the prompt.
  static final String SCHEMA = String.join("\n",
      "TABLE sales (one row per region x product x month)",
      "",
      "COLUMNS",
      "  region       text   one of: North, South, East, West, Central, NW, NE, SE",
      "  product      text   one of: Widget-A, Widget-B, Gadget-X, Gadget-Y",
      "  month        text   Jan..Dec (calendar order; monthIndex 0..11)",
      "  units        int    units sold that month",
      "  revenue      int    integer dollars (units * unit price; Gadgets cost more)",
      "  returnRate   float  fraction of units returned, 0..1",
      "",
      "BUSINESS RULES",
      "  - \"Growth\" = change in monthly revenue from Jan to Dec for a region+product.",
      "  - A return rate above 0.05 (5%) is flagged for quality review.",
      "  - Compare like-for-like: always group by region AND product, not either alone.",
      "",
      "TOOLS AVAILABLE (call them, never invent figures)",
      "  query  filter + aggregate a slice -> {matched, totalUnits, totalRevenue, avgReturnRate}",
      "  top    rank a metric (\"revenue\"|\"units\") grouped by \"product\"|\"region\" -> [{key, value}]",
      "  trend  monthly revenue series (Jan..Dec) for one region + product");

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

    GoogleGeminiClient client = new GoogleGeminiClient(Map.of(
        "api_key", apiKey, "model", "gemini-3.5-flash"));

    List<Map<String, Object>> warehouse = buildWarehouse();
    System.out.println("Warehouse: " + warehouse.size() + " rows (kept out of the prompt).");

    // --- Host tool handlers over the warehouse (the model never sees the rows) ---
    AxQuickJsHostCallable queryTool = params -> {
      Map<String, Object> p = asMap(params);
      String region = p.get("region") == null ? null : String.valueOf(p.get("region"));
      String product = p.get("product") == null ? null : String.valueOf(p.get("product"));
      String month = p.get("month") == null ? null : String.valueOf(p.get("month"));
      List<Map<String, Object>> rows = new ArrayList<>();
      for (Map<String, Object> r : warehouse) {
        if (region != null && !region.isBlank() && !region.equals(r.get("region"))) continue;
        if (product != null && !product.isBlank() && !product.equals(r.get("product"))) continue;
        if (month != null && !month.isBlank() && !month.equals(r.get("month"))) continue;
        rows.add(r);
      }
      long totalUnits = 0;
      long totalRevenue = 0;
      double sumReturn = 0;
      for (Map<String, Object> r : rows) {
        totalUnits += ((Number) r.get("units")).longValue();
        totalRevenue += ((Number) r.get("revenue")).longValue();
        sumReturn += ((Number) r.get("returnRate")).doubleValue();
      }
      double avgReturn = rows.isEmpty() ? 0 : Math.round((sumReturn / rows.size()) * 10000.0) / 10000.0;
      return Map.of(
          "matched", rows.size(), "totalUnits", totalUnits,
          "totalRevenue", totalRevenue, "avgReturnRate", avgReturn);
    };

    AxQuickJsHostCallable topTool = params -> {
      Map<String, Object> p = asMap(params);
      String metric = String.valueOf(p.getOrDefault("metric", "revenue"));
      String groupBy = String.valueOf(p.getOrDefault("groupBy", "product"));
      int limit = p.get("limit") instanceof Number n ? n.intValue() : 5;
      Map<String, Long> totals = new LinkedHashMap<>();
      for (Map<String, Object> r : warehouse) {
        String key = "region".equals(groupBy) ? String.valueOf(r.get("region")) : String.valueOf(r.get("product"));
        long value = "units".equals(metric)
            ? ((Number) r.get("units")).longValue()
            : ((Number) r.get("revenue")).longValue();
        totals.merge(key, value, Long::sum);
      }
      List<Map<String, Object>> ranked = new ArrayList<>();
      for (Map.Entry<String, Long> e : totals.entrySet()) {
        ranked.add(Map.of("key", e.getKey(), "value", e.getValue()));
      }
      ranked.sort((a, b) -> Long.compare((Long) b.get("value"), (Long) a.get("value")));
      return ranked.subList(0, Math.min(limit, ranked.size()));
    };

    AxQuickJsHostCallable trendTool = params -> {
      Map<String, Object> p = asMap(params);
      String region = String.valueOf(p.get("region"));
      String product = String.valueOf(p.get("product"));
      long[] series = new long[12];
      for (Map<String, Object> r : warehouse) {
        if (region.equals(r.get("region")) && product.equals(r.get("product"))) {
          series[((Number) r.get("monthIndex")).intValue()] = ((Number) r.get("revenue")).longValue();
        }
      }
      List<Long> out = new ArrayList<>();
      for (long v : series) out.add(v);
      return out;
    };

    AxAgent analyst = Ax.agent(
        "schema:string, question:string -> answer:string, evidence:string[] \"Concrete figures the answer is based on\"",
        Map.of(
            // Big data dictionary stays out of the prompt.
            "contextFields", List.of("schema"),
            // Tool specs advertised to the model; handlers are registered on the runtime below.
            "functions", List.of(
                Map.of(
                    "name", "query",
                    "description", "Filter the sales table and return aggregates for the matching rows.",
                    "parameters", Map.of(
                        "type", "object",
                        "properties", Map.of(
                            "region", Map.of("type", "string"),
                            "product", Map.of("type", "string"),
                            "month", Map.of("type", "string")))),
                Map.of(
                    "name", "top",
                    "description", "Rank a metric (revenue|units) grouped by product|region, highest first.",
                    "parameters", Map.of(
                        "type", "object",
                        "properties", Map.of(
                            "metric", Map.of("type", "string"),
                            "groupBy", Map.of("type", "string"),
                            "limit", Map.of("type", "number")),
                        "required", List.of("metric", "groupBy"))),
                Map.of(
                    "name", "trend",
                    "description", "Monthly revenue series (Jan..Dec) for one region and product.",
                    "parameters", Map.of(
                        "type", "object",
                        "properties", Map.of(
                            "region", Map.of("type", "string"),
                            "product", Map.of("type", "string")),
                        "required", List.of("region", "product")))),
            "contextPolicy", Map.of("preset", "lean", "budget", "balanced"),
            "runtime", Map.of("language", "JavaScript")));

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      runtime.registerCallable("query", queryTool);
      runtime.registerCallable("top", topTool);
      runtime.registerCallable("trend", trendTool);

      Map<String, Object> result = analyst.forward(
          client,
          Map.of(
              "schema", SCHEMA,
              "question", "Which region+product had the strongest Jan->Dec revenue growth, and which products have an average return rate above the 5% review threshold? Tool-calling rules: the tools are bare async functions named exactly `query`, `top`, `trend` -- call them as `await query({product:'Widget-B'})`, never as `tools.query(...)`. Do NOT wrap your turn in an IIFE like `(async()=>{...})()`; write top-level `await` and capture results in variables, then `console.log` one value to inspect, and only call `await final(task, evidence)` once you have the figures."),
          Map.of("runtime", runtime, "max_actor_steps", 40));

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

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

public final class SelfImprovingLabExample {
  @SuppressWarnings("unchecked")
  static Map<String, Object> asMap(Object value) {
    return value instanceof Map<?, ?> ? (Map<String, Object>) value : new LinkedHashMap<>();
  }

  // ---------------------------------------------------------------------------
  // The "lab": a deterministic black-box experiment. It scores an ETL config plan
  // against a hidden ideal and returns, for any failing check, the exact fix --
  // so the agent can converge by following the feedback, not by being told.
  // ---------------------------------------------------------------------------
  static final List<String> CHECKS = List.of(
      "no-nulls", "no-duplicates", "numeric-types", "trimmed-strings", "outliers-handled");
  static final Map<String, String> REMEDIES = Map.of(
      "no-nulls", "set nullPolicy=impute (or nullPolicy=drop)",
      "no-duplicates", "set dedup=on",
      "numeric-types", "set coerceTypes=on",
      "trimmed-strings", "set trim=on",
      "outliers-handled", "set outlier=clip (or outlier=winsorize)");

  static Map<String, Object> runInSandbox(String plan) {
    Map<String, String> flags = new LinkedHashMap<>();
    Matcher m = Pattern.compile("([a-z]+)\\s*=\\s*([a-z0-9]+)").matcher(plan.toLowerCase());
    while (m.find()) flags.put(m.group(1), m.group(2));

    Map<String, Boolean> ok = new LinkedHashMap<>();
    ok.put("no-nulls", List.of("impute", "drop").contains(flags.getOrDefault("nullpolicy", "")));
    ok.put("no-duplicates", "on".equals(flags.get("dedup")));
    ok.put("numeric-types", "on".equals(flags.get("coercetypes")));
    ok.put("trimmed-strings", "on".equals(flags.get("trim")));
    ok.put("outliers-handled", List.of("clip", "winsorize").contains(flags.getOrDefault("outlier", "")));

    List<String> passed = new ArrayList<>();
    List<Map<String, Object>> failed = new ArrayList<>();
    for (String c : CHECKS) {
      if (Boolean.TRUE.equals(ok.get(c))) passed.add(c);
      else failed.add(Map.of("check", c, "fix", REMEDIES.get(c)));
    }
    double score = Math.round(((double) passed.size() / CHECKS.size()) * 100.0) / 100.0;
    Map<String, Object> out = new LinkedHashMap<>();
    out.put("score", score);
    out.put("solved", passed.size() == CHECKS.size());
    out.put("passed", passed);
    out.put("failed", failed);
    out.put("logs", passed.size() + "/" + CHECKS.size() + " checks passed");
    return out;
  }

  static Map<String, Object> spec(String name, String description, Map<String, Object> props, List<String> required) {
    Map<String, Object> parameters = new LinkedHashMap<>();
    parameters.put("type", "object");
    parameters.put("properties", props);
    if (required != null && !required.isEmpty()) parameters.put("required", required);
    Map<String, Object> out = new LinkedHashMap<>();
    out.put("name", name);
    out.put("description", description);
    out.put("parameters", parameters);
    return out;
  }

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

    // An independent verifier -- a separate ax() program, not the agent grading itself.
    AxGen verifier = Ax.ax("rubric:string, evidence:json -> passed:boolean, feedback:string, missing:string[]");
    verifier.setInstruction(
        "You are an independent rubric grader, not a self-critique. Pass only when the evidence clearly satisfies every part of the rubric.");

    // In-memory rule store. Verified, reusable rules go here -- not raw failure notes.
    Map<String, String> memoryStore = new LinkedHashMap<>();

    try (AxQuickJsCodeRuntime runtime = new AxQuickJsCodeRuntime()) {
      runtime.registerCallable("runExperiment", params -> runInSandbox(String.valueOf(asMap(params).getOrDefault("plan", ""))));
      runtime.registerCallable("listChecks", params -> CHECKS);
      runtime.registerCallable("grade", params -> {
        Map<String, Object> p = asMap(params);
        return verifier.forward(client, Map.of(
            "rubric", p.getOrDefault("rubric", ""),
            "evidence", p.getOrDefault("evidence", List.of())));
      });
      runtime.registerCallable("recall", params -> {
        String topic = String.valueOf(asMap(params).getOrDefault("topic", "")).toLowerCase();
        List<String> words = topic.isBlank() ? List.of() : Arrays.asList(topic.split("\\s+"));
        List<String> out = new ArrayList<>();
        for (Map.Entry<String, String> e : memoryStore.entrySet()) {
          boolean hit = e.getKey().contains(topic);
          for (String w : words) if (!w.isBlank() && e.getKey().contains(w)) hit = true;
          if (hit) out.add(e.getValue());
        }
        return out;
      });
      runtime.registerCallable("remember", params -> {
        Map<String, Object> p = asMap(params);
        String rule = String.valueOf(p.getOrDefault("rule", ""));
        String key = rule.toLowerCase();
        key = key.substring(0, Math.min(48, key.length()));
        memoryStore.put(key, rule + " :: " + p.getOrDefault("evidence", ""));
        return Map.of("stored", true, "total", memoryStore.size());
      });

      AxAgent selfImproving = Ax.agent(
          "goal:string, rubric:string -> answer:string, experiments:string[] \"Plans tried, in order\", learnedRules:string[]",
          Map.of(
              "contextFields", List.of(),
              "functions", List.of(
                  spec("runExperiment", "Apply an ETL config plan; returns score, solved, passed[], failed[{check,fix}], logs. Pass an empty plan to discover the fixes.",
                      Map.of("plan", Map.of("type", "string")), List.of("plan")),
                  spec("listChecks", "List the data-quality checks the experiment evaluates.", Map.of(), List.of()),
                  spec("grade", "Independent rubric grader. Pass only when the evidence meets the rubric.",
                      Map.of("rubric", Map.of("type", "string"), "evidence", Map.of("type", "array", "items", Map.of("type", "string"))),
                      List.of("rubric", "evidence")),
                  spec("recall", "Recall verified rules relevant to a topic.",
                      Map.of("topic", Map.of("type", "string")), List.of("topic")),
                  spec("remember", "Store a verified, reusable rule (the rule, not raw notes).",
                      Map.of("rule", Map.of("type", "string"), "evidence", Map.of("type", "string")), List.of("rule", "evidence"))),
              "contextPolicy", Map.of("preset", "adaptive", "budget", "balanced"),
              "executorOptions", Map.of(
                  "description", String.join("\n",
                      "Use the tools -- do not answer from your own knowledge.",
                      "1. recall('etl data quality') to reuse anything already learned.",
                      "2. runExperiment('') once to see every failing check and its fix.",
                      "3. Build a plan applying all the fixes, then runExperiment again. Repeat until solved is true.",
                      "4. grade the passing evidence against the rubric.",
                      "5. For each check you fixed, remember(rule, evidence).",
                      "6. Then return the answer, the plans you tried, and the learned rules.")),
              "runtime", Map.of("language", "JavaScript")));

      Map<String, Object> result = selfImproving.forward(
          client,
          Map.of(
              "goal", "Find an ETL config plan that cleans the dirty dataset so every data-quality check passes.",
              "rubric", "All five checks (no-nulls, no-duplicates, numeric-types, trimmed-strings, outliers-handled) must pass, i.e. score 1.0. The deliverable must also record at least one verified, reusable learnedRule."),
          Map.of("runtime", runtime, "max_actor_steps", 18));

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

      // Persist the agent's verified rules so a future run's recall reuses them.
      Object learned = result.get("learnedRules");
      if (learned instanceof List<?> rules) {
        for (Object rule : rules) {
          String s = String.valueOf(rule).toLowerCase();
          memoryStore.put(s.substring(0, Math.min(48, s.length())), String.valueOf(rule));
        }
      }
      System.out.println("\nMemory now holds " + memoryStore.size() + " rule(s) for next time.");
    }
  }
}