Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Directory.Packages.props
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@
<PackageVersion Include="Azure.Monitor.OpenTelemetry.AspNetCore" Version="1.3.0" />
<PackageVersion Include="Microsoft.Extensions.Azure" Version="1.11.0" />
<PackageVersion Include="Microsoft.AspNetCore.SignalR.Client" Version="9.0.9" />
<PackageVersion Include="Azure.Search.Documents" Version="11.7.0-beta.7" />
<PackageVersion Include="Azure.Search.Documents" Version="11.8.0-beta.1" />
<PackageVersion Include="Microsoft.NET.Test.Sdk" Version="17.13.0" />
<PackageVersion Include="NSubstitute" Version="5.3.0" />
<PackageVersion Include="NSubstitute.Analyzers.CSharp" Version="1.0.17" />
Expand Down
Original file line number Diff line number Diff line change
@@ -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"
54 changes: 44 additions & 10 deletions tools/Azure.Mcp.Tools.Search/src/Services/SearchService.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -207,14 +207,14 @@ public async Task<List<KnowledgeBaseInfo>> 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))
Expand Down Expand Up @@ -246,16 +246,19 @@ public async Task<string> 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);

Expand All @@ -268,6 +271,37 @@ public async Task<string> 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<string> ProcessRetrieveResponse(Stream responseStream)
{
using var jsonDoc = await JsonDocument.ParseAsync(responseStream);
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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<KnowledgeRetrievalSemanticIntent>(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<KnowledgeRetrievalSemanticIntent>(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<KnowledgeRetrievalSemanticIntent>(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<KnowledgeRetrievalSemanticIntent>(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<KnowledgeBaseMessageTextContent>(message.Content.Single());
Assert.Equal("Show results", content.Text);
},
message =>
{
Assert.Equal("assistant", message.Role);
var content = Assert.IsType<KnowledgeBaseMessageTextContent>(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<KnowledgeBaseMessageTextContent>(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<KnowledgeBaseMessageTextContent>(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<KnowledgeBaseMessageTextContent>(message.Content.Single());
Assert.Equal("Explain search", content.Text);
}

private static async Task<string> InvokeProcessRetrieveResponse(string json)
{
using var stream = new MemoryStream(Encoding.UTF8.GetBytes(json));
Expand Down
Loading