Add model details env var

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

cli/azd/extensions/azure.foundry.ai.agents/go.mod

@@ -11,6 +11,7 @@ require (
 	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 v0.0.0-20251024053325-326f63f72d65
+	github.com/braydonk/yaml v0.9.0
 	github.com/fatih/color v1.18.0
 	github.com/google/uuid v1.6.0
 	github.com/mark3labs/mcp-go v0.41.1
@@ -20,6 +21,7 @@ require (
 )
 
 require (
+	dario.cat/mergo v1.0.2 // indirect
 	github.com/AlecAivazis/survey/v2 v2.3.7 // indirect
 	github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2 // indirect
 	github.com/AzureAD/microsoft-authentication-library-for-go v1.5.0 // indirect
@@ -29,7 +31,6 @@ require (
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/blang/semver/v4 v4.0.0 // indirect
-	github.com/braydonk/yaml v0.9.0 // indirect
 	github.com/buger/goterm v1.0.4 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
 	github.com/charmbracelet/colorprofile v0.3.2 // indirect
@@ -43,6 +44,7 @@ require (
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/dlclark/regexp2 v1.11.5 // indirect
 	github.com/drone/envsubst v1.0.3 // 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

cli/azd/extensions/azure.foundry.ai.agents/go.sum

@@ -1,4 +1,6 @@
 code.cloudfoundry.org/clock v0.0.0-20180518195852-02e53af36e6c/go.mod h1:QD9Lzhd/ux6eNQVUDVRJX/RKTigpewimNYBi7ivZKY8=
+dario.cat/mergo v1.0.2 h1:85+piFYR1tMbRrLcDwR18y4UKJ3aH1Tbzi24VRW1TK8=
+dario.cat/mergo v1.0.2/go.mod h1:E/hbnu0NxMFBjpMIE34DRGLWqDy0g5FuKDhCb31ngxA=
 github.com/AlecAivazis/survey/v2 v2.3.7 h1:6I/u8FvytdGsgonrYsVn2t8t4QiRnh6QSTqkkhIiSjQ=
 github.com/AlecAivazis/survey/v2 v2.3.7/go.mod h1:xUTIdE4KCOIjsBAE1JYsUPoCqYdZ1reCfTwbto0Fduo=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v1.19.1 h1:5YTBM8QDVIBN3sxBil89WfdAAqDZbyJTgh688DSxX5w=
@@ -90,6 +92,8 @@ 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=

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

@@ -15,13 +15,16 @@ import (
 
 	"azureaiagent/internal/pkg/agents/agent_yaml"
 	"azureaiagent/internal/pkg/agents/registry_api"
+	"azureaiagent/internal/pkg/azure/ai"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/azure/azure-dev/cli/azd/pkg/azdext"
 	"github.com/azure/azure-dev/cli/azd/pkg/exec"
 	"github.com/azure/azure-dev/cli/azd/pkg/input"
 	"github.com/azure/azure-dev/cli/azd/pkg/osutil"
 	"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"
 	"gopkg.in/yaml.v3"
@@ -41,14 +44,14 @@ type AiProjectResourceConfig struct {
 type InitAction struct {
 	azdClient *azdext.AzdClient
 	//azureClient       *azure.AzureClient
-	// azureContext *azdext.AzureContext
+	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
+	credential          azcore.TokenCredential
+	modelCatalog        map[string]*ai.AiModel
+	modelCatalogService *ai.ModelCatalogService
+	projectConfig       *azdext.ProjectConfig
+	environment         *azdext.Environment
 }
 
 // GitHubUrlInfo holds parsed information from a GitHub URL
@@ -75,8 +78,7 @@ func newInitCommand() *cobra.Command {
 			}
 			defer azdClient.Close()
 
-			//azureContext, projectConfig, err := ensureAzureContext(ctx, azdClient)
-			projectConfig, err := ensureProject(ctx, azdClient)
+			azureContext, projectConfig, err := ensureAzureContext(ctx, azdClient)
 			if err != nil {
 				return fmt.Errorf("failed to ground into a project context: %w", err)
 			}
@@ -91,13 +93,13 @@ func newInitCommand() *cobra.Command {
 			// 	return fmt.Errorf("failed to get composed resources: %w", err)
 			// }
 
-			// credential, err := azidentity.NewAzureDeveloperCLICredential(&azidentity.AzureDeveloperCLICredentialOptions{
-			// 	TenantID:                   azureContext.Scope.TenantId,
-			// 	AdditionallyAllowedTenants: []string{"*"},
-			// })
-			// if err != nil {
-			// 	return fmt.Errorf("failed to create azure credential: %w", err)
-			// }
+			credential, err := azidentity.NewAzureDeveloperCLICredential(&azidentity.AzureDeveloperCLICredentialOptions{
+				TenantID:                   azureContext.Scope.TenantId,
+				AdditionallyAllowedTenants: []string{"*"},
+			})
+			if err != nil {
+				return fmt.Errorf("failed to create azure credential: %w", err)
+			}
 
 			// console := input.NewConsole(
 			// 	false, // noPrompt
@@ -115,13 +117,13 @@ func newInitCommand() *cobra.Command {
 			action := &InitAction{
 				azdClient: azdClient,
 				// azureClient:         azure.NewAzureClient(credential),
-				// azureContext: azureContext,
+				azureContext: azureContext,
 				// composedResources:   getComposedResourcesResponse.Resources,
 				// console: console,
-				// credential:          credential,
-				// modelCatalogService: ai.NewModelCatalogService(credential),
-				projectConfig: projectConfig,
-				environment:   environment,
+				credential:          credential,
+				modelCatalogService: ai.NewModelCatalogService(credential),
+				projectConfig:       projectConfig,
+				environment:         environment,
 			}
 
 			if err := action.Run(ctx, flags); err != nil {
@@ -1176,32 +1178,32 @@ func downloadDirectoryContents(
 // 	return modelDeployment, nil
 // }
 
-// func (a *InitAction) loadAiCatalog(ctx context.Context) error {
-// 	if a.modelCatalog != nil {
-// 		return nil
-// 	}
+func (a *InitAction) loadAiCatalog(ctx context.Context) error {
+	if a.modelCatalog != nil {
+		return nil
+	}
 
-// 	spinner := ux.NewSpinner(&ux.SpinnerOptions{
-// 		Text:        "Loading AI Model Catalog",
-// 		ClearOnStop: true,
-// 	})
+	spinner := ux.NewSpinner(&ux.SpinnerOptions{
+		Text:        "Loading AI Model Catalog",
+		ClearOnStop: true,
+	})
 
-// 	if err := spinner.Start(ctx); err != nil {
-// 		return fmt.Errorf("failed to start spinner: %w", err)
-// 	}
+	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 AI model catalog: %w", err)
-// 	}
+	aiModelCatalog, err := a.modelCatalogService.ListAllModels(ctx, a.azureContext.Scope.SubscriptionId)
+	if err != nil {
+		return fmt.Errorf("failed to load AI model catalog: %w", err)
+	}
 
-// 	if err := spinner.Stop(ctx); err != nil {
-// 		return err
-// 	}
+	if err := spinner.Stop(ctx); err != nil {
+		return err
+	}
 
-// 	a.modelCatalog = aiModelCatalog
-// 	return nil
-// }
+	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 {
@@ -1290,6 +1292,7 @@ func (a *InitAction) updateEnvironment(ctx context.Context, agentManifest *agent
 	}
 
 	envName := envResponse.Environment.Name
+	deploymentDetails := []Deployment{}
 
 	// Set environment variables based on agent kind
 	switch agentDef.Kind {
@@ -1298,16 +1301,49 @@ func (a *InitAction) updateEnvironment(ctx context.Context, agentManifest *agent
 		if err := a.setEnvVar(ctx, envName, "AZURE_AI_FOUNDRY_MODEL_NAME", agentDef.Model.Id); err != nil {
 			return err
 		}
+
+		modelDeployment, err := a.getModelDeploymentDetails(ctx, agentDef.Model)
+		if err != nil {
+			return fmt.Errorf("failed to get model deployment details: %w", err)
+		}
+		deploymentDetails = append(deploymentDetails, *modelDeployment)
 	case agent_yaml.AgentKindHosted:
-		// Set environment variables for hosted agents
+		agentDef := agentManifest.Template.(agent_yaml.HostedContainerAgent)
 		if err := a.setEnvVar(ctx, envName, "ENABLE_HOSTED_AGENTS", "true"); err != nil {
 			return err
 		}
+
+		// Iterate over all models in the hosted container agent
+		for _, model := range agentDef.Models {
+			modelDeployment, err := a.getModelDeploymentDetails(ctx, model)
+			if err != nil {
+				return fmt.Errorf("failed to get model deployment details: %w", err)
+			}
+			deploymentDetails = append(deploymentDetails, *modelDeployment)
+		}
+
 	case agent_yaml.AgentKindYamlContainerApp:
-		// Set environment variables for foundry agents
+		agentDef := agentManifest.Template.(agent_yaml.ContainerAgent)
 		if err := a.setEnvVar(ctx, envName, "ENABLE_CONTAINER_AGENTS", "true"); err != nil {
 			return err
 		}
+
+		// Iterate over all models in the container agent
+		for _, model := range agentDef.Models {
+			modelDeployment, err := a.getModelDeploymentDetails(ctx, model)
+			if err != nil {
+				return fmt.Errorf("failed to get model deployment details: %w", err)
+			}
+			deploymentDetails = append(deploymentDetails, *modelDeployment)
+		}
+	}
+
+	deploymentsJson, err := json.Marshal(deploymentDetails)
+	if err != nil {
+		return fmt.Errorf("failed to marshal deployment details to JSON: %w", err)
+	}
+	if err := a.setEnvVar(ctx, envName, "AI_PROJECT_DEPLOYMENTS", string(deploymentsJson)); err != nil {
+		return err
 	}
 
 	fmt.Printf("Successfully updated environment variables for agent kind: %s\n", agentDef.Kind)
@@ -1327,3 +1363,82 @@ func (a *InitAction) setEnvVar(ctx context.Context, envName, key, value string)
 	fmt.Printf("Set environment variable: %s=%s\n", key, value)
 	return nil
 }
+
+// Deployment represents a single cognitive service account deployment
+type Deployment struct {
+	// Specify the name of cognitive service account deployment.
+	Name string `json:"name"`
+
+	// Required. Properties of Cognitive Services account deployment model.
+	Model DeploymentModel `json:"model"`
+
+	// The resource model definition representing SKU.
+	Sku DeploymentSku `json:"sku"`
+}
+
+// DeploymentModel represents the model configuration for a cognitive services deployment
+type DeploymentModel struct {
+	// Required. The name of Cognitive Services account deployment model.
+	Name string `json:"name"`
+
+	// Required. The format of Cognitive Services account deployment model.
+	Format string `json:"format"`
+
+	// Required. The version of Cognitive Services account deployment model.
+	Version string `json:"version"`
+}
+
+// DeploymentSku represents the resource model definition representing SKU
+type DeploymentSku struct {
+	// Required. The name of the resource model definition representing SKU.
+	Name string `json:"name"`
+
+	// The capacity of the resource model definition representing SKU.
+	Capacity int `json:"capacity"`
+}
+
+func (a *InitAction) getModelDeploymentDetails(ctx context.Context, model agent_yaml.Model) (*Deployment, error) {
+	modelDetails, err := a.getModelDetails(ctx, model.Id, *model.Version)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get model details: %w", err)
+	}
+
+	return &Deployment{
+		Name: *model.Deployment,
+		Model: DeploymentModel{
+			Name:    model.Id,
+			Format:  modelDetails.Format,
+			Version: *model.Version,
+		},
+		Sku: DeploymentSku{
+			Name:     "GlobalStandard",
+			Capacity: 100,
+		},
+	}, nil
+}
+
+func (a *InitAction) getModelDetails(ctx context.Context, modelName string, modelVersion 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)
+	}
+
+	deploymentOptions := ai.AiModelDeploymentOptions{
+		Versions: []string{modelVersion},
+		Skus:     []string{"GlobalStandard"},
+	}
+
+	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
+}