Add Azure AI Search support to FoundryAgentTemplate

Refactored FoundryAgentTemplate to support both Azure AI Search (raw tool) and legacy MCP tool paths, with priority given to Azure AI Search if configured. Added logic to create and clean up server-side Azure AI agents for the Azure Search path, and ensured code interpreter is only attached in the MCP path due to incompatibility. Updated MagenticAgentFactory and ReasoningAgentTemplate to use a consistent 'project_endpoint' parameter name.

Author: Fr4nc3

src/backend/v4/magentic_agents/foundry_agent.py

@@ -7,17 +7,25 @@
     ChatAgent,
     ChatMessage,
     Role,
-    HostedFileSearchTool,
-    HostedVectorStoreContent,
     HostedCodeInterpreterTool,
 )
+from azure.ai.projects.models import ConnectionType
+from agent_framework_azure_ai import AzureAIAgentClient  # Provided by agent_framework
+
+
 from v4.magentic_agents.common.lifecycle import AzureAgentBase
 from v4.magentic_agents.models.agent_models import MCPConfig, SearchConfig
 from v4.config.agent_registry import agent_registry
 
 
 class FoundryAgentTemplate(AzureAgentBase):
-    """Agent that uses Azure AI Search (RAG) and optional MCP tool via agent_framework."""
+    """Agent that uses Azure AI Search (raw tool) OR MCP tool + optional Code Interpreter.
+
+    Priority:
+      1. Azure AI Search (if search_config contains required Azure Search fields)
+      2. MCP tool (legacy path)
+    Code Interpreter is only attached on the MCP path (unless you want it also with Azure Search—currently skipped for incompatibility per request).
+    """
 
     def __init__(
         self,
@@ -37,41 +45,39 @@ def __init__(
         self.search = search_config
         self.logger = logging.getLogger(__name__)
 
+        # Decide early whether Azure Search mode should be activated
+        self._use_azure_search = self._is_azure_search_requested()
+
+        # Placeholder for server-created Azure AI agent id (if Azure Search path)
+        self._azure_server_agent_id: Optional[str] = None
+
     # -------------------------
-    # Tool construction helpers
+    # Mode detection
     # -------------------------
-    async def _make_file_search_tool(self) -> Optional[HostedFileSearchTool]:
-        """Create File Search tool (RAG capability) using vector stores."""
-        if not self.search or not self.search.vector_store_id:
-            self.logger.info("File search tool not enabled (missing vector_store_id).")
-            return None
+    def _is_azure_search_requested(self) -> bool:
+        """Determine if Azure AI Search raw tool path should be used."""
+        if not self.search:
+            return False
+        # Minimal heuristic: presence of required attributes
 
-        try:
-            # HostedFileSearchTool uses vector stores, not direct Azure AI Search indexes
-            file_search_tool = HostedFileSearchTool(
-                inputs=[HostedVectorStoreContent(vector_store_id=self.search.vector_store_id)],
-                max_results=self.search.max_results if hasattr(self.search, 'max_results') else None,
-                description="Search through indexed documents"
+        has_index = hasattr(self.search, "index_name") and bool(self.search.index_name)
+        if has_index:
+            self.logger.info(
+                "Azure AI Search requested (connection_id=%s, index=%s).",
+                getattr(self.search, "connection_name", None),
+                getattr(self.search, "index_name", None),
             )
-            self.logger.info("Created HostedFileSearchTool with vector store: %s", self.search.vector_store_id)
-            return file_search_tool
-        except Exception as ex:
-            self.logger.error("File search tool creation failed: %s", ex)
-            return None
+            return True
+        return False
+
+
 
     async def _collect_tools(self) -> List:
-        """Collect tool definitions for ChatAgent."""
+        """Collect tool definitions for ChatAgent (MCP path only)."""
         tools: List = []
 
-        # File Search tool (RAG)
-        if self.search:
-            print("Adding File Search tool.")
-            # search_tool = await self._make_file_search_tool()
-            # if search_tool:
-            #     tools.append(search_tool)
-            #     self.logger.info("Added File Search tool.")
 
-        # Code Interpreter
+        # Code Interpreter (only in MCP path per incompatibility note)
         if self.enable_code_interpreter:
             try:
                 code_tool = HostedCodeInterpreterTool()
@@ -85,28 +91,110 @@ async def _collect_tools(self) -> List:
             tools.append(self.mcp_tool)
             self.logger.info("Added MCP tool: %s", self.mcp_tool.name)
 
-        self.logger.info("Total tools collected: %d", len(tools))
+        self.logger.info("Total tools collected (MCP path): %d", len(tools))
         return tools
 
+    # -------------------------
+    # Azure Search helper
+    # -------------------------
+    async def _create_azure_search_enabled_client(self):
+        """
+        Create a server-side Azure AI agent with raw Azure AI Search tool and return an AzureAIAgentClient.
+        This mirrors your example while fitting existing lifecycle.
+
+        If these assumptions differ, adjust accordingly.
+        """
+
+        connection_id = getattr(self.search, "connection_name", "")
+        index_name = getattr(self.search, "index_name", "")
+        query_type = getattr(self.search, "search_query_type", "vector")
+        
+        # ai_search_conn_id = ""
+        # async for connection in self.client.project_client.connections.list():
+        #     if connection.type == ConnectionType.AZURE_AI_SEARCH:
+        #         ai_search_conn_id = connection.id
+        #         break
+        if not connection_id or not index_name:
+            self.logger.error(
+                "Missing azure_search_connection_id or azure_search_index_name in search_config; aborting Azure Search path."
+            )
+            return None
+
+        try:
+            azure_agent = await self.client.project_client.agents.create_agent(
+                model=self.model_deployment_name,
+                name=self.agent_name,
+                instructions=(
+                    f"{self.agent_instructions} "
+                    "Always use the Azure AI Search tool and configured index for knowledge retrieval."
+                ),
+                tools=[{"type": "azure_ai_search"}],
+                tool_resources={
+                    "azure_ai_search": {
+                        "indexes": [
+                            {
+                                "index_connection_id": connection_id,
+                                "index_name": index_name,
+                                "query_type": query_type,
+                            }
+                        ]
+                    }
+                },
+            )
+            self._azure_server_agent_id = azure_agent.id
+            self.logger.info(
+                "Created Azure server agent with Azure AI Search tool (agent_id=%s, index=%s).",
+                azure_agent.id,
+                index_name,
+            )
+
+            chat_client = AzureAIAgentClient(
+                project_client=self.client.project_client,
+                agent_id=azure_agent.id,
+            )
+            return chat_client
+        except Exception as ex:
+            self.logger.error("Failed to create Azure Search enabled agent: %s", ex)
+            return None
+
     # -------------------------
     # Agent lifecycle override
     # -------------------------
     async def _after_open(self) -> None:
         """Initialize ChatAgent after connections are established."""
         try:
-            tools = await self._collect_tools()
-            
-            self._agent = ChatAgent(
-                chat_client=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,
-            )
-            
+            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=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)
         except Exception as ex:
             self.logger.error("Failed to initialize ChatAgent: %s", ex)
@@ -128,10 +216,32 @@ async def invoke(self, prompt: str):
             raise RuntimeError("Agent not initialized; call open() first.")
 
         messages = [ChatMessage(role=Role.USER, text=prompt)]
-        
+
         async for update in self._agent.run_stream(messages=messages):
             yield update
 
+    # -------------------------
+    # Cleanup (optional override if you want to delete server-side agent)
+    # -------------------------
+    async def close(self) -> None:
+        """Extend base close to optionally delete server-side Azure agent."""
+        try:
+            if self._use_azure_search and self._azure_server_agent_id and hasattr(self, "project_client"):
+                try:
+                    await self.project_client.agents.delete_agent(self._azure_server_agent_id)
+                    self.logger.info(
+                        "Deleted Azure server agent (id=%s) during close.", self._azure_server_agent_id
+                    )
+                except Exception as ex:
+                    self.logger.warning(
+                        "Failed to delete Azure server agent (id=%s): %s",
+                        self._azure_server_agent_id,
+                        ex,
+                    )
+        finally:
+            await super().close()
+
+
 # -------------------------
 # Factory
 # -------------------------

src/backend/v4/magentic_agents/magentic_agent_factory.py

@@ -102,13 +102,12 @@ async def create_agent_from_config(self, user_id: str, agent_obj: SimpleNamespac
         if use_reasoning:
             # Get reasoning specific configuration
             azure_openai_endpoint = config.AZURE_OPENAI_ENDPOINT
-
             agent = ReasoningAgentTemplate(
                 agent_name=agent_obj.name,
                 agent_description=getattr(agent_obj, "description", ""),
                 agent_instructions=getattr(agent_obj, "system_message", ""),
                 model_deployment_name=deployment_name,
-                azure_openai_endpoint=azure_openai_endpoint,
+                project_endpoint=azure_openai_endpoint, # type: ignore
                 search_config=search_config,
                 mcp_config=mcp_config,
             )

src/backend/v4/magentic_agents/reasoning_agent.py

@@ -43,7 +43,7 @@ def __init__(
         agent_description: str,
         agent_instructions: str,
         model_deployment_name: str,
-        azure_ai_project_endpoint: str,
+        project_endpoint: str | None = "",
         search_config: SearchConfig | None = None,
         mcp_config: MCPConfig | None = None,
         max_search_docs: int = 3,
@@ -55,7 +55,7 @@ def __init__(
             agent_description: Description of the agent's purpose
             agent_instructions: System instructions for the agent
             model_deployment_name: Reasoning model deployment (e.g., "o1", "o3-mini")
-            azure_ai_project_endpoint: Azure AI Project endpoint URL
+            project_endpoint: Azure AI Project endpoint URL
             search_config: Optional search configuration for Azure AI Search
             mcp_config: Optional MCP server configuration
             max_search_docs: Maximum number of search documents to retrieve
@@ -65,7 +65,7 @@ def __init__(
         self.agent_description = agent_description
         self.base_instructions = agent_instructions
         self.model_deployment_name = model_deployment_name
-        self.project_endpoint = azure_ai_project_endpoint
+        self.project_endpoint = project_endpoint
         self.search_config = search_config
         self.max_search_docs = max_search_docs