Merge pull request #177 from microsoft/psl-agent-management

feat: Centralized agents Management

Author: Roopan-Microsoft

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 azure.identity.aio import DefaultAzureCredential
+
+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 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 = DefaultAzureCredential()
+        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/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/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")