diff --git a/Directory.Packages.props b/Directory.Packages.props index 3b1a3b4f39..3eac616895 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -91,7 +91,7 @@ - + diff --git a/servers/Azure.Mcp.Server/changelog-entries/1766191976865.yaml b/servers/Azure.Mcp.Server/changelog-entries/1766191976865.yaml new file mode 100644 index 0000000000..4804649e80 --- /dev/null +++ b/servers/Azure.Mcp.Server/changelog-entries/1766191976865.yaml @@ -0,0 +1,3 @@ +changes: + - section: "Bugs Fixed" + description: "Support for new versions of Azure AI Search knowledge bases and those set to 'minimal' reasoning effort" \ No newline at end of file diff --git a/tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs b/tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs index 3f1d33a833..78fdf1944c 100644 --- a/tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs +++ b/tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs @@ -12,10 +12,10 @@ using Azure.Mcp.Tools.Search.Models; using Azure.ResourceManager.Search; using Azure.Search.Documents; -using Azure.Search.Documents.Agents; -using Azure.Search.Documents.Agents.Models; using Azure.Search.Documents.Indexes; using Azure.Search.Documents.Indexes.Models; +using Azure.Search.Documents.KnowledgeBases; +using Azure.Search.Documents.KnowledgeBases.Models; using Azure.Search.Documents.Models; namespace Azure.Mcp.Tools.Search.Services; @@ -207,14 +207,14 @@ public async Task> ListKnowledgeBases( if (string.IsNullOrEmpty(knowledgeBaseName)) { - await foreach (var knowledgeBase in searchClient.GetKnowledgeAgentsAsync(cancellationToken: cancellationToken)) + await foreach (var knowledgeBase in searchClient.GetKnowledgeBasesAsync(cancellationToken: cancellationToken)) { bases.Add(new KnowledgeBaseInfo(knowledgeBase.Name, knowledgeBase.Description, [.. knowledgeBase.KnowledgeSources.Select(ks => ks.Name)])); } } else { - var result = await searchClient.GetKnowledgeAgentAsync(knowledgeBaseName, cancellationToken: cancellationToken); + var result = await searchClient.GetKnowledgeBaseAsync(knowledgeBaseName, cancellationToken: cancellationToken); if (result?.Value != null) { if (result.Value.Name.Equals(knowledgeBaseName, StringComparison.OrdinalIgnoreCase)) @@ -246,16 +246,19 @@ public async Task RetrieveFromKnowledgeBase( { var searchClient = await GetSearchIndexClient(serviceName, retryPolicy, cancellationToken); + var knowledgeBase = await searchClient.GetKnowledgeBaseAsync(baseName, cancellationToken: cancellationToken); + if (knowledgeBase?.Value == null) + { + throw new InvalidOperationException($"Knowledge base '{baseName}' not found in service '{serviceName}'."); + } + var clientOptions = AddDefaultPolicies(new SearchClientOptions()); clientOptions.Transport = new HttpClientTransport(TenantService.GetClient()); ConfigureRetryPolicy(clientOptions, retryPolicy); - var knowledgeBaseClient = new KnowledgeAgentRetrievalClient(searchClient.Endpoint, baseName, await GetCredential(cancellationToken: cancellationToken), clientOptions); - - var request = new KnowledgeAgentRetrievalRequest( - messages != null ? - messages.Select(m => new KnowledgeAgentMessage([new KnowledgeAgentMessageTextContent(m.message)]) { Role = m.role }) : - [new KnowledgeAgentMessage([new KnowledgeAgentMessageTextContent(query)]) { Role = "user" }]); + var knowledgeBaseClient = new KnowledgeBaseRetrievalClient(searchClient.Endpoint, baseName, await GetCredential(cancellationToken: cancellationToken), clientOptions); + var useMinimalReasoning = knowledgeBase.Value.RetrievalReasoningEffort is KnowledgeRetrievalMinimalReasoningEffort; + var request = BuildKnowledgeBaseRetrievalRequest(useMinimalReasoning, query, messages); var results = await knowledgeBaseClient.RetrieveAsync(request, cancellationToken: cancellationToken); @@ -268,6 +271,37 @@ public async Task RetrieveFromKnowledgeBase( } } + internal static KnowledgeBaseRetrievalRequest BuildKnowledgeBaseRetrievalRequest( + bool useMinimalReasoning, + string? query, + IEnumerable<(string role, string message)>? messages) + { + var request = new KnowledgeBaseRetrievalRequest(); + + if (useMinimalReasoning) + { + var intent = messages != null && messages.Any() + ? string.Join("\n", messages.Select(m => m.message)) + : query ?? string.Empty; + + request.Intents.Add(new KnowledgeRetrievalSemanticIntent(intent)); + return request; + } + + if (messages != null && messages.Any()) + { + foreach ((string role, string message) in messages) + { + request.Messages.Add(new KnowledgeBaseMessage([new KnowledgeBaseMessageTextContent(message)]) { Role = role }); + } + + return request; + } + + request.Messages.Add(new KnowledgeBaseMessage([new KnowledgeBaseMessageTextContent(query ?? string.Empty)]) { Role = "user" }); + return request; + } + internal static async Task ProcessRetrieveResponse(Stream responseStream) { using var jsonDoc = await JsonDocument.ParseAsync(responseStream); diff --git a/tools/Azure.Mcp.Tools.Search/tests/Azure.Mcp.Tools.Search.UnitTests/Service/SearchServiceTests.cs b/tools/Azure.Mcp.Tools.Search/tests/Azure.Mcp.Tools.Search.UnitTests/Service/SearchServiceTests.cs index 3e6dcc009c..0d9cd7378c 100644 --- a/tools/Azure.Mcp.Tools.Search/tests/Azure.Mcp.Tools.Search.UnitTests/Service/SearchServiceTests.cs +++ b/tools/Azure.Mcp.Tools.Search/tests/Azure.Mcp.Tools.Search.UnitTests/Service/SearchServiceTests.cs @@ -1,8 +1,11 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Collections.Generic; +using System.Linq; using System.Text; using Azure.Mcp.Tools.Search.Services; +using Azure.Search.Documents.KnowledgeBases.Models; using Xunit; namespace Azure.Mcp.Tools.Search.UnitTests.Service; @@ -100,6 +103,124 @@ public async Task ProcessRetrieveResponse_ReturnsEmptyObject_WhenNoExpectedPrope Assert.DoesNotContain("\"other\"", result); } + [Fact] + public void BuildKnowledgeBaseRetrievalRequest_UsesIntentForMinimalReasoning_WhenMessagesProvided() + { + var messages = new List<(string role, string message)> + { + ("user", "Hello"), + ("assistant", "How can I help?") + }; + + var request = SearchService.BuildKnowledgeBaseRetrievalRequest(true, null, messages); + + var intent = Assert.IsType(request.Intents.Single()); + Assert.Equal("Hello\nHow can I help?", intent.Search); + Assert.Empty(request.Messages); + } + + [Fact] + public void BuildKnowledgeBaseRetrievalRequest_UsesIntentForMinimalReasoning_WhenOnlyQueryProvided() + { + var request = SearchService.BuildKnowledgeBaseRetrievalRequest(true, "What is search?", null); + + var intent = Assert.IsType(request.Intents.Single()); + Assert.Equal("What is search?", intent.Search); + Assert.Empty(request.Messages); + } + + [Fact] + public void BuildKnowledgeBaseRetrievalRequest_UsesEmptyIntentForMinimalReasoning_WhenMessagesEmpty() + { + var messages = new List<(string role, string message)>(); + + var request = SearchService.BuildKnowledgeBaseRetrievalRequest(true, null, messages); + + var intent = Assert.IsType(request.Intents.Single()); + Assert.Equal(string.Empty, intent.Search); + Assert.Empty(request.Messages); + } + + [Fact] + public void BuildKnowledgeBaseRetrievalRequest_UsesQueryIntentForMinimalReasoning_WhenMessagesEmptyAndQueryProvided() + { + var messages = new List<(string role, string message)>(); + + var request = SearchService.BuildKnowledgeBaseRetrievalRequest(true, "Explain search", messages); + + var intent = Assert.IsType(request.Intents.Single()); + Assert.Equal("Explain search", intent.Search); + Assert.Empty(request.Messages); + } + + [Fact] + public void BuildKnowledgeBaseRetrievalRequest_UsesMessagesForStandardReasoning_WhenMessagesProvided() + { + var messages = new List<(string role, string message)> + { + ("user", "Show results"), + ("assistant", "Sure") + }; + + var request = SearchService.BuildKnowledgeBaseRetrievalRequest(false, null, messages); + + Assert.Empty(request.Intents); + Assert.Collection( + request.Messages, + message => + { + Assert.Equal("user", message.Role); + var content = Assert.IsType(message.Content.Single()); + Assert.Equal("Show results", content.Text); + }, + message => + { + Assert.Equal("assistant", message.Role); + var content = Assert.IsType(message.Content.Single()); + Assert.Equal("Sure", content.Text); + }); + } + + [Fact] + public void BuildKnowledgeBaseRetrievalRequest_UsesQueryMessageForStandardReasoning_WhenNoMessagesProvided() + { + var request = SearchService.BuildKnowledgeBaseRetrievalRequest(false, "Explain indexing", null); + + Assert.Empty(request.Intents); + var message = Assert.Single(request.Messages); + Assert.Equal("user", message.Role); + var content = Assert.IsType(message.Content.Single()); + Assert.Equal("Explain indexing", content.Text); + } + + [Fact] + public void BuildKnowledgeBaseRetrievalRequest_UsesEmptyQueryMessageForStandardReasoning_WhenMessagesEmpty() + { + var messages = new List<(string role, string message)>(); + + var request = SearchService.BuildKnowledgeBaseRetrievalRequest(false, null, messages); + + Assert.Empty(request.Intents); + var message = Assert.Single(request.Messages); + Assert.Equal("user", message.Role); + var content = Assert.IsType(message.Content.Single()); + Assert.Equal(string.Empty, content.Text); + } + + [Fact] + public void BuildKnowledgeBaseRetrievalRequest_UsesQueryMessageForStandardReasoning_WhenMessagesEmptyAndQueryProvided() + { + var messages = new List<(string role, string message)>(); + + var request = SearchService.BuildKnowledgeBaseRetrievalRequest(false, "Explain search", messages); + + Assert.Empty(request.Intents); + var message = Assert.Single(request.Messages); + Assert.Equal("user", message.Role); + var content = Assert.IsType(message.Content.Single()); + Assert.Equal("Explain search", content.Text); + } + private static async Task InvokeProcessRetrieveResponse(string json) { using var stream = new MemoryStream(Encoding.UTF8.GetBytes(json));