[AI Agents Extension] Add check for missing sku capacity (#6053)

* Add check for missing sku capacity

Signed-off-by: trangevi <[email protected]>

* Update cli/azd/extensions/azure.ai.agents/internal/cmd/init.go

Co-authored-by: JeffreyCA <[email protected]>

---------

Signed-off-by: trangevi <[email protected]>
Co-authored-by: JeffreyCA <[email protected]>

Author: trangevi

cli/azd/extensions/azure.ai.agents/internal/cmd/init.go

@@ -12,6 +12,7 @@ import (
 	"path/filepath"
 	"regexp"
 	"slices"
+	"strconv"
 	"strings"
 
 	"azureaiagent/internal/pkg/agents/agent_yaml"
@@ -1415,5 +1416,24 @@ func (a *InitAction) getModelDetails(ctx context.Context, modelName string, mode
 		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,
+			},
+		})
+		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
 }

cli/azd/extensions/azure.ai.agents/internal/pkg/azure/ai/model_catalog.go

@@ -347,10 +347,15 @@ func (c *ModelCatalogService) GetModelDeployment(
 				Sku: AiModelDeploymentSku{
 					Name:      *sku.Name,
 					UsageName: *sku.UsageName,
-					Capacity:  *sku.Capacity.Default,
 				},
 			}
 
+			if sku.Capacity.Default != nil {
+				modelDeployment.Sku.Capacity = *sku.Capacity.Default
+			} else {
+				modelDeployment.Sku.Capacity = -1
+			}
+
 			break
 		}
 	}