Skip to content

Commit c814856

Browse files
scottleppscottlepp
authored
support multiple results formats (#11)
1 parent 35d6f32 commit c814856

1 file changed

Lines changed: 187 additions & 51 deletions

File tree

‎internal/handlers/loki.go‎

Lines changed: 187 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,10 @@ func NewLokiQueryTool() mcp.Tool {
121121
mcp.WithString("org",
122122
mcp.Description(fmt.Sprintf("Organization ID for the query (default: %s from %s env var)", orgID, EnvLokiOrgID)),
123123
),
124+
mcp.WithString("format",
125+
mcp.Description("Output format: raw, json, or text (default: raw)"),
126+
mcp.DefaultString("raw"),
127+
),
124128
)
125129
}
126130

@@ -195,6 +199,12 @@ func HandleLokiQuery(ctx context.Context, request mcp.CallToolRequest) (*mcp.Cal
195199
limit = int(limitVal)
196200
}
197201

202+
// Extract format parameter
203+
format := "raw" // default
204+
if formatArg, ok := args["format"].(string); ok && formatArg != "" {
205+
format = formatArg
206+
}
207+
198208
// Build query URL
199209
queryURL, err := buildLokiQueryURL(lokiURL, queryString, start, end, limit)
200210
if err != nil {
@@ -208,7 +218,7 @@ func HandleLokiQuery(ctx context.Context, request mcp.CallToolRequest) (*mcp.Cal
208218
}
209219

210220
// Format results
211-
formattedResult, err := formatLokiResults(result)
221+
formattedResult, err := formatLokiResults(result, format)
212222
if err != nil {
213223
return nil, fmt.Errorf("failed to format results: %v", err)
214224
}
@@ -356,50 +366,102 @@ func executeLokiQuery(ctx context.Context, queryURL string, username, password,
356366
}
357367

358368
// formatLokiResults formats the Loki query results into a readable string
359-
func formatLokiResults(result *LokiResult) (string, error) {
369+
func formatLokiResults(result *LokiResult, format string) (string, error) {
360370
if len(result.Data.Result) == 0 {
361-
return "No logs found matching the query", nil
371+
switch format {
372+
case "json":
373+
return "{\"message\": \"No logs found matching the query\"}", nil
374+
default:
375+
return "No logs found matching the query", nil
376+
}
362377
}
363378

364-
var output string
365-
output = fmt.Sprintf("Found %d streams:\n\n", len(result.Data.Result))
379+
switch format {
380+
case "json":
381+
// Return raw JSON response
382+
jsonBytes, err := json.MarshalIndent(result, "", " ")
383+
if err != nil {
384+
return "", fmt.Errorf("failed to marshal JSON: %v", err)
385+
}
386+
return string(jsonBytes), nil
387+
388+
case "raw":
389+
// Return raw log lines with timestamps and labels in simple format
390+
var output string
391+
for _, entry := range result.Data.Result {
392+
// Build labels string
393+
var labels string
394+
if len(entry.Stream) > 0 {
395+
labelParts := make([]string, 0, len(entry.Stream))
396+
for k, v := range entry.Stream {
397+
labelParts = append(labelParts, fmt.Sprintf("%s=%s", k, v))
398+
}
399+
labels = "{" + strings.Join(labelParts, ",") + "} "
400+
}
366401

367-
for i, entry := range result.Data.Result {
368-
// Format stream labels
369-
streamInfo := "Stream "
370-
if len(entry.Stream) > 0 {
371-
streamInfo += "("
372-
first := true
373-
for k, v := range entry.Stream {
374-
if !first {
375-
streamInfo += ", "
402+
for _, val := range entry.Values {
403+
if len(val) >= 2 {
404+
// Parse timestamp and convert to readable format
405+
ts, err := strconv.ParseFloat(val[0], 64)
406+
var timestamp string
407+
if err == nil {
408+
// Convert to time - Loki returns timestamps in nanoseconds
409+
t := time.Unix(0, int64(ts))
410+
timestamp = t.Format(time.RFC3339)
411+
} else {
412+
timestamp = val[0]
413+
}
414+
415+
output += fmt.Sprintf("%s %s%s\n", timestamp, labels, val[1])
376416
}
377-
streamInfo += fmt.Sprintf("%s=%s", k, v)
378-
first = false
379417
}
380-
streamInfo += ")"
381418
}
419+
return output, nil
420+
421+
case "text":
422+
// Return formatted text with timestamps and stream info (original behavior)
423+
var output string
424+
output = fmt.Sprintf("Found %d streams:\n\n", len(result.Data.Result))
425+
426+
for i, entry := range result.Data.Result {
427+
// Format stream labels
428+
streamInfo := "Stream "
429+
if len(entry.Stream) > 0 {
430+
streamInfo += "("
431+
first := true
432+
for k, v := range entry.Stream {
433+
if !first {
434+
streamInfo += ", "
435+
}
436+
streamInfo += fmt.Sprintf("%s=%s", k, v)
437+
first = false
438+
}
439+
streamInfo += ")"
440+
}
382441

383-
output += fmt.Sprintf("%s %d:\n", streamInfo, i+1)
384-
385-
// Format log entries
386-
for _, val := range entry.Values {
387-
if len(val) >= 2 {
388-
// Parse timestamp
389-
ts, err := strconv.ParseFloat(val[0], 64)
390-
if err == nil {
391-
// Convert to time - Loki returns timestamps in nanoseconds already
392-
timestamp := time.Unix(0, int64(ts))
393-
output += fmt.Sprintf("[%s] %s\n", timestamp.Format(time.RFC3339), val[1])
394-
} else {
395-
output += fmt.Sprintf("[%s] %s\n", val[0], val[1])
442+
output += fmt.Sprintf("%s %d:\n", streamInfo, i+1)
443+
444+
// Format log entries
445+
for _, val := range entry.Values {
446+
if len(val) >= 2 {
447+
// Parse timestamp
448+
ts, err := strconv.ParseFloat(val[0], 64)
449+
if err == nil {
450+
// Convert to time - Loki returns timestamps in nanoseconds already
451+
timestamp := time.Unix(0, int64(ts))
452+
output += fmt.Sprintf("[%s] %s\n", timestamp.Format(time.RFC3339), val[1])
453+
} else {
454+
output += fmt.Sprintf("[%s] %s\n", val[0], val[1])
455+
}
396456
}
397457
}
458+
output += "\n"
398459
}
399-
output += "\n"
400-
}
460+
return output, nil
401461

402-
return output, nil
462+
default:
463+
return "", fmt.Errorf("unsupported format: %s. Supported formats: raw, json, text", format)
464+
}
403465
}
404466

405467
// NewLokiLabelNamesTool creates and returns a tool for getting all label names from Grafana Loki
@@ -440,6 +502,10 @@ func NewLokiLabelNamesTool() mcp.Tool {
440502
mcp.WithString("org",
441503
mcp.Description(fmt.Sprintf("Organization ID for the query (default: %s from %s env var)", orgID, EnvLokiOrgID)),
442504
),
505+
mcp.WithString("format",
506+
mcp.Description("Output format: raw, json, or text (default: raw)"),
507+
mcp.DefaultString("raw"),
508+
),
443509
)
444510
}
445511

@@ -485,6 +551,10 @@ func NewLokiLabelValuesTool() mcp.Tool {
485551
mcp.WithString("org",
486552
mcp.Description(fmt.Sprintf("Organization ID for the query (default: %s from %s env var)", orgID, EnvLokiOrgID)),
487553
),
554+
mcp.WithString("format",
555+
mcp.Description("Output format: raw, json, or text (default: raw)"),
556+
mcp.DefaultString("raw"),
557+
),
488558
)
489559
}
490560

@@ -549,6 +619,12 @@ func HandleLokiLabelNames(ctx context.Context, request mcp.CallToolRequest) (*mc
549619
end = endTime.Unix()
550620
}
551621

622+
// Extract format parameter
623+
format := "raw" // default
624+
if formatArg, ok := args["format"].(string); ok && formatArg != "" {
625+
format = formatArg
626+
}
627+
552628
// Build labels URL
553629
labelsURL, err := buildLokiLabelsURL(lokiURL, start, end)
554630
if err != nil {
@@ -562,7 +638,7 @@ func HandleLokiLabelNames(ctx context.Context, request mcp.CallToolRequest) (*mc
562638
}
563639

564640
// Format results
565-
formattedResult, err := formatLokiLabelsResults(result)
641+
formattedResult, err := formatLokiLabelsResults(result, format)
566642
if err != nil {
567643
return nil, fmt.Errorf("failed to format results: %v", err)
568644
}
@@ -632,6 +708,12 @@ func HandleLokiLabelValues(ctx context.Context, request mcp.CallToolRequest) (*m
632708
end = endTime.Unix()
633709
}
634710

711+
// Extract format parameter
712+
format := "raw" // default
713+
if formatArg, ok := args["format"].(string); ok && formatArg != "" {
714+
format = formatArg
715+
}
716+
635717
// Build label values URL
636718
labelValuesURL, err := buildLokiLabelValuesURL(lokiURL, labelName, start, end)
637719
if err != nil {
@@ -645,7 +727,7 @@ func HandleLokiLabelValues(ctx context.Context, request mcp.CallToolRequest) (*m
645727
}
646728

647729
// Format results
648-
formattedResult, err := formatLokiLabelValuesResults(labelName, result)
730+
formattedResult, err := formatLokiLabelValuesResults(labelName, result, format)
649731
if err != nil {
650732
return nil, fmt.Errorf("failed to format results: %v", err)
651733
}
@@ -824,33 +906,87 @@ func executeLokiLabelValuesQuery(ctx context.Context, queryURL string, username,
824906
}
825907

826908
// formatLokiLabelsResults formats the Loki labels results into a readable string
827-
func formatLokiLabelsResults(result *LokiLabelsResult) (string, error) {
909+
func formatLokiLabelsResults(result *LokiLabelsResult, format string) (string, error) {
828910
if len(result.Data) == 0 {
829-
return "No labels found", nil
911+
switch format {
912+
case "json":
913+
return "{\"message\": \"No labels found\"}", nil
914+
default:
915+
return "No labels found", nil
916+
}
830917
}
831918

832-
var output string
833-
output = fmt.Sprintf("Found %d labels:\n\n", len(result.Data))
919+
switch format {
920+
case "json":
921+
// Return raw JSON response
922+
jsonBytes, err := json.MarshalIndent(result, "", " ")
923+
if err != nil {
924+
return "", fmt.Errorf("failed to marshal JSON: %v", err)
925+
}
926+
return string(jsonBytes), nil
834927

835-
for i, label := range result.Data {
836-
output += fmt.Sprintf("%d. %s\n", i+1, label)
837-
}
928+
case "raw":
929+
// Return raw label names only, one per line
930+
var output string
931+
for _, label := range result.Data {
932+
output += label + "\n"
933+
}
934+
return output, nil
838935

839-
return output, nil
936+
case "text":
937+
// Return formatted text with numbering (original behavior)
938+
var output string
939+
output = fmt.Sprintf("Found %d labels:\n\n", len(result.Data))
940+
941+
for i, label := range result.Data {
942+
output += fmt.Sprintf("%d. %s\n", i+1, label)
943+
}
944+
return output, nil
945+
946+
default:
947+
return "", fmt.Errorf("unsupported format: %s. Supported formats: raw, json, text", format)
948+
}
840949
}
841950

842951
// formatLokiLabelValuesResults formats the Loki label values results into a readable string
843-
func formatLokiLabelValuesResults(labelName string, result *LokiLabelValuesResult) (string, error) {
952+
func formatLokiLabelValuesResults(labelName string, result *LokiLabelValuesResult, format string) (string, error) {
844953
if len(result.Data) == 0 {
845-
return fmt.Sprintf("No values found for label '%s'", labelName), nil
954+
switch format {
955+
case "json":
956+
return fmt.Sprintf("{\"message\": \"No values found for label '%s'\"}", labelName), nil
957+
default:
958+
return fmt.Sprintf("No values found for label '%s'", labelName), nil
959+
}
846960
}
847961

848-
var output string
849-
output = fmt.Sprintf("Found %d values for label '%s':\n\n", len(result.Data), labelName)
962+
switch format {
963+
case "json":
964+
// Return raw JSON response
965+
jsonBytes, err := json.MarshalIndent(result, "", " ")
966+
if err != nil {
967+
return "", fmt.Errorf("failed to marshal JSON: %v", err)
968+
}
969+
return string(jsonBytes), nil
850970

851-
for i, value := range result.Data {
852-
output += fmt.Sprintf("%d. %s\n", i+1, value)
853-
}
971+
case "raw":
972+
// Return raw label values only, one per line
973+
var output string
974+
for _, value := range result.Data {
975+
output += value + "\n"
976+
}
977+
return output, nil
854978

855-
return output, nil
979+
case "text":
980+
// Return formatted text with numbering (original behavior)
981+
var output string
982+
output = fmt.Sprintf("Found %d values for label '%s':\n\n", len(result.Data), labelName)
983+
984+
for i, value := range result.Data {
985+
output += fmt.Sprintf("%d. %s\n", i+1, value)
986+
}
987+
return output, nil
988+
989+
default:
990+
return "", fmt.Errorf("unsupported format: %s. Supported formats: raw, json, text", format)
991+
}
856992
}

0 commit comments

Comments
 (0)