Merge pull request #605 from Fr4nc3/macae-rfp-af-101725

Add Azure AI Search support to FoundryAgentTemplate

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