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"
	"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)
}