From 875c0427068d9cb062880e5bfc3d4a91884638fd Mon Sep 17 00:00:00 2001 From: Yu Yi Date: Tue, 13 Oct 2026 14:52:49 -0400 Subject: [PATCH] docs: align model examples with supported providers The model client implements the Google provider while the guide also advertises an unsupported Anthropic configuration. Correct the manifest guidance and exercise its actual model YAML through the API, store, and model provider client. --- docs/manifests.md | 34 ++---- internal/server/manifests_test.go | 177 ++++++++++++++++++++++++++++++ 2 files changed, 185 insertions(+), 26 deletions(-) create mode 100644 internal/server/manifests_test.go diff --git a/docs/manifests.md b/docs/manifests.md index 16f32e35..3a0222df 100644 --- a/docs/manifests.md +++ b/docs/manifests.md @@ -1,6 +1,6 @@ # Writing manifests -All four kinds can live in one multi-document YAML file. The examples below show every kind; [`examples/simple.yaml`](../examples/simple.yaml) is the smallest working Task. [`examples/claude-code/task.yaml`](../examples/claude-code/task.yaml) shows a Task that runs Claude Code headlessly against a cloned repository. +The three supported kinds, `Task`, `Workspace`, and `Model`, can live in one multi-document YAML file. The examples below show every kind; [`examples/simple.yaml`](../examples/simple.yaml) is the smallest working Task. [`examples/claude-code/task.yaml`](../examples/claude-code/task.yaml) shows a Task that runs Claude Code headlessly against a cloned repository. `metadata.name` and `metadata.atespace` become Substrate resource names, so they must be lowercase RFC 1123 labels: at most 63 lowercase alphanumeric characters or `-`, starting and ending with an alphanumeric character. `ax apply` rejects anything else up front rather than letting the task fail later with `ActorCreationFailed`. @@ -15,7 +15,7 @@ metadata: name: task123 atespace: default spec: - image: "ghcr.io/my-org/my-agent-image" + image: "ghcr.io/my-org/my-agent-image@sha256:" command: ["python", "agent.py"] env: - name: ENVIRONMENT @@ -101,10 +101,14 @@ spec: ## Model -For Google models, store the API key and set `provider: google`. +Only `google` is implemented as a model provider. Other providers are accepted when applying a `Model`, but generation fails with an `unsupported provider` error. + +`Model` resources configure AX's model client (`default-model` in the atespace by default). These resources do not configure the provider or model used by arbitrary agent frameworks inside a Task image. + +Store the Google API key in a Kubernetes Secret in the same namespace as the Model's `metadata.atespace`, and set `provider: google`. ```bash -kubectl create secret generic gemini-api-secret --from-literal=GEMINI_API_KEY="AIzaSy..." +kubectl create secret generic gemini-api-secret --namespace default --from-literal=GEMINI_API_KEY="" ``` Then reference it from the `Model`: @@ -122,25 +126,3 @@ spec: name: gemini-api-secret key: GEMINI_API_KEY ``` - -For Anthropic models, store the key the same way and set `provider: anthropic`. - -```bash -kubectl create secret generic anthropic-api-secret --from-literal=ANTHROPIC_API_KEY="sk-ant-..." -``` - -```yaml -apiVersion: ax.io/v1alpha1 -kind: Model -metadata: - name: claude-model - atespace: default -spec: - provider: anthropic - model: claude-opus-5 - secretKey: - name: anthropic-api-secret - key: ANTHROPIC_API_KEY - parameters: - maxTokens: 16000 -``` diff --git a/internal/server/manifests_test.go b/internal/server/manifests_test.go new file mode 100644 index 00000000..85e3283e --- /dev/null +++ b/internal/server/manifests_test.go @@ -0,0 +1,177 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server_test + +import ( + "context" + "encoding/json" + "fmt" + "net" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/google/ax/internal/model" + "github.com/google/ax/internal/server" + "github.com/google/ax/internal/store/memory" + "github.com/google/ax/internal/substrate/substratetest" + "github.com/google/ax/pkg/apis/v1alpha1" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/protobuf/proto" + "gopkg.in/yaml.v3" +) + +func TestManifestGuideModels(t *testing.T) { + guide, err := os.ReadFile("../../docs/manifests.md") + if err != nil { + t.Fatal(err) + } + var examples []*v1alpha1.Model + for _, block := range strings.Split(string(guide), "```yaml\n")[1:] { + source := strings.SplitN(block, "```", 2)[0] + var meta struct { + Kind string `yaml:"kind"` + } + if err := yaml.Unmarshal([]byte(source), &meta); err != nil { + t.Fatal(err) + } + if meta.Kind != v1alpha1.KindModel { + continue + } + var m v1alpha1.Model + if err := yaml.Unmarshal([]byte(source), &m); err != nil { + t.Fatal(err) + } + if m.Metadata == nil || m.Spec == nil || m.Spec.SecretKey == nil { + t.Fatal("Model example requires metadata, spec and secretKey") + } + examples = append(examples, &m) + } + if len(examples) == 0 { + t.Fatal("guide has no Model examples") + } + + for _, example := range examples { + t.Run(example.Metadata.Name, func(t *testing.T) { + for _, unsupported := range []bool{false, true} { + name := "documented" + if unsupported { + name = "unsupported-provider" + } + t.Run(name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + t.Cleanup(cancel) + st := memory.NewStore() + srv := server.NewServer(st, newTestReconciler(substratetest.NewServer(t))) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + done := make(chan error, 1) + go func() { done <- srv.GRPCServer().Serve(ln) }() + t.Cleanup(func() { + srv.GRPCServer().Stop() + if err := <-done; err != nil { + t.Errorf("gRPC server cleanup: %v", err) + } + t.Log("gRPC fixture stopped") + }) + conn, err := grpc.NewClient(ln.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := conn.Close(); err != nil { + t.Error(err) + } + }) + api := v1alpha1.NewAXClient(conn) + m := proto.Clone(example).(*v1alpha1.Model) + if unsupported { + m.Spec.Provider = "anthropic" + } + if _, err := api.UpdateModel(ctx, &v1alpha1.UpdateModelRequest{Model: m}); err != nil { + t.Fatal(err) + } + saved, err := api.GetModel(ctx, &v1alpha1.GetModelRequest{Atespace: m.Metadata.Atespace, Name: m.Metadata.Name}) + if err != nil { + t.Fatal(err) + } + if !proto.Equal(saved.Spec, m.Spec) { + t.Fatal("API/store changed Model spec") + } + t.Logf("YAML -> gRPC UpdateModel/GetModel: provider=%s model=%s", saved.Spec.Provider, saved.Spec.Model) + + var calls atomic.Int32 + fixture := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + if r.Method != http.MethodPost || r.URL.Path != "/v1beta/models/"+m.Spec.Model+":generateContent" || r.URL.Query().Get("key") != "fixture-key" { + t.Errorf("unexpected provider request method/path or fixture authentication") + } + var body struct { + Contents []struct{ Parts []struct{ Text string } } + SystemInstruction struct{ Parts []struct{ Text string } } + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + t.Error(err) + } + if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 1 || !strings.Contains(body.Contents[0].Parts[0].Text, "Install dependencies") || len(body.SystemInstruction.Parts) != 1 { + t.Error("request did not contain prompt and system instruction") + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"candidates":[{"content":{"parts":[{"text":"fixture bootstrap plan"}]}}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":3,"totalTokenCount":10}}`) + })) + t.Cleanup(func() { fixture.Close(); t.Log("HTTP provider fixture stopped") }) + client, err := model.NewClientFromStore(ctx, st, m.Metadata.Atespace, m.Metadata.Name, + model.WithBaseURL(fixture.URL), model.WithSecretResolver(func(secretName, key string) (string, error) { + if secretName != m.Spec.SecretKey.Name || key != m.Spec.SecretKey.Key { + return "", fmt.Errorf("unexpected secret reference") + } + return "fixture-key", nil + })) + if err != nil { + t.Fatal(err) + } + resp, err := client.Generate(ctx, &model.GenerateRequest{ + Prompt: "Install dependencies", + SystemInstruction: "System workspace planner", + }) + if unsupported { + if err == nil || !strings.Contains(err.Error(), `unsupported provider "anthropic"`) { + t.Fatalf("expected unsupported provider error, got %v", err) + } + if calls.Load() != 0 { + t.Fatal("unsupported provider contacted HTTP fixture") + } + t.Logf("negative control: %v; HTTP requests=0", err) + return + } + if err != nil { + t.Fatal(err) + } + if resp.Content != "fixture bootstrap plan" || resp.Model != m.Spec.Model || calls.Load() != 1 { + t.Fatalf("expected fixture plan and one HTTP request, got %+v; requests=%d", resp, calls.Load()) + } + t.Logf("store -> model client -> Google HTTP: model=%s content=%q; requests=%d", resp.Model, resp.Content, calls.Load()) + }) + } + }) + } +}