package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"time"

	ax "github.com/ax-llm/ax/packages/go"
	axgoja "github.com/ax-llm/ax/packages/go/runtime/goja"
)

func geminiClient() *ax.GoogleGeminiClient {
	apiKey := os.Getenv("GOOGLE_APIKEY")
	if apiKey == "" {
		panic("Set GOOGLE_APIKEY to run this example.")
	}
	model := os.Getenv("AX_GEMINI_MODEL")
	if model == "" {
		model = "gemini-3.5-flash"
	}
	return ax.NewGoogleGeminiClient(map[string]ax.Value{"api_key": apiKey, "model": model})
}

func printJSON(value ax.Value) {
	data, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

// ---------------------------------------------------------------------------
// 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.
// ---------------------------------------------------------------------------
func buildLogDump() []ax.Value {
	start := time.Date(2026, 3, 2, 13, 0, 0, 0, time.UTC)
	events := []ax.Value{}

	push := func(i int, event map[string]ax.Value) {
		event["timestamp"] = start.Add(time.Duration(i*2) * time.Second).Format("2006-01-02T15:04:05Z")
		event["requestId"] = fmt.Sprintf("req-%d", 100000+i)
		events = append(events, ax.Value(event))
	}

	for i := 0; i < 1600; i++ {
		// Routine, healthy traffic across the fleet.
		push(i, ax.Object("level", "INFO", "service", "gateway", "statusCode", 200, "latencyMs", 40+(i%30), "message", "route ok GET /checkout"))
		push(i, ax.Object("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(i, ax.Object("level", "ERROR", "service", "payments-gw", "statusCode", 504, "latencyMs", 10000, "tenantTier", "enterprise", "message", "upstream timeout calling acquirer (10s)"))
			push(i, ax.Object("level", "ERROR", "service", "checkout-api", "statusCode", 502, "tenantTier", "enterprise", "message", "bad gateway from svc-payments-gw"))
			if i%3 == 0 {
				push(i, ax.Object("level", "WARN", "service", "payments-gw", "message", "connection pool exhausted (max=64) waiting=200+"))
				push(i, ax.Object("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(i, ax.Object("level", "WARN", "service", "catalog-cron", "latencyMs", 0, "message", "rebuild step pinning CPU at 95% on shared node"))
			push(i, ax.Object("level", "ERROR", "service", "search-api", "statusCode", 429, "message", "rate limited: downstream catalog unavailable"))
		}
	}

	return events
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	client := geminiClient()

	logs := buildLogDump()
	fmt.Printf("Generated %d log events (kept out of the prompt).\n", len(logs))

	logRLM := ax.NewAgent(
		`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[string]ax.Value{
			// The export stays in the runtime; only extracted evidence reaches the model.
			"contextFields":   ax.Array("logs"),
			"contextPolicy":   ax.Object("preset", "lean", "budget", "balanced"),
			"maxRuntimeChars": 12000,
			"runtime":         ax.Object("language", "JavaScript"),
		},
	)

	report, err := logRLM.Forward(
		ctx,
		client,
		map[string]ax.Value{
			"logs": ax.Array(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[string]ax.Value{"runtime": axgoja.NewRuntime(), "max_actor_steps": 40},
	)
	if err != nil {
		panic(err)
	}

	fmt.Println("\n=== Report ===")
	printJSON(report)
	fmt.Println("\n=== Usage ===")
	printJSON(logRLM.GetUsage())
}

package main

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	ax "github.com/ax-llm/ax/packages/go"
	axgoja "github.com/ax-llm/ax/packages/go/runtime/goja"
)

func geminiClient() *ax.GoogleGeminiClient {
	apiKey := os.Getenv("GOOGLE_APIKEY")
	if apiKey == "" {
		panic("Set GOOGLE_APIKEY to run this example.")
	}
	model := os.Getenv("AX_GEMINI_MODEL")
	if model == "" {
		model = "gemini-3.5-flash"
	}
	return ax.NewGoogleGeminiClient(map[string]ax.Value{"api_key": apiKey, "model": model})
}

type module struct {
	path    string
	imports []string
	writes  string
}

// ---------------------------------------------------------------------------
// 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.
// ---------------------------------------------------------------------------
func buildModuleIndex() []module {
	core := []module{
		{"packages/api/middleware/auth.ts", []string{"packages/shared"}, "-"},
		{"packages/api/middleware/rateLimit.ts", []string{"packages/db"}, "-"},
		{"packages/api/routes/checkout.ts", []string{"packages/api/middleware/auth.ts", "packages/services/orders/createOrder.ts", "packages/services/payments/charge.ts"}, "-"},
		{"packages/api/routes/search.ts", []string{"packages/api/middleware/auth.ts", "packages/services/catalog/searchCatalog.ts"}, "-"},
		{"packages/services/orders/createOrder.ts", []string{"packages/db", "packages/clients/bus"}, "orders"},
		{"packages/services/orders/orderRepo.ts", []string{"packages/db"}, "orders"},
		{"packages/services/payments/charge.ts", []string{"packages/clients/acquirer", "packages/db"}, "payments"},
		{"packages/services/payments/refund.ts", []string{"packages/clients/acquirer", "packages/db"}, "refunds"},
		{"packages/services/catalog/searchCatalog.ts", []string{"packages/db"}, "-"},
		{"packages/clients/acquirer/index.ts", []string{"packages/shared"}, "-"},
		{"packages/clients/bus/index.ts", []string{"packages/shared"}, "-"},
	}
	// Filler modules so the index is genuinely large; some also depend on the acquirer.
	filler := []module{}
	for i := 0; i < 110; i++ {
		dep := "packages/db"
		if i%4 == 0 {
			dep = "packages/clients/acquirer"
		}
		writes := "-"
		if i%6 == 0 {
			writes = "audit"
		}
		filler = append(filler, module{
			path:    fmt.Sprintf("packages/services/feature%d/handler.ts", i),
			imports: []string{dep, "packages/shared"},
			writes:  writes,
		})
	}
	return append(core, filler...)
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	client := geminiClient()

	modules := buildModuleIndex()
	records := make([]string, 0, len(modules))
	for _, m := range modules {
		records = append(records, fmt.Sprintf("PATH: %s\nIMPORTS: %s\nWRITES: %s", m.path, strings.Join(m.imports, ", "), m.writes))
	}
	codebaseIndex := strings.Join(records, "\n\n")
	fmt.Printf("Module index: %d records (kept out of the prompt).\n", len(modules))

	analyst := ax.NewAgent(
		`context:string, question:string -> answer:string, paths:string[] "Exact PATH values from the index that answer the question"`,
		map[string]ax.Value{
			"contextFields": ax.Array("context"),
			"contextPolicy": ax.Object("preset", "adaptive", "budget", "balanced"),
			"contextOptions": ax.Object(
				"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": ax.Object("maxChars", 1800, "infiniteEvolve", false, "evolveSteps", 1),
			"runtime":    ax.Object("language", "JavaScript"),
		},
	)

	questions := []string{
		"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?",
	}

	runtime := axgoja.NewRuntime()
	for _, question := range questions {
		output, err := analyst.Forward(
			ctx,
			client,
			map[string]ax.Value{"context": codebaseIndex, "question": question},
			map[string]ax.Value{"runtime": runtime, "max_actor_steps": 24},
		)
		if err != nil {
			panic(err)
		}
		result, _ := output.(map[string]ax.Value)
		answer, _ := result["answer"].(string)
		paths := []string{}
		if raw, ok := result["paths"].([]ax.Value); ok {
			for _, p := range raw {
				paths = append(paths, fmt.Sprint(p))
			}
		}
		fmt.Println("\nQ:", question)
		fmt.Println("A:", answer)
		fmt.Println("Paths:", strings.Join(paths, ", "))
	}

	fmt.Println("\nThe context map evolved on the first query and was reused for the rest.")
}

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"strings"
	"time"

	ax "github.com/ax-llm/ax/packages/go"
	axgoja "github.com/ax-llm/ax/packages/go/runtime/goja"
)

func geminiClient() *ax.GoogleGeminiClient {
	apiKey := os.Getenv("GOOGLE_APIKEY")
	if apiKey == "" {
		panic("Set GOOGLE_APIKEY to run this example.")
	}
	model := os.Getenv("AX_GEMINI_MODEL")
	if model == "" {
		model = "gemini-3.5-flash"
	}
	return ax.NewGoogleGeminiClient(map[string]ax.Value{"api_key": apiKey, "model": model})
}

func printJSON(value ax.Value) {
	data, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

func roundTo(v float64, places int) float64 {
	scale := math.Pow(10, float64(places))
	return math.Round(v*scale) / scale
}

// ---------------------------------------------------------------------------
// 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.
// ---------------------------------------------------------------------------
var months = []string{"Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"}

type row struct {
	region     string
	product    string
	monthIndex int
	month      string
	units      int
	revenue    int
	returnRate float64
}

func buildWarehouse() []row {
	regions := []string{"North", "South", "East", "West", "Central", "NW", "NE", "SE"}
	products := []string{"Widget-A", "Widget-B", "Gadget-X", "Gadget-Y"}
	rows := []row{}
	seed := int64(7)
	rand := func() float64 {
		seed = (seed*1103515245 + 12345) & 0x7FFFFFFF
		return float64(seed) / float64(0x7FFFFFFF)
	}

	for _, region := range regions {
		for _, product := range products {
			trend := 25 // a planted winner
			if product == "Gadget-X" && region == "East" {
				trend = 90
			}
			for m := 0; m < len(months); m++ {
				units := int(math.Round(400 + rand()*1200 + float64(m*trend)))
				price := 38
				if strings.HasPrefix(product, "Gadget") {
					price = 60
				}
				extra := 0.0
				if product == "Widget-B" {
					extra = 0.03
				}
				returnRate := roundTo(0.01+rand()*0.05+extra, 3)
				rows = append(rows, row{
					region: region, product: product, monthIndex: m, month: months[m],
					units: units, revenue: units * price, returnRate: returnRate,
				})
			}
		}
	}
	return rows
}

var warehouse = buildWarehouse()

// 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.
var schema = strings.TrimSpace(`
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
`)

func asMap(value ax.Value) map[string]ax.Value {
	if out, ok := value.(map[string]ax.Value); ok {
		return out
	}
	return map[string]ax.Value{}
}

func asString(value ax.Value) string {
	if value == nil {
		return ""
	}
	if s, ok := value.(string); ok {
		return s
	}
	return fmt.Sprint(value)
}

// --- Host tool handlers over the warehouse (the model never sees the rows) ---
func queryTool(params ax.Value) (ax.Value, error) {
	p := asMap(params)
	region, product, month := asString(p["region"]), asString(p["product"]), asString(p["month"])
	matched := 0
	totalUnits := 0
	totalRevenue := 0
	sumReturn := 0.0
	for _, r := range warehouse {
		if (region != "" && r.region != region) || (product != "" && r.product != product) || (month != "" && r.month != month) {
			continue
		}
		matched++
		totalUnits += r.units
		totalRevenue += r.revenue
		sumReturn += r.returnRate
	}
	avgReturn := 0.0
	if matched > 0 {
		avgReturn = roundTo(sumReturn/float64(matched), 4)
	}
	return ax.Object("matched", matched, "totalUnits", totalUnits, "totalRevenue", totalRevenue, "avgReturnRate", avgReturn), nil
}

func topTool(params ax.Value) (ax.Value, error) {
	p := asMap(params)
	metric := asString(p["metric"])
	if metric == "" {
		metric = "revenue"
	}
	groupBy := asString(p["groupBy"])
	if groupBy == "" {
		groupBy = "product"
	}
	limit := 5
	if raw, ok := p["limit"]; ok {
		switch v := raw.(type) {
		case int:
			limit = v
		case int64:
			limit = int(v)
		case float64:
			limit = int(v)
		}
	}
	totals := map[string]int{}
	order := []string{}
	for _, r := range warehouse {
		key := r.product
		if groupBy == "region" {
			key = r.region
		}
		if _, seen := totals[key]; !seen {
			order = append(order, key)
		}
		if metric == "units" {
			totals[key] += r.units
		} else {
			totals[key] += r.revenue
		}
	}
	// Stable rank by value descending.
	for i := 0; i < len(order); i++ {
		for j := i + 1; j < len(order); j++ {
			if totals[order[j]] > totals[order[i]] {
				order[i], order[j] = order[j], order[i]
			}
		}
	}
	ranked := []ax.Value{}
	for i, key := range order {
		if i >= limit {
			break
		}
		ranked = append(ranked, ax.Object("key", key, "value", totals[key]))
	}
	return ax.Array(ranked...), nil
}

func trendTool(params ax.Value) (ax.Value, error) {
	p := asMap(params)
	region, product := asString(p["region"]), asString(p["product"])
	series := make([]ax.Value, 12)
	for i := range series {
		series[i] = 0
	}
	for _, r := range warehouse {
		if r.region == region && r.product == product {
			series[r.monthIndex] = r.revenue
		}
	}
	return ax.Array(series...), nil
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	client := geminiClient()

	fmt.Printf("Warehouse: %d rows (kept out of the prompt; reachable only via tools).\n", len(warehouse))

	runtime := axgoja.NewRuntime(
		axgoja.WithCallable("query", queryTool),
		axgoja.WithCallable("top", topTool),
		axgoja.WithCallable("trend", trendTool),
	)

	analyst := ax.NewAgent(
		`schema:string, question:string -> answer:string, evidence:string[] "Concrete figures the answer is based on"`,
		map[string]ax.Value{
			// Big data dictionary stays out of the prompt.
			"contextFields": ax.Array("schema"),
			// Tool specs advertised to the model; handlers are registered on the runtime above.
			"functions": ax.Array(
				ax.Object(
					"name", "query",
					"description", "Filter the sales table and return aggregates for the matching rows.",
					"parameters", ax.Object(
						"type", "object",
						"properties", ax.Object(
							"region", ax.Object("type", "string"),
							"product", ax.Object("type", "string"),
							"month", ax.Object("type", "string"),
						),
					),
				),
				ax.Object(
					"name", "top",
					"description", "Rank a metric (revenue|units) grouped by product|region, highest first.",
					"parameters", ax.Object(
						"type", "object",
						"properties", ax.Object(
							"metric", ax.Object("type", "string"),
							"groupBy", ax.Object("type", "string"),
							"limit", ax.Object("type", "number"),
						),
						"required", ax.Array("metric", "groupBy"),
					),
				),
				ax.Object(
					"name", "trend",
					"description", "Monthly revenue series (Jan..Dec) for one region and product.",
					"parameters", ax.Object(
						"type", "object",
						"properties", ax.Object(
							"region", ax.Object("type", "string"),
							"product", ax.Object("type", "string"),
						),
						"required", ax.Array("region", "product"),
					),
				),
			),
			"contextPolicy": ax.Object("preset", "lean", "budget", "balanced"),
			"executorOptions": ax.Object("description", strings.Join([]string{
				"Consult the schema for column meaning and business rules.",
				"Answer using the warehouse tools -- never invent figures.",
				"Return rates: call query({product}) for each of Widget-A, Widget-B, Gadget-X, Gadget-Y and read avgReturnRate; any product with avgReturnRate > 0.05 is above the review threshold.",
				"Growth: call trend({region, product}); the returned array is revenue Jan..Dec, so growth = last element minus first element. Compare a few region+product pairs and report the largest.",
				"Cite the concrete numbers you observed as evidence, then call final(...).",
			}, "\n")),
			"runtime": ax.Object("language", "JavaScript"),
		},
	)

	result, err := analyst.Forward(
		ctx,
		client,
		map[string]ax.Value{
			"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?",
		},
		map[string]ax.Value{"runtime": runtime, "max_actor_steps": 40},
	)
	if err != nil {
		panic(err)
	}

	printJSON(result)
}

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"regexp"
	"strings"
	"time"

	ax "github.com/ax-llm/ax/packages/go"
	axgoja "github.com/ax-llm/ax/packages/go/runtime/goja"
)

func openAIClient() *ax.OpenAICompatibleClient {
	apiKey := os.Getenv("OPENAI_API_KEY")
	if apiKey == "" {
		apiKey = os.Getenv("OPENAI_APIKEY")
	}
	if apiKey == "" {
		panic("Set OPENAI_API_KEY or OPENAI_APIKEY to run this example.")
	}
	model := os.Getenv("AX_OPENAI_MODEL")
	if model == "" {
		model = "gpt-5.4-mini"
	}
	return ax.NewOpenAICompatibleClient(map[string]ax.Value{"api_key": apiKey, "model": model, "model_config": ax.Object("temperature", 0)})
}

func printJSON(value ax.Value) {
	data, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

func asMap(value ax.Value) map[string]ax.Value {
	if out, ok := value.(map[string]ax.Value); ok {
		return out
	}
	return map[string]ax.Value{}
}

func asString(value ax.Value) string {
	if value == nil {
		return ""
	}
	if s, ok := value.(string); ok {
		return s
	}
	return fmt.Sprint(value)
}

// ---------------------------------------------------------------------------
// 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.
// ---------------------------------------------------------------------------
var checks = []string{"no-nulls", "no-duplicates", "numeric-types", "trimmed-strings", "outliers-handled"}

var remedies = map[string]string{
	"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)",
}

var flagRe = regexp.MustCompile(`([a-z]+)\s*=\s*([a-z0-9]+)`)

func runInSandbox(plan string) ax.Value {
	flags := map[string]string{}
	for _, m := range flagRe.FindAllStringSubmatch(strings.ToLower(plan), -1) {
		flags[m[1]] = m[2]
	}
	ok := map[string]bool{
		"no-nulls":         flags["nullpolicy"] == "impute" || flags["nullpolicy"] == "drop",
		"no-duplicates":    flags["dedup"] == "on",
		"numeric-types":    flags["coercetypes"] == "on",
		"trimmed-strings":  flags["trim"] == "on",
		"outliers-handled": flags["outlier"] == "clip" || flags["outlier"] == "winsorize",
	}
	passed := []ax.Value{}
	failed := []ax.Value{}
	for _, c := range checks {
		if ok[c] {
			passed = append(passed, c)
		} else {
			failed = append(failed, ax.Object("check", c, "fix", remedies[c]))
		}
	}
	score := math.Round(float64(len(passed))/float64(len(checks))*100) / 100
	return ax.Object(
		"score", score,
		"solved", len(passed) == len(checks),
		"passed", ax.Array(passed...),
		"failed", ax.Array(failed...),
		"logs", fmt.Sprintf("%d/%d checks passed", len(passed), len(checks)),
	)
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	client := openAIClient()

	// An independent verifier -- a separate ax() program, not the agent grading itself.
	verifier := ax.NewAx(
		"rubric:string, evidence:json -> passed:boolean, feedback:string, missing:string[]",
		map[string]ax.Value{"instruction": "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.
	memoryStore := map[string]string{}
	memoryOrder := []string{}
	storeRule := func(key, value string) {
		if _, seen := memoryStore[key]; !seen {
			memoryOrder = append(memoryOrder, key)
		}
		memoryStore[key] = value
	}

	runtime := axgoja.NewRuntime(
		axgoja.WithCallable("runExperiment", func(p ax.Value) (ax.Value, error) {
			return runInSandbox(asString(asMap(p)["plan"])), nil
		}),
		axgoja.WithCallable("listChecks", func(p ax.Value) (ax.Value, error) {
			out := []ax.Value{}
			for _, c := range checks {
				out = append(out, c)
			}
			return ax.Array(out...), nil
		}),
		axgoja.WithCallable("grade", func(p ax.Value) (ax.Value, error) {
			m := asMap(p)
			evidence := m["evidence"]
			if evidence == nil {
				evidence = ax.Array()
			}
			return verifier.Forward(ctx, client, map[string]ax.Value{"rubric": asString(m["rubric"]), "evidence": evidence}, nil)
		}),
		axgoja.WithCallable("recallRules", func(p ax.Value) (ax.Value, error) {
			t := strings.ToLower(asString(asMap(p)["topic"]))
			words := strings.Fields(t)
			out := []ax.Value{}
			for _, k := range memoryOrder {
				match := strings.Contains(k, t)
				if !match {
					for _, w := range words {
						if strings.Contains(k, w) {
							match = true
							break
						}
					}
				}
				if match {
					out = append(out, memoryStore[k])
				}
			}
			return ax.Array(out...), nil
		}),
		axgoja.WithCallable("remember", func(p ax.Value) (ax.Value, error) {
			m := asMap(p)
			rule := asString(m["rule"])
			key := strings.ToLower(rule)
			if len(key) > 48 {
				key = key[:48]
			}
			storeRule(key, fmt.Sprintf("%s :: %s", rule, asString(m["evidence"])))
			return ax.Object("stored", true, "total", len(memoryStore)), nil
		}),
	)

	spec := func(name, description string, props map[string]ax.Value, required ...string) ax.Value {
		parameters := ax.Object("type", "object", "properties", props)
		if len(required) > 0 {
			reqd := []ax.Value{}
			for _, r := range required {
				reqd = append(reqd, r)
			}
			parameters["required"] = ax.Array(reqd...)
		}
		return ax.Object("name", name, "description", description, "parameters", parameters)
	}

	selfImproving := ax.NewAgent(
		`goal:string, rubric:string -> answer:string, experiments:string[] "Plans tried, in order", learnedRules:string[]`,
		map[string]ax.Value{
			"contextFields": ax.Array(),
			"functions": ax.Array(
				spec("runExperiment", "Apply an ETL config plan; returns score, solved, passed[], failed[{check,fix}], logs. Pass an empty plan to discover the fixes.", map[string]ax.Value{"plan": ax.Object("type", "string")}, "plan"),
				spec("listChecks", "List the data-quality checks the experiment evaluates.", map[string]ax.Value{}),
				spec("grade", "Independent rubric grader. Pass only when the evidence meets the rubric.", map[string]ax.Value{"rubric": ax.Object("type", "string"), "evidence": ax.Object("type", "array", "items", ax.Object("type", "string"))}, "rubric", "evidence"),
				spec("recallRules", "Recall verified rules relevant to a topic.", map[string]ax.Value{"topic": ax.Object("type", "string")}, "topic"),
				spec("remember", "Store a verified, reusable rule (the rule, not raw notes).", map[string]ax.Value{"rule": ax.Object("type", "string"), "evidence": ax.Object("type", "string")}, "rule", "evidence"),
			),
			"contextPolicy": ax.Object("preset", "adaptive", "budget", "balanced"),
			"executorOptions": ax.Object("description", strings.Join([]string{
				"Use the tools -- do not answer from your own knowledge.",
				"1. recallRules('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.",
			}, "\n")),
			"runtime": ax.Object("language", "JavaScript"),
		},
	)

	result, err := selfImproving.Forward(
		ctx,
		client,
		map[string]ax.Value{
			"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.",
		},
		map[string]ax.Value{"runtime": runtime, "max_actor_steps": 18},
	)
	if err != nil {
		panic(err)
	}

	printJSON(result)

	// Persist the agent's verified rules so a future run's recall reuses them.
	if learned, ok := asMap(result)["learnedRules"].([]ax.Value); ok {
		for _, rule := range learned {
			r := asString(rule)
			key := strings.ToLower(r)
			if len(key) > 48 {
				key = key[:48]
			}
			storeRule(key, r)
		}
	}
	fmt.Printf("\nMemory now holds %d rule(s) for next time.\n", len(memoryStore))
}