Centralized agents Management

Author: Shreyas-Microsoft

src/backend/app.py

@@ -17,6 +17,7 @@
 
 from sql_agents.agents.agent_config import AgentBaseConfig
 from sql_agents.helpers.agents_manager import SqlAgents
+from sql_agents.agent_manager import set_sql_agents, clear_sql_agents
 
 import uvicorn
 # from agent_services.agents_routes import router as agents_router
@@ -32,22 +33,6 @@
 azure_client = None
 
 
-def get_sql_agents() -> SqlAgents:
-    """Get the global SQL agents instance."""
-    return sql_agents
-
-
-async def update_agent_config(convert_from: str, convert_to: str):
-    """Update the global agent configuration for different SQL conversion types."""
-    global sql_agents
-    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.logger.info(f"Updated agent configuration: {convert_from} -> {convert_to}")
-    else:
-        logger.logger.warning("SQL agents not initialized, cannot update configuration")
-
-
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Manage application lifespan - startup and shutdown."""
@@ -73,6 +58,9 @@ async def lifespan(app: FastAPI):
 
         # 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:
@@ -84,9 +72,12 @@ async def lifespan(app: FastAPI):
     # Shutdown
     try:
         if sql_agents:
-            logger.logger.info("Cleaning up 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()

src/backend/sql_agents/agent_manager.py

@@ -0,0 +1,45 @@
+"""
+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."""
+    global _sql_agents
+    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")