adding the multiple env mapping feature

Author: RoxyFarhad

src/agentex/lib/cli/handlers/deploy_handlers.py

@@ -292,9 +292,8 @@ def deploy_agent(
     check_and_switch_cluster_context(cluster_name)
 
     manifest = AgentManifest.from_yaml(file_path=manifest_path)
-
-    # Load agent environment configuration
     agent_env_config = None
+    # Load agent environment configuration
     if environment_name:
         manifest_dir = Path(manifest_path).parent
         environments_config = manifest.load_environments_config(manifest_dir)

src/agentex/lib/sdk/config/environment_config.py

@@ -7,7 +7,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, override
+from typing import Any, Dict, override, List
 from pathlib import Path
 
 import yaml
@@ -18,13 +18,12 @@
 
 class AgentAuthConfig(BaseModel):
     """Authentication configuration for an agent in a specific environment."""
-    
+
     principal: Dict[str, Any] = Field(
-        ..., 
-        description="Principal configuration for agent authorization and registration"
+        ..., description="Principal configuration for agent authorization and registration"
     )
-    
-    @field_validator('principal')
+
+    @field_validator("principal")
     @classmethod
     def validate_principal_required_fields(cls, v: Any) -> Dict[str, Any]:
         """Ensure principal has required fields for agent registration."""
@@ -35,142 +34,177 @@ def validate_principal_required_fields(cls, v: Any) -> Dict[str, Any]:
 
 class AgentKubernetesConfig(BaseModel):
     """Kubernetes configuration for an agent in a specific environment."""
-    
-    namespace: str = Field(
-        ..., 
-        description="Kubernetes namespace where the agent will be deployed"
-    )
-    
-    @field_validator('namespace')
+
+    namespace: str = Field(..., description="Kubernetes namespace where the agent will be deployed")
+
+    @field_validator("namespace")
     @classmethod
     def validate_namespace_format(cls, v: str) -> str:
         """Ensure namespace follows Kubernetes naming conventions."""
         if not v or not v.strip():
             raise ValueError("Namespace cannot be empty")
-        
+
         # Basic Kubernetes namespace validation
         namespace = v.strip().lower()
-        if not namespace.replace('-', '').replace('.', '').isalnum():
-            raise ValueError(
-                f"Namespace '{v}' must contain only lowercase letters, numbers, "
-                "hyphens, and periods"
-            )
-        
+        if not namespace.replace("-", "").replace(".", "").isalnum():
+            raise ValueError(f"Namespace '{v}' must contain only lowercase letters, numbers, hyphens, and periods")
+
         if len(namespace) > 63:
             raise ValueError(f"Namespace '{v}' cannot exceed 63 characters")
-        
+
         return namespace
 
 
 class AgentEnvironmentConfig(BaseModel):
     """Complete configuration for an agent in a specific environment."""
-    
-    kubernetes: AgentKubernetesConfig | None = Field(
-        default=None, 
-        description="Kubernetes deployment configuration"
-    )
-    auth: AgentAuthConfig = Field(
-        ..., 
-        description="Authentication and authorization configuration"
-    )
-    helm_repository_name: str = Field(
-        default="scale-egp", 
-        description="Helm repository name for the environment"
+
+    kubernetes: AgentKubernetesConfig | None = Field(default=None, description="Kubernetes deployment configuration")
+    environment: str | None = Field(
+        default=None,
+        description="The environment keyword that this specific environment maps to: either dev, staging, prod",
     )
+    auth: AgentAuthConfig = Field(..., description="Authentication and authorization configuration")
+    helm_repository_name: str = Field(default="scale-egp", description="Helm repository name for the environment")
     helm_repository_url: str = Field(
-        default="https://scale-egp-helm-charts-us-west-2.s3.amazonaws.com/charts", 
-        description="Helm repository url for the environment"
+        default="https://scale-egp-helm-charts-us-west-2.s3.amazonaws.com/charts",
+        description="Helm repository url for the environment",
     )
     helm_overrides: Dict[str, Any] = Field(
-        default_factory=dict, 
-        description="Helm chart value overrides for environment-specific tuning"
+        default_factory=dict, description="Helm chart value overrides for environment-specific tuning"
     )
 
 
 class AgentEnvironmentsConfig(UtilsBaseModel):
     """All environment configurations for an agent."""
-    
-    schema_version: str = Field(
-        default="v1", 
-        description="Schema version for validation and compatibility"
-    )
+
+    schema_version: str = Field(default="v1", description="Schema version for validation and compatibility")
     environments: Dict[str, AgentEnvironmentConfig] = Field(
-        ..., 
-        description="Environment-specific configurations (dev, prod, etc.)"
+        ..., description="Environment-specific configurations (dev, prod, etc.)"
     )
-    
-    @field_validator('schema_version')
+
+    @field_validator("schema_version")
     @classmethod
     def validate_schema_version(cls, v: str) -> str:
         """Ensure schema version is supported."""
-        supported_versions = ['v1']
+        supported_versions = ["v1"]
         if v not in supported_versions:
-            raise ValueError(
-                f"Schema version '{v}' not supported. "
-                f"Supported versions: {', '.join(supported_versions)}"
-            )
+            raise ValueError(f"Schema version '{v}' not supported. Supported versions: {', '.join(supported_versions)}")
         return v
-    
-    @field_validator('environments')
+
+    @field_validator("environments")
     @classmethod
     def validate_environments_not_empty(cls, v: Dict[str, AgentEnvironmentConfig]) -> Dict[str, AgentEnvironmentConfig]:
         """Ensure at least one environment is defined."""
         if not v:
             raise ValueError("At least one environment must be defined")
         return v
-    
+
     def get_config_for_env(self, env_name: str) -> AgentEnvironmentConfig:
         """Get configuration for a specific environment.
-        
+
         Args:
             env_name: Name of the environment (e.g., 'dev', 'prod')
-            
+
         Returns:
             AgentEnvironmentConfig for the specified environment
-            
+
         Raises:
             ValueError: If environment is not found
         """
         if env_name not in self.environments:
-            available_envs = ', '.join(self.environments.keys())
+            available_envs = ", ".join(self.environments.keys())
             raise ValueError(
-                f"Environment '{env_name}' not found in environments.yaml. "
-                f"Available environments: {available_envs}"
+                f"Environment '{env_name}' not found in environments.yaml. Available environments: {available_envs}"
             )
         return self.environments[env_name]
-    
+
+    def get_configs_for_env(self, env: str) -> List[AgentEnvironmentConfig]:
+        """Get configuration for a specific environment based on the expected mapping.
+        The environment is either:
+        1. explicitly specified like so using a key-map in the environments conifg:
+        environments:
+        dev-aws:
+            environment: "dev"
+            kubernetes:
+            namespace: "sgp-000-hello-acp"
+            auth:
+            principal:
+                user_id: 73d0c8bd-4726-434c-9686-eb627d89f078
+                account_id: 6887f093600ecd59bbbd3095
+            helm_overrides:
+
+        or: it it can be defined at the top level:
+        dev:
+            kubernetes:
+            namespace: "sgp-000-hello-acp"
+            auth:
+            principal:
+                user_id: 73d0c8bd-4726-434c-9686-eb627d89f078
+                account_id: 6887f093600ecd59bbbd3095
+            helm_overrides:
+
+        if its not explicitly defined, we assumed the level key (environment name) is the environment
+        Args:
+            env_name: Name of the environment (e.g., 'dev', 'prod')
+
+        Returns:
+            AgentEnvironmentConfig for the specified environment
+
+        Raises:
+            ValueError: If environment is not found
+        """
+        envs_to_deploy = []
+        if env in self.environments:
+            # this supports if the top-level key is just "dev, staging, etc" and matches
+            # the environment name exactly without any explicit mapping
+            envs_to_deploy.append(self.environments[env])
+
+        for _, config in self.environments.items():
+            if config.environment == env:
+                envs_to_deploy.append(config)
+
+        if len(envs_to_deploy) == 0:
+            available_envs = [env.environment for _, env in self.environments.items()] + [
+                env_name for env_name in self.environments
+            ]
+            raise ValueError(
+                f"Environment '{envs_to_deploy}' not found in environments.yaml. Available environments: {available_envs}"
+            )
+
+        return envs_to_deploy
+
     def list_environments(self) -> list[str]:
         """Get list of all configured environment names."""
         return list(self.environments.keys())
-    
+
     @classmethod
     @override
     def from_yaml(cls, file_path: str) -> "AgentEnvironmentsConfig":
         """Load configuration from environments.yaml file.
-        
+
         Args:
             file_path: Path to environments.yaml file
-            
+
         Returns:
             Parsed and validated AgentEnvironmentsConfig
-            
+
         Raises:
             FileNotFoundError: If file doesn't exist
             ValueError: If file is invalid or doesn't validate
         """
         path = Path(file_path)
         if not path.exists():
             raise FileNotFoundError(f"environments.yaml not found: {file_path}")
-        
+
         try:
-            with open(path, 'r') as f:
+            with open(path, "r") as f:
                 data = yaml.safe_load(f)
-            
+
             if not data:
                 raise ValueError("environments.yaml file is empty")
-            
+
             return cls.model_validate(data)
-            
+
         except yaml.YAMLError as e:
             raise ValueError(f"Invalid YAML format in {file_path}: {e}") from e
         except Exception as e:
@@ -179,18 +213,18 @@ def from_yaml(cls, file_path: str) -> "AgentEnvironmentsConfig":
 
 def load_environments_config_from_manifest_dir(manifest_dir: Path) -> AgentEnvironmentsConfig | None:
     """Helper function to load environments.yaml from same directory as manifest.yaml.
-    
+
     Args:
         manifest_dir: Directory containing manifest.yaml
-        
+
     Returns:
         AgentEnvironmentsConfig if environments.yaml exists, None otherwise
-        
+
     Raises:
         ValueError: If environments.yaml exists but is invalid
     """
     environments_file = manifest_dir / "environments.yaml"
     if not environments_file.exists():
         return None
-    
+
     return AgentEnvironmentsConfig.from_yaml(str(environments_file))