From c77ad306f6e77ed3fc1f451a7ec25d1b2e032d2e Mon Sep 17 00:00:00 2001 From: Alex Jones Date: Fri, 25 Apr 2025 09:12:48 +0100 Subject: [PATCH] feat: consolidating code duplication Signed-off-by: Alex Jones --- pkg/server/config/config.go | 37 ++++++----- pkg/server/mcp.go | 122 ++++++++++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 15 deletions(-) diff --git a/pkg/server/config/config.go b/pkg/server/config/config.go index c728f63..c9dc4cf 100644 --- a/pkg/server/config/config.go +++ b/pkg/server/config/config.go @@ -1,8 +1,9 @@ package config import ( - schemav1 "buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go/schema/v1" "context" + + schemav1 "buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go/schema/v1" "github.com/k8sgpt-ai/k8sgpt/pkg/cache" "github.com/k8sgpt-ai/k8sgpt/pkg/custom" "github.com/spf13/viper" @@ -20,19 +21,13 @@ const ( notUsedInsecure = false ) -func (h *Handler) AddConfig(ctx context.Context, i *schemav1.AddConfigRequest) (*schemav1.AddConfigResponse, error, -) { - - resp, err := h.syncIntegration(ctx, i) - if err != nil { - return resp, err - } - +// ApplyConfig applies the configuration changes from the request +func (h *Handler) ApplyConfig(ctx context.Context, i *schemav1.AddConfigRequest) error { if i.CustomAnalyzers != nil { // We need to add the custom analyzers to the viper config and save them var customAnalyzers = make([]custom.CustomAnalyzer, 0) if err := viper.UnmarshalKey("custom_analyzers", &customAnalyzers); err != nil { - return resp, err + return err } else { // If there are analyzers are already in the config we will append the ones with new names for _, ca := range i.CustomAnalyzers { @@ -56,7 +51,7 @@ func (h *Handler) AddConfig(ctx context.Context, i *schemav1.AddConfigRequest) ( // save the config viper.Set("custom_analyzers", customAnalyzers) if err := viper.WriteConfig(); err != nil { - return resp, err + return err } } } @@ -74,18 +69,30 @@ func (h *Handler) AddConfig(ctx context.Context, i *schemav1.AddConfigRequest) ( case *schemav1.Cache_InterplexCache: remoteCache, err = cache.NewCacheProvider("interplex", notUsedBucket, notUsedRegion, i.Cache.GetInterplexCache().Endpoint, notUsedStorageAcc, notUsedContainerName, notUsedProjectId, notUsedInsecure) default: - return resp, status.Error(codes.InvalidArgument, "Invalid cache configuration") + return status.Error(codes.InvalidArgument, "Invalid cache configuration") } if err != nil { - return resp, err + return err } err = cache.AddRemoteCache(remoteCache) if err != nil { - return resp, err + return err } - } + return nil +} + +func (h *Handler) AddConfig(ctx context.Context, i *schemav1.AddConfigRequest) (*schemav1.AddConfigResponse, error) { + resp, err := h.syncIntegration(ctx, i) + if err != nil { + return resp, err + } + + if err := h.ApplyConfig(ctx, i); err != nil { + return resp, err + } + return resp, nil } diff --git a/pkg/server/mcp.go b/pkg/server/mcp.go index 77a1dd3..e11ced9 100644 --- a/pkg/server/mcp.go +++ b/pkg/server/mcp.go @@ -17,10 +17,13 @@ import ( "context" "fmt" + schemav1 "buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go/schema/v1" "github.com/k8sgpt-ai/k8sgpt/pkg/ai" "github.com/k8sgpt-ai/k8sgpt/pkg/analysis" + "github.com/k8sgpt-ai/k8sgpt/pkg/server/config" mcp_golang "github.com/metoro-io/mcp-golang" "github.com/metoro-io/mcp-golang/transport/stdio" + "github.com/spf13/viper" ) // MCPServer represents an MCP server for k8sgpt @@ -81,6 +84,11 @@ func (s *MCPServer) Start() error { return fmt.Errorf("failed to register cluster-info tool: %v", err) } + // Register config tool + if err := s.server.RegisterTool("config", "Configure K8sGPT settings", s.handleConfig); err != nil { + return fmt.Errorf("failed to register config tool: %v", err) + } + // Register resources if err := s.registerResources(); err != nil { return fmt.Errorf("failed to register resources: %v", err) @@ -126,8 +134,59 @@ type ClusterInfoResponse struct { Info string `json:"info"` } +// ConfigRequest represents the input parameters for the config tool +type ConfigRequest struct { + CustomAnalyzers []struct { + Name string `json:"name"` + Connection struct { + Url string `json:"url"` + Port int `json:"port"` + } `json:"connection"` + } `json:"customAnalyzers,omitempty"` + Cache struct { + Type string `json:"type"` + // S3 specific fields + BucketName string `json:"bucketName,omitempty"` + Region string `json:"region,omitempty"` + Endpoint string `json:"endpoint,omitempty"` + Insecure bool `json:"insecure,omitempty"` + // Azure specific fields + StorageAccount string `json:"storageAccount,omitempty"` + ContainerName string `json:"containerName,omitempty"` + // GCS specific fields + ProjectId string `json:"projectId,omitempty"` + } `json:"cache,omitempty"` +} + +// ConfigResponse represents the output of the config tool +type ConfigResponse struct { + Status string `json:"status"` +} + // handleAnalyze handles the analyze tool func (s *MCPServer) handleAnalyze(ctx context.Context, request *AnalyzeRequest) (*mcp_golang.ToolResponse, error) { + // Get stored configuration + var configAI ai.AIConfiguration + if err := viper.UnmarshalKey("ai", &configAI); err != nil { + return mcp_golang.NewToolResponse(mcp_golang.NewTextContent(fmt.Sprintf("Failed to load AI configuration: %v", err))), nil + } + + // Use stored configuration if not specified in request + if request.Backend == "" { + if configAI.DefaultProvider != "" { + request.Backend = configAI.DefaultProvider + } else if len(configAI.Providers) > 0 { + request.Backend = configAI.Providers[0].Name + } else { + request.Backend = "openai" // fallback default + } + } + + // Get stored filters if not specified + if len(request.Filters) == 0 { + request.Filters = viper.GetStringSlice("active_filters") + } + // Create a new analysis with the request parameters analysis, err := analysis.NewAnalysis( request.Backend, @@ -172,6 +231,69 @@ func (s *MCPServer) handleClusterInfo(ctx context.Context, request *ClusterInfoR return mcp_golang.NewToolResponse(mcp_golang.NewTextContent(info)), nil } +// handleConfig handles the config tool +func (s *MCPServer) handleConfig(ctx context.Context, request *ConfigRequest) (*mcp_golang.ToolResponse, error) { + // Create a new config handler + handler := &config.Handler{} + + // Convert request to AddConfigRequest + addConfigReq := &schemav1.AddConfigRequest{ + CustomAnalyzers: make([]*schemav1.CustomAnalyzer, 0), + } + + // Add custom analyzers if present + if len(request.CustomAnalyzers) > 0 { + for _, ca := range request.CustomAnalyzers { + addConfigReq.CustomAnalyzers = append(addConfigReq.CustomAnalyzers, &schemav1.CustomAnalyzer{ + Name: ca.Name, + Connection: &schemav1.Connection{ + Url: ca.Connection.Url, + Port: fmt.Sprintf("%d", ca.Connection.Port), + }, + }) + } + } + + // Add cache configuration if present + if request.Cache.Type != "" { + cacheConfig := &schemav1.Cache{} + switch request.Cache.Type { + case "s3": + cacheConfig.CacheType = &schemav1.Cache_S3Cache{ + S3Cache: &schemav1.S3Cache{ + BucketName: request.Cache.BucketName, + Region: request.Cache.Region, + Endpoint: request.Cache.Endpoint, + Insecure: request.Cache.Insecure, + }, + } + case "azure": + cacheConfig.CacheType = &schemav1.Cache_AzureCache{ + AzureCache: &schemav1.AzureCache{ + StorageAccount: request.Cache.StorageAccount, + ContainerName: request.Cache.ContainerName, + }, + } + case "gcs": + cacheConfig.CacheType = &schemav1.Cache_GcsCache{ + GcsCache: &schemav1.GCSCache{ + BucketName: request.Cache.BucketName, + Region: request.Cache.Region, + ProjectId: request.Cache.ProjectId, + }, + } + } + addConfigReq.Cache = cacheConfig + } + + // Apply the configuration using the shared function + if err := handler.ApplyConfig(ctx, addConfigReq); err != nil { + return mcp_golang.NewToolResponse(mcp_golang.NewTextContent(fmt.Sprintf("Failed to add config: %v", err))), nil + } + + return mcp_golang.NewToolResponse(mcp_golang.NewTextContent("Successfully added configuration")), nil +} + // registerPrompts registers the prompts for the MCP server func (s *MCPServer) registerPrompts() error { // Register any prompts needed for the MCP server