Enhance agent orchestration and configuration for Azure AI Search integration

Author: Abdul-Microsoft

src/backend/v4/api/router.py

@@ -303,6 +303,17 @@ async def process_request(
         )
         await memory_store.add_plan(plan)
 
+        # Ensure orchestration is initialized before running
+        # Force rebuild for each new task since Magentic workflows cannot be reused after completion
+        team_service = TeamService(memory_store)
+        await OrchestrationManager.get_current_or_new_orchestration(
+            user_id=user_id,
+            team_config=team,
+            team_switched=False,
+            team_service=team_service,
+            force_rebuild=True,  # Always rebuild workflow for new tasks
+        )
+
         track_event_if_configured(
             "PlanCreated",
             {

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

@@ -51,7 +51,7 @@ def __init__(
         self._agent: ChatAgent | None = None
         self.team_service: TeamService | None = team_service
         self.team_config: TeamConfiguration | None = team_config
-        self.client: Optional[AzureAIClient] = None
+        self.client: Optional[AgentsClient] = None
         self.project_endpoint = project_endpoint
         self.creds: Optional[DefaultAzureCredential] = None
         self.memory_store: Optional[DatabaseBase] = memory_store
@@ -228,7 +228,25 @@ def get_agent_id(self, chat_client) -> str:
         return id
 
     async def get_database_team_agent(self) -> Optional[AzureAIClient]:
-        """Retrieve existing team agent from database, if any."""
+        """Retrieve existing team agent from database, if any.
+        
+        NOTE: Agent reuse is currently DISABLED to ensure fresh agents are created
+        with the correct Azure AI Search configuration.
+        This prevents issues with stale agents that may not have the search tool configured.
+        
+        To re-enable agent reuse, set ENABLE_AGENT_REUSE=true in environment.
+        """
+        import os
+        
+        # DISABLED: Always create fresh agents to ensure Azure AI Search tool is configured
+        enable_reuse = os.environ.get("ENABLE_AGENT_REUSE", "false").lower() == "true"
+        if not enable_reuse:
+            self.logger.info(
+                "Agent reuse DISABLED: Creating fresh agent with search tools (agent_name=%s)",
+                self.agent_name,
+            )
+            return None
+        
         chat_client = None
         try:
             agent_id = await get_database_team_agent_id(
@@ -251,15 +269,15 @@ async def get_database_team_agent(self) -> Optional[AzureAIClient]:
                 )
                 return None
 
-            # Create client with resolved ID, preferring project_client for RAI agents
+            # Create client with resolved ID
             if self.agent_name == "RAIAgent" and self.project_client:
                 chat_client = AzureAIClient(
-                    project_client=self.project_client,
+                    project_endpoint=self.project_endpoint,
                     agent_id=resolved,
                     credential=self.creds,
                 )
                 self.logger.info(
-                    "RAI.AgentReuseSuccess: Created AzureAIClient via Projects SDK (id=%s)",
+                    "RAI.AgentReuseSuccess: Created AzureAIClient (id=%s)",
                     resolved,
                 )
             else:
@@ -284,33 +302,36 @@ async def get_database_team_agent(self) -> Optional[AzureAIClient]:
     async def save_database_team_agent(self) -> None:
         """Save current team agent to database (only if truly new or changed)."""
         try:
-            if self._agent.id is None:
-                self.logger.error("Cannot save database team agent: agent_id is None")
+            if self._agent is None or self._agent.id is None:
+                self.logger.error("Cannot save database team agent: agent or agent_id is None")
                 return
 
+            # Use the agent ID from ChatAgent (set during creation)
+            agent_id = self._agent.id
+
             # Check if stored ID matches current ID
             stored_id = await get_database_team_agent_id(
                 self.memory_store, self.team_config, self.agent_name
             )
-            if stored_id == self._agent.chat_client.agent_id:
+            if stored_id == agent_id:
                 self.logger.info(
-                    "RAI reuse: id unchanged (id=%s); skip save.", self._agent.id
+                    "RAI reuse: id unchanged (id=%s); skip save.", agent_id
                 )
                 return
 
             currentAgent = CurrentTeamAgent(
                 team_id=self.team_config.team_id,
                 team_name=self.team_config.name,
                 agent_name=self.agent_name,
-                agent_foundry_id=self._agent.chat_client.agent_id,
+                agent_foundry_id=agent_id,
                 agent_description=self.agent_description,
                 agent_instructions=self.agent_instructions,
             )
             await self.memory_store.add_team_agent(currentAgent)
             self.logger.info(
                 "Saved team agent to database (agent_name=%s, id=%s)",
                 self.agent_name,
-                self._agent.id,
+                agent_id,
             )
 
         except Exception as ex:

src/backend/v4/magentic_agents/foundry_agent.py

@@ -7,7 +7,12 @@
                              Role)
 from agent_framework_azure_ai import \
     AzureAIClient  # Provided by agent_framework
-from azure.ai.projects.models import ConnectionType
+from azure.ai.projects.models import (
+    PromptAgentDefinition,
+    AzureAISearchAgentTool,
+    AzureAISearchToolResource,
+    AISearchIndexResource,
+)
 from common.config.app_config import config
 from common.database.database_base import DatabaseBase
 from common.models.messages_af import TeamConfiguration
@@ -65,19 +70,23 @@ def __init__(
         self._use_azure_search = self._is_azure_search_requested()
         self.use_reasoning = use_reasoning
 
-        # Placeholder for server-created Azure AI agent id (if Azure Search path)
+        # Placeholder for server-created Azure AI agent id/version (if Azure Search path)
         self._azure_server_agent_id: Optional[str] = None
+        self._azure_server_agent_version: Optional[str] = None
 
     # -------------------------
     # Mode detection
     # -------------------------
     def _is_azure_search_requested(self) -> bool:
         """Determine if Azure AI Search raw tool path should be used."""
+        print(f"[DEBUG _is_azure_search_requested] Agent={self.agent_name}, search={self.search}")
         if not self.search:
+            print(f"[DEBUG _is_azure_search_requested] Agent={self.agent_name}: No search config, returning False")
             return False
         # Minimal heuristic: presence of required attributes
 
         has_index = hasattr(self.search, "index_name") and bool(self.search.index_name)
+        print(f"[DEBUG _is_azure_search_requested] Agent={self.agent_name}: has_index={has_index}, index_name={getattr(self.search, 'index_name', None)}")
         if has_index:
             self.logger.info(
                 "Azure AI Search requested (connection_id=%s, index=%s).",
@@ -113,14 +122,17 @@ async def _collect_tools(self) -> List:
     # -------------------------
     async def _create_azure_search_enabled_client(self, chatClient=None) -> Optional[AzureAIClient]:
         """
-        Create a server-side Azure AI agent with Azure AI Search raw tool.
+        Create a server-side Azure AI agent with Azure AI Search tool using create_version.
+
+        This uses the AIProjectClient.agents.create_version() approach with:
+        - PromptAgentDefinition for agent configuration
+        - AzureAISearchAgentTool with AzureAISearchToolResource for search capability
+        - AISearchIndexResource for index configuration with project_connection_id
 
         Requirements:
-          - An Azure AI Project Connection (type=AZURE_AI_SEARCH) that contains either:
-              a) API key + endpoint, OR
-              b) Managed Identity (RBAC enabled on the Search service with Search Service Contributor + Search Index Data Reader).
+          - An Azure AI Project Connection for Azure AI Search
           - search_config.index_name must exist in the Search service.
-
+          - search_config.connection_name should match the AI Project connection name
 
         Returns:
             AzureAIClient | None
@@ -134,92 +146,106 @@ async def _create_azure_search_enabled_client(self, chatClient=None) -> Optional
             self.logger.error("Search configuration missing.")
             return None
 
-        desired_connection_name = getattr(self.search, "connection_name", None)
+        # Get connection name - this is used as project_connection_id in create_version
+        connection_name = getattr(self.search, "connection_name", None)
+        if not connection_name:
+            # Fallback to environment variable
+            connection_name = config.AZURE_AI_SEARCH_CONNECTION_NAME
+            self.logger.info("Using connection_name from environment: %s", connection_name)
+
         index_name = getattr(self.search, "index_name", "")
         query_type = getattr(self.search, "search_query_type", "simple")
+        top_k = getattr(self.search, "top_k", 5)
 
         if not index_name:
             self.logger.error(
                 "index_name not provided in search_config; aborting Azure Search path."
             )
             return None
 
-        resolved_connection_id = None
-
-        try:
-            async for connection in self.project_client.connections.list():
-                if connection.type == ConnectionType.AZURE_AI_SEARCH:
-
-                    if (
-                        desired_connection_name
-                        and connection.name == desired_connection_name
-                    ):
-                        resolved_connection_id = connection.id
-                        break
-                    # Fallback: if no specific connection requested and none resolved yet, take the first
-                    if not desired_connection_name and not resolved_connection_id:
-                        resolved_connection_id = connection.id
-                        # Do not break yet; we log but allow chance to find a name match later. If not, this stays.
-
-            if not resolved_connection_id:
-                self.logger.error(
-                    "No Azure AI Search connection resolved. " "connection_name=%s",
-                    desired_connection_name,
-                )
-            #  return None
-
-            self.logger.info(
-                "Using Azure AI Search connection (id=%s, requested_name=%s).",
-                resolved_connection_id,
-                desired_connection_name,
+        if not connection_name:
+            self.logger.error(
+                "connection_name not provided; aborting Azure Search path."
             )
-        except Exception as ex:
-            self.logger.error("Failed to enumerate connections: %s", ex)
             return None
 
-        # Create agent with raw tool
+        self.logger.info(
+            "Creating Azure AI Search agent with create_version: connection_name=%s, index=%s, query_type=%s, top_k=%s",
+            connection_name,
+            index_name,
+            query_type,
+            top_k,
+        )
+
+        # Create agent using create_version with PromptAgentDefinition and AzureAISearchAgentTool
+        # This approach matches the Knowledge Mining Solution Accelerator pattern
         try:
-            azure_agent = await self.client.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."
+            enhanced_instructions = (
+                f"{self.agent_instructions} "
+                "Always use the Azure AI Search tool and configured index for knowledge retrieval."
+            )
+            
+            print(f"[DEBUG] Creating agent '{self.agent_name}' with instructions (first 200 chars): {enhanced_instructions[:200]}...")
+            print(f"[DEBUG] Agent model: {self.model_deployment_name}")
+            print(f"[DEBUG] Search config: connection={connection_name}, index={index_name}, query_type={query_type}, top_k={top_k}")
+
+            azure_agent = await self.project_client.agents.create_version(
+                agent_name=self.agent_name,
+                definition=PromptAgentDefinition(
+                    model=self.model_deployment_name,
+                    instructions=enhanced_instructions,
+                    tools=[
+                        AzureAISearchAgentTool(
+                            azure_ai_search=AzureAISearchToolResource(
+                                indexes=[
+                                    AISearchIndexResource(
+                                        project_connection_id=connection_name,
+                                        index_name=index_name,
+                                        query_type=query_type,
+                                        top_k=top_k,
+                                    )
+                                ]
+                            )
+                        )
+                    ],
                 ),
-                tools=[{"type": "azure_ai_search"}],
-                tool_resources={
-                    "azure_ai_search": {
-                        "indexes": [
-                            {
-                                "index_connection_id": resolved_connection_id,
-                                "index_name": index_name,
-                                "query_type": query_type,
-                            }
-                        ]
-                    }
-                },
             )
+
             self._azure_server_agent_id = azure_agent.id
+            self._azure_server_agent_version = azure_agent.version
             self.logger.info(
-                "Created Azure server agent with Azure AI Search tool (agent_id=%s, index=%s, query_type=%s).",
+                "Created Azure AI Search agent via create_version (name=%s, id=%s, version=%s, connection=%s, index=%s, query_type=%s, top_k=%s).",
+                azure_agent.name,
                 azure_agent.id,
+                azure_agent.version,
+                connection_name,
                 index_name,
                 query_type,
+                top_k,
             )
+            print(f"[DEBUG] Created agent via create_version: name={azure_agent.name}, id={azure_agent.id}, version={azure_agent.version}")
+            print(f"[DEBUG] Agent definition: {azure_agent.definition}")
+            print(f"[DEBUG] Agent instructions from definition: {getattr(azure_agent.definition, 'instructions', 'N/A')}")
 
+            # Wrap in AzureAIClient using agent_name and agent_version (NOT agent_id)
+            # AzureAIClient constructor: agent_name, agent_version, project_endpoint, credential
             chat_client = AzureAIClient(
-                project_client=self.project_client,
-                agent_id=azure_agent.id,
+                project_endpoint=self.project_endpoint,
+                agent_name=azure_agent.name,
+                agent_version=azure_agent.version,  # Use the specific version we just created
                 credential=self.creds,
             )
             return chat_client
+
         except Exception as ex:
             self.logger.error(
-                "Failed to create Azure Search enabled agent (connection_id=%s, index=%s): %s",
-                resolved_connection_id,
+                "Failed to create Azure Search enabled agent via create_version (connection=%s, index=%s): %s",
+                connection_name,
                 index_name,
                 ex,
             )
+            import traceback
+            traceback.print_exc()
             return None
 
     # -------------------------

src/backend/v4/magentic_agents/models/agent_models.py

@@ -43,6 +43,8 @@ class SearchConfig:
     connection_name: str | None = None
     endpoint: str | None = None
     index_name: str | None = None
+    search_query_type: str = "simple"  # Options: "simple", "vector_simple", "vector", "semantic", "hybrid"
+    top_k: int = 5  # Number of results to return
 
     @classmethod
     def from_env(cls, index_name: str) -> "SearchConfig":
@@ -58,5 +60,7 @@ def from_env(cls, index_name: str) -> "SearchConfig":
         return cls(
             connection_name=connection_name,
             endpoint=endpoint,
-            index_name=index_name
+            index_name=index_name,
+            search_query_type="simple",  # Use simple query type (keyword search)
+            top_k=5
         )

src/backend/v4/orchestration/orchestration_manager.py

@@ -231,17 +231,22 @@ async def get_current_or_new_orchestration(
         team_config: TeamConfiguration,
         team_switched: bool,
         team_service: TeamService = None,
+        force_rebuild: bool = False,
     ):
         """
         Return an existing workflow for the user or create a new one if:
           - None exists
           - Team switched flag is True
+          - force_rebuild is True (for new tasks after workflow completion)
         """
         current = orchestration_config.get_current_orchestration(user_id)
-        if current is None or team_switched:
-            if current is not None and team_switched:
+        needs_rebuild = current is None or team_switched or force_rebuild
+        
+        if needs_rebuild:
+            if current is not None and (team_switched or force_rebuild):
+                reason = "team switched" if team_switched else "force rebuild for new task"
                 cls.logger.info(
-                    "Team switched, closing previous agents for user '%s'", user_id
+                    "Rebuilding orchestration for user '%s' (reason: %s)", user_id, reason
                 )
                 # Close prior agents (same logic as old version)
                 for agent in getattr(current, "_participants", {}).values():