package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"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))
}

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

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

	// Keep the handbook in the runtime, out of the prompt.
	assistant := ax.NewAgent(
		`question:string, handbook:string -> answer:string, citations:string[] "Handbook sections the answer relies on"`,
		map[string]ax.Value{"contextFields": ax.Array("handbook"), "runtime": ax.Object("language", "JavaScript")},
	)

	output, err := assistant.Forward(
		ctx,
		client,
		map[string]ax.Value{
			"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[string]ax.Value{"runtime": axgoja.NewRuntime(), "max_actor_steps": 12},
	)
	if err != nil {
		panic(err)
	}
	printJSON(output)
}

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"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))
}

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

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

	triage := ax.NewAgent(
		`report:string, question:string -> severity:class "low, medium, high, critical", rootCause:string, nextSteps:string[], evidence:string[] "Quoted log lines that support the assessment"`,
		map[string]ax.Value{
			"contextFields": ax.Array("report"),
			"contextPolicy": ax.Object("preset", "lean", "budget", "balanced"),
			"runtime":       ax.Object("language", "JavaScript"),
		},
	)

	output, err := triage.Forward(
		ctx,
		client,
		map[string]ax.Value{
			"report":   report,
			"question": "What happened, how bad was it, and what should the on-call do next? Cite the lines you relied on.",
		},
		map[string]ax.Value{"runtime": axgoja.NewRuntime(), "max_actor_steps": 12},
	)
	if err != nil {
		panic(err)
	}
	printJSON(output)
}

package main

import (
	"context"
	"fmt"
	ax "github.com/ax-llm/ax/packages/go"
	axgoja "github.com/ax-llm/ax/packages/go/runtime/goja"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

func main() {
	key := os.Getenv("OPENAI_API_KEY")
	if key == "" {
		key = os.Getenv("OPENAI_APIKEY")
	}
	if key == "" {
		panic("Set OPENAI_API_KEY or OPENAI_APIKEY.")
	}
	client := ax.NewAI("openai", ax.Object("api_key", key, "model", "gpt-6-astra", "model_config", ax.Object("thinkingTokenBudget", "low", "max_tokens", 4096)))
	pending := make(chan struct{})
	var finished, overlap atomic.Bool
	var queued, pendingOnce sync.Once
	var applied atomic.Int32
	control := ax.RunControl()
	control.OnEvent(func(event map[string]ax.Value) {
		if event["type"] == "tool.started" {
			queued.Do(func() {
				if err := control.Steer("Include the word VERIFIED in the final answer."); err != nil {
					panic(err)
				}
				if err := control.SetThinkingTokenBudget("medium"); err != nil {
					panic(err)
				}
			})
		}
		if event["type"] == "applied" {
			applied.Add(1)
		}
	})
	slow := ax.Fn("slow_reference").Execution("background").WithContextHandler(func(ctx context.Context, _ map[string]ax.Value) (ax.Value, error) {
		pendingOnce.Do(func() { close(pending) })
		select {
		case <-time.After(6 * time.Second):
			finished.Store(true)
			return "REF-42", nil
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	})
	label := ax.Fn("local_label").Execution("background").WithHandler(func(_ map[string]ax.Value) (ax.Value, error) {
		select {
		case <-pending:
			if !finished.Load() {
				overlap.Store(true)
			}
		case <-time.After(3 * time.Second):
		}
		return "LAUNCH", nil
	})
	program := ax.NewAgent("question -> answer", ax.Object("runtime", ax.Object("language", "JavaScript"), "directResponse", "off", "functions", ax.Array(slow, label)))
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	result, err := program.Forward(ctx, client, ax.Object("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."), ax.Object("runtime", axgoja.NewRuntime(), "max_actor_steps", 12, "control", control, "serviceTier", "standard", "maxSteps", 6))
	if err != nil {
		panic(err)
	}
	answer := fmt.Sprint(result)
	for _, word := range []string{"REF-42", "LAUNCH", "VERIFIED"} {
		if !strings.Contains(answer, word) {
			panic("Missing final result: " + answer)
		}
	}
	if !overlap.Load() {
		panic("No independent model work while background tool was pending")
	}
	if applied.Load() < 2 {
		panic("Control updates were not applied")
	}
	fmt.Println(answer)
	fmt.Println("Background overlap verified; steering and reasoning applied at the next response.")
}

package main

import (
	"context"
	"encoding/json"
	"fmt"
	ax "github.com/ax-llm/ax/packages/go"
	axgoja "github.com/ax-llm/ax/packages/go/runtime/goja"
	"os"
	"strings"
	"time"
)

func main() {
	key := os.Getenv("OPENAI_API_KEY")
	if key == "" {
		key = os.Getenv("OPENAI_APIKEY")
	}
	if key == "" {
		panic("Set OPENAI_API_KEY or OPENAI_APIKEY.")
	}
	client := ax.NewAI("openai", ax.Object("api_key", key, "model", "gpt-6-astra", "model_config", ax.Object("thinkingTokenBudget", "low", "max_tokens", 4096)))
	control := ax.RunControl()
	applied := false
	control.OnEvent(func(event map[string]ax.Value) {
		if event["type"] == "applied" && event["path"] == "root/team.researcher/executor" {
			applied = true
		}
	})
	if err := control.Steer("Include VERIFIED in the final answer."); err != nil {
		panic(err)
	}
	if err := control.Steer("Include CHILD-CHECK in the final answer.", "root/team.researcher"); err != nil {
		panic(err)
	}
	if err := control.SetThinkingTokenBudget("medium", "root/team.researcher/executor"); err != nil {
		panic(err)
	}
	runtime := func() *axgoja.Runtime {
		return axgoja.NewRuntime(axgoja.WithRuntimePolicy(ax.Object("timeoutMs", 180000)))
	}
	child := ax.NewAgent("question -> answer", ax.Object("directResponse", "off", "runtime", runtime()))
	parent := ax.NewAgent("question -> answer", ax.Object("directResponse", "off", "runtime", runtime())).AddChildAgent("team", "researcher", child)
	ctx, cancel := context.WithTimeout(context.Background(), 180*time.Second)
	defer cancel()
	result, err := parent.Forward(ctx, client, ax.Object("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."), ax.Object("control", control, "serviceTier", "standard", "max_actor_steps", 8))
	if err != nil {
		panic(err)
	}
	for _, word := range []string{"42", "VERIFIED", "CHILD-CHECK"} {
		if !strings.Contains(fmt.Sprint(result), word) {
			panic(fmt.Sprint(result))
		}
	}
	if !applied {
		panic("Child control did not apply")
	}
	fmt.Println(result)
	usage, err := json.Marshal(parent.GetUsage())
	if err != nil {
		panic(err)
	}
	fmt.Println(string(usage))
}

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"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))
}

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

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

	specialist := ax.NewAgent(
		`brief:string, goal:string -> plan:string[] "Ordered, concrete steps", answer:string, risks:string[]`,
		map[string]ax.Value{
			"contextFields":   ax.Array("brief"),
			"contextPolicy":   ax.Object("preset", "checkpointed", "budget", "balanced"),
			"maxRuntimeChars": 3000,
			"runtime":         ax.Object("language", "JavaScript"),
		},
	)

	output, err := specialist.Forward(
		ctx,
		client,
		map[string]ax.Value{
			"brief": brief,
			"goal":  "Propose a safe, incremental 2-quarter plan to split checkout out first, respecting the hard constraints.",
		},
		map[string]ax.Value{"runtime": axgoja.NewRuntime(), "max_actor_steps": 12},
	)
	if err != nil {
		panic(err)
	}
	printJSON(output)
}

package main

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

	ax "github.com/ax-llm/ax/packages/go"
)

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

type panelist struct {
	model  string
	client ax.AIClient
}

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

	openaiKey := os.Getenv("OPENAI_API_KEY")
	if openaiKey == "" {
		openaiKey = os.Getenv("OPENAI_APIKEY")
	}
	googleKey := os.Getenv("GOOGLE_APIKEY")
	if googleKey == "" {
		googleKey = os.Getenv("GOOGLE_API_KEY")
	}
	anthropicKey := os.Getenv("ANTHROPIC_APIKEY")
	if anthropicKey == "" {
		anthropicKey = os.Getenv("ANTHROPIC_API_KEY")
	}
	if openaiKey == "" || googleKey == "" || anthropicKey == "" {
		panic("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.
	panel := []panelist{
		{"openai/gpt-5.4-mini", ax.NewOpenAICompatibleClient(map[string]ax.Value{"api_key": openaiKey, "model": "gpt-5.4-mini", "model_config": ax.Object("temperature", 0)})},
		{"google/gemini-3.5-flash", ax.NewGoogleGeminiClient(map[string]ax.Value{"api_key": googleKey, "model": "gemini-3.5-flash"})},
		{"anthropic/claude-haiku-4.5", ax.NewAnthropicClient(map[string]ax.Value{"api_key": anthropicKey, "model": "claude-haiku-4-5"})},
	}

	researcher := ax.NewAx(
		"question:string -> answer:string, keyFindings:string[], citations:string[], confidence:number",
		map[string]ax.Value{"instruction": "Answer independently. Use evidence. Call out uncertainty. Do not optimize for consensus."},
	)

	judge := ax.NewAx(
		"question:string, candidates:json -> consensus:string[], contradictions:string[], uniqueInsights:string[], blindSpots:string[]",
		map[string]ax.Value{"instruction": "Compare the candidates. Find agreement, conflicts, missing coverage, and unique useful points."},
	)

	synthesizer := ax.NewAx(
		"question:string, candidates:json, review:json -> answer:string, citations:string[], caveats:string[]",
		map[string]ax.Value{"instruction": "Write one final answer grounded in the candidates and review. Resolve conflicts explicitly."},
	)

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

	candidates := []ax.Value{}
	for _, p := range panel {
		response, err := researcher.Forward(ctx, p.client, map[string]ax.Value{"question": question}, nil)
		if err != nil {
			panic(err)
		}
		candidate := map[string]ax.Value{"model": p.model}
		if fields, ok := response.(map[string]ax.Value); ok {
			for k, v := range fields {
				candidate[k] = v
			}
		}
		candidates = append(candidates, candidate)
	}

	// The judge + synthesizer run on one of the panel clients (OpenAI here).
	orchestrator := panel[0].client
	review, err := judge.Forward(ctx, orchestrator, map[string]ax.Value{"question": question, "candidates": ax.Array(candidates...)}, nil)
	if err != nil {
		panic(err)
	}
	final, err := synthesizer.Forward(ctx, orchestrator, map[string]ax.Value{"question": question, "candidates": ax.Array(candidates...), "review": review}, nil)
	if err != nil {
		panic(err)
	}

	printJSON(final)
}