Refactor agent initialization to support persistent clients

Agent initialization now checks for existing database-backed AzureAIAgentClient instances and reuses them if available, improving persistence and resource management. The constructors for agent base classes were updated to accept and store model_deployment_name. FoundryAgentTemplate and ReasoningAgentTemplate now consistently use the database client if present, and only save new agents when necessary.

Author: Fr4nc3

src/backend/v4/magentic_agents/common/lifecycle.py

@@ -49,6 +49,7 @@ def __init__(
         agent_name: str | None = None,
         agent_description: str | None = None,
         agent_instructions: str | None = None,
+        model_deployment_name: str | None = None,
     ) -> None:
         self._stack: AsyncExitStack | None = None
         self.mcp_cfg: MCPConfig | None = mcp
@@ -63,6 +64,7 @@ def __init__(
         self.agent_name: str | None = agent_name
         self.agent_description: str | None = agent_description
         self.agent_instructions: str | None = agent_instructions
+        self.model_deployment_name: str | None = model_deployment_name
         self.logger = logging.getLogger(__name__)
 
     async def open(self) -> "MCPEnabledBase":
@@ -133,22 +135,27 @@ async def _after_open(self) -> None:
         """Subclasses must build self._agent here."""
         raise NotImplementedError
 
-    async def get_database_team_agent(self) -> Optional[CurrentTeamAgent]:
+    async def get_database_team_agent(self) -> Optional[AzureAIAgentClient]:
         """Retrieve existing team agent from database, if any."""
-        agent = None
+        chat_client = None
         try:
             currentAgent = await self.memory_store.get_team_agent(
                 team_id=self.team_config.team_id,
                 agent_name=self.agent_name
             )
             if currentAgent and currentAgent.agent_foundry_id:
-                agent = self.client.get_agent(
-                    agent_id=currentAgent.agent_foundry_id
-                )
+                agent = await self.client.get_agent(agent_id=currentAgent.agent_foundry_id)
+                if agent:
+                    chat_client = AzureAIAgentClient(
+                        project_endpoint=self.project_endpoint,
+                        agent_id=currentAgent.agent_foundry_id,
+                        model_deployment_name=self.model_deployment_name,
+                        async_credential=self.creds,
+                    )
 
         except Exception as ex:  # Consider narrowing this to specific exceptions if possible
             self.logger.error("Failed to initialize ReasoningAgentTemplate: %s", ex)
-        return agent
+        return chat_client
 
     async def save_database_team_agent(self) -> None:
         """Save current team agent to database."""
@@ -213,13 +220,14 @@ def __init__(
             agent_name=agent_name,
             agent_description=agent_description,
             agent_instructions=agent_instructions,
+            model_deployment_name=model_deployment_name,
         )
 
         self._created_ephemeral: bool = (
             False  # reserved if you add ephemeral agent cleanup
         )
 
-        self.model_deployment_name = model_deployment_name
+
 
     async def open(self) -> "AzureAgentBase":
         if self._stack is not None:

src/backend/v4/magentic_agents/foundry_agent.py

@@ -104,7 +104,7 @@ async def _collect_tools(self) -> List:
     # -------------------------
     # Azure Search helper
     # -------------------------
-    async def _create_azure_search_enabled_client(self):
+    async def _create_azure_search_enabled_client(self, chatClient=None) -> Optional[AzureAIAgentClient]:
         """
         Create a server-side Azure AI agent with Azure AI Search raw tool.
 
@@ -118,6 +118,9 @@ async def _create_azure_search_enabled_client(self):
         Returns:
             AzureAIAgentClient | None
         """
+        if chatClient:
+            return chatClient
+        
         if not self.search:
             self.logger.error("Search configuration missing.")
             return None
@@ -217,55 +220,48 @@ async def _after_open(self) -> None:
         """Initialize ChatAgent after connections are established."""
 
         try:
-            agent = await self.get_database_team_agent()
-            if agent is None:
-                if self._use_azure_search:
-                    # Azure Search mode (skip MCP + Code Interpreter due to incompatibility)
-                    self.logger.info(
-                        "Initializing agent in Azure AI Search mode (exclusive)."
-                    )
-                    chat_client = await self._create_azure_search_enabled_client()
-                    if not chat_client:
-                        raise RuntimeError(
-                            "Azure AI Search mode requested but setup failed."
-                        )
-
-                    # In Azure Search raw tool path, tools/tool_choice are handled server-side.
-                    self._agent = ChatAgent(
-                        chat_client=chat_client,
-                        instructions=self.agent_instructions,
-                        name=self.agent_name,
-                        description=self.agent_description,
-                        tool_choice="required",  # Force usage
-                        temperature=0.7,
-                        model_id=self.model_deployment_name,
-                    )
-                else:
-                    # use MCP path
-                    self.logger.info("Initializing agent in MCP mode.")
-                    tools = await self._collect_tools()
-                    self._agent = ChatAgent(
-                        chat_client=AzureAIAgentClient(
-                            project_endpoint=self.project_endpoint,
-                            model_deployment_name=self.model_deployment_name,
-                            async_credential=self.creds,
-                        ),
-                        instructions=self.agent_instructions,
-                        name=self.agent_name,
-                        description=self.agent_description,
-                        tools=tools if tools else None,
-                        tool_choice="auto" if tools else "none",
-                        temperature=0.7,
-                        model_id=self.model_deployment_name,
+            chatClient= await self.get_database_team_agent()
+            
+            if self._use_azure_search:
+                # Azure Search mode (skip MCP + Code Interpreter due to incompatibility)
+                self.logger.info(
+                    "Initializing agent in Azure AI Search mode (exclusive)."
+                )
+                chat_client = await self._create_azure_search_enabled_client(chatClient)
+                if not chat_client:
+                    raise RuntimeError(
+                        "Azure AI Search mode requested but setup failed."
                     )
 
-                self.logger.info("Initialized ChatAgent '%s'", self.agent_name)
+                # In Azure Search raw tool path, tools/tool_choice are handled server-side.
+                self._agent = ChatAgent(
+                    chat_client=chat_client,
+                    instructions=self.agent_instructions,
+                    name=self.agent_name,
+                    description=self.agent_description,
+                    tool_choice="required",  # Force usage
+                    temperature=0.7,
+                    model_id=self.model_deployment_name,
+                )
+            else:
+                # use MCP path
+                self.logger.info("Initializing agent in MCP mode.")
+                tools = await self._collect_tools()
+                self._agent = ChatAgent(
+                    chat_client=chatClient or self.client,
+                    instructions=self.agent_instructions,
+                    name=self.agent_name,
+                    description=self.agent_description,
+                    tools=tools if tools else None,
+                    tool_choice="auto" if tools else "none",
+                    temperature=0.7,
+                    model_id=self.model_deployment_name,
+                )
 
+            self.logger.info("Initialized ChatAgent '%s'", self.agent_name)
+            if not chatClient: # Only save if we didn't load from DB
                 await self.save_database_team_agent()
 
-            else:
-                self.logger.info("Using existing ChatAgent '%s'", self.agent_name)
-                self._agent = agent
         except Exception as ex:
             self.logger.error("Failed to initialize ChatAgent: %s", ex)
             raise

src/backend/v4/magentic_agents/reasoning_agent.py

@@ -72,8 +72,9 @@ def __init__(
             agent_name=agent_name,
             agent_description=agent_description,    
             agent_instructions=agent_instructions,
+            model_deployment_name=model_deployment_name,
         )
-        self.model_deployment_name = model_deployment_name
+
 
         self.search_config = search_config
         self.max_search_docs = max_search_docs
@@ -106,22 +107,20 @@ async def _after_open(self) -> None:
             # Prepare tools for the agent
             tools = self._prepare_tools()
 
-            agent = await self.get_database_team_agent()
-            if not agent:
-                self._agent = ChatAgent(
-                    chat_client=self.client,
-                    instructions=self.base_instructions,
-                    name=self.agent_name,
-                    description=self.agent_description,
-                    tools=tools if tools else None,
-                    tool_choice="auto" if tools else "none",
-                    temperature=1.0,  # Reasoning models use fixed temperature
-                    model_id=self.model_deployment_name,
-                )
+            chatClient = await self.get_database_team_agent()
+            
+            self._agent = ChatAgent(
+                chat_client = chatClient or self.client,
+                instructions=self.agent_instructions,
+                name=self.agent_name,
+                description=self.agent_description,
+                tools=tools if tools else None,
+                tool_choice="auto" if tools else "none",
+                temperature=1.0,  # Reasoning models use fixed temperature
+                model_id=self.model_deployment_name,
+            )
+            if not chatClient:
                 await self.save_database_team_agent()
-                    
-            else:
-                self._agent = agent
                 self.logger.info("Using existing ReasoningAgent '%s'", self.agent_name)
             # Register agent globally
             try: