Generation Generation — Go examples backed by real provider calls. go examples examples/generation src/examples/go/generation example Generation

These Go examples are real runnable files. Edit the source file first; this page is rebuilt from the checked-in example and its metadata header.

Go Prompt-Cached Generation

Runs GPT-6 structured generation with stable OpenAI prompt-cache affinity.

Go
package main

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

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

func openAIClient() ax.AIClient {
	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-6-luna"
	}
	return ax.NewAI("openai", 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 main() {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	client := openAIClient()
	program := ax.NewAx("question:string -> answer:string", nil)
	output, err := program.Forward(ctx, client, map[string]ax.Value{"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."}, map[string]ax.Value{"timeout": 30000, "promptCacheKey": "ax-openai-example", "contextCache": ax.Object()})
	if err != nil {
		panic(err)
	}
	printJSON(output)
}

Go Astra Generation

Runs Astra through the standard generator with automatic Responses routing and prompt caching.

Go
package main

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

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

func openAIClient() ax.AIClient {
	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-6-astra"
	}
	return ax.NewAI("openai", map[string]ax.Value{"api_key": apiKey, "model": model, "model_config": ax.Object("thinkingTokenBudget", "low", "max_tokens", 2048)})
}

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

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	client := openAIClient()
	program := ax.NewAx("question:string -> answer:string", nil)
	output, err := program.Forward(ctx, client, map[string]ax.Value{"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."}, map[string]ax.Value{"serviceTier": "standard", "promptCacheKey": "ax-openai-example", "contextCache": ax.Object()})
	if err != nil {
		panic(err)
	}
	printJSON(output)
}

Go GPT-6.1 Sol Generation

Runs GPT-6.1 Sol through the standard generator with automatic Responses routing and prompt caching.

Go
package main

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

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

func openAIClient() ax.AIClient {
	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-6.1-sol"
	}
	return ax.NewAI("openai", map[string]ax.Value{"api_key": apiKey, "model": model, "model_config": ax.Object("thinkingTokenBudget", "low", "max_tokens", 2048)})
}

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

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	client := openAIClient()
	program := ax.NewAx("question:string -> answer:string", nil)
	output, err := program.Forward(ctx, client, map[string]ax.Value{"question": "In one sentence, explain Ax as a language-agnostic LLM programming library."}, map[string]ax.Value{"serviceTier": "standard", "promptCacheKey": "ax-gpt61-sol-example", "contextCache": ax.Object()})
	if err != nil {
		panic(err)
	}
	printJSON(output)
}

Go Model Catalog

Lists static models and named OpenAI-compatible profiles with portable thinking levels and service tiers.

Go
package main

import (
	"fmt"

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

func values(value ax.Value) []ax.Value {
	switch items := value.(type) {
	case []ax.Value:
		return items
	case *ax.AxArray:
		return items.Items
	default:
		panic(fmt.Sprintf("expected catalog array, got %T", value))
	}
}

func provider(catalog ax.Value, name string) map[string]ax.Value {
	for _, item := range values(catalog) {
		entry := item.(map[string]ax.Value)
		if entry["name"] == name {
			return entry
		}
	}
	panic("missing provider " + name)
}

func contains(items ax.Value, expected string) bool {
	for _, item := range values(items) {
		if item == expected {
			return true
		}
	}
	return false
}

func main() {
	catalog := ax.GetSupportedAIModels(map[string]ax.Value{})
	azure := provider(catalog, "azure-openai")
	openrouter := provider(catalog, "openrouter")
	azureCapabilities := azure["capabilities"].(map[string]ax.Value)
	openrouterCapabilities := openrouter["capabilities"].(map[string]ax.Value)

	if azure["isDynamic"] != true || len(values(azure["models"])) != 0 {
		panic("unexpected Azure profile")
	}
	if !contains(azureCapabilities["thinkingLevels"], "high") {
		panic("missing Azure thinking levels")
	}
	if !contains(azureCapabilities["serviceTiers"], "priority") {
		panic("missing Azure service tier")
	}
	if !contains(openrouterCapabilities["serviceTiers"], "flex") {
		panic("missing OpenRouter service tier")
	}

	fmt.Printf("%d providers; Azure and OpenRouter named profiles are available\n", len(values(catalog)))
}

Go Jev Signature Decisions

Converts Jev probabilities into boolean and class outputs with a provider threshold.

Go
package main

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

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	model := ax.NewAI("typesafe", map[string]ax.Value{"api_key": os.Getenv("TYPESAFE_APIKEY"), "trueThreshold": 0.9})
	triage := ax.NewAx("ticket:string -> urgent:boolean(true \"Customers cannot complete a core task\", false \"Routine request\") \"Needs immediate attention?\", team:class \"support, billing, engineering\"", nil)
	decisionValue, err := triage.Forward(ctx, model, map[string]ax.Value{"ticket": "Checkout is unavailable for all customers after the latest deployment."}, nil)
	if err != nil {
		panic(err)
	}
	values := decisionValue.(map[string]ax.Value)
	decision := map[string]ax.Value{"urgent": values["urgent"], "team": values["team"]}
	if _, ok := decision["urgent"].(bool); !ok {
		panic("Invalid boolean")
	}
	output := decision
	data, err := json.MarshalIndent(output, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

Go OpenAI Signature Decisions

Converts OpenAI probabilities into boolean and class outputs with a provider threshold.

Go
package main

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

func main() {
	apiKey := os.Getenv("OPENAI_API_KEY")
	if apiKey == "" {
		apiKey = os.Getenv("OPENAI_APIKEY")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	model := ax.NewAI("openai-decisions", map[string]ax.Value{"api_key": apiKey, "trueThreshold": 0.9})
	triage := ax.NewAx("ticket:string -> urgent:boolean(true \"Customers cannot complete a core task\", false \"Routine request\") \"Needs immediate attention?\", team:class \"support, billing, engineering\"", nil)
	decisionValue, err := triage.Forward(ctx, model, map[string]ax.Value{"ticket": "Checkout is unavailable for all customers after the latest deployment."}, nil)
	if err != nil {
		panic(err)
	}
	values := decisionValue.(map[string]ax.Value)
	decision := map[string]ax.Value{"urgent": values["urgent"], "team": values["team"]}
	if _, ok := decision["urgent"].(bool); !ok {
		panic("Invalid boolean")
	}
	output := decision
	data, err := json.MarshalIndent(output, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

Go Meta Muse Spark

Selects any of Meta’s three protocols through the existing chat API.

Go
package main

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

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

func main() {
	key := os.Getenv("MODEL_API_KEY")
	if key == "" {
		panic("Set MODEL_API_KEY to run this example.")
	}
	for _, profile := range []string{"meta", "meta-chat", "meta-messages"} {
		client := ax.NewAI(profile, map[string]ax.Value{"api_key": key, "model": "muse-spark-1.3"})
		response, err := client.Chat(context.Background(), map[string]ax.Value{
			"chat_prompt":  ax.Array(ax.Object("role", "user", "content", "Name a solar-powered sailboat.")),
			"model_config": ax.Object("thinking_token_budget", "highest"),
		}, nil)
		if err != nil {
			panic(err)
		}
		data, _ := json.Marshal(response)
		fmt.Println(profile, string(data))
	}
}

Go Structured Extraction

Extracts structured fields and labels from support text with OpenAI.

Go
package main

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

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


func openAIClient() ax.AIClient {
	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.NewAI("openai", 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 main() {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	client := openAIClient()
	program := ax.NewAx("ticket:string -> priority:class \"high, normal, low\", summary:string, labels:string[]", nil)
	output, err := program.Forward(ctx, client, map[string]ax.Value{"ticket": "Checkout has failed for enterprise customers since 09:00. Support wants a concise summary and tags."}, nil)
	if err != nil { panic(err) }
	printJSON(output)
}

Go Vertex MaaS Renewable Credentials

Calls a Vertex MaaS OpenAI-compatible endpoint with a fresh bearer token per request.

Go
package main

import (
	"context"
	"fmt"
	"os"

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

func required(name string) string {
	value := os.Getenv(name)
	if value == "" {
		panic("Set " + name + " to run this example.")
	}
	return value
}

func vertexCredentialProvider() ax.AxCredentialProvider {
	return ax.AxCredentialProviderFunc(
		func(_ context.Context, _ ax.AxCredentialRequest) (map[string]string, error) {
			// Replace this environment lookup with the host application's ADC token
			// source. Ax calls the hook again for every request attempt and retry.
			token := os.Getenv("GOOGLE_VERTEX_ACCESS_TOKEN")
			if token == "" {
				return nil, fmt.Errorf("set GOOGLE_VERTEX_ACCESS_TOKEN to run this example")
			}
			return map[string]string{"Authorization": "Bearer " + token}, nil
		},
	)
}

func main() {
	model := os.Getenv("AX_VERTEX_MODEL")
	if model == "" {
		model = "google/gemma-4-26b-a4b-it-maas"
	}
	client := ax.NewAI("vertex-ai", map[string]ax.Value{
		"base_url":            required("VERTEX_AI_API_URL"),
		"model":               model,
		"credential_provider": vertexCredentialProvider(),
	})
	out, err := client.Chat(context.Background(), map[string]ax.Value{
		"chat_prompt":     ax.Array(ax.Object("role", "user", "content", "Reply with the word ready.")),
		"response_format": ax.Object("type", "json_object"),
	}, nil)
	if err != nil {
		panic(err)
	}
	fmt.Println(out)
}

Go Jev Hybrid Reply

Passes Jev decisions to a second Ax program to generate a customer reply.

Go
package main

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

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	model := ax.NewAI("typesafe", map[string]ax.Value{"api_key": os.Getenv("TYPESAFE_APIKEY"), "trueThreshold": 0.9})
	triage := ax.NewAx("ticket:string -> urgent:boolean(true \"Customers cannot complete a core task\", false \"Routine request\") \"Needs immediate attention?\", team:class \"support, billing, engineering\"", nil)
	decisionValue, err := triage.Forward(ctx, model, map[string]ax.Value{"ticket": "Checkout is unavailable for all customers after the latest deployment."}, nil)
	if err != nil {
		panic(err)
	}
	values := decisionValue.(map[string]ax.Value)
	decision := map[string]ax.Value{"urgent": values["urgent"], "team": values["team"]}
	if _, ok := decision["urgent"].(bool); !ok {
		panic("Invalid boolean")
	}
	key := os.Getenv("OPENAI_API_KEY")
	if key == "" {
		key = os.Getenv("OPENAI_APIKEY")
	}
	writer := ax.NewAI("openai", map[string]ax.Value{"api_key": key, "model": "gpt-5.6-luna"})
	inputs := map[string]ax.Value{"ticket": "Checkout is unavailable for all customers after the latest deployment.", "urgent": decision["urgent"], "team": decision["team"]}
	replyValue, err := ax.NewAx("ticket:string, urgent:boolean, team:string -> reply:string", nil).Forward(ctx, writer, inputs, nil)
	if err != nil {
		panic(err)
	}
	replyValues := replyValue.(map[string]ax.Value)
	reply := map[string]ax.Value{"reply": replyValues["reply"]}
	if text, ok := reply["reply"].(string); !ok || text == "" {
		panic("Empty reply")
	}
	output := map[string]any{"decision": decision, "reply": reply}
	data, err := json.MarshalIndent(output, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

Go Signature Constraints

Builds native constrained fields and runs the signature with OpenAI.

Go
package main

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

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

func openAIClient() ax.AIClient {
	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.NewAI("openai", map[string]ax.Value{
		"api_key":      apiKey,
		"model":        model,
		"model_config": ax.Object("temperature", 0),
	})
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	signature := ax.AxSignature{
		Description: "Extract a constrained restaurant booking",
		Inputs: []ax.Field{
			{
				Name: "requestText", Title: "Request Text", Description: "Booking request",
				Type: ax.FieldType{Name: "string", MinLength: 10, MaxLength: 500},
			},
			{
				Name: "contactEmail", Title: "Contact Email", Description: "Contact email",
				Type: ax.FieldType{Name: "string", Format: "email"},
			},
		},
		Outputs: []ax.Field{
			{
				Name: "partySize", Title: "Party Size", Description: "Guests",
				Type: ax.FieldType{Name: "number", Minimum: 1, Maximum: 12},
			},
			{
				Name: "bookingCode", Title: "Booking Code", Description: "Must look like ABC-1234",
				Type: ax.FieldType{
					Name: "string", Pattern: `^[A-Z]{3}-\d{4}$`,
					PatternDescription: "Three letters, a dash, and four digits",
				},
			},
		},
	}
	program := ax.NewAx("requestText:string -> partySize:number, bookingCode:string", nil)
	program.Signature = signature
	output, err := program.Forward(
		ctx,
		openAIClient(),
		map[string]ax.Value{
			"requestText":  "Book dinner for four people under the name Ada Lovelace.",
			"contactEmail": "ada@example.com",
		},
		nil,
	)
	if err != nil {
		panic(err)
	}
	data, err := json.MarshalIndent(output, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

Centralized Usage Observer

Attributes every completed model call to a tenant, user, and request from one global observer.

Go
package main

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

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

func main() {
	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"
	}

	events := []ax.AxUsageEvent{}
	ax.SetUsageObserver(func(event ax.AxUsageEvent) {
		events = append(events, event)
	})
	defer ax.SetUsageObserver(nil)

	client := ax.NewAI("openai", map[string]ax.Value{
		"api_key": apiKey,
		"model":   model,
		"usageContext": ax.Object(
			"tenantId", "tenant-42",
			"feature", "support-chat",
			"attributes", ax.Object("environment", "example"),
		),
	})
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	_, err := client.Chat(
		ctx,
		map[string]ax.Value{
			"chat_prompt": ax.Array(
				ax.Object("role", "user", "content", "Reply with one short greeting."),
			),
		},
		ax.Object(
			"usageContext",
			ax.Object("userId", "user-7", "requestId", fmt.Sprintf("request-%d", time.Now().UnixNano())),
		),
	)
	if err != nil {
		panic(err)
	}
	output, err := json.MarshalIndent(events, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(output))
}

Go Incremental Provider Stream

Pulls OpenAI SSE events incrementally and closes the response body.

Go
package main

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

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

func main() {
	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.6-luna"
	}
	client, ok := ax.NewAI("openai", map[string]ax.Value{"api_key": apiKey, "model": model}).(ax.StreamingAIClient)
	if !ok {
		panic("OpenAI provider does not expose incremental streaming")
	}
	started := time.Now()
	stream, err := client.StreamEvents(context.Background(), map[string]ax.Value{
		"chat_prompt": ax.Array(ax.Object("role", "user", "content", "Reply with exactly: streaming works")),
	}, nil)
	if err != nil {
		panic(err)
	}
	defer stream.Close()
	for stream.Next() {
		event := stream.Value().(map[string]ax.Value)
		var results []ax.Value
		switch value := event["results"].(type) {
		case *ax.AxArray:
			results = value.Items
		case []ax.Value:
			results = value
		}
		if len(results) == 0 {
			continue
		}
		result := results[0].(map[string]ax.Value)
		if content, ok := result["content"].(string); ok && content != "" {
			fmt.Printf("[%d ms] %s", time.Since(started).Milliseconds(), content)
		}
	}
	if err := stream.Err(); err != nil {
		panic(err)
	}
	fmt.Println()
}

Go Decision Evidence Reranking

Ranks documents through Ax signatures while retaining decision evidence and candidate IDs.

Go
package main

import (
	"context"
	"encoding/json"
	"fmt"
	ax "github.com/ax-llm/ax/packages/go"
	"os"
	"sort"
	"time"
)

func main() {
	key := os.Getenv("OPENAI_API_KEY")
	if key == "" {
		key = os.Getenv("OPENAI_APIKEY")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	model := ax.NewAI("openai-decisions", map[string]ax.Value{"api_key": key})
	judge := ax.NewAx("query:string, document:string -> relevant:boolean(evidence, threshold 0.7), relevance:score \"Irrelevant, Related, Partial answer, Direct answer\"(evidence)", nil)
	candidates := []struct{ id, text string }{{"refund", "Request a refund within 30 days through your account billing page."}, {"password", "Reset your password using the sign-in page."}}
	ranked := []map[string]ax.Value{}
	for _, item := range candidates {
		result, err := judge.Forward(ctx, model, map[string]ax.Value{"query": "How do I request a refund?", "document": item.text}, nil)
		if err != nil {
			panic(err)
		}
		values := result.(map[string]ax.Value)
		values["id"] = item.id
		ranked = append(ranked, values)
	}
	score := func(i int) float64 { return ranked[i]["relevance"].(map[string]ax.Value)["value"].(float64) }
	sort.SliceStable(ranked, func(i, j int) bool { return score(i) > score(j) })
	data, err := json.MarshalIndent(ranked, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

Go Portable Cancellation

Uses context cancellation to stop a provider request before transport and preserve its reason.

Go
package main

import (
	"context"
	"errors"
	"fmt"
	"strings"

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

func main() {
	transport := ax.NewScriptedTransport([]ax.Value{ax.Object("status", 200, "json", ax.Object())})
	client := ax.NewOpenAICompatibleClient(map[string]ax.Value{
		"api_key": "test-key", "model": "gpt-5.6-luna", "transport": transport,
	})
	ctx, cancel := context.WithCancelCause(context.Background())
	cancel(errors.New("user stopped"))

	_, err := client.Chat(ctx, map[string]ax.Value{
		"chat_prompt": ax.Array(ax.Object("role", "user", "content", "This must not be sent.")),
	}, nil)
	var aborted ax.AxAIServiceAbortedError
	if !errors.As(err, &aborted) || aborted.Retryable || !strings.Contains(err.Error(), "user stopped") {
		panic(fmt.Sprintf("wrong cancellation error: %v", err))
	}
	if len(transport.Requests) != 0 {
		panic("pre-cancelled request reached transport")
	}
	fmt.Println("cancelled before transport: user stopped")
}

Go Sonnet 5.5 Between-Tool Thinking

Disables pre-response thinking while retaining signed updates for replay.

Go
package main

import (
	"context"
	"fmt"
	ax "github.com/ax-llm/ax/packages/go"
	"os"
)

func main() {
	apiKey := os.Getenv("ANTHROPIC_API_KEY")
	if apiKey == "" {
		apiKey = os.Getenv("ANTHROPIC_APIKEY")
	}
	if apiKey == "" {
		panic("Set ANTHROPIC_API_KEY or ANTHROPIC_APIKEY to run this example.")
	}
	client := ax.NewAI("anthropic", ax.Object("api_key", apiKey, "model", "claude-sonnet-5-5"))
	response, err := client.Chat(context.Background(), ax.Object(
		"chat_prompt", ax.Array(ax.Object("role", "user", "content", "Reply with exactly: Sonnet 5.5 works")),
		"model_config", ax.Object("thinkingTokenBudget", "none", "effort", "high", "maxTokens", 128),
	), nil)
	if err != nil {
		panic(err)
	}
	results := response.(map[string]ax.Value)["results"].(*ax.AxArray).Items
	fmt.Println(results[0].(map[string]ax.Value)["content"])
}

Go Streaming Field Deltas

Streams AxGen output as TypeScript-style {Version, Index, Delta} field deltas and merges them as they arrive.

Go
package main

import (
	"context"
	"fmt"
	"os"

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

func main() {
	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"
	}
	client := ax.NewAI("openai", map[string]ax.Value{"api_key": apiKey, "model": model})
	story := ax.NewAx(`topic:string -> title:string, story:string "Three short sentences"`, nil)

	// Each delta holds the new text of one field. Merge a sample's deltas
	// (strings append, other values replace) and start over when the version
	// changes: a retry or a replaced step starts a new version.
	merged := map[string]ax.Value{}
	version := 0
	values := map[string]ax.Value{"topic": "a lighthouse keeper's cat"}
	for delta, err := range story.StreamingForward(context.Background(), client, values, nil) {
		if err != nil {
			panic(err)
		}
		if delta.Version != version {
			merged, version = map[string]ax.Value{}, delta.Version
			fmt.Println("\n[retry: starting over]")
		}
		for field, value := range delta.Delta {
			text, isText := value.(string)
			if previous, ok := merged[field].(string); ok && isText {
				merged[field] = previous + text
			} else {
				merged[field] = value
			}
			if field == "story" && isText {
				fmt.Print(text)
			}
		}
	}
	fmt.Printf("\nTitle: %v\n", merged["title"])
}

Go Date Fields

Parses date, datetime and range outputs into ISO 8601 by default as TypeScript does, and passes a time.Time as an input.

Go
package main

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

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

func main() {
	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-6-luna"
	}
	client := ax.NewAI("openai", map[string]ax.Value{"api_key": apiKey, "model": model})

	// Date outputs parse by default; set parseDates (or parse_dates) to false to keep text.
	planner := ax.NewAx(
		`emailText:string, sentAt:datetime -> meetingStartsAt:datetime "Start time with its time zone", meetingDay:date, travelWindow:dateRange "First and last day away"`,
        nil,
	)
	out, err := planner.Forward(context.Background(), client, map[string]ax.Value{
		"emailText": "Can we meet next Tuesday at 3pm New York time? I'm travelling from the 8th to the 12th.",
		// A time.Time is rendered as TypeScript renders a Date.
		"sentAt": time.Date(2024, 5, 2, 16, 30, 0, 0, time.UTC),
	}, nil)
	if err != nil {
		panic(err)
	}
	fields := out.(map[string]ax.Value)
	fmt.Println(fields["meetingStartsAt"]) // e.g. 2024-05-07T19:00:00.000Z
	fmt.Println(fields["meetingDay"])      // e.g. 2024-05-07T00:00:00.000Z
	fmt.Println(fields["travelWindow"])    // map[end:... start:2024-05-08T00:00:00.000Z]
}

Go Field Processor Feedback

Sends a field processor’s note back to the model for another step, as TypeScript does, and trims the final answer with a field transform.

Go
package main

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

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

func main() {
	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"
	}
	client := ax.NewAI("openai", map[string]ax.Value{"api_key": apiKey, "model": model})
	summarize := ax.NewAx("text:string -> summary:string", nil)

	// A non-nil result goes back to the model as a user message, and the next
	// step's answer replaces this one.
	summarize.AddFieldProcessor("summary", func(value ax.Value, _ ax.AxFieldProcessorContext) (ax.Value, error) {
		summary, _ := value.(string)
		if words := len(strings.Fields(summary)); words > 12 {
			return fmt.Sprintf("That summary has %d words; answer again in at most 12 words.", words), nil
		}
		return nil, nil
	})
	summarize.AddFieldTransform("summary", "trim")

	text := "The committee met on Tuesday to review the budget. After a long debate about " +
		"the new library wing, they approved the plan and asked staff to find a builder " +
		"who can start in spring, while keeping the reading room open during the work."
	output, err := summarize.Forward(context.Background(), client, map[string]ax.Value{"text": text}, nil)
	if err != nil {
		panic(err)
	}
	fmt.Println(output.(map[string]ax.Value)["summary"])
}

Go Haiku 5.5 Adaptive Thinking

Uses Haiku 5.5 adaptive thinking at low effort for a short response.

  • Provider: anthropic
  • Env: ANTHROPIC_API_KEY, ANTHROPIC_APIKEY
  • Level: intermediate
  • Run: npm run example -- go src/examples/go/generation/haiku55.go
  • Source: src/examples/go/generation/haiku55.go
Go
package main

import (
	"context"
	"fmt"
	ax "github.com/ax-llm/ax/packages/go"
	"os"
)

func main() {
	apiKey := os.Getenv("ANTHROPIC_API_KEY")
	if apiKey == "" {
		apiKey = os.Getenv("ANTHROPIC_APIKEY")
	}
	if apiKey == "" {
		panic("Set ANTHROPIC_API_KEY or ANTHROPIC_APIKEY to run this example.")
	}
	client := ax.NewAI("anthropic", ax.Object("api_key", apiKey, "model", "claude-haiku-5-5"))
	response, err := client.Chat(context.Background(), ax.Object(
		"chat_prompt", ax.Array(ax.Object("role", "user", "content", "Reply with exactly: Haiku 5.5 works")),
		"model_config", ax.Object("thinkingTokenBudget", "low", "showThoughts", false, "maxTokens", 2048),
	), nil)
	if err != nil {
		panic(err)
	}
	results := response.(map[string]ax.Value)["results"].(*ax.AxArray).Items
	fmt.Println(results[0].(map[string]ax.Value)["content"])
}

Go Tool Result Formatting

Formats a structured inventory tool result as concise text for the model.

Go
package main

import (
	"context"
	"fmt"
	ax "github.com/ax-llm/ax/packages/go"
	"os"
)

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.")
	}
	model := os.Getenv("AX_OPENAI_MODEL")
	if model == "" {
		model = "gpt-6-luna"
	}
	client := ax.NewAI("openai", ax.Object("api_key", key, "model", model))
	inventory := ax.Fn("inventory").WithHandler(func(_ map[string]ax.Value) (ax.Value, error) {
		return ax.Object("available", 12, "warehouse", "A"), nil
	})
	inventory.Description = "Read the current stock count."
	program := ax.NewAx("question:string -> answer:string", nil)
	program.Functions = []ax.Tool{inventory}
	// The model receives this text; tool traces retain the original object.
	program.SetFunctionResultFormatter(func(result ax.Value) (string, error) {
		return fmt.Sprintf("%v units available", result.(map[string]ax.Value)["available"]), nil
	})
	result, err := program.Forward(context.Background(), client, ax.Object("question", "Call inventory and report how many units are available."), nil)
	if err != nil {
		panic(err)
	}
	fmt.Println(result.(map[string]ax.Value)["answer"])
}

Go Gemini Flex Inference

Sends latency-tolerant work through Gemini Flex and reports the applied tier.

Go
package main

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

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

func geminiAPIKey() string {
	value := os.Getenv("GOOGLE_API_KEY")
	if value == "" {
		value = os.Getenv("GOOGLE_APIKEY")
	}
	if value == "" {
		panic("Set GOOGLE_API_KEY or GOOGLE_APIKEY to run this example.")
	}
	return value
}

func main() {
	model := os.Getenv("AX_GEMINI_MODEL")
	if model == "" {
		model = "gemini-3.8-flash"
	}
	client := ax.NewAI("google-gemini", map[string]ax.Value{
		"api_key": geminiAPIKey(),
		"model":   model,
	})
	out, err := client.Chat(context.Background(), map[string]ax.Value{
		"chat_prompt": ax.Array(ax.Object(
			"role", "user",
			"content", "Explain in one sentence why batch evaluations save time.",
		)),
	}, map[string]ax.Value{"service_tier": "flex"})
	if err != nil {
		panic(err)
	}
	data, err := json.MarshalIndent(out, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

Go Native File Routing

Summarizes a PDF through a provider router without replacing the native file with extracted text.

Go
package main

import (
 "context"
 "fmt"
 "os"
 ax "github.com/ax-llm/ax/packages/go"
)
func main() {
 key:=os.Getenv("OPENAI_API_KEY");if key==""{key=os.Getenv("OPENAI_APIKEY")}
 client := ax.NewAI("openai", ax.Object("api_key",key,"model","gpt-6-astra","model_config",ax.Object("thinkingTokenBudget","low")))
 router := ax.NewProviderRouter(ax.Object("providers",ax.Object("primary",client)))
 program := ax.NewAx("document:file -> summary:string",nil)
 result,err := program.Forward(context.Background(),router,ax.Object("document",ax.Object("filename","report.pdf","mimeType","application/pdf","data",os.Getenv("AX_PDF_BASE64"))),ax.Object("serviceTier","standard"))
 if err != nil {panic(err)}
 fmt.Println(result)
}

Go Automatic Background Tools

Uses ordinary generation with background tools, steering, and a reasoning update.

Go
package main

import (
	"context"
	"fmt"
	ax "github.com/ax-llm/ax/packages/go"
	"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").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.NewAx("question -> answer", nil)
	program.Functions = []ax.Tool{slow, label}
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	result, err := program.Forward(ctx, client, ax.Object("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."), ax.Object("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.")
}

Go Cancel Background Work

Cancels a live Astra run through the high-level controller and observes cooperative tool cancellation.

Go
package main
import (
 "context"
 "fmt"
 "os"
 "strings"
 "sync/atomic"
 "time"
 ax "github.com/ax-llm/ax/packages/go"
)
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.")}
 control:=ax.RunControl();settled:=make(chan struct{});var started atomic.Int64
 lookup:=ax.Fn("lookup").Execution("background").WithContextHandler(func(ctx context.Context,_ map[string]ax.Value)(ax.Value,error){started.Store(time.Now().UnixNano());control.Abort();select {case <-ctx.Done():close(settled);return "LATE: discard this result",nil;case <-time.After(2*time.Second):return nil,fmt.Errorf("Tool missed cancellation")}})
 program:=ax.NewAx("question -> answer",nil);program.Functions=[]ax.Tool{lookup}
 client:=ax.NewAI("openai",ax.Object("api_key",key,"model","gpt-6-astra","model_config",ax.Object("thinkingTokenBudget","low","max_tokens",2048)))
 ctx,cancel:=context.WithTimeout(context.Background(),90*time.Second);defer cancel()
 _,err:=program.Forward(ctx,client,ax.Object("question","Call lookup once and return its result."),ax.Object("control",control,"serviceTier","standard"))
 if err==nil||started.Load()==0||!strings.Contains(err.Error(),"unresolved calls") {panic(fmt.Sprintf("Expected cancelled pending work: %v",err))}
 elapsed:=time.Since(time.Unix(0,started.Load()));if elapsed>2*time.Second {panic("Cancellation blocked caller")};select {case <-settled:case <-time.After(2*time.Second):panic("Tool missed cancellation")}
 fmt.Printf("Cancelled in %s; %v\n",elapsed,err)
}

Go Contextual Generation

Answers from supplied context and returns compact citations with OpenAI.

Go
package main

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

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


func openAIClient() ax.AIClient {
	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.NewAI("openai", 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 main() {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	client := openAIClient()
	program := ax.NewAx("context:string, question:string -> answer:string, citations:string[]", nil)
	output, err := program.Forward(ctx, client, map[string]ax.Value{"context": "Ax uses signatures, ai(), ax(), agent(), flow(), and optimize().", "question": "How should a new developer think about Ax?"}, nil)
	if err != nil { panic(err) }
	printJSON(output)
}

Go Jev Native Questions

Uses structured criteria, native scoring, model discovery, and probability-based decisions.

Go
package main

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

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	client := ax.Typesafe(map[string]ax.Value{"api_key": os.Getenv("TYPESAFE_APIKEY")})
	models, err := client.ListModels(ctx, nil)
	if err != nil {
		panic(err)
	}
	if len(models) == 0 {
		panic("Empty model catalog")
	}
	request := ax.TypesafeRequest{
		State: map[string]ax.Value{
			"ticket":  "Checkout is unavailable for all customers after the latest deployment.",
			"account": map[string]ax.Value{"tier": "enterprise", "notes": nil},
		},
		Questions: map[string]ax.TypesafeQuestion{
			"urgent": {Type: "noul",
				Instructions: map[string]ax.Value{"question": "Does this need immediate attention?"},
				Criteria:     map[string]ax.Value{"true": "Customers cannot complete a core task", "false": "Routine request"}},
			"team": {Type: "choice", Instructions: "Who should handle the ticket?",
				Criteria: map[string]ax.Value{"support": "Usage guidance",
					"billing": map[string]ax.Value{"scope": "Invoices and payments"}, "engineering": "Product failures"}},
			"severity": {Type: "score", Instructions: "Rate customer impact",
				Criteria: []ax.Value{"Minor inconvenience", "One task blocked", "Core task unavailable", "Widespread outage"}},
		},
	}
	response, err := client.SystemOne(ctx, request, nil)
	if err != nil {
		panic(err)
	}
	probability := response.Answers["urgent"].Noul
	score := response.Answers["severity"].Score
	if probability < 0 || probability > 1 || score < 0 || score > 3 {
		panic("Invalid answer bounds")
	}
	// Apply thresholds and custom score scales in application code.
	output := map[string]any{"page_on_call": probability >= 0.9, "severity_1_to_5": 1 + 4*score/3, "response": response}
	data, err := json.MarshalIndent(output, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(data))
}

Go Adaptive Provider Balancing

Routes equivalent chat traffic using shared reliability, latency, and cost statistics.

Go
package main

import (
	"context"
	"fmt"
	"os"

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

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 to run this example.")
	}
	model := os.Getenv("AX_OPENAI_MODEL")
	if model == "" {
		model = "gpt-5.4-mini"
	}
	services := []ax.AxAIService{
		ax.NewAI("openai", map[string]ax.Value{"api_key": key, "model": model}).(ax.AxAIService),
		ax.NewAI("openai", map[string]ax.Value{"api_key": key, "model": model}).(ax.AxAIService),
	}

	store := ax.NewAxInMemoryBalancerStatsStore()
	routeKeys := []string{"openai-primary", "openai-backup"}
	events := []string{}
	strategy := &ax.AxBalancerAdaptiveStrategy{
		DeadlineMs: 6_000, BadOutcomeCost: 0.02,
		ExpectedTokens: map[string]ax.Value{"promptTokens": 1_200.0, "completionTokens": 300.0},
		Namespace:      "support-summary-v1", StatsStore: store,
		RouteKey: func(_ ax.AxAIService, index int) string { return routeKeys[index] },
		Slice: func(context map[string]ax.Value) string {
			if options, ok := context["options"].(map[string]ax.Value); ok && options["stream"] == true {
				return "streaming"
			}
			return "interactive"
		},
		OnRoutingEvent: func(event ax.AxBalancerRoutingEvent) { events = append(events, fmt.Sprint(event["type"])) },
	}
	balancer, err := ax.NewAxBalancerWithOptions(services, ax.AxBalancerOptions{Strategy: strategy})
	if err != nil {
		panic(err)
	}
	response, err := balancer.Chat(context.Background(), map[string]ax.Value{"model": model, "chat_prompt": ax.Array(ax.Object("role", "user", "content", "Summarize why shared routing state matters."))}, nil)
	if err != nil {
		panic(err)
	}
	fmt.Println(response)
	fmt.Println(events)
}

Go OpenAI Native Decisions

Uses ordered predicate, choice, and score questions with explicit rubrics and raw probabilities.

Go
package main

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

func main() {
	apiKey := os.Getenv("OPENAI_API_KEY")
	if apiKey == "" {
		apiKey = os.Getenv("OPENAI_APIKEY")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	client := ax.OpenAIDecisions(map[string]ax.Value{"api_key": apiKey})
	var request map[string]ax.Value
	if err := json.Unmarshal([]byte(`{"input": "Checkout is unavailable for all customers after the latest deployment.", "questions": [{"type": "predicate", "name": "urgent", "instructions": "Are customers unable to complete a core task?"}, {"type": "choice", "name": "team", "instructions": "Who should handle the ticket?", "choices": [{"value": "support", "description": "Usage guidance"}, {"value": "billing", "description": "Invoices and payments"}, {"value": "engineering", "description": "Product failures"}]}, {"type": "score", "name": "severity", "instructions": "Rate customer impact", "levels": [{"label": "Minor inconvenience"}, {"label": "One task blocked"}, {"label": "Core task unavailable"}, {"label": "Widespread outage"}]}]}`), &request); err != nil {
		panic(err)
	}
	response, err := client.Create(ctx, request, nil)
	if err != nil {
		panic(err)
	}
	// Handle per-question refusals before using probability, choice, or score.
	output, err := json.MarshalIndent(response, "", "  ")
	if err != nil {
		panic(err)
	}
	fmt.Println(string(output))
}

Portable Runtime Hooks

Applies global and forward-scoped rate limiting, tracing, and metrics to AxGen, AxAgent, and AxFlow.

Go
package main

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

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

type logSpan struct{ name string }
func (s *logSpan) SetAttributes(map[string]ax.Value) {}
func (s *logSpan) AddEvent(name string, _ map[string]ax.Value) { fmt.Printf("[span:event] %s %s\n", s.name, name) }
func (s *logSpan) RecordException(err error) { fmt.Printf("[span:error] %s %v\n", s.name, err) }
func (s *logSpan) SetStatus(string, string) {}
func (s *logSpan) End() { fmt.Printf("[span:end] %s\n", s.name) }

type logTracer struct{}
func (logTracer) StartSpan(start ax.AxSpanStart) ax.AxSpan {
	fmt.Printf("[span:start] %s\n", start.Name)
	return &logSpan{name: start.Name}
}

type logInstrument struct{ name string }
func (i logInstrument) Add(value float64, _ map[string]ax.Value) { fmt.Printf("[metric] %s += %g\n", i.name, value) }
func (i logInstrument) Record(value float64, _ map[string]ax.Value) { fmt.Printf("[metric] %s = %g\n", i.name, value) }

type logMeter struct{}
func (logMeter) CreateCounter(name string, _ ax.AxMetricInstrumentOptions) ax.AxCounter { return logInstrument{name} }
func (logMeter) CreateHistogram(name string, _ ax.AxMetricInstrumentOptions) ax.AxHistogram { return logInstrument{name} }
func (logMeter) CreateGauge(name string, _ ax.AxMetricInstrumentOptions) ax.AxGauge { return logInstrument{name} }

func limiter(label string) ax.AxRateLimiter {
	return ax.AxRateLimiterFunc(func(next ax.AxRequestExecutor, info ax.AxRateLimitInfo) (ax.Value, error) {
		fmt.Printf("[limit:%s] %s %s/%s stream=%t\n", label, info.Operation, info.Provider, info.Model, info.Streaming)
		return next()
	})
}

func main() {
	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" }
	client := ax.NewAI("openai", map[string]ax.Value{
		"api_key": apiKey, "model": model, "model_config": ax.Object("temperature", 0),
	})
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	defer cancel()
	tracer, meter := logTracer{}, logMeter{}
	override := ax.AxRuntimeHooks{RateLimiter: limiter("forward"), Tracer: tracer, Meter: meter}

	ax.SetRateLimiter(limiter("global"))
	ax.SetTracer(tracer)
	ax.SetMeter(meter)
	defer func() { ax.SetRateLimiter(nil); ax.SetTracer(nil); ax.SetMeter(nil) }()

	direct := ax.NewAx("topic:string -> summary:string", nil)
	result, err := direct.Forward(ctx, client, map[string]ax.Value{"topic": "portable Ax runtime hooks"}, nil)
	if err != nil { panic(err) }
	fmt.Println(result)

	helper := ax.NewAgent("question:string -> answer:string", nil)
	result, err = helper.ForwardWithHooks(ctx, client, map[string]ax.Value{"question": "What does a rate limiter wrap?"}, map[string]ax.Value{"runtime": axgoja.NewRuntime(), "max_actor_steps": 12}, override)
	if err != nil { panic(err) }
	fmt.Println(result)

	workflow := ax.NewFlow(map[string]ax.Value{"id": "examples.runtimeHooks"}).
		Execute("outline", ax.NewAx("topic:string -> outline:string", nil), nil).
		Execute("polish", ax.NewAx("outline:string -> answer:string", nil), nil).
		Returns(map[string]ax.Value{"answer": "answer"})
	result, err = workflow.ForwardWithHooks(ctx, client, map[string]ax.Value{"topic": "Ax runtime hooks"}, nil, override)
	if err != nil { panic(err) }
	for _, key := range []string{"answer"} {
		switch fmt.Sprint(result.(map[string]ax.Value)[key]) {
		case "", "<nil>", "[]":
			panic(fmt.Sprintf("flow output field %s is empty: %v", key, result))
		}
	}
	fmt.Println(result)
}
Docs