refactor: Agent framework v2 changes.#794
Conversation
… scripts and configuration
…emove unused agent factory classes, streamline ChatService and HistoryService, and enhance SQLTool for better database interaction. Update API routes and configuration management for improved clarity and performance. Upgrade dependencies in requirements.txt for compatibility with new Azure SDK versions.
…gin, ChatService, HistoryService, and App - Removed test cases for deprecated agent factories in test_app.py. - Updated test cases in test_chat_with_data_plugin.py to reflect changes in the ChatWithDataPlugin implementation. - Refactored test cases in test_chat_service.py to remove dependencies on mock requests and streamline the setup. - Enhanced test_generate_title in test_history_service.py to utilize the new v2 agent framework. - Improved error handling tests in test_chat_service.py for rate limits and general exceptions. - Consolidated mock setups and assertions for clarity and maintainability.
…lation and access
…security in ChatService and HistoryService
…edge-Mining-Solution-Accelerator into km-agentframework-v2
…edge-Mining-Solution-Accelerator into km-agentframework-v2
…steps in Deployment Guide
…edge-Mining-Solution-Accelerator into km-agentframework-v2
…edge-Mining-Solution-Accelerator into km-agentframework-v2
…edge-Mining-Solution-Accelerator into km-agentframework-v2
…edge-Mining-Solution-Accelerator into km-agentframework-v2
…edge-Mining-Solution-Accelerator into km-agentframework-v2
- Updated 00_create_sample_data_files.py to improve CSV and JSON export functions, ensuring better error handling and code readability. - Modified 01_create_search_index.py to include additional whitespace for consistency. - Enhanced 03_cu_process_data_text.py by implementing asynchronous processing for embeddings and agent creation, improving performance and scalability. - Updated 04_cu_process_custom_data.py to streamline the search index creation process and improve error handling. - Adjusted requirements.txt to include new agent framework dependencies and ensure compatibility. - Enhanced process_sample_data.sh and run_create_index_scripts.sh to support new solution_name parameter for better configuration management.
…processing scripts
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 39 out of 39 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| projectEndpoint=$(extract_value "azureAiAgentEndpoint" "azurE_AI_AGENT_ENDPOINT") | ||
| solutionName=$(extract_value "solutionName" "solutioN_NAME") | ||
| gptModelName=$(extract_value "azureAIAgentModelDeploymentName" "azurE_AI_AGENT_MODEL_DEPLOYMENT_NAME") | ||
| aiFoundryResourceId=$(extract_value "aiFoundryResourceId" "aI_FOUNDRY_RESOURCE_ID") | ||
| apiAppName=$(extract_value "apiAppName" "apI_APP_NAME") | ||
| aiSearchConnectionName=$(extract_value "azureAISearchConnectionName" "azurE_AI_SEARCH_CONNECTION_NAME") | ||
| aiSearchIndex=$(extract_value "azureAISearchIndex" "azurE_AI_SEARCH_INDEX") |
There was a problem hiding this comment.
The fallback keys passed to extract_value (e.g., "azurE_AI_AGENT_ENDPOINT", "solutioN_NAME") don’t match the actual output names used elsewhere in this repo (typically all-caps like AZURE_AI_AGENT_ENDPOINT). Since the grep is case-sensitive, these fallbacks will never be found and get_values_from_az_deployment may fail unexpectedly. Use the real output keys (e.g., AZURE_AI_AGENT_ENDPOINT, SOLUTION_NAME, etc.) or normalize output parsing (e.g., with jq).
| projectEndpoint=$(extract_value "azureAiAgentEndpoint" "azurE_AI_AGENT_ENDPOINT") | |
| solutionName=$(extract_value "solutionName" "solutioN_NAME") | |
| gptModelName=$(extract_value "azureAIAgentModelDeploymentName" "azurE_AI_AGENT_MODEL_DEPLOYMENT_NAME") | |
| aiFoundryResourceId=$(extract_value "aiFoundryResourceId" "aI_FOUNDRY_RESOURCE_ID") | |
| apiAppName=$(extract_value "apiAppName" "apI_APP_NAME") | |
| aiSearchConnectionName=$(extract_value "azureAISearchConnectionName" "azurE_AI_SEARCH_CONNECTION_NAME") | |
| aiSearchIndex=$(extract_value "azureAISearchIndex" "azurE_AI_SEARCH_INDEX") | |
| projectEndpoint=$(extract_value "azureAiAgentEndpoint" "AZURE_AI_AGENT_ENDPOINT") | |
| solutionName=$(extract_value "solutionName" "SOLUTION_NAME") | |
| gptModelName=$(extract_value "azureAIAgentModelDeploymentName" "AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME") | |
| aiFoundryResourceId=$(extract_value "aiFoundryResourceId" "AZURE_AI_FOUNDRY_RESOURCE_ID") | |
| apiAppName=$(extract_value "apiAppName" "API_APP_NAME") | |
| aiSearchConnectionName=$(extract_value "azureAISearchConnectionName" "AZURE_AI_SEARCH_CONNECTION_NAME") | |
| aiSearchIndex=$(extract_value "azureAISearchIndex" "AZURE_AI_SEARCH_INDEX") |
| thread_cache = ExpCache(maxsize=1000, ttl=3600.0) | ||
| return thread_cache | ||
|
|
||
| async def stream_openai_text(self, conversation_id: str, query: str) -> StreamingResponse: |
There was a problem hiding this comment.
stream_openai_text is annotated to return StreamingResponse, but the function uses yield and therefore returns an async generator. This mismatch can confuse callers/type-checkers (and StreamingResponse is constructed in the route instead). Update the return annotation to an async-iterator type (or refactor to build/return a StreamingResponse here).
…prove title extraction; update fetch_azure_search_content_endpoint to return structured content with title.
…ng and streamline title extraction
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 40 out of 40 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| citation_list.append(f"{{\"url\": \"{url}\", \"title\": \"{title}\"}}") | ||
| yield ", \"citations\": [" + ",".join(citation_list) + "]}" |
There was a problem hiding this comment.
citation_list.append(f"{\"url\": \"{url}\", \"title\": \"{title}\"}") builds a JSON fragment without escaping url/title. If either contains quotes, backslashes, or newlines, parseCitationFromMessage() on the frontend will fail to JSON.parse the citations block. Build citations as Python dicts and serialize with json.dumps (or at least JSON-escape the interpolated values) before yielding the citations string.
| citation_list.append(f"{{\"url\": \"{url}\", \"title\": \"{title}\"}}") | |
| yield ", \"citations\": [" + ",".join(citation_list) + "]}" | |
| citation_list.append({"url": url, "title": title}) | |
| yield ', "citations": ' + json.dumps(citation_list) + "}" |
| async def stream_openai_text(self, conversation_id: str, query: str) -> StreamingResponse: | ||
| """ | ||
| Get a streaming text response from OpenAI. |
There was a problem hiding this comment.
The function signature advertises -> StreamingResponse, but this method is an async generator yielding string chunks (it’s iterated in stream_chat_request). Consider updating the return annotation to AsyncIterator[str]/AsyncGenerator[str, None] (or returning an actual StreamingResponse) to keep typing accurate and avoid confusing API consumers.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 42 out of 42 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| async def stream_openai_text(self, conversation_id: str, query: str) -> StreamingResponse: | ||
| """ | ||
| Get a streaming text response from OpenAI. | ||
| """ |
There was a problem hiding this comment.
stream_openai_text is annotated as returning StreamingResponse, but it is implemented as an async generator that yields chunks and is iterated with async for by stream_chat_request. This mismatch can confuse readers and break type checking. Consider updating the return annotation to an AsyncIterator[str]/AsyncGenerator[str, None] (or return an actual StreamingResponse here and move the generator out).
…onment variables; update .env.sample for Azure resource ID and configure logging level in ChatService
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 42 out of 42 changed files in this pull request and generated 7 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME=<agent-model-deployment> | ||
|
|
||
| # Agent Framework v2 Configuration (Set by deployment) | ||
| AI_FOUNDRY_RESOURCE_ID=<ai-foundry-resource-id> |
There was a problem hiding this comment.
The new env var name here (AI_FOUNDRY_RESOURCE_ID) doesn’t match the infrastructure outputs / sample env (AZURE_AI_FOUNDRY_RESOURCE_ID). This inconsistency will cause misconfiguration for local dev and contradicts the note below. Align the docs to the actual variable name used by the deployment outputs and .env.sample.
| AI_FOUNDRY_RESOURCE_ID=<ai-foundry-resource-id> | |
| AZURE_AI_FOUNDRY_RESOURCE_ID=<ai-foundry-resource-id> |
| @patch("services.chat_service.AIProjectClient") | ||
| @patch("services.chat_service.get_azure_credential_async") | ||
| @patch("services.chat_service.Config") |
There was a problem hiding this comment.
The @patch("services.chat_service.get_azure_credential_async") decorator here patches an awaited async function but will supply a non-awaitable MagicMock by default. These tests will fail when _delete_thread_async does await get_azure_credential_async(...). Use new_callable=AsyncMock for the patch (and set its return_value to the credential mock).
| def expire(self, time=None): | ||
| """Remove expired items and delete associated Azure AI threads.""" | ||
| items = super().expire(time) | ||
| for key, thread_id in items: | ||
| for key, thread_conversation_id in items: | ||
| try: | ||
| if self.agent: | ||
| thread = AzureAIAgentThread(client=self.agent.client, thread_id=thread_id) | ||
| asyncio.create_task(thread.delete()) | ||
| print(f"Thread deleted : {thread_id}") | ||
| # Create task for async deletion with proper session management | ||
| asyncio.create_task(self._delete_thread_async(thread_conversation_id)) | ||
| logger.info("Scheduled thread deletion: %s", thread_conversation_id) |
There was a problem hiding this comment.
asyncio.create_task(...) is called from synchronous cache methods (expire/popitem). If these methods ever run without a running event loop (e.g., called from a threadpool or during sync code paths), it will raise RuntimeError: no running event loop and thread cleanup won’t happen. Consider using asyncio.get_running_loop() and falling back to asyncio.run(...) / a background task queue / or logging+skipping when no loop is available.
| async def stream_openai_text(self, conversation_id: str, query: str) -> StreamingResponse: | ||
| """ | ||
| Get a streaming text response from OpenAI. |
There was a problem hiding this comment.
stream_openai_text is annotated as returning StreamingResponse, but it’s implemented as an async generator yielding string chunks. This mismatch can confuse callers/type-checkers and makes the public API of the service unclear. Update the return type to an async-iterator type (e.g., AsyncIterator[str]) or wrap it in a StreamingResponse at this layer consistently.
| seen_doc_ids.add(doc_id) | ||
|
|
||
| citation_list.append(f"{{\"url\": \"{url}\", \"title\": \"{title}\"}}") | ||
| yield ", \"citations\": [" + ",".join(citation_list) + "]}" |
There was a problem hiding this comment.
Citations are assembled via string interpolation into a JSON-like fragment. If url or title contains quotes/backslashes/newlines, it can break the downstream parser and potentially enable content injection. Prefer building citation objects and serializing with json.dumps (or at least escape url/title).
| @pytest.mark.asyncio | ||
| @patch('services.chat_service.AzureAIAgentThread') | ||
| @patch('services.chat_service.TruncationObject') | ||
| async def test_stream_openai_text_empty_query(self, mock_truncation_class, mock_thread_class, chat_service): | ||
| """Test streaming with empty query.""" | ||
| mock_response = MagicMock() | ||
| mock_response.content = "Please provide a query." | ||
| mock_response.thread.id = "thread_id" | ||
| @patch("services.chat_service.SQLTool") | ||
| @patch("services.chat_service.get_sqldb_connection") | ||
| @patch("services.chat_service.AzureAIProjectAgentProvider") | ||
| @patch("services.chat_service.AIProjectClient") | ||
| @patch("services.chat_service.get_azure_credential_async") | ||
| async def test_stream_openai_text_success( |
There was a problem hiding this comment.
These patches target async functions (get_azure_credential_async, get_sqldb_connection) but use the default MagicMock. Since the production code awaits these functions, the tests will raise TypeError when trying to await a non-awaitable mock. Use new_callable=AsyncMock (or patch with an AsyncMock explicitly) for awaited callables.
| with patch("services.history_service.get_azure_credential_async", return_value=mock_credential): | ||
| with patch("services.history_service.AIProjectClient", return_value=mock_project_client): | ||
| with patch("services.history_service.AzureAIProjectAgentProvider", return_value=mock_provider): |
There was a problem hiding this comment.
get_azure_credential_async is an async function and is awaited in HistoryService.generate_title. Patching it with return_value=... creates a non-awaitable mock, which will fail at runtime. Patch with new_callable=AsyncMock (or set AsyncMock(return_value=mock_credential)), then configure __aenter__/__aexit__ on the returned credential object.
…s; add solution name parameter and improve parameter validation
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 43 out of 43 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| async def call_topic_mapping_agent(input_text, list_of_topics): | ||
| """Use Topic Mapping Agent with Agent Framework to map topic to category.""" | ||
| async with ( | ||
| AsyncAzureCliCredential(process_timeout=30) as async_cred, | ||
| AIProjectClient(endpoint=AI_PROJECT_ENDPOINT, credential=async_cred) as project_client, | ||
| ): | ||
| # Create provider for agent management | ||
| provider = AzureAIProjectAgentProvider(project_client=project_client) | ||
|
|
||
| # Get agent using provider | ||
| agent = await provider.get_agent(name=TOPIC_MAPPING_AGENT_NAME) | ||
|
|
||
| query = f"""Find the closest topic for this text: '{input_text}' from this list of topics: {list_of_topics}""" | ||
|
|
||
| result = await agent.run(query) | ||
| return result.text.strip() | ||
|
|
||
| cursor.execute('SELECT * FROM processed_data') | ||
| rows = [tuple(row) for row in cursor.fetchall()] | ||
| column_names = [i[0] for i in cursor.description] | ||
| df_processed_data = pd.DataFrame(rows, columns=column_names) | ||
| df_processed_data = df_processed_data[df_processed_data['ConversationId'].isin(conversationIds)] | ||
|
|
||
| # Map topics using agent asynchronously | ||
| async def map_all_topics(): | ||
| """Map all topics to categories using agent.""" | ||
| for _, row in df_processed_data.iterrows(): | ||
| mined_topic_str = await call_topic_mapping_agent(row['topic'], str(mined_topics_list)) | ||
| cursor.execute("UPDATE processed_data SET mined_topic = ? WHERE ConversationId = ?", (mined_topic_str, row['ConversationId'])) |
There was a problem hiding this comment.
call_topic_mapping_agent creates a new Azure credential + AIProjectClient + provider for every single row in map_all_topics(). For custom datasets this will be very slow and may trigger throttling. Reuse one AIProjectClient/provider/agent for the whole loop (and optionally use bounded concurrency).
| async def call_topic_mapping_agent(input_text, list_of_topics): | |
| """Use Topic Mapping Agent with Agent Framework to map topic to category.""" | |
| async with ( | |
| AsyncAzureCliCredential(process_timeout=30) as async_cred, | |
| AIProjectClient(endpoint=AI_PROJECT_ENDPOINT, credential=async_cred) as project_client, | |
| ): | |
| # Create provider for agent management | |
| provider = AzureAIProjectAgentProvider(project_client=project_client) | |
| # Get agent using provider | |
| agent = await provider.get_agent(name=TOPIC_MAPPING_AGENT_NAME) | |
| query = f"""Find the closest topic for this text: '{input_text}' from this list of topics: {list_of_topics}""" | |
| result = await agent.run(query) | |
| return result.text.strip() | |
| cursor.execute('SELECT * FROM processed_data') | |
| rows = [tuple(row) for row in cursor.fetchall()] | |
| column_names = [i[0] for i in cursor.description] | |
| df_processed_data = pd.DataFrame(rows, columns=column_names) | |
| df_processed_data = df_processed_data[df_processed_data['ConversationId'].isin(conversationIds)] | |
| # Map topics using agent asynchronously | |
| async def map_all_topics(): | |
| """Map all topics to categories using agent.""" | |
| for _, row in df_processed_data.iterrows(): | |
| mined_topic_str = await call_topic_mapping_agent(row['topic'], str(mined_topics_list)) | |
| cursor.execute("UPDATE processed_data SET mined_topic = ? WHERE ConversationId = ?", (mined_topic_str, row['ConversationId'])) | |
| cursor.execute('SELECT * FROM processed_data') | |
| rows = [tuple(row) for row in cursor.fetchall()] | |
| column_names = [i[0] for i in cursor.description] | |
| df_processed_data = pd.DataFrame(rows, columns=column_names) | |
| df_processed_data = df_processed_data[df_processed_data['ConversationId'].isin(conversationIds)] | |
| # Map topics using agent asynchronously | |
| async def map_all_topics(): | |
| """Map all topics to categories using a single, reused agent instance.""" | |
| async with ( | |
| AsyncAzureCliCredential(process_timeout=30) as async_cred, | |
| AIProjectClient(endpoint=AI_PROJECT_ENDPOINT, credential=async_cred) as project_client, | |
| ): | |
| # Create provider for agent management | |
| provider = AzureAIProjectAgentProvider(project_client=project_client) | |
| # Get agent once using provider | |
| agent = await provider.get_agent(name=TOPIC_MAPPING_AGENT_NAME) | |
| for _, row in df_processed_data.iterrows(): | |
| query = ( | |
| f"Find the closest topic for this text: '{row['topic']}' " | |
| f"from this list of topics: {mined_topics_list}" | |
| ) | |
| result = await agent.run(query) | |
| mined_topic_str = result.text.strip() | |
| cursor.execute( | |
| "UPDATE processed_data SET mined_topic = ? WHERE ConversationId = ?", | |
| (mined_topic_str, row['ConversationId']), | |
| ) |
| url = request_json.get("url") | ||
|
|
||
| if not url: | ||
| return JSONResponse(content={"error": "URL is required"}, status_code=400) | ||
|
|
||
| # Get Azure AD token | ||
| config = Config() | ||
| credential = get_azure_credential(client_id=config.azure_client_id) | ||
| token = credential.get_token("https://search.azure.com/.default") | ||
| access_token = token.token | ||
|
|
||
| # Define blocking request call | ||
| def fetch_content(): | ||
| try: | ||
| response = requests.get( | ||
| url, | ||
| headers={ | ||
| "Authorization": f"Bearer {access_token}", | ||
| "Content-Type": "application/json" | ||
| }, | ||
| timeout=10 | ||
| ) | ||
| if response.status_code == 200: | ||
| data = response.json() | ||
| content = data.get("content", "") | ||
| return content | ||
| title = data.get("sourceurl", "") | ||
| return {"content": content, "title": title} |
There was a problem hiding this comment.
This endpoint makes an authenticated request to a caller-provided url using a bearer token. Without validating that url is an expected Azure AI Search endpoint, this can leak the access token to arbitrary hosts (SSRF/token exfiltration). Also, title is being populated from sourceurl, which looks like the document URL, not a human-readable title (e.g., chunk_id/title). Restrict url to the known search service domain (or build it server-side from safe inputs) and return a real title field.
| async def call_topic_mapping_agent(input_text, list_of_topics): | ||
| """Use Topic Mapping Agent with Agent Framework to map topic to category.""" | ||
| async with ( | ||
| AsyncAzureCliCredential(process_timeout=30) as async_cred, | ||
| AIProjectClient(endpoint=AI_PROJECT_ENDPOINT, credential=async_cred) as project_client, | ||
| ): | ||
| # Create provider for agent management | ||
| provider = AzureAIProjectAgentProvider(project_client=project_client) | ||
|
|
||
| # Get agent using provider | ||
| agent = await provider.get_agent(name=TOPIC_MAPPING_AGENT_NAME) | ||
|
|
||
| query = f"""Find the closest topic for this text: '{input_text}' from this list of topics: {list_of_topics}""" | ||
|
|
||
| result = await agent.run(query) | ||
| return result.text.strip() | ||
|
|
||
| cursor.execute('SELECT * FROM processed_data') | ||
| rows = [tuple(row) for row in cursor.fetchall()] | ||
| column_names = [i[0] for i in cursor.description] | ||
| df_processed_data = pd.DataFrame(rows, columns=column_names) | ||
| df_processed_data = df_processed_data[df_processed_data['ConversationId'].isin(conversationIds)] | ||
|
|
||
| # Map topics using agent asynchronously | ||
| async def map_all_topics(): | ||
| """Map all topics to categories using agent.""" | ||
| for _, row in df_processed_data.iterrows(): | ||
| mined_topic_str = await call_topic_mapping_agent(row['topic'], str(mined_topics_list)) | ||
| cursor.execute("UPDATE processed_data SET mined_topic = ? WHERE ConversationId = ?", (mined_topic_str, row['ConversationId'])) |
There was a problem hiding this comment.
call_topic_mapping_agent creates a new Azure credential + AIProjectClient + provider for every single row in map_all_topics(). For non-trivial datasets this will be extremely slow and may hit rate limits. Reuse a single AIProjectClient/provider/agent for the whole mapping loop (and optionally run mappings with bounded concurrency).
Purpose
This pull request updates the Azure deployment scripts, infrastructure templates, and documentation to streamline agent creation, clarify parameter usage, and improve naming consistency for Azure AI Search connections. The changes impact both the deployment process and the configuration outputs, making it easier to set up agents and reference required resources.
Deployment and Agent Creation Process:
run_create_agents_scripts.sh) in both PowerShell and Bash deployment hooks, as well as in the deployment guide documentation. This helps users understand when and how to create agents during setup. [1] [2] [3]Azure AI Search Connection Naming and Outputs:
aiSearchConnectionNamevariable (instead of reusingaiSearchName) for Azure AI Search connections throughout the Bicep and ARM templates. Updated resource definitions, outputs, and module parameters to use this new variable, ensuring consistent and clear naming. [1] [2] [3] [4] [5] [6] [7]API_APP_NAME,AZURE_AI_FOUNDRY_RESOURCE_ID,AGENT_NAME_CONVERSATION, andAGENT_NAME_TITLE, making these available for downstream services and documentation. [1] [2] [3]Parameter Order Consistency:
Infrastructure Template Updates:
Dependency and Output Adjustments:
These changes collectively improve the deployment experience, clarify resource configuration, and make the setup process more robust and user-friendly.
Does this introduce a breaking change?
Golden Path Validation
Deployment Validation