1+ package prometheus
2+
3+ import (
4+ "context"
5+ "fmt"
6+ "time"
7+
8+ "github.com/google/jsonschema-go/jsonschema"
9+ p8s_api "github.com/prometheus/client_golang/api"
10+ "github.com/prometheus/client_golang/api/prometheus/v1"
11+ "k8s.io/utils/ptr"
12+
13+ "github.com/containers/kubernetes-mcp-server/pkg/api"
14+ "github.com/containers/kubernetes-mcp-server/pkg/config"
15+ "github.com/containers/kubernetes-mcp-server/pkg/kubernetes"
16+ )
17+
18+ // NewToolset returns a new toolset for Prometheus.
19+ func NewToolset (cfg * config.PrometheusConfig ) (api.Toolset , error ) {
20+ if cfg == nil || cfg .URL == "" {
21+ return & disabledToolset {}, nil
22+ }
23+
24+ client , err := p8s_api .NewClient (p8s_api.Config {
25+ Address : cfg .URL ,
26+ })
27+ if err != nil {
28+ return nil , fmt .Errorf ("failed to create prometheus client: %w" , err )
29+ }
30+
31+ return & prometheusToolset {
32+ api : v1 .NewAPI (client ),
33+ }, nil
34+ }
35+
36+ type prometheusToolset struct {
37+ api v1.API
38+ }
39+
40+ func (t * prometheusToolset ) GetName () string {
41+ return "prometheus"
42+ }
43+
44+ func (t * prometheusToolset ) GetDescription () string {
45+ return "Tools for interacting with Prometheus"
46+ }
47+
48+ func (t * prometheusToolset ) GetTools (_ kubernetes.Openshift ) []api.ServerTool {
49+ return []api.ServerTool {
50+ {
51+ Tool : api.Tool {
52+ Name : "prometheus.runQuery" ,
53+ Description : "Run a PromQL query." ,
54+ InputSchema : & jsonschema.Schema {
55+ Type : "object" ,
56+ Properties : map [string ]* jsonschema.Schema {
57+ "query" : {
58+ Type : "string" ,
59+ Description : "The PromQL query to run." ,
60+ },
61+ },
62+ Required : []string {"query" },
63+ },
64+ Annotations : api.ToolAnnotations {
65+ ReadOnlyHint : ptr .To (true ),
66+ },
67+ },
68+ Handler : runQueryHandler (t .api ),
69+ },
70+ }
71+ }
72+
73+ func runQueryHandler (p8sAPI v1.API ) api.ToolHandlerFunc {
74+ return func (params api.ToolHandlerParams ) (* api.ToolCallResult , error ) {
75+ query , _ := params .GetArguments ()["query" ].(string )
76+
77+ ctx , cancel := context .WithTimeout (context .Background (), 10 * time .Second )
78+ defer cancel ()
79+
80+ result , warnings , err := p8sAPI .Query (ctx , query , time .Now ())
81+ if err != nil {
82+ return api .NewToolCallResult ("" , fmt .Errorf ("failed to run query: %w" , err )), nil
83+ }
84+ if len (warnings ) > 0 {
85+ // Not treating warnings as errors for now
86+ }
87+
88+ return api .NewToolCallResult (result .String (), nil ), nil
89+ }
90+ }
91+
92+ type disabledToolset struct {}
93+
94+ func (t * disabledToolset ) GetName () string {
95+ return "prometheus"
96+ }
97+
98+ func (t * disabledToolset ) GetDescription () string {
99+ return "Prometheus toolset is disabled. Please configure it in the server settings."
100+ }
101+
102+ func (t * disabledToolset ) GetTools (_ kubernetes.Openshift ) []api.ServerTool {
103+ return nil
104+ }
0 commit comments