From 9ebe41415950f9cd5e532cb9c7f9bb6a56df20e9 Mon Sep 17 00:00:00 2001 From: Jeffrey Chen Date: Wed, 11 Feb 2026 00:49:06 +0000 Subject: [PATCH 1/4] Update azure.ai.agents to use gRPC model catalog and quota functionality --- .../extensions/azure.ai.agents/extension.yaml | 1 + cli/azd/extensions/azure.ai.agents/go.mod | 19 +- cli/azd/extensions/azure.ai.agents/go.sum | 28 +- .../azure.ai.agents/internal/cmd/init.go | 850 +----------------- .../internal/cmd/init_models.go | 825 +++++++++++++++++ .../azure.ai.agents/internal/cmd/init_test.go | 69 ++ .../azure.ai.agents/internal/cmd/root.go | 2 +- .../internal/pkg/azure/ai/model_catalog.go | 429 --------- 8 files changed, 935 insertions(+), 1288 deletions(-) create mode 100644 cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go delete mode 100644 cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go diff --git a/cli/azd/extensions/azure.ai.agents/extension.yaml b/cli/azd/extensions/azure.ai.agents/extension.yaml index acfe4c89bc6..e467bef2fe8 100644 --- a/cli/azd/extensions/azure.ai.agents/extension.yaml +++ b/cli/azd/extensions/azure.ai.agents/extension.yaml @@ -6,6 +6,7 @@ description: Extension for the Foundry Agent Service. (Preview) usage: azd ai agent [options] # NOTE: Make sure version.txt is in sync with this version. version: 0.1.9-preview +requiredAzdVersion: ">1.23.4" language: go capabilities: - custom-commands diff --git a/cli/azd/extensions/azure.ai.agents/go.mod b/cli/azd/extensions/azure.ai.agents/go.mod index 667288ea1d3..a1c79f24d61 100644 --- a/cli/azd/extensions/azure.ai.agents/go.mod +++ b/cli/azd/extensions/azure.ai.agents/go.mod @@ -10,7 +10,7 @@ require ( github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices v1.8.0 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armresources v1.2.0 github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armsubscriptions v1.3.0 - github.com/azure/azure-dev/cli/azd v0.0.0-20260122173819-89795b295491 + github.com/azure/azure-dev/cli/azd v0.0.0-20260210235925-af2fc5ae4ef3 github.com/braydonk/yaml v0.9.0 github.com/drone/envsubst v1.0.3 github.com/fatih/color v1.18.0 @@ -19,6 +19,8 @@ require ( github.com/spf13/cobra v1.10.1 github.com/spf13/pflag v1.0.10 go.yaml.in/yaml/v3 v3.0.4 + google.golang.org/genproto/googleapis/rpc v0.0.0-20251007200510-49b9836ed3ff + google.golang.org/grpc v1.76.0 google.golang.org/protobuf v1.36.10 gopkg.in/yaml.v3 v3.0.1 ) @@ -48,7 +50,6 @@ require ( github.com/clipperhouse/uax29/v2 v2.2.0 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/dlclark/regexp2 v1.11.5 // indirect - github.com/eiannone/keyboard v0.0.0-20220611211555-0d226195f203 // indirect github.com/go-logr/logr v1.4.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect github.com/golang-jwt/jwt/v5 v5.3.0 // indirect @@ -56,6 +57,7 @@ require ( github.com/gorilla/css v1.0.1 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/invopop/jsonschema v0.13.0 // indirect + github.com/jmespath-community/go-jmespath v1.1.1 // indirect github.com/joho/godotenv v1.5.1 // indirect github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect github.com/kylelemons/godebug v1.1.0 // indirect @@ -90,11 +92,10 @@ require ( go.opentelemetry.io/otel/sdk v1.38.0 // indirect go.opentelemetry.io/otel/trace v1.38.0 // indirect go.uber.org/atomic v1.11.0 // indirect - golang.org/x/crypto v0.45.0 // indirect - golang.org/x/net v0.47.0 // indirect - golang.org/x/sys v0.38.0 // indirect - golang.org/x/term v0.37.0 // indirect - golang.org/x/text v0.31.0 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20251007200510-49b9836ed3ff // indirect - google.golang.org/grpc v1.76.0 // indirect + golang.org/x/crypto v0.46.0 // indirect + golang.org/x/exp v0.0.0-20250911091902-df9299821621 // indirect + golang.org/x/net v0.48.0 // indirect + golang.org/x/sys v0.39.0 // indirect + golang.org/x/term v0.38.0 // indirect + golang.org/x/text v0.32.0 // indirect ) diff --git a/cli/azd/extensions/azure.ai.agents/go.sum b/cli/azd/extensions/azure.ai.agents/go.sum index 7f43fb64537..bcb74bdf033 100644 --- a/cli/azd/extensions/azure.ai.agents/go.sum +++ b/cli/azd/extensions/azure.ai.agents/go.sum @@ -51,8 +51,8 @@ github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWp github.com/aymanbagabas/go-udiff v0.2.0/go.mod h1:RE4Ex0qsGkTAJoQdQQCA0uG+nAzJO/pI/QwceO5fgrA= github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= -github.com/azure/azure-dev/cli/azd v0.0.0-20260122173819-89795b295491 h1:k3ymt8L+5ovZCsXbqWzmY8rGqiGh89GXnJ+ojKbCxk8= -github.com/azure/azure-dev/cli/azd v0.0.0-20260122173819-89795b295491/go.mod h1:j+bdvNwQPdYtSfFe/xbfWqYr8Guw9hiP1JOVpIBERj0= +github.com/azure/azure-dev/cli/azd v0.0.0-20260210235925-af2fc5ae4ef3 h1:c9y1yqY7GiVG+yLqemPCXh+jm2rQuDOBUEUR1OKuIE4= +github.com/azure/azure-dev/cli/azd v0.0.0-20260210235925-af2fc5ae4ef3/go.mod h1:4gxguvyC2AMKdookMDH1H6fYgtH9ogNb1BIx9WfeINM= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/benbjohnson/clock v1.3.5 h1:VvXlSJBzZpA/zum6Sj74hxwYI2DIxRWuNIoXAzHZz5o= @@ -96,8 +96,6 @@ github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZ github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/drone/envsubst v1.0.3 h1:PCIBwNDYjs50AsLZPYdfhSATKaRg/FJmDc2D6+C2x8g= github.com/drone/envsubst v1.0.3/go.mod h1:N2jZmlMufstn1KEqvbHjw40h1KyTmnVzHcSc9bFiJ2g= -github.com/eiannone/keyboard v0.0.0-20220611211555-0d226195f203 h1:XBBHcIb256gUJtLmY22n99HaZTz+r2Z51xUPi01m3wg= -github.com/eiannone/keyboard v0.0.0-20220611211555-0d226195f203/go.mod h1:E1jcSv8FaEny+OP/5k9UxZVw9YFWGj7eI4KR/iOBqCg= github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM= github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= @@ -132,6 +130,8 @@ github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= +github.com/jmespath-community/go-jmespath v1.1.1 h1:bFikPhsi/FdmlZhVgSCd2jj1e7G/rw+zyQfyg5UF+L4= +github.com/jmespath-community/go-jmespath v1.1.1/go.mod h1:4gOyFJsR/Gk+05RgTKYrifT7tBPWD8Lubtb5jRrfy9I= github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs= @@ -247,8 +247,8 @@ go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= -golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q= -golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4= +golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU= +golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0= golang.org/x/exp v0.0.0-20250911091902-df9299821621 h1:2id6c1/gto0kaHYyrixvknJ8tUK/Qs5IsmBtrc+FtgU= golang.org/x/exp v0.0.0-20250911091902-df9299821621/go.mod h1:TwQYMMnGpvZyc+JpB/UAuTNIsVJifOlSkrZkhcvpVUk= golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= @@ -256,8 +256,8 @@ golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73r golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= -golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= -golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= +golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU= +golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY= golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= @@ -272,18 +272,18 @@ golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc= -golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk= +golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= -golang.org/x/term v0.37.0 h1:8EGAD0qCmHYZg6J17DvsMy9/wJ7/D/4pV/wfnld5lTU= -golang.org/x/term v0.37.0/go.mod h1:5pB4lxRNYYVZuTLmy8oR2BH8dflOR+IbTYFD8fi3254= +golang.org/x/term v0.38.0 h1:PQ5pkm/rLO6HnxFR7N2lJHOZX6Kez5Y1gDSJla6jo7Q= +golang.org/x/term v0.38.0/go.mod h1:bSEAKrOT1W+VSu9TSCMtoGEOUcKxOKgl3LE5QEF/xVg= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= -golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM= -golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM= +golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU= +golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go index 48f3ef5f83f..2384864f19d 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go @@ -15,14 +15,12 @@ import ( "os" "path/filepath" "regexp" - "slices" "strconv" "strings" "azureaiagent/internal/pkg/agents/agent_yaml" "azureaiagent/internal/pkg/agents/registry_api" "azureaiagent/internal/pkg/azure" - "azureaiagent/internal/pkg/azure/ai" "azureaiagent/internal/project" "github.com/Azure/azure-sdk-for-go/sdk/azcore" @@ -35,7 +33,6 @@ import ( "github.com/azure/azure-dev/cli/azd/pkg/osutil" "github.com/azure/azure-dev/cli/azd/pkg/output" "github.com/azure/azure-dev/cli/azd/pkg/tools/github" - "github.com/azure/azure-dev/cli/azd/pkg/ux" "github.com/fatih/color" "github.com/spf13/cobra" "google.golang.org/protobuf/types/known/structpb" @@ -43,7 +40,7 @@ import ( ) type initFlags struct { - rootFlagsDefinition + *rootFlagsDefinition projectResourceId string manifestPointer string src string @@ -61,14 +58,14 @@ type InitAction struct { //azureClient *azure.AzureClient azureContext *azdext.AzureContext //composedResources []*azdext.ComposedResource - console input.Console - credential azcore.TokenCredential - modelCatalog map[string]*ai.AiModel - modelCatalogService *ai.ModelCatalogService - projectConfig *azdext.ProjectConfig - environment *azdext.Environment - flags *initFlags - deploymentDetails []project.Deployment + console input.Console + credential azcore.TokenCredential + modelCatalog map[string]*azdext.AiModel + locationWarningShown bool + projectConfig *azdext.ProjectConfig + environment *azdext.Environment + flags *initFlags + deploymentDetails []project.Deployment } // GitHubUrlInfo holds parsed information from a GitHub URL @@ -82,7 +79,7 @@ type GitHubUrlInfo struct { const AiAgentHost = "azure.ai.agent" const ContainerAppHost = "containerapp" -func newInitCommand(rootFlags rootFlagsDefinition) *cobra.Command { +func newInitCommand(rootFlags *rootFlagsDefinition) *cobra.Command { flags := &initFlags{ rootFlagsDefinition: rootFlags, } @@ -141,12 +138,11 @@ func newInitCommand(rootFlags rootFlagsDefinition) *cobra.Command { // azureClient: azure.NewAzureClient(credential), azureContext: azureContext, // composedResources: getComposedResourcesResponse.Resources, - console: console, - credential: credential, - modelCatalogService: ai.NewModelCatalogService(credential), - projectConfig: projectConfig, - environment: environment, - flags: flags, + console: console, + credential: credential, + projectConfig: projectConfig, + environment: environment, + flags: flags, } if err := action.Run(ctx); err != nil { @@ -1709,819 +1705,3 @@ func downloadDirectoryContentsWithoutGhCli( return nil } - -// func (a *InitAction) validateResources(ctx context.Context, agentYaml map[string]interface{}) error { -// fmt.Println("Reading model name from agent.yaml...") - -// // Extract the model name from agentYaml -// agentModelName, ok := agentYaml["model"].(string) -// if !ok || agentModelName == "" { -// return fmt.Errorf("extracting model name from agent YAML: model name missing or empty") -// } - -// fmt.Println("Reading current azd project resources...") - -// // Check if the ai.project resource already exists and has the required model -// existingResourceName, err := a.checkResourceExistsAndHasModel(agentModelName) -// if err != nil { -// return fmt.Errorf("checking if ai.project resource has model '%s': %w", agentModelName, err) -// } - -// if existingResourceName == "" { -// return a.addResource(ctx, agentModelName) -// } - -// fmt.Printf("Validated: ai.project resource '%s' has required model '%s'\n", existingResourceName, agentModelName) -// return nil -// } - -// // checkResourceExistsAndHasModel checks if the given ai.project resource has the specified model -// func (a *InitAction) checkResourceExistsAndHasModel(modelName string) (string, error) { -// // Look for ai.project resource -// var aiProjectResource *azdext.ComposedResource -// for _, resource := range a.composedResources { -// if resource.Type == "ai.project" { -// aiProjectResource = resource -// break -// } -// } - -// if aiProjectResource == nil { -// fmt.Println("No 'ai.project' resource found in current azd project.") -// return "", nil -// } - -// fmt.Println("'ai.project' resource found in current azd project. Checking for required model...") - -// // Parse the resource config to check for models -// if len(aiProjectResource.Config) > 0 { -// var config map[string]interface{} -// if err := yaml.Unmarshal(aiProjectResource.Config, &config); err != nil { -// return "", fmt.Errorf("parsing resource config: %w", err) -// } - -// // Check the models array - based on azure.yaml format: models[].name -// if models, ok := config["Models"].([]interface{}); ok { -// for _, model := range models { -// if modelObj, ok := model.(map[string]interface{}); ok { -// if name, ok := modelObj["Name"].(string); ok { -// if name == modelName { -// fmt.Printf("Found matching model: %s\n", name) -// return aiProjectResource.Name, nil -// } -// } -// } -// } -// } -// } - -// fmt.Printf("Model '%s' not found in resource '%s'\n", modelName, aiProjectResource.Name) -// return "", nil -// } - -// func (a *InitAction) addResource(ctx context.Context, agentModelName string) error { -// // Look for existing ai.project resource -// var aiProject *azdext.ComposedResource -// var aiProjectConfig *AiProjectResourceConfig - -// for _, resource := range a.composedResources { -// if resource.Type == "ai.project" { -// aiProject = resource - -// // Parse existing config if it exists -// if len(resource.Config) > 0 { -// if err := yaml.Unmarshal(resource.Config, &aiProjectConfig); err != nil { -// return fmt.Errorf("failed to unmarshal AI project config: %w", err) -// } -// } -// break -// } -// } - -// // Create new ai.project resource if it doesn't exist -// if aiProject == nil { -// fmt.Println("Adding new 'ai.project' resource to azd project.") -// aiProject = &azdext.ComposedResource{ -// Name: generateResourceName("ai-project", a.composedResources), -// Type: "ai.project", -// } -// aiProjectConfig = &AiProjectResourceConfig{} -// } - -// // Prompt user for model details -// modelDetails, err := a.promptForModelDetails(ctx, agentModelName) -// if err != nil { -// return fmt.Errorf("failed to get model details: %w", err) -// } - -// fmt.Println("Got model details, adding to ai.project resource.") -// // Convert the ai.AiModelDeployment to the map format expected by AiProjectResourceConfig -// defaultModel := map[string]interface{}{ -// "name": modelDetails.Name, -// "format": modelDetails.Format, -// "version": modelDetails.Version, -// "sku": map[string]interface{}{ -// "name": modelDetails.Sku.Name, -// "usageName": modelDetails.Sku.UsageName, -// "capacity": modelDetails.Sku.Capacity, -// }, -// } -// aiProjectConfig.Models = append(aiProjectConfig.Models, defaultModel) - -// // Marshal the config as JSON (since the struct has json tags) -// configJson, err := json.Marshal(aiProjectConfig) -// if err != nil { -// return fmt.Errorf("failed to marshal AI project config: %w", err) -// } - -// // Update the resource config -// aiProject.Config = configJson - -// // Add the resource to the project -// _, err = a.azdClient.Compose().AddResource(ctx, &azdext.AddResourceRequest{ -// Resource: aiProject, -// }) -// if err != nil { -// return fmt.Errorf("failed to add resource %s: %w", aiProject.Name, err) -// } - -// fmt.Printf("Added AI project resource '%s' to azure.yaml\n", aiProject.Name) -// return nil -// } - -// func (a *InitAction) promptForModelDetails(ctx context.Context, modelName string) (*ai.AiModelDeployment, error) { -// // Load the AI model catalog if not already loaded -// if err := a.loadAiCatalog(ctx); err != nil { -// return nil, err -// } - -// // Check if the model exists in the catalog -// var model *ai.AiModel -// model, exists := a.modelCatalog[modelName] -// if !exists { -// return nil, fmt.Errorf("model '%s' not found in AI model catalog", modelName) -// } - -// availableVersions, err := a.modelCatalogService.ListModelVersions(ctx, model) -// if err != nil { -// return nil, fmt.Errorf("listing versions for model '%s': %w", modelName, err) -// } - -// availableSkus, err := a.modelCatalogService.ListModelSkus(ctx, model) -// if err != nil { -// return nil, fmt.Errorf("listing SKUs for model '%s': %w", modelName, err) -// } - -// modelVersionSelection, err := selectFromList( -// ctx, a.console, "Which model version do you want to use?", availableVersions, nil) -// if err != nil { -// return nil, err -// } - -// skuSelection, err := selectFromList(ctx, a.console, "Select model SKU", availableSkus, nil) -// if err != nil { -// return nil, err -// } - -// deploymentOptions := ai.AiModelDeploymentOptions{ -// Versions: []string{modelVersionSelection}, -// Skus: []string{skuSelection}, -// } - -// modelDeployment, err := a.modelCatalogService.GetModelDeployment(ctx, model, &deploymentOptions) -// if err != nil { -// return nil, fmt.Errorf("failed to get model deployment: %w", err) -// } - -// return modelDeployment, nil -// } - -func (a *InitAction) loadAiCatalog(ctx context.Context) error { - if a.modelCatalog != nil { - return nil - } - - spinner := ux.NewSpinner(&ux.SpinnerOptions{ - Text: "Loading the model catalog", - ClearOnStop: true, - }) - - if err := spinner.Start(ctx); err != nil { - return fmt.Errorf("failed to start spinner: %w", err) - } - - aiModelCatalog, err := a.modelCatalogService.ListAllModels(ctx, a.azureContext.Scope.SubscriptionId, "") - - if err != nil { - return fmt.Errorf("failed to load the model catalog: %w", err) - } - - if err := spinner.Stop(ctx); err != nil { - return err - } - - a.modelCatalog = aiModelCatalog - return nil -} - -// // generateResourceName generates a unique resource name, similar to the AI builder pattern -// func generateResourceName(desiredName string, existingResources []*azdext.ComposedResource) string { -// resourceMap := map[string]struct{}{} -// for _, resource := range existingResources { -// resourceMap[resource.Name] = struct{}{} -// } - -// if _, exists := resourceMap[desiredName]; !exists { -// return desiredName -// } -// // If the desired name already exists, append a number (always 2 digits) to the name -// nextIndex := 1 -// for { -// newName := fmt.Sprintf("%s-%02d", desiredName, nextIndex) -// if _, exists := resourceMap[newName]; !exists { -// return newName -// } -// nextIndex++ -// } -// } - -func (a *InitAction) selectFromList( - ctx context.Context, property string, options []string, defaultOpt string) (string, error) { - - if len(options) == 1 { - fmt.Printf("Only one %s available: %s\n", property, options[0]) - return options[0], nil - } - - slices.Sort(options) - - // Convert default value to string for comparison - defaultStr := options[0] - if defaultOpt != "" { - defaultStr = defaultOpt - } - - if a.flags.NoPrompt { - fmt.Printf("No prompt mode enabled, selecting default %s: %s\n", property, defaultStr) - return defaultStr, nil - } - - // Create choices for the select prompt - choices := make([]*azdext.SelectChoice, len(options)) - defaultIndex := int32(0) - for i, val := range options { - choices[i] = &azdext.SelectChoice{ - Value: val, - Label: val, - } - if val == defaultStr { - defaultIndex = int32(i) - } - } - resp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ - Options: &azdext.SelectOptions{ - Message: fmt.Sprintf("Select %s", property), - Choices: choices, - SelectedIndex: &defaultIndex, - }, - }) - if err != nil { - return "", fmt.Errorf("failed to prompt for enum value: %w", err) - } - - return options[*resp.Value], nil -} - -func (a *InitAction) setEnvVar(ctx context.Context, key, value string) error { - _, err := a.azdClient.Environment().SetValue(ctx, &azdext.SetEnvRequest{ - EnvName: a.environment.Name, - Key: key, - Value: value, - }) - if err != nil { - return fmt.Errorf("failed to set environment variable %s=%s: %w", key, value, err) - } - - fmt.Printf("Set environment variable: %s=%s\n", key, value) - return nil -} - -func (a *InitAction) getModelDeploymentDetails(ctx context.Context, model agent_yaml.Model) (*project.Deployment, error) { - resp, err := a.azdClient.Environment().GetValue(ctx, &azdext.GetEnvRequest{ - EnvName: a.environment.Name, - Key: "AZURE_AI_PROJECT_ID", - }) - if err != nil { - return nil, fmt.Errorf("failed to get the environment variable AZURE_AI_PROJECT_ID from your azd environment: %w", err) - } - - foundryProjectId := resp.Value - if foundryProjectId != "" { - // Extract subscription and account name from foundry project ID - // Format: /subscriptions/{subscription}/resourceGroups/{rg}/providers/Microsoft.CognitiveServices/accounts/{account}/projects/{project} - parts := strings.Split(foundryProjectId, "/") - var subscription, resourceGroup, accountName string - - if len(parts) >= 9 { - subscription = parts[2] // subscriptions/{subscription} - resourceGroup = parts[4] // resourceGroups/{rg} - accountName = parts[8] // accounts/{account} - } - - deploymentsClient, err := armcognitiveservices.NewDeploymentsClient(subscription, a.credential, azure.NewArmClientOptions()) - if err != nil { - return nil, fmt.Errorf("failed to create deployments client: %w", err) - } - - pager := deploymentsClient.NewListPager(resourceGroup, accountName, nil) - var deployments []*armcognitiveservices.Deployment - for pager.More() { - page, err := pager.NextPage(ctx) - if err != nil { - return nil, fmt.Errorf("failed to list deployments: %w", err) - } - deployments = append(deployments, page.Value...) - } - - // Check for existing deployments that match the requested model - matchingDeployments := make(map[string]*armcognitiveservices.Deployment) - for _, deployment := range deployments { - if deployment.Properties != nil && deployment.Properties.Model != nil { - deployedModel := deployment.Properties.Model - if deployedModel.Name != nil { - if *deployedModel.Name == model.Id { - matchingDeployments[*deployment.Name] = deployment - } - } - } - } - - // If we found matching deployments, prompt the user - if len(matchingDeployments) > 0 { - fmt.Printf("In your Microsoft Foundry project, found %d existing model deployment(s) matching your model %s.\n", len(matchingDeployments), model.Id) - - // Build options list with existing deployments plus "Create new deployment" option - var options []string - for deploymentName := range matchingDeployments { - options = append(options, deploymentName) - } - options = append(options, "Create new model deployment") - - // Use selectFromList to choose between existing deployments or creating new one - selection, err := a.selectFromList(ctx, "deployment", options, options[0]) - if err != nil { - return nil, fmt.Errorf("failed to select deployment: %w", err) - } - - // Check if user chose to create new deployment - if selection != "Create new model deployment" { - // User chose an existing deployment by name - fmt.Printf("Using existing model deployment: %s\n", selection) - - // Get the selected deployment from the map and return its details - if deployment, exists := matchingDeployments[selection]; exists { - return &project.Deployment{ - Name: selection, - Model: project.DeploymentModel{ - Name: model.Id, - Format: *deployment.Properties.Model.Format, - Version: *deployment.Properties.Model.Version, - }, - Sku: project.DeploymentSku{ - Name: *deployment.SKU.Name, - Capacity: int(*deployment.SKU.Capacity), - }, - }, nil - } - } - } - } - - modelDetails, err := a.getModelDetails(ctx, model.Id) - if err != nil { - return nil, fmt.Errorf("failed to get model details: %w", err) - } - - message := fmt.Sprintf("Enter model deployment name for model '%s' (defaults to model name)", modelDetails.Name) - - modelDeploymentInput, err := a.azdClient.Prompt().Prompt(ctx, &azdext.PromptRequest{ - Options: &azdext.PromptOptions{ - Message: message, - IgnoreHintKeys: true, - DefaultValue: modelDetails.Name, - }, - }) - if err != nil { - return nil, fmt.Errorf("failed to prompt for text value: %w", err) - } - - modelDeployment := modelDeploymentInput.Value - - return &project.Deployment{ - Name: modelDeployment, - Model: project.DeploymentModel{ - Name: modelDetails.Name, - Format: modelDetails.Format, - Version: modelDetails.Version, - }, - Sku: project.DeploymentSku{ - Name: modelDetails.Sku.Name, - Capacity: int(modelDetails.Sku.Capacity), - }, - }, nil -} - -var defaultSkuPriority = []string{"GlobalStandard", "DataZoneStandard", "Standard"} - -func (a *InitAction) getModelDetails(ctx context.Context, modelName string) (*ai.AiModelDeployment, error) { - // Load the AI model catalog if not already loaded - if err := a.loadAiCatalog(ctx); err != nil { - return nil, err - } - - // Check if the model exists in the catalog - var model *ai.AiModel - model, exists := a.modelCatalog[modelName] - if !exists { - // Model not found - prompt user for alternative - selectedModel, err := a.promptForAlternativeModel(ctx, modelName) - if err != nil { - return nil, err - } - if selectedModel == nil { - return nil, fmt.Errorf("no model selected, exiting") - } - model = selectedModel - modelName = model.Name - } - - // Check if the model is available in the current location - currentLocation := a.azureContext.Scope.Location - if _, hasLocation := model.ModelDetailsByLocation[currentLocation]; !hasLocation { - // Model not available in current location - prompt user for action - resolvedModel, resolvedLocation, err := a.promptForModelLocationMismatch(ctx, model, currentLocation) - if err != nil { - return nil, err - } - if resolvedModel == nil { - return nil, fmt.Errorf("model unavailable in current location and no alternative selected, exiting") - } - model = resolvedModel - modelName = model.Name - currentLocation = resolvedLocation - } - - availableVersions, defaultVersion, err := a.modelCatalogService.ListModelVersions(ctx, model, currentLocation) - if err != nil { - return nil, fmt.Errorf("listing versions for model '%s': %w", modelName, err) - } - - modelVersion, err := a.selectFromList(ctx, "model version", availableVersions, defaultVersion) - if err != nil { - return nil, err - } - - availableSkus, err := a.modelCatalogService.ListModelSkus(ctx, model, currentLocation, modelVersion) - if err != nil { - return nil, fmt.Errorf("listing SKUs for model '%s': %w", modelName, err) - } - - // Determine default SKU based on priority list - defaultSku := "" - for _, sku := range defaultSkuPriority { - if slices.Contains(availableSkus, sku) { - defaultSku = sku - break - } - } - - skuSelection, err := a.selectFromList(ctx, "model SKU", availableSkus, defaultSku) - if err != nil { - return nil, err - } - - deploymentOptions := ai.AiModelDeploymentOptions{ - Versions: []string{modelVersion}, - Skus: []string{skuSelection}, - } - - modelDeployment, err := a.modelCatalogService.GetModelDeployment(ctx, model, &deploymentOptions) - if err != nil { - return nil, fmt.Errorf("failed to get model deployment: %w", err) - } - - if modelDeployment.Sku.Capacity == -1 { - skuCapacity, err := a.azdClient.Prompt().Prompt(ctx, &azdext.PromptRequest{ - Options: &azdext.PromptOptions{ - Message: "Selected model SKU has no default capacity. Please enter desired capacity", - IgnoreHintKeys: true, - Required: true, - DefaultValue: "10", - }, - }) - if err != nil { - return nil, fmt.Errorf("failed to prompt for text value: %w", err) - } - - capacity, err := strconv.Atoi(skuCapacity.Value) - if err != nil { - return nil, fmt.Errorf("invalid capacity value: %w", err) - } - modelDeployment.Sku.Capacity = int32(capacity) - } - - return modelDeployment, nil -} - -func (a *InitAction) promptForAlternativeModel(ctx context.Context, originalModelName string) (*ai.AiModel, error) { - fmt.Println(output.WithErrorFormat("The model '%s' could not be found in the model catalog for your subscription in any region.\n", originalModelName)) - - // Ask if they want to select a different model or exit - choices := []*azdext.SelectChoice{ - {Label: "Select a different model", Value: "select"}, - {Label: "Exit", Value: "exit"}, - } - - defaultIndex := int32(1) - selectResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ - Options: &azdext.SelectOptions{ - Message: "Would you like to select a different model or exit?", - Choices: choices, - SelectedIndex: &defaultIndex, - }, - }) - if err != nil { - return nil, fmt.Errorf("failed to prompt for model selection choice: %w", err) - } - - if choices[*selectResp.Value].Value == "exit" { - return nil, nil - } - - // Ask if they want all models or region-specific models - regionChoices := []*azdext.SelectChoice{ - {Label: fmt.Sprintf("Models available in my current region (%s)", a.azureContext.Scope.Location), Value: "region"}, - {Label: "All available models", Value: "all"}, - } - - regionResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ - Options: &azdext.SelectOptions{ - Message: "Which models would you like to explore?", - Choices: regionChoices, - }, - }) - if err != nil { - return nil, fmt.Errorf("failed to prompt for region choice: %w", err) - } - - // Get the list of model names based on the choice - var modelNames []string - if regionChoices[*regionResp.Value].Value == "region" { - // Filter models that have the current region - for name, model := range a.modelCatalog { - if _, hasLocation := model.ModelDetailsByLocation[a.azureContext.Scope.Location]; hasLocation { - modelNames = append(modelNames, name) - } - } - } else { - // All models - for name := range a.modelCatalog { - modelNames = append(modelNames, name) - } - } - - if len(modelNames) == 0 { - return nil, fmt.Errorf("no models available for selection") - } - - // Sort the model names - slices.Sort(modelNames) - - // Create choices for the model selection - modelChoices := make([]*azdext.SelectChoice, len(modelNames)) - for i, name := range modelNames { - modelChoices[i] = &azdext.SelectChoice{ - Label: name, - Value: name, - } - } - - modelResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ - Options: &azdext.SelectOptions{ - Message: "Select a model", - Choices: modelChoices, - }, - }) - if err != nil { - return nil, fmt.Errorf("failed to prompt for model selection: %w", err) - } - - selectedModelName := modelNames[*modelResp.Value] - return a.modelCatalog[selectedModelName], nil -} - -func (a *InitAction) promptForModelLocationMismatch(ctx context.Context, model *ai.AiModel, currentLocation string) (*ai.AiModel, string, error) { - fmt.Println(output.WithErrorFormat("The model '%s' is not available in your current location '%s'.", model.Name, currentLocation)) - fmt.Println("Would you like to use a different model, or select a different location?") - fmt.Println(output.WithWarningFormat( - "WARNING: If you switch locations:\n" + - "• Your AZD environment will use a new default region.\n" + - "• Any existing Azure AI Foundry project created in your current region may fail.\n\n" + - "Recommended options:\n" + - "1) Select a different model in this region (safe), or\n" + - "2) Create a new Foundry project after changing regions.")) - - // Ask what they want to do - choices := []*azdext.SelectChoice{ - {Label: "Select a different model available in this location", Value: "model"}, - {Label: "Select a different location for this model", Value: "location"}, - {Label: "Exit", Value: "exit"}, - } - - defaultIndex := int32(2) - selectResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ - Options: &azdext.SelectOptions{ - Message: "What would you like to do?", - Choices: choices, - SelectedIndex: &defaultIndex, - }, - }) - if err != nil { - return nil, "", fmt.Errorf("failed to prompt for action choice: %w", err) - } - - selectedChoice := choices[*selectResp.Value].Value - - if selectedChoice == "exit" { - return nil, "", nil - } - - if selectedChoice == "location" { - // Get available locations for this model - var locationNames []string - for locationName := range model.ModelDetailsByLocation { - locationNames = append(locationNames, locationName) - } - - if len(locationNames) == 0 { - return nil, "", fmt.Errorf("no locations available for model '%s'", model.Name) - } - - slices.Sort(locationNames) - - // Create choices for location selection - locationChoices := make([]*azdext.SelectChoice, len(locationNames)) - for i, name := range locationNames { - locationChoices[i] = &azdext.SelectChoice{ - Label: name, - Value: name, - } - } - - locationResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ - Options: &azdext.SelectOptions{ - Message: fmt.Sprintf("Select a location for model '%s'", model.Name), - Choices: locationChoices, - }, - }) - if err != nil { - return nil, "", fmt.Errorf("failed to prompt for location selection: %w", err) - } - - selectedLocation := locationNames[*locationResp.Value] - - // Update the azd environment with the new location - envResponse, err := a.azdClient.Environment().GetCurrent(ctx, &azdext.EmptyRequest{}) - if err != nil { - return nil, "", fmt.Errorf("failed to get current azd environment: %w", err) - } - - a.azureContext.Scope.Location = selectedLocation - _, err = a.azdClient.Environment().SetValue(ctx, &azdext.SetEnvRequest{ - EnvName: envResponse.Environment.Name, - Key: "AZURE_LOCATION", - Value: a.azureContext.Scope.Location, - }) - if err != nil { - return nil, "", fmt.Errorf("failed to update AZURE_LOCATION in azd environment: %w", err) - } - - fmt.Printf("Updated AZURE_LOCATION to '%s' in your azd environment.\n", selectedLocation) - - return model, selectedLocation, nil - } - - // selectedChoice == "model" - // Get models available in the current location - var modelNames []string - for name, m := range a.modelCatalog { - if _, hasLocation := m.ModelDetailsByLocation[currentLocation]; hasLocation { - modelNames = append(modelNames, name) - } - } - - if len(modelNames) == 0 { - return nil, "", fmt.Errorf("no models available in location '%s'", currentLocation) - } - - slices.Sort(modelNames) - - // Create choices for model selection - modelChoices := make([]*azdext.SelectChoice, len(modelNames)) - for i, name := range modelNames { - modelChoices[i] = &azdext.SelectChoice{ - Label: name, - Value: name, - } - } - - modelResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ - Options: &azdext.SelectOptions{ - Message: fmt.Sprintf("Select a model available in '%s'", currentLocation), - Choices: modelChoices, - }, - }) - if err != nil { - return nil, "", fmt.Errorf("failed to prompt for model selection: %w", err) - } - - selectedModelName := modelNames[*modelResp.Value] - return a.modelCatalog[selectedModelName], currentLocation, nil -} - -func (a *InitAction) ProcessModels(ctx context.Context, manifest *agent_yaml.AgentManifest) (*agent_yaml.AgentManifest, []project.Deployment, error) { - // Convert the template to bytes - templateBytes, err := yaml.Marshal(manifest.Template) - if err != nil { - return nil, nil, fmt.Errorf("failed to marshal agent template to YAML: %w", err) - } - - // Convert the bytes to a dictionary - var templateDict map[string]interface{} - if err := yaml.Unmarshal(templateBytes, &templateDict); err != nil { - return nil, nil, fmt.Errorf("failed to unmarshal agent template from YAML: %w", err) - } - - // Convert the dictionary to bytes - dictJsonBytes, err := yaml.Marshal(templateDict) - if err != nil { - return nil, nil, fmt.Errorf("failed to marshal templateDict to YAML: %w", err) - } - - // Convert the bytes to an Agent Definition - var agentDef agent_yaml.AgentDefinition - if err := yaml.Unmarshal(dictJsonBytes, &agentDef); err != nil { - return nil, nil, fmt.Errorf("failed to unmarshal YAML to AgentDefinition: %w", err) - } - - deploymentDetails := []project.Deployment{} - paramValues := registry_api.ParameterValues{} - switch agentDef.Kind { - case agent_yaml.AgentKindPrompt: - agentDef := manifest.Template.(agent_yaml.PromptAgent) - - modelDeployment, err := a.getModelDeploymentDetails(ctx, agentDef.Model) - if err != nil { - return nil, nil, fmt.Errorf("failed to get model deployment details: %w", err) - } - deploymentDetails = append(deploymentDetails, *modelDeployment) - paramValues["deploymentName"] = modelDeployment.Name - case agent_yaml.AgentKindHosted: - // Iterate over all models in the manifest for the container agent - for _, resource := range manifest.Resources { - // Convert the resource to bytes - resourceBytes, err := yaml.Marshal(resource) - if err != nil { - return nil, nil, fmt.Errorf("failed to marshal resource to YAML: %w", err) - } - - // Convert the bytes to an Agent Definition - var resourceDef agent_yaml.Resource - if err := yaml.Unmarshal(resourceBytes, &resourceDef); err != nil { - return nil, nil, fmt.Errorf("failed to unmarshal YAML to Resource: %w", err) - } - - if resourceDef.Kind == agent_yaml.ResourceKindModel { - resource := resource.(agent_yaml.ModelResource) - model := agent_yaml.Model{ - Id: resource.Id, - } - modelDeployment, err := a.getModelDeploymentDetails(ctx, model) - if err != nil { - return nil, nil, fmt.Errorf("failed to get model deployment details: %w", err) - } - deploymentDetails = append(deploymentDetails, *modelDeployment) - paramValues[resource.Name] = modelDeployment.Name - } - } - } - - updatedManifest, err := registry_api.InjectParameterValuesIntoManifest(manifest, paramValues) - if err != nil { - return nil, nil, fmt.Errorf("failed to inject deployment names into manifest: %w", err) - } - - fmt.Println("Model deployment details processed and injected into agent definition. Deployment details can also be found in the JSON formatted AI_PROJECT_DEPLOYMENTS environment variable.") - - return updatedManifest, deploymentDetails, nil -} diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go new file mode 100644 index 00000000000..140d960c91a --- /dev/null +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go @@ -0,0 +1,825 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package cmd + +import ( + "context" + "fmt" + "slices" + "strings" + + "azureaiagent/internal/pkg/agents/agent_yaml" + "azureaiagent/internal/pkg/agents/registry_api" + "azureaiagent/internal/pkg/azure" + "azureaiagent/internal/project" + + "github.com/Azure/azure-sdk-for-go/sdk/azcore/to" + "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices" + "github.com/azure/azure-dev/cli/azd/pkg/azdext" + "github.com/azure/azure-dev/cli/azd/pkg/output" + "github.com/azure/azure-dev/cli/azd/pkg/ux" + "google.golang.org/genproto/googleapis/rpc/errdetails" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + "gopkg.in/yaml.v3" +) + +var defaultSkuPriority = []string{"GlobalStandard", "DataZoneStandard", "Standard"} + +func (a *InitAction) loadAiCatalog(ctx context.Context) error { + if a.modelCatalog != nil { + return nil + } + + spinner := ux.NewSpinner(&ux.SpinnerOptions{ + Text: "Loading the model catalog", + ClearOnStop: true, + }) + + if err := spinner.Start(ctx); err != nil { + return fmt.Errorf("failed to start spinner: %w", err) + } + + modelResp, err := a.azdClient.Ai().ListModels(ctx, &azdext.ListModelsRequest{ + AzureContext: a.azureContext, + }) + stopErr := spinner.Stop(ctx) + if err != nil { + return fmt.Errorf("failed to load the model catalog: %w", err) + } + if stopErr != nil { + return stopErr + } + + a.modelCatalog = mapModelsByName(modelResp.Models) + + return nil +} + +func mapModelsByName(models []*azdext.AiModel) map[string]*azdext.AiModel { + modelMap := make(map[string]*azdext.AiModel, len(models)) + for _, model := range models { + if model == nil { + continue + } + modelMap[model.Name] = model + } + + return modelMap +} + +func (a *InitAction) updateEnvLocation(ctx context.Context, selectedLocation string) error { + envResponse, err := a.azdClient.Environment().GetCurrent(ctx, &azdext.EmptyRequest{}) + if err != nil { + return fmt.Errorf("failed to get current azd environment: %w", err) + } + + if a.azureContext == nil { + a.azureContext = &azdext.AzureContext{} + } + if a.azureContext.Scope == nil { + a.azureContext.Scope = &azdext.AzureScope{} + } + a.azureContext.Scope.Location = selectedLocation + + _, err = a.azdClient.Environment().SetValue(ctx, &azdext.SetEnvRequest{ + EnvName: envResponse.Environment.Name, + Key: "AZURE_LOCATION", + Value: a.azureContext.Scope.Location, + }) + if err != nil { + return fmt.Errorf("failed to update AZURE_LOCATION in azd environment: %w", err) + } + + fmt.Println(output.WithSuccessFormat("Updated AZURE_LOCATION to '%s' in your azd environment.", selectedLocation)) + return nil +} + +func (a *InitAction) selectFromList( + ctx context.Context, property string, options []string, defaultOpt string) (string, error) { + + if len(options) == 1 { + fmt.Printf("Only one %s available: %s\n", property, options[0]) + return options[0], nil + } + + slices.Sort(options) + + defaultStr := options[0] + if defaultOpt != "" { + defaultStr = defaultOpt + } + + if a.flags.NoPrompt { + fmt.Printf("No prompt mode enabled, selecting default %s: %s\n", property, defaultStr) + return defaultStr, nil + } + + choices := make([]*azdext.SelectChoice, len(options)) + defaultIndex := int32(0) + for i, val := range options { + choices[i] = &azdext.SelectChoice{ + Value: val, + Label: val, + } + if val == defaultStr { + defaultIndex = int32(i) + } + } + resp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ + Options: &azdext.SelectOptions{ + Message: fmt.Sprintf("Select %s", property), + Choices: choices, + SelectedIndex: &defaultIndex, + }, + }) + if err != nil { + return "", fmt.Errorf("failed to prompt for enum value: %w", err) + } + + return options[*resp.Value], nil +} + +func (a *InitAction) setEnvVar(ctx context.Context, key, value string) error { + _, err := a.azdClient.Environment().SetValue(ctx, &azdext.SetEnvRequest{ + EnvName: a.environment.Name, + Key: key, + Value: value, + }) + if err != nil { + return fmt.Errorf("failed to set environment variable %s=%s: %w", key, value, err) + } + + fmt.Printf("Set environment variable: %s=%s\n", key, value) + return nil +} + +func (a *InitAction) getModelDeploymentDetails(ctx context.Context, model agent_yaml.Model) (*project.Deployment, error) { + resp, err := a.azdClient.Environment().GetValue(ctx, &azdext.GetEnvRequest{ + EnvName: a.environment.Name, + Key: "AZURE_AI_PROJECT_ID", + }) + if err != nil { + return nil, fmt.Errorf("failed to get the environment variable AZURE_AI_PROJECT_ID from your azd environment: %w", err) + } + + foundryProjectId := resp.Value + if foundryProjectId != "" { + parts := strings.Split(foundryProjectId, "/") + var subscription, resourceGroup, accountName string + + if len(parts) >= 9 { + subscription = parts[2] + resourceGroup = parts[4] + accountName = parts[8] + } + + deploymentsClient, err := armcognitiveservices.NewDeploymentsClient(subscription, a.credential, azure.NewArmClientOptions()) + if err != nil { + return nil, fmt.Errorf("failed to create deployments client: %w", err) + } + + pager := deploymentsClient.NewListPager(resourceGroup, accountName, nil) + var deployments []*armcognitiveservices.Deployment + for pager.More() { + page, err := pager.NextPage(ctx) + if err != nil { + return nil, fmt.Errorf("failed to list deployments: %w", err) + } + deployments = append(deployments, page.Value...) + } + + matchingDeployments := make(map[string]*armcognitiveservices.Deployment) + for _, deployment := range deployments { + if deployment.Properties != nil && deployment.Properties.Model != nil { + deployedModel := deployment.Properties.Model + if deployedModel.Name != nil { + if *deployedModel.Name == model.Id { + matchingDeployments[*deployment.Name] = deployment + } + } + } + } + + if len(matchingDeployments) > 0 { + fmt.Printf("In your Microsoft Foundry project, found %d existing model deployment(s) matching your model %s.\n", len(matchingDeployments), model.Id) + + var options []string + for deploymentName := range matchingDeployments { + options = append(options, deploymentName) + } + options = append(options, "Create new model deployment") + + selection, err := a.selectFromList(ctx, "deployment", options, options[0]) + if err != nil { + return nil, fmt.Errorf("failed to select deployment: %w", err) + } + + if selection != "Create new model deployment" { + fmt.Printf("Using existing model deployment: %s\n", selection) + + if deployment, exists := matchingDeployments[selection]; exists { + return &project.Deployment{ + Name: selection, + Model: project.DeploymentModel{ + Name: model.Id, + Format: *deployment.Properties.Model.Format, + Version: *deployment.Properties.Model.Version, + }, + Sku: project.DeploymentSku{ + Name: *deployment.SKU.Name, + Capacity: int(*deployment.SKU.Capacity), + }, + }, nil + } + } + } + } + + modelDetails, err := a.getModelDetails(ctx, model.Id) + if err != nil { + return nil, fmt.Errorf("failed to get model details: %w", err) + } + + message := fmt.Sprintf("Enter model deployment name for model '%s' (defaults to model name)", modelDetails.ModelName) + + modelDeploymentInput, err := a.azdClient.Prompt().Prompt(ctx, &azdext.PromptRequest{ + Options: &azdext.PromptOptions{ + Message: message, + IgnoreHintKeys: true, + DefaultValue: modelDetails.ModelName, + }, + }) + if err != nil { + return nil, fmt.Errorf("failed to prompt for text value: %w", err) + } + + modelDeployment := modelDeploymentInput.Value + + return &project.Deployment{ + Name: modelDeployment, + Model: project.DeploymentModel{ + Name: modelDetails.ModelName, + Format: modelDetails.Format, + Version: modelDetails.Version, + }, + Sku: project.DeploymentSku{ + Name: modelDetails.Sku.Name, + Capacity: int(modelDetails.Capacity), + }, + }, nil +} + +func (a *InitAction) getModelDetails(ctx context.Context, modelName string) (*azdext.AiModelDeployment, error) { + if err := a.loadAiCatalog(ctx); err != nil { + return nil, err + } + + model, exists := a.modelCatalog[modelName] + if !exists { + selectedModel, err := a.promptForAlternativeModel(ctx, modelName) + if err != nil { + return nil, err + } + if selectedModel == nil { + return nil, fmt.Errorf("no model selected, exiting") + } + model = selectedModel + } + + currentLocation := a.azureContext.Scope.Location + if !slices.Contains(model.Locations, currentLocation) { + resolvedModel, resolvedLocation, err := a.promptForModelLocationMismatch( + ctx, + model, + currentLocation, + fmt.Sprintf("The model '%s' is not available in your current location '%s'.", model.Name, currentLocation), + modelRecoveryReasonAvailability, + ) + if err != nil { + return nil, err + } + if resolvedModel == nil { + return nil, fmt.Errorf("model unavailable in current location and no alternative selected, exiting") + } + model = resolvedModel + currentLocation = resolvedLocation + } + + if a.flags.NoPrompt { + fmt.Println("No prompt mode enabled, automatically selecting a model deployment based on availability and quota...") + return a.resolveModelDeploymentNoPrompt(ctx, model, currentLocation) + } + + for { + deploymentResp, err := a.azdClient.Prompt().PromptAiDeployment(ctx, &azdext.PromptAiDeploymentRequest{ + AzureContext: a.azureContext, + ModelName: model.Name, + Options: &azdext.AiModelDeploymentOptions{ + Locations: []string{currentLocation}, + }, + Quota: &azdext.QuotaCheckOptions{ + MinRemainingCapacity: 1, + }, + }) + if err == nil { + return deploymentResp.Deployment, nil + } + + if !isRecoverableDeploymentSelectionError(err) { + return nil, fmt.Errorf("failed to prompt for model deployment: %w", err) + } + + resolvedModel, resolvedLocation, resolveErr := a.promptForModelLocationMismatch( + ctx, + model, + currentLocation, + fmt.Sprintf( + "Not enough available quota to deploy model '%s' in '%s'.", + model.Name, + currentLocation, + ), + modelRecoveryReasonQuota, + ) + if resolveErr != nil { + return nil, resolveErr + } + if resolvedModel == nil { + return nil, fmt.Errorf("model unavailable due to quota constraints and no alternative selected, exiting") + } + + model = resolvedModel + currentLocation = resolvedLocation + } +} + +func (a *InitAction) resolveModelDeploymentNoPrompt( + ctx context.Context, + model *azdext.AiModel, + location string, +) (*azdext.AiModelDeployment, error) { + resolveResp, err := a.azdClient.Ai().ResolveModelDeployments(ctx, &azdext.ResolveModelDeploymentsRequest{ + AzureContext: a.azureContext, + ModelName: model.Name, + Options: &azdext.AiModelDeploymentOptions{ + Locations: []string{location}, + }, + Quota: &azdext.QuotaCheckOptions{ + MinRemainingCapacity: 1, + }, + }) + if err != nil { + return nil, fmt.Errorf("failed to resolve model deployment: %w", err) + } + + if len(resolveResp.Deployments) == 0 { + return nil, fmt.Errorf("no deployment candidates found for model '%s' in location '%s'", model.Name, location) + } + + orderedCandidates := slices.Clone(resolveResp.Deployments) + defaultVersions := make(map[string]struct{}, len(model.Versions)) + for _, version := range model.Versions { + if version.IsDefault { + defaultVersions[version.Version] = struct{}{} + } + } + + slices.SortFunc(orderedCandidates, func(a, b *azdext.AiModelDeployment) int { + _, aDefault := defaultVersions[a.Version] + _, bDefault := defaultVersions[b.Version] + if aDefault != bDefault { + if aDefault { + return -1 + } + return 1 + } + + aSkuPriority := skuPriority(a.Sku.Name) + bSkuPriority := skuPriority(b.Sku.Name) + if aSkuPriority != bSkuPriority { + if aSkuPriority < bSkuPriority { + return -1 + } + return 1 + } + + if cmp := strings.Compare(a.Version, b.Version); cmp != 0 { + return cmp + } + + if cmp := strings.Compare(a.Sku.Name, b.Sku.Name); cmp != 0 { + return cmp + } + + return strings.Compare(a.Sku.UsageName, b.Sku.UsageName) + }) + + for _, candidate := range orderedCandidates { + capacity, ok := resolveNoPromptCapacity(candidate) + if !ok { + continue + } + + return cloneDeploymentWithCapacity(candidate, capacity), nil + } + + return nil, fmt.Errorf("no deployment candidates found for model '%s' with a valid non-interactive capacity", model.Name) +} + +func resolveNoPromptCapacity(candidate *azdext.AiModelDeployment) (int32, bool) { + capacity := candidate.Capacity + if capacity <= 0 { + capacity = max(candidate.Sku.MinCapacity, int32(1)) + } + + if candidate.Sku.CapacityStep > 0 && capacity%candidate.Sku.CapacityStep != 0 { + step := candidate.Sku.CapacityStep + capacity = ((capacity + step - 1) / step) * step + } + + if candidate.Sku.MinCapacity > 0 && capacity < candidate.Sku.MinCapacity { + capacity = candidate.Sku.MinCapacity + } + if candidate.Sku.MaxCapacity > 0 && capacity > candidate.Sku.MaxCapacity { + return 0, false + } + + if candidate.RemainingQuota != nil && float64(capacity) > *candidate.RemainingQuota { + return 0, false + } + + return capacity, true +} + +func cloneDeploymentWithCapacity(candidate *azdext.AiModelDeployment, capacity int32) *azdext.AiModelDeployment { + if candidate == nil { + return nil + } + + cloned := proto.Clone(candidate).(*azdext.AiModelDeployment) + cloned.Capacity = capacity + return cloned +} + +func skuPriority(skuName string) int { + for i, preferred := range defaultSkuPriority { + if preferred == skuName { + return i + } + } + + return len(defaultSkuPriority) +} + +func isRecoverableDeploymentSelectionError(err error) bool { + if err == nil { + return false + } + + return hasAiErrorReason(err, + azdext.AiErrorReasonNoValidSkus, + azdext.AiErrorReasonNoDeploymentMatch, + azdext.AiErrorReasonModelNotFound, + azdext.AiErrorReasonNoModelsMatch, + ) +} + +func hasAiErrorReason(err error, reasons ...string) bool { + if err == nil { + return false + } + + st, ok := status.FromError(err) + if !ok { + return false + } + + for _, detail := range st.Details() { + info, ok := detail.(*errdetails.ErrorInfo) + if !ok || info.Domain != azdext.AiErrorDomain { + continue + } + + if slices.Contains(reasons, info.Reason) { + return true + } + } + + return false +} + +type modelRecoveryReason string + +const ( + modelRecoveryReasonAvailability modelRecoveryReason = "availability" + modelRecoveryReasonQuota modelRecoveryReason = "quota" + modelLocationSwitchWarning = "WARNING: If you switch locations:\n" + + "• Your AZD environment will use a new default region.\n" + + "• Any existing Azure AI Foundry project created in your current region may fail.\n" + + "• Quota availability varies by region and model.\n\n" + + "Recommended options:\n" + + "1) Select a different model in this region (safe), or\n" + + "2) Create a new Foundry project after changing regions." +) + +func (a *InitAction) promptForAlternativeModel( + ctx context.Context, + originalModelName string, +) (*azdext.AiModel, error) { + fmt.Println(output.WithErrorFormat("The model '%s' could not be found in the model catalog for your subscription in any region.\n", originalModelName)) + + choices := []*azdext.SelectChoice{ + {Label: "Select a different model", Value: "select"}, + {Label: "Exit", Value: "exit"}, + } + + defaultIndex := int32(1) + selectResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ + Options: &azdext.SelectOptions{ + Message: "Would you like to select a different model or exit?", + Choices: choices, + SelectedIndex: &defaultIndex, + EnableFiltering: to.Ptr(false), + }, + }) + if err != nil { + return nil, fmt.Errorf("failed to prompt for model selection choice: %w", err) + } + + if choices[*selectResp.Value].Value == "exit" { + return nil, nil + } + + regionChoices := []*azdext.SelectChoice{ + {Label: fmt.Sprintf("Models available in my current region (%s)", a.azureContext.Scope.Location), Value: "region"}, + {Label: "All available models", Value: "all"}, + } + + regionResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ + Options: &azdext.SelectOptions{ + Message: "Which models would you like to explore?", + Choices: regionChoices, + EnableFiltering: to.Ptr(false), + }, + }) + if err != nil { + return nil, fmt.Errorf("failed to prompt for region choice: %w", err) + } + + promptReq := &azdext.PromptAiModelRequest{ + AzureContext: a.azureContext, + SelectOptions: &azdext.SelectOptions{ + Message: "Select a model", + }, + } + + if regionChoices[*regionResp.Value].Value == "region" { + promptReq.Filter = &azdext.AiModelFilterOptions{ + Locations: []string{a.azureContext.Scope.Location}, + } + } + + modelResp, err := a.azdClient.Prompt().PromptAiModel(ctx, promptReq) + if err != nil { + return nil, fmt.Errorf("failed to prompt for model selection: %w", err) + } + + return modelResp.Model, nil +} + +func (a *InitAction) promptForModelLocationMismatch( + ctx context.Context, + model *azdext.AiModel, + currentLocation string, + reasonMessage string, + reasonKind modelRecoveryReason, +) (*azdext.AiModel, string, error) { + currentModel := model + message := reasonMessage + + for { + if message == "" { + message = fmt.Sprintf( + "The model '%s' is not available in your current location '%s'.", + currentModel.Name, + currentLocation, + ) + } + + fmt.Println(output.WithErrorFormat(message)) + + modelChoiceLabel := fmt.Sprintf("Choose a different model in %s", currentLocation) + + choices := []*azdext.SelectChoice{ + {Label: modelChoiceLabel, Value: "model"}, + {Label: "Choose a different model (all regions)", Value: "model_all_regions"}, + {Label: fmt.Sprintf("Choose a different location for %s", currentModel.Name), Value: "location"}, + {Label: "Exit setup", Value: "exit"}, + } + + if !a.locationWarningShown { + fmt.Println() + fmt.Println(output.WithWarningFormat(modelLocationSwitchWarning)) + a.locationWarningShown = true + fmt.Println() + } + + defaultIndex := int32(3) + selectResp, err := a.azdClient.Prompt().Select(ctx, &azdext.SelectRequest{ + Options: &azdext.SelectOptions{ + Message: "What would you like to do?", + Choices: choices, + SelectedIndex: &defaultIndex, + EnableFiltering: to.Ptr(false), + }, + }) + if err != nil { + return nil, "", fmt.Errorf("failed to prompt for action choice: %w", err) + } + + selectedChoice := choices[*selectResp.Value].Value + + if selectedChoice == "exit" { + return nil, "", nil + } + + if selectedChoice == "location" { + locationResp, err := a.azdClient.Prompt().PromptAiModelLocationWithQuota(ctx, + &azdext.PromptAiModelLocationWithQuotaRequest{ + AzureContext: a.azureContext, + ModelName: currentModel.Name, + AllowedLocations: currentModel.Locations, + Quota: &azdext.QuotaCheckOptions{ + MinRemainingCapacity: 1, + }, + SelectOptions: &azdext.SelectOptions{ + Message: fmt.Sprintf("Select a location for model '%s'", currentModel.Name), + }, + }, + ) + if err != nil { + if hasAiErrorReason(err, azdext.AiErrorReasonNoLocationsWithQuota) { + message = fmt.Sprintf("No locations have sufficient quota for model '%s'.", currentModel.Name) + continue + } + + return nil, "", fmt.Errorf("failed to prompt for location selection: %w", err) + } + + selectedLocation := locationResp.Location.Name + if err := a.updateEnvLocation(ctx, selectedLocation); err != nil { + return nil, "", err + } + + return currentModel, selectedLocation, nil + } + + if selectedChoice == "model_all_regions" { + modelResp, err := a.azdClient.Prompt().PromptAiModel(ctx, &azdext.PromptAiModelRequest{ + AzureContext: a.azureContext, + Filter: &azdext.AiModelFilterOptions{ + ExcludeModelNames: []string{currentModel.Name}, + }, + Quota: &azdext.QuotaCheckOptions{ + MinRemainingCapacity: 1, + }, + SelectOptions: &azdext.SelectOptions{ + Message: "Select a model from all regions", + }, + }) + if err != nil { + if hasAiErrorReason(err, azdext.AiErrorReasonNoModelsMatch) { + message = "No alternative models were found across all regions." + continue + } + + return nil, "", fmt.Errorf("failed to prompt for model selection across all regions: %w", err) + } + + selectedModel := modelResp.Model + locationResp, err := a.azdClient.Prompt().PromptAiModelLocationWithQuota(ctx, + &azdext.PromptAiModelLocationWithQuotaRequest{ + AzureContext: a.azureContext, + ModelName: selectedModel.Name, + AllowedLocations: selectedModel.Locations, + Quota: &azdext.QuotaCheckOptions{ + MinRemainingCapacity: 1, + }, + SelectOptions: &azdext.SelectOptions{ + Message: fmt.Sprintf("Select a location for model '%s'", selectedModel.Name), + }, + }, + ) + if err != nil { + if hasAiErrorReason(err, azdext.AiErrorReasonNoLocationsWithQuota) { + currentModel = selectedModel + message = fmt.Sprintf("No locations have sufficient quota for model '%s'.", selectedModel.Name) + continue + } + + return nil, "", fmt.Errorf("failed to prompt for location selection: %w", err) + } + + selectedLocation := locationResp.Location.Name + if err := a.updateEnvLocation(ctx, selectedLocation); err != nil { + return nil, "", err + } + + return selectedModel, selectedLocation, nil + } + + promptReq := &azdext.PromptAiModelRequest{ + AzureContext: a.azureContext, + Filter: &azdext.AiModelFilterOptions{ + Locations: []string{currentLocation}, + }, + SelectOptions: &azdext.SelectOptions{ + Message: fmt.Sprintf("Select a model available in '%s'", currentLocation), + }, + Quota: &azdext.QuotaCheckOptions{ + MinRemainingCapacity: 1, + }, + } + + modelResp, err := a.azdClient.Prompt().PromptAiModel(ctx, promptReq) + if err != nil { + if hasAiErrorReason(err, azdext.AiErrorReasonNoModelsMatch) { + message = fmt.Sprintf("No models are available in your current location '%s'.", currentLocation) + continue + } + + return nil, "", fmt.Errorf("failed to prompt for model selection: %w", err) + } + + return modelResp.Model, currentLocation, nil + } +} + +func (a *InitAction) ProcessModels(ctx context.Context, manifest *agent_yaml.AgentManifest) (*agent_yaml.AgentManifest, []project.Deployment, error) { + templateBytes, err := yaml.Marshal(manifest.Template) + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal agent template to YAML: %w", err) + } + + var templateDict map[string]interface{} + if err := yaml.Unmarshal(templateBytes, &templateDict); err != nil { + return nil, nil, fmt.Errorf("failed to unmarshal agent template from YAML: %w", err) + } + + dictJsonBytes, err := yaml.Marshal(templateDict) + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal templateDict to YAML: %w", err) + } + + var agentDef agent_yaml.AgentDefinition + if err := yaml.Unmarshal(dictJsonBytes, &agentDef); err != nil { + return nil, nil, fmt.Errorf("failed to unmarshal YAML to AgentDefinition: %w", err) + } + + deploymentDetails := []project.Deployment{} + paramValues := registry_api.ParameterValues{} + switch agentDef.Kind { + case agent_yaml.AgentKindPrompt: + agentDef := manifest.Template.(agent_yaml.PromptAgent) + + modelDeployment, err := a.getModelDeploymentDetails(ctx, agentDef.Model) + if err != nil { + return nil, nil, fmt.Errorf("failed to get model deployment details: %w", err) + } + deploymentDetails = append(deploymentDetails, *modelDeployment) + paramValues["deploymentName"] = modelDeployment.Name + case agent_yaml.AgentKindHosted: + for _, resource := range manifest.Resources { + resourceBytes, err := yaml.Marshal(resource) + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal resource to YAML: %w", err) + } + + var resourceDef agent_yaml.Resource + if err := yaml.Unmarshal(resourceBytes, &resourceDef); err != nil { + return nil, nil, fmt.Errorf("failed to unmarshal YAML to Resource: %w", err) + } + + if resourceDef.Kind == agent_yaml.ResourceKindModel { + resource := resource.(agent_yaml.ModelResource) + model := agent_yaml.Model{Id: resource.Id} + modelDeployment, err := a.getModelDeploymentDetails(ctx, model) + if err != nil { + return nil, nil, fmt.Errorf("failed to get model deployment details: %w", err) + } + deploymentDetails = append(deploymentDetails, *modelDeployment) + paramValues[resource.Name] = modelDeployment.Name + } + } + } + + updatedManifest, err := registry_api.InjectParameterValuesIntoManifest(manifest, paramValues) + if err != nil { + return nil, nil, fmt.Errorf("failed to inject deployment names into manifest: %w", err) + } + + fmt.Println("Model deployment details processed and injected into agent definition. Deployment details can also be found in the JSON formatted AI_PROJECT_DEPLOYMENTS environment variable.") + + return updatedManifest, deploymentDetails, nil +} diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_test.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_test.go index 3305b663dda..7ec7503be9f 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_test.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_test.go @@ -8,8 +8,77 @@ import ( "os" "path/filepath" "testing" + + "github.com/azure/azure-dev/cli/azd/pkg/azdext" + "google.golang.org/genproto/googleapis/rpc/errdetails" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) +func TestIsRecoverableDeploymentSelectionError_StructuredReason(t *testing.T) { + t.Parallel() + + st := status.New(codes.FailedPrecondition, "no valid SKUs for selected model") + withDetails, err := st.WithDetails(&errdetails.ErrorInfo{ + Reason: azdext.AiErrorReasonNoValidSkus, + Domain: azdext.AiErrorDomain, + }) + if err != nil { + t.Fatalf("failed to attach grpc error details: %v", err) + } + + if !isRecoverableDeploymentSelectionError(withDetails.Err()) { + t.Fatalf("expected structured AI reason to be recoverable") + } +} + +func TestIsRecoverableDeploymentSelectionError_NonRecoverableStructuredReason(t *testing.T) { + t.Parallel() + + st := status.New(codes.InvalidArgument, "quota location is required") + withDetails, err := st.WithDetails(&errdetails.ErrorInfo{ + Reason: azdext.AiErrorReasonQuotaLocation, + Domain: azdext.AiErrorDomain, + }) + if err != nil { + t.Fatalf("failed to attach grpc error details: %v", err) + } + + if isRecoverableDeploymentSelectionError(withDetails.Err()) { + t.Fatalf("expected structured quota-location error to be non-recoverable") + } +} + +func TestIsRecoverableDeploymentSelectionError_UnstructuredError(t *testing.T) { + t.Parallel() + + if isRecoverableDeploymentSelectionError( + status.Error(codes.Internal, "no deployment found for model \"foo\" with the specified options"), + ) { + t.Fatalf("expected unstructured error to be non-recoverable") + } +} + +func TestHasAiErrorReason(t *testing.T) { + t.Parallel() + + st := status.New(codes.NotFound, "no locations with sufficient quota") + withDetails, err := st.WithDetails(&errdetails.ErrorInfo{ + Reason: azdext.AiErrorReasonNoLocationsWithQuota, + Domain: azdext.AiErrorDomain, + }) + if err != nil { + t.Fatalf("failed to attach grpc error details: %v", err) + } + + if !hasAiErrorReason(withDetails.Err(), azdext.AiErrorReasonNoLocationsWithQuota) { + t.Fatalf("expected reason to be detected") + } + if hasAiErrorReason(withDetails.Err(), azdext.AiErrorReasonNoValidSkus) { + t.Fatalf("expected non-matching reason to be false") + } +} + func TestCopyDirectory_RefusesToCopyIntoSubtree(t *testing.T) { t.Parallel() diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/root.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/root.go index b5616220d4c..b035d4f18b1 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/root.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/root.go @@ -48,7 +48,7 @@ func NewRootCommand() *cobra.Command { rootCmd.AddCommand(newListenCommand()) rootCmd.AddCommand(newVersionCommand()) - rootCmd.AddCommand(newInitCommand(rootFlags)) + rootCmd.AddCommand(newInitCommand(&rootFlags)) rootCmd.AddCommand(newMcpCommand()) rootCmd.AddCommand(newMetadataCommand()) diff --git a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go b/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go deleted file mode 100644 index 9c593e4f1d6..00000000000 --- a/cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go +++ /dev/null @@ -1,429 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package ai - -import ( - "context" - "errors" - "fmt" - "slices" - "strings" - "sync" - - "azureaiagent/internal/pkg/azure" - - "github.com/Azure/azure-sdk-for-go/sdk/azcore" - "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/cognitiveservices/armcognitiveservices" - "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armsubscriptions" -) - -type AiModel struct { - Name string - ModelDetailsByLocation map[string][]*armcognitiveservices.Model -} - -type ModelCatalogService struct { - credential azcore.TokenCredential - azureClient *azure.AzureClient -} - -func NewModelCatalogService(credential azcore.TokenCredential) *ModelCatalogService { - return &ModelCatalogService{ - credential: credential, - azureClient: azure.NewAzureClient(credential), - } -} - -func (c *ModelCatalogService) ListAllCapabilities(ctx context.Context, allModels map[string]*AiModel) []string { - return filterDistinctModelData(allModels, func(m *armcognitiveservices.Model) []string { - capabilities := []string{} - for key := range m.Model.Capabilities { - capabilities = append(capabilities, key) - } - - return capabilities - }) -} - -func (c *ModelCatalogService) ListAllStatuses(ctx context.Context, allModels map[string]*AiModel) []string { - return filterDistinctModelData(allModels, func(m *armcognitiveservices.Model) []string { - return []string{string(*m.Model.LifecycleStatus)} - }) -} - -func (c *ModelCatalogService) ListAllFormats(ctx context.Context, allModels map[string]*AiModel) []string { - return filterDistinctModelData(allModels, func(m *armcognitiveservices.Model) []string { - return []string{*m.Model.Format} - }) -} - -func (c *ModelCatalogService) ListAllKinds(ctx context.Context, allModels map[string]*AiModel) []string { - return filterDistinctModelData(allModels, func(m *armcognitiveservices.Model) []string { - return []string{*m.Kind} - }) -} - -func (c *ModelCatalogService) ListModelVersions(ctx context.Context, model *AiModel, location string) ([]string, string, error) { - versions := make(map[string]struct{}) - defaultVersion := "" - - models, exists := model.ModelDetailsByLocation[location] - if !exists { - return nil, "", fmt.Errorf("no model details found for location '%s'", location) - } - - for _, m := range models { - versions[*m.Model.Version] = struct{}{} - if m.Model.IsDefaultVersion != nil && *m.Model.IsDefaultVersion { - defaultVersion = *m.Model.Version - } - } - - versionList := make([]string, 0, len(versions)) - for version := range versions { - versionList = append(versionList, version) - } - - slices.Sort(versionList) - - return versionList, defaultVersion, nil -} - -func (c *ModelCatalogService) ListModelSkus(ctx context.Context, model *AiModel, location string, modelVersion string) ([]string, error) { - skus := make(map[string]struct{}) - - models, exists := model.ModelDetailsByLocation[location] - if !exists { - return nil, fmt.Errorf("no model details found for location '%s'", location) - } - - for _, m := range models { - if *m.Model.Version == modelVersion { - for _, sku := range m.Model.SKUs { - skus[*sku.Name] = struct{}{} - } - } - } - - skuList := make([]string, 0, len(skus)) // Create with capacity, not length - for sku := range skus { - skuList = append(skuList, sku) - } - - slices.Sort(skuList) - - return skuList, nil -} - -type FilterOptions struct { - Capabilities []string - Statuses []string - Formats []string - Kinds []string - Locations []string -} - -func (c *ModelCatalogService) ListFilteredModels( - ctx context.Context, - allModels map[string]*AiModel, - options *FilterOptions, -) []*AiModel { - if options == nil { - options = &FilterOptions{} - } - - filteredModels := []*AiModel{} - - for _, model := range allModels { - // Initialize flags to true if the corresponding filter is not provided. - isCapabilityMatch := len(options.Capabilities) == 0 - isLocationMatch := len(options.Locations) == 0 - isStatusMatch := len(options.Statuses) == 0 - isFormatMatch := len(options.Formats) == 0 - isKindMatch := len(options.Kinds) == 0 - - for locationName, models := range model.ModelDetailsByLocation { - for _, m := range models { - if !isCapabilityMatch && len(options.Capabilities) > 0 { - for modelCapability := range m.Model.Capabilities { - if slices.Contains(options.Capabilities, modelCapability) { - isCapabilityMatch = true - break - } - } - } - - if !isLocationMatch && len(options.Locations) > 0 && - slices.Contains(options.Locations, locationName) { - isLocationMatch = true - } - - if !isStatusMatch && len(options.Statuses) > 0 && - slices.Contains(options.Statuses, string(*m.Model.LifecycleStatus)) { - isStatusMatch = true - } - - if !isFormatMatch && len(options.Formats) > 0 && - slices.Contains(options.Formats, *m.Model.Format) { - isFormatMatch = true - } - - if !isKindMatch && len(options.Kinds) > 0 && - slices.Contains(options.Kinds, *m.Kind) { - isKindMatch = true - } - } - } - - if isLocationMatch && isCapabilityMatch && isFormatMatch && isStatusMatch && isKindMatch { - filteredModels = append(filteredModels, model) - } - } - - // Sort the filtered models by name - slices.SortFunc(filteredModels, func(a, b *AiModel) int { - return strings.Compare(a.Name, b.Name) - }) - - return filteredModels -} - -func (c *ModelCatalogService) ListAllModels(ctx context.Context, subscriptionId string, location string) (map[string]*AiModel, error) { - var locations []*armsubscriptions.Location - var err error - - if location == "" { - // If no specific location provided, get all locations - locations, err = c.azureClient.ListLocations(ctx, subscriptionId) - if err != nil { - return nil, err - } - } else { - // If specific location provided, create a single-item slice with that location - locations = []*armsubscriptions.Location{ - { - Name: &location, - }, - } - } - - modelsClient, err := createModelsClient(subscriptionId, c.credential) - if err != nil { - return nil, err - } - - var locationResults sync.Map - var wg sync.WaitGroup - - for _, location := range locations { - wg.Add(1) - go func(location *armsubscriptions.Location) { - defer wg.Done() - pager := modelsClient.NewListPager(*location.Name, nil) - - results := []*armcognitiveservices.Model{} - - for pager.More() { - page, err := pager.NextPage(ctx) - if err != nil { - return - } - - results = append(results, page.Value...) - } - - locationResults.Store(location, results) - }(location) - } - - wg.Wait() - - modelMap := map[string]*AiModel{} - - locationResults.Range(func(key, value any) bool { - location := key.(*armsubscriptions.Location) - modelsList := value.([]*armcognitiveservices.Model) - - for _, model := range modelsList { - - modelName := *model.Model.Name - existingModel, exists := modelMap[modelName] - if !exists { - existingModel = &AiModel{ - Name: modelName, - ModelDetailsByLocation: make(map[string][]*armcognitiveservices.Model), - } - } - - locationName := *location.Name - existingModel.ModelDetailsByLocation[locationName] = append(existingModel.ModelDetailsByLocation[locationName], model) - - modelMap[modelName] = existingModel - } - - return true - }) - - return modelMap, nil - // allModels := []*AiModel{} - // for _, model := range modelMap { - // allModels = append(allModels, model) - // } - - // slices.SortFunc(allModels, func(a, b *AiModel) int { - // return strings.Compare(a.Name, b.Name) - // }) - - // return allModels, nil -} - -type AiModelDeployment struct { - Name string - Format string - Version string - Sku AiModelDeploymentSku -} - -type AiModelDeploymentSku struct { - Name string - UsageName string - Capacity int32 -} - -type AiModelDeploymentOptions struct { - Locations []string - Versions []string - Skus []string -} - -func (c *ModelCatalogService) GetModelDeployment( - ctx context.Context, - model *AiModel, - options *AiModelDeploymentOptions, -) (*AiModelDeployment, error) { - if options == nil { - options = &AiModelDeploymentOptions{ - Skus: []string{ - "GlobalStandard", - "Standard", - }, - } - } - - var modelDeployment *AiModelDeployment - hasDefaultVersion := c.hasDefaultVersion(model) - - for locationName, models := range model.ModelDetailsByLocation { - if modelDeployment != nil { - break - } - - // Check for location match if specified - if len(options.Locations) > 0 && !slices.Contains(options.Locations, locationName) { - continue - } - - for _, m := range models { - if modelDeployment != nil { - break - } - - // Check for default version if no version is specified - if len(options.Versions) > 0 { - if !slices.Contains(options.Versions, *m.Model.Version) { - continue - } - } else if hasDefaultVersion { - // Not all models have a default version - if m.Model.IsDefaultVersion != nil && !*m.Model.IsDefaultVersion { - continue - } - } - - // Check for SKU match if specified - for _, sku := range m.Model.SKUs { - - if !slices.Contains(options.Skus, *sku.Name) { - continue - } - - modelDeployment = &AiModelDeployment{ - Name: *m.Model.Name, - Format: *m.Model.Format, - Version: *m.Model.Version, - Sku: AiModelDeploymentSku{ - Name: *sku.Name, - UsageName: *sku.UsageName, - }, - } - - if sku.Capacity.Default != nil { - modelDeployment.Sku.Capacity = *sku.Capacity.Default - } else { - modelDeployment.Sku.Capacity = -1 - } - - break - } - } - } - - if modelDeployment == nil { - return nil, errors.New("No model deployment found for the specified options") - } - - return modelDeployment, nil -} - -func (c *ModelCatalogService) hasDefaultVersion(model *AiModel) bool { - for _, models := range model.ModelDetailsByLocation { - for _, m := range models { - if m.Model.IsDefaultVersion != nil && *m.Model.IsDefaultVersion { - return true - } - } - } - return false -} - -func filterDistinctModelData( - models map[string]*AiModel, - filterFunc func(*armcognitiveservices.Model) []string, -) []string { - filtered := make(map[string]struct{}) - for _, model := range models { - for _, locationModels := range model.ModelDetailsByLocation { - for _, m := range locationModels { - value := filterFunc(m) - for _, v := range value { - if v != "" { - filtered[v] = struct{}{} - } - } - } - } - } - - list := make([]string, len(filtered)) - i := 0 - for value := range filtered { - list[i] = value - i++ - } - - slices.Sort(list) - return list -} - -func createModelsClient( - subscriptionId string, - credential azcore.TokenCredential, -) (*armcognitiveservices.ModelsClient, error) { - client, err := armcognitiveservices.NewModelsClient(subscriptionId, credential, azure.NewArmClientOptions()) - if err != nil { - return nil, err - } - - return client, nil -} From aa73bb74568ac8c335af1e673c564d364123de65 Mon Sep 17 00:00:00 2001 From: Jeffrey Chen Date: Wed, 11 Feb 2026 17:58:13 +0000 Subject: [PATCH 2/4] Address feedback --- .../internal/cmd/init_models.go | 30 +++++++++++-------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go index 140d960c91a..d217d47c556 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go @@ -167,14 +167,15 @@ func (a *InitAction) getModelDeploymentDetails(ctx context.Context, model agent_ foundryProjectId := resp.Value if foundryProjectId != "" { parts := strings.Split(foundryProjectId, "/") - var subscription, resourceGroup, accountName string - - if len(parts) >= 9 { - subscription = parts[2] - resourceGroup = parts[4] - accountName = parts[8] + if len(parts) < 9 { + return nil, fmt.Errorf( + "invalid AZURE_AI_PROJECT_ID format: expected at least 9 path segments, got %d", len(parts)) } + subscription := parts[2] + resourceGroup := parts[4] + accountName := parts[8] + deploymentsClient, err := armcognitiveservices.NewDeploymentsClient(subscription, a.credential, azure.NewArmClientOptions()) if err != nil { return nil, fmt.Errorf("failed to create deployments client: %w", err) @@ -192,13 +193,16 @@ func (a *InitAction) getModelDeploymentDetails(ctx context.Context, model agent_ matchingDeployments := make(map[string]*armcognitiveservices.Deployment) for _, deployment := range deployments { - if deployment.Properties != nil && deployment.Properties.Model != nil { - deployedModel := deployment.Properties.Model - if deployedModel.Name != nil { - if *deployedModel.Name == model.Id { - matchingDeployments[*deployment.Name] = deployment - } - } + if deployment.Name == nil || + deployment.Properties == nil || deployment.Properties.Model == nil || + deployment.Properties.Model.Name == nil || deployment.Properties.Model.Format == nil || + deployment.Properties.Model.Version == nil || + deployment.SKU == nil || deployment.SKU.Name == nil || deployment.SKU.Capacity == nil { + continue + } + + if *deployment.Properties.Model.Name == model.Id { + matchingDeployments[*deployment.Name] = deployment } } From a64ba9308d48d82e9754e90ea50a7799d2239ef7 Mon Sep 17 00:00:00 2001 From: Jeffrey Chen Date: Wed, 11 Feb 2026 18:31:23 +0000 Subject: [PATCH 3/4] Add simple compatibility check --- .../azure.ai.agents/internal/cmd/init.go | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go index 2384864f19d..edb9f65deaa 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init.go @@ -35,6 +35,8 @@ import ( "github.com/azure/azure-dev/cli/azd/pkg/tools/github" "github.com/fatih/color" "github.com/spf13/cobra" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/structpb" "gopkg.in/yaml.v3" ) @@ -79,6 +81,23 @@ type GitHubUrlInfo struct { const AiAgentHost = "azure.ai.agent" const ContainerAppHost = "containerapp" +// checkAiModelServiceAvailable is a temporary check to ensure the azd host supports +// required gRPC services. Remove once azd core enforces requiredAzdVersion. +func checkAiModelServiceAvailable(ctx context.Context, azdClient *azdext.AzdClient) error { + _, err := azdClient.Ai().ListModels(ctx, &azdext.ListModelsRequest{}) + if err == nil { + return nil + } + + if st, ok := status.FromError(err); ok && st.Code() == codes.Unimplemented { + return fmt.Errorf( + "this version of the azure.ai.agents extension is incompatible with your installed version of azd. " + + "Please upgrade azd to the latest version") + } + + return nil +} + func newInitCommand(rootFlags *rootFlagsDefinition) *cobra.Command { flags := &initFlags{ rootFlagsDefinition: rootFlags, @@ -99,6 +118,10 @@ func newInitCommand(rootFlags *rootFlagsDefinition) *cobra.Command { } defer azdClient.Close() + if err := checkAiModelServiceAvailable(ctx, azdClient); err != nil { + return err + } + // Wait for debugger if AZD_EXT_DEBUG is set if err := azdext.WaitForDebugger(ctx, azdClient); err != nil { if errors.Is(err, context.Canceled) || errors.Is(err, azdext.ErrDebuggerAborted) { From c60b45b9b90ca9fbea6166b3cbdf4fa60095f892 Mon Sep 17 00:00:00 2001 From: Jeffrey Chen Date: Thu, 12 Feb 2026 21:27:57 +0000 Subject: [PATCH 4/4] Address comments --- .../internal/cmd/init_models.go | 18 +- .../internal/cmd/init_models_test.go | 174 ++++++++++++++++++ 2 files changed, 183 insertions(+), 9 deletions(-) create mode 100644 cli/azd/extensions/azure.ai.agents/internal/cmd/init_models_test.go diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go index d217d47c556..e6382c3e5c5 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models.go @@ -75,23 +75,23 @@ func (a *InitAction) updateEnvLocation(ctx context.Context, selectedLocation str return fmt.Errorf("failed to get current azd environment: %w", err) } - if a.azureContext == nil { - a.azureContext = &azdext.AzureContext{} - } - if a.azureContext.Scope == nil { - a.azureContext.Scope = &azdext.AzureScope{} - } - a.azureContext.Scope.Location = selectedLocation - _, err = a.azdClient.Environment().SetValue(ctx, &azdext.SetEnvRequest{ EnvName: envResponse.Environment.Name, Key: "AZURE_LOCATION", - Value: a.azureContext.Scope.Location, + Value: selectedLocation, }) if err != nil { return fmt.Errorf("failed to update AZURE_LOCATION in azd environment: %w", err) } + if a.azureContext == nil { + a.azureContext = &azdext.AzureContext{} + } + if a.azureContext.Scope == nil { + a.azureContext.Scope = &azdext.AzureScope{} + } + a.azureContext.Scope.Location = selectedLocation + fmt.Println(output.WithSuccessFormat("Updated AZURE_LOCATION to '%s' in your azd environment.", selectedLocation)) return nil } diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models_test.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models_test.go new file mode 100644 index 00000000000..1c3b3993913 --- /dev/null +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/init_models_test.go @@ -0,0 +1,174 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package cmd + +import ( + "testing" + + "github.com/azure/azure-dev/cli/azd/pkg/azdext" + "github.com/stretchr/testify/assert" +) + +func TestResolveNoPromptCapacity(t *testing.T) { + floatPtr := func(v float64) *float64 { return &v } + + tests := []struct { + name string + candidate *azdext.AiModelDeployment + wantCapacity int32 + wantOk bool + }{ + { + name: "uses existing positive capacity", + candidate: &azdext.AiModelDeployment{ + Capacity: 10, + Sku: &azdext.AiModelSku{}, + }, + wantCapacity: 10, + wantOk: true, + }, + { + name: "zero capacity defaults to max(minCapacity, 1)", + candidate: &azdext.AiModelDeployment{ + Capacity: 0, + Sku: &azdext.AiModelSku{MinCapacity: 5}, + }, + wantCapacity: 5, + wantOk: true, + }, + { + name: "zero capacity with zero minCapacity defaults to 1", + candidate: &azdext.AiModelDeployment{ + Capacity: 0, + Sku: &azdext.AiModelSku{MinCapacity: 0}, + }, + wantCapacity: 1, + wantOk: true, + }, + { + name: "negative capacity defaults to max(minCapacity, 1)", + candidate: &azdext.AiModelDeployment{ + Capacity: -3, + Sku: &azdext.AiModelSku{MinCapacity: 2}, + }, + wantCapacity: 2, + wantOk: true, + }, + { + name: "rounds up to capacity step", + candidate: &azdext.AiModelDeployment{ + Capacity: 7, + Sku: &azdext.AiModelSku{CapacityStep: 5}, + }, + wantCapacity: 10, + wantOk: true, + }, + { + name: "already aligned to step", + candidate: &azdext.AiModelDeployment{ + Capacity: 10, + Sku: &azdext.AiModelSku{CapacityStep: 5}, + }, + wantCapacity: 10, + wantOk: true, + }, + { + name: "enforces minCapacity then step alignment", + candidate: &azdext.AiModelDeployment{ + Capacity: 0, + Sku: &azdext.AiModelSku{MinCapacity: 10, CapacityStep: 3}, + }, + wantCapacity: 12, // min=10, rounded up to next step of 3 + wantOk: true, + }, + { + name: "exceeds maxCapacity returns false", + candidate: &azdext.AiModelDeployment{ + Capacity: 20, + Sku: &azdext.AiModelSku{MaxCapacity: 10}, + }, + wantCapacity: 0, + wantOk: false, + }, + { + name: "exceeds remaining quota returns false", + candidate: &azdext.AiModelDeployment{ + Capacity: 10, + Sku: &azdext.AiModelSku{}, + RemainingQuota: floatPtr(5), + }, + wantCapacity: 0, + wantOk: false, + }, + { + name: "within remaining quota returns true", + candidate: &azdext.AiModelDeployment{ + Capacity: 5, + Sku: &azdext.AiModelSku{}, + RemainingQuota: floatPtr(10), + }, + wantCapacity: 5, + wantOk: true, + }, + { + name: "nil remaining quota is not checked", + candidate: &azdext.AiModelDeployment{ + Capacity: 100, + Sku: &azdext.AiModelSku{MaxCapacity: 200}, + RemainingQuota: nil, + }, + wantCapacity: 100, + wantOk: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + capacity, ok := resolveNoPromptCapacity(tt.candidate) + assert.Equal(t, tt.wantOk, ok) + assert.Equal(t, tt.wantCapacity, capacity) + }) + } +} + +func TestSkuPriority(t *testing.T) { + tests := []struct { + name string + skuName string + wantPrio int + }{ + { + name: "GlobalStandard is highest priority", + skuName: "GlobalStandard", + wantPrio: 0, + }, + { + name: "DataZoneStandard is second priority", + skuName: "DataZoneStandard", + wantPrio: 1, + }, + { + name: "Standard is third priority", + skuName: "Standard", + wantPrio: 2, + }, + { + name: "unknown SKU returns fallback priority", + skuName: "UnknownSku", + wantPrio: len(defaultSkuPriority), + }, + { + name: "empty string returns fallback priority", + skuName: "", + wantPrio: len(defaultSkuPriority), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := skuPriority(tt.skuName) + assert.Equal(t, tt.wantPrio, got) + }) + } +}