Skip to content

Migrate Search tools to new tool design #2954

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
alzimmermsft wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Migrate Search tools to new tool design #2954

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
3 changes: 3 additions & 0 deletions servers/Azure.Mcp.Server/changelog-entries/1782416793886.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
changes:
- section: "Breaking Changes"
description: "Removed unused parameters from Search tools."
3 changes: 1 addition & 2 deletions tools/Azure.Mcp.Tools.Search/src/Azure.Mcp.Tools.Search.csproj
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
<EmbeddedResource Include="**\Resources\*.json" />
</ItemGroup>
<ItemGroup>
<ProjectReference Include="..\..\..\core\Azure.Mcp.Core\src\Azure.Mcp.Core.csproj" />
<ProjectReference Include="$(RepoRoot)core\Azure.Mcp.Core\src\Azure.Mcp.Core.csproj" />
</ItemGroup>
<ItemGroup />
<ItemGroup>
Expand All @@ -16,6 +16,5 @@
<PackageReference Include="Azure.Search.Documents" />
<PackageReference Include="Microsoft.Extensions.DependencyInjection" />
<PackageReference Include="ModelContextProtocol" />
<PackageReference Include="System.CommandLine" />
</ItemGroup>
</Project>
34 changes: 5 additions & 29 deletions tools/Azure.Mcp.Tools.Search/src/Commands/Index/IndexGetCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,11 @@
// Licensed under the MIT License.

using Azure.Mcp.Tools.Search.Models;
using Azure.Mcp.Tools.Search.Options;
using Azure.Mcp.Tools.Search.Options.Index;
using Azure.Mcp.Tools.Search.Services;
using Microsoft.Extensions.Logging;
using Microsoft.Mcp.Core.Commands;
using Microsoft.Mcp.Core.Extensions;
using Microsoft.Mcp.Core.Models.Command;
using Microsoft.Mcp.Core.Models.Option;

namespace Azure.Mcp.Tools.Search.Commands.Index;

Expand All @@ -28,39 +25,18 @@ namespace Azure.Mcp.Tools.Search.Commands.Index;
ReadOnly = true,
Secret = false,
LocalRequired = false)]
public sealed class IndexGetCommand(ILogger<IndexGetCommand> logger, ISearchService searchService) : GlobalCommand<IndexGetOptions>()
public sealed class IndexGetCommand(ILogger<IndexGetCommand> logger, ISearchService searchService)
: AuthenticatedCommand<IndexGetOptions, IndexGetCommand.IndexGetCommandResult>
{
private readonly ILogger<IndexGetCommand> _logger = logger;
private readonly ISearchService _searchService = searchService;

protected override void RegisterOptions(Command command)
public override async Task<CommandResponse> ExecuteAsync(CommandContext context, IndexGetOptions options, CancellationToken cancellationToken)
{
base.RegisterOptions(command);
command.Options.Add(SearchOptionDefinitions.Service);
command.Options.Add(SearchOptionDefinitions.Index.AsOptional());
}

protected override IndexGetOptions BindOptions(ParseResult parseResult)
{
var options = base.BindOptions(parseResult);
options.Service = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.Service.Name);
options.Index = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.Index.Name);
return options;
}

public override async Task<CommandResponse> ExecuteAsync(CommandContext context, ParseResult parseResult, CancellationToken cancellationToken)
{
if (!Validate(parseResult.CommandResult, context.Response).IsValid)
{
return context.Response;
}

var options = BindOptions(parseResult);

try
{
var indexes = await _searchService.GetIndexDetails(
options.Service!,
options.Service,
options.Index,
options.RetryPolicy,
cancellationToken);
Expand All @@ -76,5 +52,5 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
return context.Response;
}

internal sealed record IndexGetCommandResult(List<IndexInfo> Indexes);
public sealed record IndexGetCommandResult(List<IndexInfo> Indexes);
}
38 changes: 7 additions & 31 deletions tools/Azure.Mcp.Tools.Search/src/Commands/Index/IndexQueryCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,11 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using Azure.Mcp.Tools.Search.Options;
using System.Text.Json;
using Azure.Mcp.Tools.Search.Options.Index;
using Azure.Mcp.Tools.Search.Services;
using Microsoft.Extensions.Logging;
using Microsoft.Mcp.Core.Commands;
using Microsoft.Mcp.Core.Extensions;
using Microsoft.Mcp.Core.Models.Command;

namespace Azure.Mcp.Tools.Search.Commands.Index;
Expand All @@ -25,43 +24,20 @@ namespace Azure.Mcp.Tools.Search.Commands.Index;
ReadOnly = true,
Secret = false,
LocalRequired = false)]
public sealed class IndexQueryCommand(ILogger<IndexQueryCommand> logger, ISearchService searchService) : GlobalCommand<IndexQueryOptions>()
public sealed class IndexQueryCommand(ILogger<IndexQueryCommand> logger, ISearchService searchService)
: AuthenticatedCommand<IndexQueryOptions, List<JsonElement>>
{
private readonly ILogger<IndexQueryCommand> _logger = logger;
private readonly ISearchService _searchService = searchService;

protected override void RegisterOptions(Command command)
public override async Task<CommandResponse> ExecuteAsync(CommandContext context, IndexQueryOptions options, CancellationToken cancellationToken)
{
base.RegisterOptions(command);
command.Options.Add(SearchOptionDefinitions.Service);
command.Options.Add(SearchOptionDefinitions.Index);
command.Options.Add(SearchOptionDefinitions.Query);
}

protected override IndexQueryOptions BindOptions(ParseResult parseResult)
{
var options = base.BindOptions(parseResult);
options.Service = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.Service.Name);
options.Index = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.Index.Name);
options.Query = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.Query.Name);
return options;
}

public override async Task<CommandResponse> ExecuteAsync(CommandContext context, ParseResult parseResult, CancellationToken cancellationToken)
{
if (!Validate(parseResult.CommandResult, context.Response).IsValid)
{
return context.Response;
}

var options = BindOptions(parseResult);

try
{
var results = await _searchService.QueryIndex(
options.Service!,
options.Index!,
options.Query!,
options.Service,
options.Index,
options.Query,
options.RetryPolicy,
cancellationToken);

Expand Down
34 changes: 5 additions & 29 deletions tools/Azure.Mcp.Tools.Search/src/Commands/Knowledge/KnowledgeBaseGetCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,11 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using Azure.Mcp.Tools.Search.Options;
using Azure.Mcp.Tools.Search.Options.Knowledge;
using Azure.Mcp.Tools.Search.Services;
using Microsoft.Extensions.Logging;
using Microsoft.Mcp.Core.Commands;
using Microsoft.Mcp.Core.Extensions;
using Microsoft.Mcp.Core.Models.Command;
using Microsoft.Mcp.Core.Models.Option;

namespace Azure.Mcp.Tools.Search.Commands.Knowledge;

Expand All @@ -30,38 +27,17 @@ Gets the details of Azure AI Search knowledge bases. Knowledge bases encapsulate
ReadOnly = true,
Secret = false,
LocalRequired = false)]
public sealed class KnowledgeBaseGetCommand(ILogger<KnowledgeBaseGetCommand> logger, ISearchService searchService) : GlobalCommand<KnowledgeBaseGetOptions>()
public sealed class KnowledgeBaseGetCommand(ILogger<KnowledgeBaseGetCommand> logger, ISearchService searchService)
: AuthenticatedCommand<KnowledgeBaseGetOptions, KnowledgeBaseGetCommand.KnowledgeBaseGetCommandResult>
{
private readonly ILogger<KnowledgeBaseGetCommand> _logger = logger;
private readonly ISearchService _searchService = searchService;

protected override void RegisterOptions(Command command)
public override async Task<CommandResponse> ExecuteAsync(CommandContext context, KnowledgeBaseGetOptions options, CancellationToken cancellationToken)
{
base.RegisterOptions(command);
command.Options.Add(SearchOptionDefinitions.Service);
command.Options.Add(SearchOptionDefinitions.KnowledgeBase.AsOptional());
}

protected override KnowledgeBaseGetOptions BindOptions(ParseResult parseResult)
{
var options = base.BindOptions(parseResult);
options.Service = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.Service.Name);
options.KnowledgeBase = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.KnowledgeBase.Name);
return options;
}

public override async Task<CommandResponse> ExecuteAsync(CommandContext context, ParseResult parseResult, CancellationToken cancellationToken)
{
if (!Validate(parseResult.CommandResult, context.Response).IsValid)
{
return context.Response;
}

var options = BindOptions(parseResult);

try
{
var bases = await _searchService.ListKnowledgeBases(options.Service!, options.KnowledgeBase, options.RetryPolicy, cancellationToken);
var bases = await _searchService.ListKnowledgeBases(options.Service, options.KnowledgeBase, options.RetryPolicy, cancellationToken);
context.Response.Results = ResponseResult.Create(new(bases ?? []), SearchJsonContext.Default.KnowledgeBaseGetCommandResult);
}
catch (Exception ex)
Expand All @@ -73,5 +49,5 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,
return context.Response;
}

internal sealed record KnowledgeBaseGetCommandResult(List<Models.KnowledgeBaseInfo> KnowledgeBases);
public sealed record KnowledgeBaseGetCommandResult(List<Models.KnowledgeBaseInfo> KnowledgeBases);
}
82 changes: 27 additions & 55 deletions tools/Azure.Mcp.Tools.Search/src/Commands/Knowledge/KnowledgeBaseRetrieveCommand.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,11 @@
// Licensed under the MIT License.

using System.Net;
using Azure.Mcp.Core.Commands;
using Azure.Mcp.Tools.Search.Options;
using Azure.Mcp.Tools.Search.Options.Knowledge;
using Azure.Mcp.Tools.Search.Services;
using Microsoft.Extensions.Logging;
using Microsoft.Mcp.Core.Commands;
using Microsoft.Mcp.Core.Extensions;
using Microsoft.Mcp.Core.Models.Command;
using Microsoft.Mcp.Core.Models.Option;

namespace Azure.Mcp.Tools.Search.Commands.Knowledge;

Expand All @@ -34,70 +30,46 @@ namespace Azure.Mcp.Tools.Search.Commands.Knowledge;
ReadOnly = true,
Secret = false,
LocalRequired = false)]
public sealed class KnowledgeBaseRetrieveCommand(ILogger<KnowledgeBaseRetrieveCommand> logger, ISearchService searchService) : GlobalCommand<KnowledgeBaseRetrieveOptions>()
public sealed class KnowledgeBaseRetrieveCommand(ILogger<KnowledgeBaseRetrieveCommand> logger, ISearchService searchService)
: AuthenticatedCommand<KnowledgeBaseRetrieveOptions, KnowledgeBaseRetrieveCommand.KnowledgeBaseRetrieveCommandResult>
{
private readonly ILogger<KnowledgeBaseRetrieveCommand> _logger = logger;
private readonly ISearchService _searchService = searchService;

protected override void RegisterOptions(Command command)
public override void ValidateOptions(KnowledgeBaseRetrieveOptions options, ValidationResult validationResult)
{
base.RegisterOptions(command);
command.Options.Add(SearchOptionDefinitions.Service);
command.Options.Add(SearchOptionDefinitions.KnowledgeBase);
command.Options.Add(SearchOptionDefinitions.KnowledgeQuery.AsOptional());
command.Options.Add(SearchOptionDefinitions.Messages.AsOptional());
command.Validators.Add(commandResult =>
base.ValidateOptions(options, validationResult);

if (string.IsNullOrEmpty(options.Query) && (options.Messages == null || options.Messages.Length == 0))
{
var query = commandResult.GetValueOrDefault<string>(SearchOptionDefinitions.KnowledgeQuery.Name);
var messages = commandResult.GetValueOrDefault<string[]>(SearchOptionDefinitions.Messages.Name) ?? [];
if (string.IsNullOrEmpty(query) && messages.Length == 0)
{
commandResult.AddError("Either --query or at least one --messages entry must be provided.");
}
else if (!string.IsNullOrEmpty(query) && messages.Length > 0)
{
commandResult.AddError("Specifying both --query and --messages is not allowed.");
}
validationResult.Errors.Add("Either --query or at least one --messages entry must be provided.");
}
else if (!string.IsNullOrEmpty(options.Query) && options.Messages is { Length: > 0 })
{
validationResult.Errors.Add("Specifying both --query and --messages is not allowed.");
}

if (messages.Length > 0)
if (options.Messages is { Length: > 0 })
{
foreach ((var index, var message) in options.Messages.Index())
{
foreach ((var index, var message) in messages.Index())
try
Comment thread
alzimmermsft marked this conversation as resolved.
{
try
{
ParseMessage(message);
}
catch (ArgumentException ex)
{
commandResult.AddError($"Message {index}: {ex.Message}");
continue;
}
ParseMessage(message);
}
catch (ArgumentException ex)
{
validationResult.Errors.Add($"Message {index}: {ex.Message}");
continue;
}
}
});
}

protected override KnowledgeBaseRetrieveOptions BindOptions(ParseResult parseResult)
{
var options = base.BindOptions(parseResult);
options.Service = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.Service.Name);
options.KnowledgeBase = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.KnowledgeBase.Name);
options.Query = parseResult.GetValueOrDefault<string>(SearchOptionDefinitions.KnowledgeQuery.Name);
options.Messages = parseResult.GetValueOrDefault<string[]>(SearchOptionDefinitions.Messages.Name) ?? [];
return options;
}
}

public override async Task<CommandResponse> ExecuteAsync(CommandContext context, ParseResult parseResult, CancellationToken cancellationToken)
public override async Task<CommandResponse> ExecuteAsync(CommandContext context, KnowledgeBaseRetrieveOptions options, CancellationToken cancellationToken)
{
if (!Validate(parseResult.CommandResult, context.Response).IsValid)
{
return context.Response;
}

var options = BindOptions(parseResult);

List<(string role, string message)>? parsedMessages = null;
if (options.Messages.Length > 0)
if (options.Messages is { Length: > 0 })
{
try
{
Expand All @@ -113,7 +85,7 @@ public override async Task<CommandResponse> ExecuteAsync(CommandContext context,

try
{
var result = await _searchService.RetrieveFromKnowledgeBase(options.Service!, options.KnowledgeBase!, options.Query, parsedMessages, options.RetryPolicy, cancellationToken);
var result = await _searchService.RetrieveFromKnowledgeBase(options.Service, options.KnowledgeBase, options.Query, parsedMessages, options.RetryPolicy, cancellationToken);
context.Response.Results = ResponseResult.Create(new(result), SearchJsonContext.Default.KnowledgeBaseRetrieveCommandResult);
}
catch (Exception ex)
Expand Down Expand Up @@ -145,5 +117,5 @@ internal static (string role, string content) ParseMessage(string message)
return (role, content);
}

internal sealed record KnowledgeBaseRetrieveCommandResult(string RetrievalResult);
public sealed record KnowledgeBaseRetrieveCommandResult(string RetrievalResult);
}
Loading