Merge pull request #175 from microsoft/dev

feat: Enhancements to SQL Migration: Agent Management, Retry Logic, and SFI Changes

Author: Prajwal-Microsoft

infra/main.bicep

@@ -530,6 +530,10 @@ module containerAppBackend 'br/public:avm/res/app/container-app:0.17.0' = {
               name: 'AZURE_CLIENT_ID'
               value: appIdentity.outputs.clientId // NOTE: This is the client ID of the managed identity, not the Entra application, and is needed for the App Service to access the Cosmos DB account.
             }
+            {
+              name: 'APP_ENV'
+              value: 'prod'
+            }
           ],
           enableMonitoring
             ? [
@@ -604,6 +608,10 @@ module containerAppFrontend 'br/public:avm/res/app/container-app:0.17.0' = {
             name: 'API_URL'
             value: 'https://${containerAppBackend.outputs.fqdn}'
           }
+          {
+            name: 'APP_ENV'
+            value: 'prod'
+          }
         ]
         image: 'cmsacontainerreg.azurecr.io/cmsafrontend:${imageVersion}'
         name: 'cmsafrontend'

infra/main.json

@@ -6,7 +6,7 @@
     "_generator": {
       "name": "bicep",
       "version": "0.36.177.2456",
-      "templateHash": "9967677186411813072"
+      "templateHash": "12716753818171697569"
     },
     "name": "Modernize Your Code Solution Accelerator",
     "description": "CSA CTO Gold Standard Solution Accelerator for Modernize Your Code. \r\n"
@@ -41832,7 +41832,7 @@
               {
                 "name": "cmsabackend",
                 "image": "[format('cmsacontainerreg.azurecr.io/cmsabackend:{0}', parameters('imageVersion'))]",
-                "env": "[concat(createArray(createObject('name', 'COSMOSDB_ENDPOINT', 'value', reference('cosmosDb').outputs.endpoint.value), createObject('name', 'COSMOSDB_DATABASE', 'value', reference('cosmosDb').outputs.databaseName.value), createObject('name', 'COSMOSDB_BATCH_CONTAINER', 'value', reference('cosmosDb').outputs.containerNames.value.batch), createObject('name', 'COSMOSDB_FILE_CONTAINER', 'value', reference('cosmosDb').outputs.containerNames.value.file), createObject('name', 'COSMOSDB_LOG_CONTAINER', 'value', reference('cosmosDb').outputs.containerNames.value.log), createObject('name', 'AZURE_BLOB_ACCOUNT_NAME', 'value', reference('storageAccount').outputs.name.value), createObject('name', 'AZURE_BLOB_CONTAINER_NAME', 'value', variables('appStorageContainerName')), createObject('name', 'AZURE_OPENAI_ENDPOINT', 'value', format('https://{0}.openai.azure.com/', reference('aiServices').outputs.name.value)), createObject('name', 'MIGRATOR_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'PICKER_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'FIXER_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'SEMANTIC_VERIFIER_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'SYNTAX_CHECKER_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'SELECTION_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'TERMINATION_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME', 'value', variables('modelDeployment').name), createObject('name', 'AI_PROJECT_ENDPOINT', 'value', reference('aiServices').outputs.aiProjectInfo.value.apiEndpoint), createObject('name', 'AZURE_AI_AGENT_PROJECT_CONNECTION_STRING', 'value', reference('aiServices').outputs.aiProjectInfo.value.apiEndpoint), createObject('name', 'AZURE_AI_AGENT_PROJECT_NAME', 'value', reference('aiServices').outputs.aiProjectInfo.value.name), createObject('name', 'AZURE_AI_AGENT_RESOURCE_GROUP_NAME', 'value', resourceGroup().name), createObject('name', 'AZURE_AI_AGENT_SUBSCRIPTION_ID', 'value', subscription().subscriptionId), createObject('name', 'AZURE_AI_AGENT_ENDPOINT', 'value', reference('aiServices').outputs.aiProjectInfo.value.apiEndpoint), createObject('name', 'AZURE_CLIENT_ID', 'value', reference('appIdentity').outputs.clientId.value)), if(parameters('enableMonitoring'), createArray(createObject('name', 'APPLICATIONINSIGHTS_INSTRUMENTATION_KEY', 'value', reference('applicationInsights').outputs.instrumentationKey.value), createObject('name', 'APPLICATIONINSIGHTS_CONNECTION_STRING', 'value', reference('applicationInsights').outputs.connectionString.value)), createArray()))]",
+                "env": "[concat(createArray(createObject('name', 'COSMOSDB_ENDPOINT', 'value', reference('cosmosDb').outputs.endpoint.value), createObject('name', 'COSMOSDB_DATABASE', 'value', reference('cosmosDb').outputs.databaseName.value), createObject('name', 'COSMOSDB_BATCH_CONTAINER', 'value', reference('cosmosDb').outputs.containerNames.value.batch), createObject('name', 'COSMOSDB_FILE_CONTAINER', 'value', reference('cosmosDb').outputs.containerNames.value.file), createObject('name', 'COSMOSDB_LOG_CONTAINER', 'value', reference('cosmosDb').outputs.containerNames.value.log), createObject('name', 'AZURE_BLOB_ACCOUNT_NAME', 'value', reference('storageAccount').outputs.name.value), createObject('name', 'AZURE_BLOB_CONTAINER_NAME', 'value', variables('appStorageContainerName')), createObject('name', 'AZURE_OPENAI_ENDPOINT', 'value', format('https://{0}.openai.azure.com/', reference('aiServices').outputs.name.value)), createObject('name', 'MIGRATOR_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'PICKER_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'FIXER_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'SEMANTIC_VERIFIER_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'SYNTAX_CHECKER_AGENT_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'SELECTION_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'TERMINATION_MODEL_DEPLOY', 'value', variables('modelDeployment').name), createObject('name', 'AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME', 'value', variables('modelDeployment').name), createObject('name', 'AI_PROJECT_ENDPOINT', 'value', reference('aiServices').outputs.aiProjectInfo.value.apiEndpoint), createObject('name', 'AZURE_AI_AGENT_PROJECT_CONNECTION_STRING', 'value', reference('aiServices').outputs.aiProjectInfo.value.apiEndpoint), createObject('name', 'AZURE_AI_AGENT_PROJECT_NAME', 'value', reference('aiServices').outputs.aiProjectInfo.value.name), createObject('name', 'AZURE_AI_AGENT_RESOURCE_GROUP_NAME', 'value', resourceGroup().name), createObject('name', 'AZURE_AI_AGENT_SUBSCRIPTION_ID', 'value', subscription().subscriptionId), createObject('name', 'AZURE_AI_AGENT_ENDPOINT', 'value', reference('aiServices').outputs.aiProjectInfo.value.apiEndpoint), createObject('name', 'AZURE_CLIENT_ID', 'value', reference('appIdentity').outputs.clientId.value), createObject('name', 'APP_ENV', 'value', 'prod')), if(parameters('enableMonitoring'), createArray(createObject('name', 'APPLICATIONINSIGHTS_INSTRUMENTATION_KEY', 'value', reference('applicationInsights').outputs.instrumentationKey.value), createObject('name', 'APPLICATIONINSIGHTS_CONNECTION_STRING', 'value', reference('applicationInsights').outputs.connectionString.value)), createArray()))]",
                 "resources": {
                   "cpu": 1,
                   "memory": "2.0Gi"
@@ -43334,6 +43334,10 @@
                   {
                     "name": "API_URL",
                     "value": "[format('https://{0}', reference('containerAppBackend').outputs.fqdn.value)]"
+                  },
+                  {
+                    "name": "APP_ENV",
+                    "value": "prod"
                   }
                 ],
                 "image": "[format('cmsacontainerreg.azurecr.io/cmsafrontend:{0}', parameters('imageVersion'))]",

src/backend/.env.sample

@@ -23,4 +23,6 @@ AZURE_AI_AGENT_PROJECT_CONNECTION_STRING = ""
 AZURE_AI_AGENT_SUBSCRIPTION_ID = ""
 AZURE_AI_AGENT_RESOURCE_GROUP_NAME = ""
 AZURE_AI_AGENT_PROJECT_NAME = ""
-AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME = ""
+AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME = ""
+
+APP_ENV = "dev"

src/backend/app.py

@@ -1,13 +1,24 @@
 """Create and configure the FastAPI application."""
+from contextlib import asynccontextmanager
+
 from api.api_routes import router as backend_router
 
+from common.config.config import app_config
 from common.logger.app_logger import AppLogger
 
 from dotenv import load_dotenv
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
+from helper.azure_credential_utils import get_azure_credential
+
+from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent  # pylint: disable=E0611
+
+from sql_agents.agent_manager import clear_sql_agents, set_sql_agents
+from sql_agents.agents.agent_config import AgentBaseConfig
+from sql_agents.helpers.agents_manager import SqlAgents
+
 import uvicorn
 # from agent_services.agents_routes import router as agents_router
 
@@ -17,10 +28,67 @@
 # Configure logging
 logger = AppLogger("app")
 
+# Global variables for agents
+sql_agents: SqlAgents = None
+azure_client = None
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manage application lifespan - startup and shutdown."""
+    global sql_agents, azure_client
+
+    # Startup
+    try:
+        logger.logger.info("Initializing SQL agents...")
+
+        # Create Azure credentials and client
+        creds = get_azure_credential(app_config.azure_client_id)
+        azure_client = AzureAIAgent.create_client(
+            credential=creds,
+            endpoint=app_config.ai_project_endpoint
+        )
+
+        # Setup agent configuration with default conversion settings
+        agent_config = AgentBaseConfig(
+            project_client=azure_client,
+            sql_from="informix",  # Default source dialect
+            sql_to="tsql"         # Default target dialect
+        )
+
+        # Create SQL agents
+        sql_agents = await SqlAgents.create(agent_config)
+
+        # Set the global agents instance
+        set_sql_agents(sql_agents)
+        logger.logger.info("SQL agents initialized successfully.")
+
+    except Exception as exc:
+        logger.logger.error("Failed to initialize SQL agents: %s", exc)
+        # Don't raise the exception to allow the app to start even if agents fail
+
+    yield  # Application runs here
+
+    # Shutdown
+    try:
+        if sql_agents:
+            logger.logger.info("Application shutting down - cleaning up SQL agents...")
+            await sql_agents.delete_agents()
+            logger.logger.info("SQL agents cleaned up successfully.")
+
+            # Clear the global agents instance
+            await clear_sql_agents()
+
+        if azure_client:
+            await azure_client.close()
+
+    except Exception as exc:
+        logger.logger.error("Error during agent cleanup: %s", exc)
+
 
 def create_app() -> FastAPI:
     """Create and return the FastAPI application instance."""
-    app = FastAPI(title="Code Gen Accelerator", version="1.0.0")
+    app = FastAPI(title="Code Gen Accelerator", version="1.0.0", lifespan=lifespan)
 
     # Configure CORS
     app.add_middleware(

src/backend/common/config/config.py

@@ -14,7 +14,9 @@
 
 import os
 
-from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
+from azure.identity.aio import ClientSecretCredential
+
+from helper.azure_credential_utils import get_azure_credential
 
 
 class Config:
@@ -49,7 +51,7 @@ def __init__(self):
             "SYNTAX_CHECKER_AGENT_MODEL_DEPLOY"
         )
 
-        self.__azure_credentials = DefaultAzureCredential()
+        self.__azure_credentials = get_azure_credential(self.azure_client_id)
 
     def get_azure_credentials(self):
         """Retrieve Azure credentials, either from environment variables or managed identity."""

src/backend/common/services/batch_service.py

@@ -61,6 +61,14 @@ async def get_file_report(self, file_id: str) -> Optional[Dict]:
             storage = await BlobStorageFactory.get_storage()
             if file_record.translated_path not in ["", None]:
                 translated_content = await storage.get_file(file_record.translated_path)
+            else:
+                # If translated_path is empty, try to get translated content from logs
+                # Look for the final success log with the translated result
+                if logs:
+                    for log in logs:
+                        if (log.get("log_type") == "success" and log.get("agent_type") == "agents" and log.get("last_candidate")):
+                            translated_content = log["last_candidate"]
+                            break
         except IOError as e:
             self.logger.error(f"Error downloading file content: {str(e)}")
 
@@ -79,6 +87,14 @@ async def get_file_translated(self, file: dict):
             storage = await BlobStorageFactory.get_storage()
             if file["translated_path"] not in ["", None]:
                 translated_content = await storage.get_file(file["translated_path"])
+            else:
+                # If translated_path is empty, try to get translated content from logs
+                # Look for the final success log with the translated result
+                if "logs" in file and file["logs"]:
+                    for log in file["logs"]:
+                        if (log.get("log_type") == "success" and log.get("agent_type") == "agents" and log.get("last_candidate")):
+                            translated_content = log["last_candidate"]
+                            break
         except IOError as e:
             self.logger.error(f"Error downloading file content: {str(e)}")
 
@@ -126,19 +142,22 @@ async def get_batch_summary(self, batch_id: str, user_id: str) -> Optional[Dict]
                 try:
                     logs = await self.database.get_file_logs(file["file_id"])
                     file["logs"] = logs
-                    if file["translated_path"]:
-                        try:
-                            translated_content = await self.get_file_translated(file)
-                            file["translated_content"] = translated_content
-                        except Exception as e:
-                            self.logger.error(
-                                f"Error retrieving translated content for file {file['file_id']}: {str(e)}"
-                            )
+                    # Try to get translated content for all files, but prioritize completed files
+                    try:
+                        translated_content = await self.get_file_translated(file)
+                        file["translated_content"] = translated_content
+                    except Exception as e:
+                        self.logger.error(
+                            f"Error retrieving translated content for file {file['file_id']}: {str(e)}"
+                        )
+                        # Ensure translated_content field exists even if empty
+                        file["translated_content"] = ""
                 except Exception as e:
                     self.logger.error(
                         f"Error retrieving logs for file {file['file_id']}: {str(e)}"
                     )
                     file["logs"] = []  # Set empty logs on error
+                    file["translated_content"] = ""  # Ensure field exists
 
             return {
                 "files": files,

src/backend/common/storage/blob_azure.py

@@ -1,11 +1,13 @@
 from typing import Any, BinaryIO, Dict, Optional
 
-from azure.identity import DefaultAzureCredential
 from azure.storage.blob import BlobServiceClient
 
+from common.config.config import app_config
 from common.logger.app_logger import AppLogger
 from common.storage.blob_base import BlobStorageBase
 
+from helper.azure_credential_utils import get_azure_credential
+
 
 class AzureBlobStorage(BlobStorageBase):
     def __init__(self, account_name: str, container_name: Optional[str] = None):
@@ -15,7 +17,7 @@ def __init__(self, account_name: str, container_name: Optional[str] = None):
             self.container_name = container_name
             self.service_client = None
             self.container_client = None
-            credential = DefaultAzureCredential()  # Using Entra Authentication
+            credential = get_azure_credential(app_config.azure_client_id)  # Using Entra Authentication
             self.service_client = BlobServiceClient(
                 account_url=f"https://{self.account_name}.blob.core.windows.net/",
                 credential=credential,

src/backend/helper/azure_credential_utils.py

@@ -0,0 +1,36 @@
+import os
+
+from azure.identity import DefaultAzureCredential, ManagedIdentityCredential
+from azure.identity.aio import DefaultAzureCredential as AioDefaultAzureCredential, ManagedIdentityCredential as AioManagedIdentityCredential
+
+
+async def get_azure_credential_async(client_id=None):
+    """
+    Returns an Azure credential asynchronously based on the application environment.
+    If the environment is 'dev', it uses AioDefaultAzureCredential.
+    Otherwise, it uses AioManagedIdentityCredential.
+    Args:
+        client_id (str, optional): The client ID for the Managed Identity Credential.
+    Returns:
+        Credential object: Either AioDefaultAzureCredential or AioManagedIdentityCredential.
+    """
+    if os.getenv("APP_ENV", "prod").lower() == 'dev':
+        return AioDefaultAzureCredential()  # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development
+    else:
+        return AioManagedIdentityCredential(client_id=client_id)
+
+
+def get_azure_credential(client_id=None):
+    """
+    Returns an Azure credential based on the application environment.
+    If the environment is 'dev', it uses DefaultAzureCredential.
+    Otherwise, it uses ManagedIdentityCredential.
+    Args:
+        client_id (str, optional): The client ID for the Managed Identity Credential.
+    Returns:
+        Credential object: Either DefaultAzureCredential or ManagedIdentityCredential.
+    """
+    if os.getenv("APP_ENV", "prod").lower() == 'dev':
+        return DefaultAzureCredential()  # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development
+    else:
+        return ManagedIdentityCredential(client_id=client_id)

src/backend/sql_agents/agent_manager.py

@@ -0,0 +1,44 @@
+"""
+Global agent manager for SQL agents.
+This module manages the global SQL agents instance to avoid circular imports.
+"""
+
+import logging
+from typing import Optional
+
+from sql_agents.helpers.agents_manager import SqlAgents
+
+logger = logging.getLogger(__name__)
+
+# Global variable to store the SQL agents instance
+_sql_agents: Optional[SqlAgents] = None
+
+
+def set_sql_agents(agents: SqlAgents) -> None:
+    """Set the global SQL agents instance."""
+    global _sql_agents
+    _sql_agents = agents
+    logger.info("Global SQL agents instance has been set")
+
+
+def get_sql_agents() -> Optional[SqlAgents]:
+    """Get the global SQL agents instance."""
+    return _sql_agents
+
+
+async def update_agent_config(convert_from: str, convert_to: str) -> None:
+    """Update the global agent configuration for different SQL conversion types."""
+    if _sql_agents and _sql_agents.agent_config:
+        _sql_agents.agent_config.sql_from = convert_from
+        _sql_agents.agent_config.sql_to = convert_to
+        logger.info(f"Updated agent configuration: {convert_from} -> {convert_to}")
+    else:
+        logger.warning("SQL agents not initialized, cannot update configuration")
+
+
+async def clear_sql_agents() -> None:
+    """Clear the global SQL agents instance."""
+    global _sql_agents
+    await _sql_agents.delete_agents()
+    _sql_agents = None
+    logger.info("Global SQL agents instance has been cleared")