SearchTools.cs

// Licensed to Elasticsearch B.V under one or more agreements.
// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
// See the LICENSE file in the project root for more information

using System.ComponentModel;
using System.Diagnostics;
using System.Text.Json;
using Elastic.Documentation.Assembler.Mcp;
using Elastic.Documentation.Mcp.Remote.Responses;
using Elastic.Documentation.Mcp.Remote.Telemetry;
using Elastic.Documentation.Search;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;

namespace Elastic.Documentation.Mcp.Remote.Tools;

/// <summary>
/// MCP tools for semantic search operations on Elastic documentation.
/// </summary>
[McpServerToolType]
public class SearchTools(IFullSearchService fullSearchGateway, ILogger<SearchTools> logger)
{
	/// <summary>
	/// Performs semantic search across all Elastic documentation.
	/// </summary>
	[McpServerTool, McpToolName("search_{resource}"), Description(
		"Searches all published {docs} by meaning. " +
		"Use when the user asks about Elastic product features, needs to find existing docs pages, " +
		"verify published content, or research what documentation exists on a topic. " +
		"Returns relevant documents with AI summaries, relevance scores, and navigation context.")]
	public async Task<CallToolResult> SemanticSearch(
		[Description("The search query - can be a question or keywords")] string query,
		[Description("Page number (1-based, default: 1)")] int pageNumber = 1,
		[Description("Number of results per page (default: 10, max: 50)")] int pageSize = 10,
		[Description("Filter by product ID (e.g., 'elasticsearch', 'kibana')")] string? productFilter = null,
		[Description("Filter by navigation section (e.g., 'reference', 'getting-started')")] string? sectionFilter = null,
		CancellationToken cancellationToken = default)
	{
		var toolName = McpToolTelemetry.ResolveToolName("search_{resource}");
		using var activity = McpToolTelemetry.StartActivity(toolName);
		var payload = McpToolTelemetry.SetPayloadMetadata(activity, new Dictionary<string, object?>
		{
			["query"] = query,
			["pageNumber"] = pageNumber,
			["pageSize"] = pageSize,
			["productFilter"] = productFilter,
			["sectionFilter"] = sectionFilter
		});
		McpToolTelemetry.LogStart(logger, toolName, payload);
		var duration = Stopwatch.StartNew();
		var outcome = "failure";

		try
		{
			pageSize = Math.Clamp(pageSize, 1, 50);
			pageNumber = Math.Max(1, pageNumber);

			var request = new FullSearchRequest
			{
				Query = query,
				PageNumber = pageNumber,
				PageSize = pageSize,
				ProductFilter = productFilter != null ? [productFilter] : null,
				SectionFilter = sectionFilter != null ? [sectionFilter] : null,
				IncludeHighlighting = false
			};

			var result = await fullSearchGateway.SearchAsync(request, cancellationToken);

			var response = new SemanticSearchResponse
			{
				Query = query,
				TotalHits = result.TotalResults,
				IsSemanticQuery = result.IsSemanticQuery,
				Results = result.Results.Select(r => new SearchResultDto
				{
					Url = r.Url,
					Title = r.Title,
					Description = r.Description,
					Score = r.Score,
					AiShortSummary = r.AiShortSummary,
					NavigationSection = r.NavigationSection,
					Product = r.Product?.DisplayName,
					LastUpdated = r.LastUpdated
				}).ToList()
			};

			McpToolTelemetry.MarkSuccess(activity);
			outcome = "success";
			return McpToolResults.Ok(JsonSerializer.Serialize(response, McpJsonContext.Default.SemanticSearchResponse));
		}
		catch (OperationCanceledException)
		{
			McpToolTelemetry.MarkCancelled(activity);
			outcome = "cancelled";
			throw;
		}
		catch (Exception ex) when (ex is not OutOfMemoryException and not StackOverflowException)
		{
			McpToolTelemetry.MarkFailure(activity, ex);
			logger.LogError(ex, "SemanticSearch failed for query '{Query}'", query);
			return McpToolResults.Error(JsonSerializer.Serialize(new ErrorResponse(ex.Message), McpJsonContext.Default.ErrorResponse));
		}
		finally
		{
			duration.Stop();
			McpToolTelemetry.LogCompletion(logger, toolName, duration.ElapsedMilliseconds, outcome);
		}
	}

	/// <summary>
	/// Finds documents related to a given topic or document URL.
	/// </summary>
	[McpServerTool, McpToolName("find_related_{resource}"), Description(
		"Finds {docs} pages related to a given topic. " +
		"Use when exploring what documentation exists around a subject, building context for writing, " +
		"or discovering related content the user should be aware of.")]
	public async Task<CallToolResult> FindRelatedDocs(
		[Description("Topic or search terms to find related documents for")] string topic,
		[Description("Maximum number of related documents to return (default: 10)")] int limit = 10,
		[Description("Filter by product ID (e.g., 'elasticsearch', 'kibana')")] string? productFilter = null,
		CancellationToken cancellationToken = default)
	{
		var toolName = McpToolTelemetry.ResolveToolName("find_related_{resource}");
		using var activity = McpToolTelemetry.StartActivity(toolName);
		var payload = McpToolTelemetry.SetPayloadMetadata(activity, new Dictionary<string, object?>
		{
			["topic"] = topic,
			["limit"] = limit,
			["productFilter"] = productFilter
		});
		McpToolTelemetry.LogStart(logger, toolName, payload);
		var duration = Stopwatch.StartNew();
		var outcome = "failure";

		try
		{
			limit = Math.Clamp(limit, 1, 20);

			var request = new FullSearchRequest
			{
				Query = topic,
				PageNumber = 1,
				PageSize = limit,
				ProductFilter = productFilter != null ? [productFilter] : null,
				IncludeHighlighting = false
			};

			var result = await fullSearchGateway.SearchAsync(request, cancellationToken);

			var response = new RelatedDocsResponse
			{
				Topic = topic,
				Count = result.Results.Count,
				RelatedDocs = result.Results.Select(r => new RelatedDocDto
				{
					Url = r.Url,
					Title = r.Title,
					Description = r.Description,
					Score = r.Score,
					AiShortSummary = r.AiShortSummary,
					Product = r.Product?.DisplayName
				}).ToList()
			};

			McpToolTelemetry.MarkSuccess(activity);
			outcome = "success";
			return McpToolResults.Ok(JsonSerializer.Serialize(response, McpJsonContext.Default.RelatedDocsResponse));
		}
		catch (OperationCanceledException)
		{
			McpToolTelemetry.MarkCancelled(activity);
			outcome = "cancelled";
			throw;
		}
		catch (Exception ex) when (ex is not OutOfMemoryException and not StackOverflowException)
		{
			McpToolTelemetry.MarkFailure(activity, ex);
			logger.LogError(ex, "FindRelatedDocs failed for topic '{Topic}'", topic);
			return McpToolResults.Error(JsonSerializer.Serialize(new ErrorResponse(ex.Message), McpJsonContext.Default.ErrorResponse));
		}
		finally
		{
			duration.Stop();
			McpToolTelemetry.LogCompletion(logger, toolName, duration.ElapsedMilliseconds, outcome);
		}
	}
}