feat(tools): add tool to query exemplars

Signed-off-by: TJ Hoplock <[email protected]>

Author: tjhop

pkg/mcp/promapi.go

@@ -143,6 +143,20 @@ func rangeQueryApiCall(ctx context.Context, query string, start, end time.Time,
 	return string(jsonBytes), nil
 }
 
+func exemplarQueryApiCall(ctx context.Context, query string, start, end time.Time) (string, error) {
+	res, err := apiV1Client.QueryExemplars(ctx, query, start, end)
+	if err != nil {
+		return "", fmt.Errorf("error executing exemplar query: %w", err)
+	}
+
+	jsonBytes, err := json.Marshal(res)
+	if err != nil {
+		return "", fmt.Errorf("error converting query response to JSON: %w", err)
+	}
+
+	return string(jsonBytes), nil
+}
+
 func seriesApiCall(ctx context.Context, matchers []string, start, end time.Time) (string, error) {
 	result, warnings, err := apiV1Client.Series(ctx, matchers, start, end)
 	if err != nil {

pkg/mcp/server.go

@@ -27,6 +27,7 @@ func NewServer(logger *slog.Logger) *server.MCPServer {
 	// add tools
 	mcpServer.AddTool(queryTool, queryToolHandler)
 	mcpServer.AddTool(rangeQueryTool, rangeQueryToolHandler)
+	mcpServer.AddTool(exemplarQueryTool, exemplarQueryToolHandler)
 	mcpServer.AddTool(seriesTool, seriesToolHandler)
 	mcpServer.AddTool(labelNamesTool, labelNamesToolHandler)
 	mcpServer.AddTool(labelValuesTool, labelValuesToolHandler)

pkg/mcp/tools.go

@@ -45,6 +45,22 @@ var (
 		),
 	)
 
+	exemplarQueryTool = mcp.NewTool("exemplar_query",
+		mcp.WithDescription("Execute a exemplar query against the Prometheus datasource"),
+		mcp.WithString("query",
+			mcp.Required(),
+			mcp.Description("Query to be executed"),
+		),
+		mcp.WithString("start_time",
+			mcp.Description("[Optional] Start timestamp for the query to be executed at."+
+				" Must be either Unix timestamp or RFC3339. Defaults to 5m ago."),
+		),
+		mcp.WithString("end_time",
+			mcp.Description("[Optional] End timestamp for the query to be executed at."+
+				" Must be either Unix timestamp or RFC3339. Defaults to current time."),
+		),
+	)
+
 	seriesTool = mcp.NewTool("series",
 		mcp.WithDescription("Finds series by label matchers"),
 		mcp.WithArray("matchers",
@@ -227,6 +243,38 @@ func rangeQueryToolHandler(ctx context.Context, request mcp.CallToolRequest) (*m
 	return mcp.NewToolResultText(data), err
 }
 
+func exemplarQueryToolHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+	args := request.Params.Arguments
+	query, ok := args["query"].(string)
+	if !ok {
+		return nil, errors.New("query must be a string")
+	}
+
+	endTs := time.Now()
+	startTs := endTs.Add(DefaultLookbackDelta)
+
+	if argEndTime, ok := args["end_time"].(string); ok {
+		parsedEndTime, err := mcpProm.ParseTimestamp(argEndTime)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse end_time %s from args: %w", argEndTime, err)
+		}
+
+		endTs = parsedEndTime
+	}
+
+	if argStartTime, ok := args["start_time"].(string); ok {
+		parsedStartTime, err := mcpProm.ParseTimestamp(argStartTime)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse start_time %s from args: %w", argStartTime, err)
+		}
+
+		startTs = parsedStartTime
+	}
+
+	data, err := exemplarQueryApiCall(ctx, query, startTs, endTs)
+	return mcp.NewToolResultText(data), err
+}
+
 func seriesToolHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
 	args := request.Params.Arguments
 	argMatchers, ok := args["matchers"].([]any)