feat: adding a query mode for the schednex scheduler (#1257)

Signed-off-by: AlexsJones <alexsimonjones@gmail.com>

go.mod

@@ -28,9 +28,9 @@ require (
 require github.com/adrg/xdg v0.4.0
 
 require (
-	buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 v2.22.0-20240807134501-ea98c104104d.1
+	buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 v2.22.0-20240920204244-7a91c8620515.1
-	buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go v1.5.1-20240807134501-ea98c104104d.1
+	buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go v1.5.1-20240920204244-7a91c8620515.1
-	buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240807134501-ea98c104104d.2
+	buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240920204244-7a91c8620515.2
 	cloud.google.com/go/storage v1.43.0
 	cloud.google.com/go/vertexai v0.7.1
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.6.0
@@ -57,6 +57,7 @@ require (
 	atomicgo.dev/cursor v0.2.0 // indirect
 	atomicgo.dev/keyboard v0.2.9 // indirect
 	atomicgo.dev/schedule v0.1.0 // indirect
+	buf.build/gen/go/k8sgpt-ai/k8sgpt/connectrpc/go v1.17.0-20240920204244-7a91c8620515.1 // indirect
 	cloud.google.com/go v0.115.0 // indirect
 	cloud.google.com/go/ai v0.3.5-0.20240409161017-ce55ad694f21 // indirect
 	cloud.google.com/go/aiplatform v1.68.0 // indirect

go.sum

@@ -6,12 +6,30 @@ atomicgo.dev/keyboard v0.2.9 h1:tOsIid3nlPLZ3lwgG8KZMp/SFmr7P0ssEN5JUsm78K8=
 atomicgo.dev/keyboard v0.2.9/go.mod h1:BC4w9g00XkxH/f1HXhW2sXmJFOCWbKn9xrOunSFtExQ=
 atomicgo.dev/schedule v0.1.0 h1:nTthAbhZS5YZmgYbb2+DH8uQIZcTlIrd4eYr3UQxEjs=
 atomicgo.dev/schedule v0.1.0/go.mod h1:xeUa3oAkiuHYh8bKiQBRojqAMq3PXXbJujjb0hw8pEU=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/connectrpc/go v1.17.0-20240807134501-ea98c104104d.1 h1:m0wXoD4sYVtjBIm5Foxjlul5GTemO6d3vHnRgRae7co=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/connectrpc/go v1.17.0-20240807134501-ea98c104104d.1/go.mod h1:37t+BcJa0Wfn61FR99AwyHueonghoWpx7NuAkwJL50o=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/connectrpc/go v1.17.0-20240920203222-070d6b308f49.1 h1:z7KWoJGA7pGfp49KaqKX2CA73OssaimZavN19E35bTI=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/connectrpc/go v1.17.0-20240920203222-070d6b308f49.1/go.mod h1:6Pn118mP1faw8jG/EoARsHLyBr9hYPEuJLySsNAsMPo=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/connectrpc/go v1.17.0-20240920204244-7a91c8620515.1 h1:fvowVes7HyfrXBxJY7f9V5EW/Y/LNr4aZRv1GrMrjKw=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/connectrpc/go v1.17.0-20240920204244-7a91c8620515.1/go.mod h1:vY5/6MV/yEDzj4w3qjNX+qEcYwPrUQk26N6slfQJRUY=
 buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 v2.22.0-20240807134501-ea98c104104d.1 h1:H93Xb0qlnrjrw/vAgtRfvrn/doNpVLa2PC8LGXtdg+4=
 buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 v2.22.0-20240807134501-ea98c104104d.1/go.mod h1:aJlkczECoZ6cLkWvQKJnfwSfRbRAn93fHT1S/ynKA0o=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 v2.22.0-20240920203222-070d6b308f49.1 h1:3C1ibEuf2I2uzUCKc6bYHj2E2Jmd1rUZSSFBWwdaScY=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 v2.22.0-20240920203222-070d6b308f49.1/go.mod h1:xOsrq27xOxf0dStHMAZ2rgWFVUVHUPYdJsuMJtqetLc=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 v2.22.0-20240920204244-7a91c8620515.1 h1:qgcrhhBtW1KVxZ0izTXciQ/KR9RmE2V2EdkY+0pqE5g=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc-ecosystem/gateway/v2 v2.22.0-20240920204244-7a91c8620515.1/go.mod h1:qVT3VFVrkD5nFUD/KEzjLbTyLwIubUsrc/TyG0X35EU=
 buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go v1.5.1-20240807134501-ea98c104104d.1 h1:W/4o/wizyw9oBry0PRZFDH9i55u8MLxmODUSZHW7krU=
 buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go v1.5.1-20240807134501-ea98c104104d.1/go.mod h1:cJGWJGXgaTTqiPaZAuQPvkjkxQCWoAlCWte3/ZrpU8s=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go v1.5.1-20240920203222-070d6b308f49.1 h1:ukk/DWLRwDymW6NvHnx1KQOMILBtr/0ceqkd46aXtCc=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go v1.5.1-20240920203222-070d6b308f49.1/go.mod h1:VKNx/1hd1SlIkbD2qIDBvGH8AebitAGtil+RhYij/5g=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go v1.5.1-20240920204244-7a91c8620515.1 h1:zrgiD9fwDXvwXw5P+lG5O+JEKvrUxsNs2xNS+c41xOA=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go v1.5.1-20240920204244-7a91c8620515.1/go.mod h1:UWLDF7rVrPUgIhYUkufxK2MQoCt8c2L2Mr2XKGepJdI=
 buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240807134501-ea98c104104d.2 h1:D3HPkRDfG1F4/yjNYTYkVgnh6aa1Fj14tiP5sBJ0pfA=
 buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240807134501-ea98c104104d.2/go.mod h1:1wq1qVxvJkTEUQsF5/XjmhQYXYhbVoLSGhKnzS3ie54=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240920203222-070d6b308f49.2 h1:IHb6o9t61M3MY9nSSjo9RSG9o3xB6Fv2/AEc6aaF/Es=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240920203222-070d6b308f49.2/go.mod h1:1wq1qVxvJkTEUQsF5/XjmhQYXYhbVoLSGhKnzS3ie54=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240920204244-7a91c8620515.2 h1:KNevBMEFHuPv6Z1Uhp9V3UNW06xnvp6hA2Jf6HPLdas=
+buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240920204244-7a91c8620515.2/go.mod h1:1wq1qVxvJkTEUQsF5/XjmhQYXYhbVoLSGhKnzS3ie54=
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 cloud.google.com/go v0.38.0/go.mod h1:990N+gfupTy94rShfmMCWGDn0LpTmnzTp2qbd1dvSRU=

pkg/server/query/handler.go

@@ -0,0 +1,7 @@
+package query
+
+import rpc "buf.build/gen/go/k8sgpt-ai/k8sgpt/grpc/go/schema/v1/schemav1grpc"
+
+type Handler struct {
+	rpc.UnimplementedServerQueryServiceServer
+}

pkg/server/query/query.go

@@ -0,0 +1,27 @@
+package query
+
+import (
+	schemav1 "buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go/schema/v1"
+	"context"
+	"github.com/k8sgpt-ai/k8sgpt/pkg/ai"
+)
+
+func (h *Handler) Query(ctx context.Context, i *schemav1.QueryRequest) (
+	*schemav1.QueryResponse,
+	error,
+) {
+	aiClient := ai.NewClient(i.Backend)
+	defer aiClient.Close()
+
+	resp, err := aiClient.GetCompletion(ctx, i.Query)
+	var errMessage string = ""
+	if err != nil {
+		errMessage = err.Error()
+	}
+	return &schemav1.QueryResponse{
+		Response: resp,
+		Error: &schemav1.QueryError{
+			Message: errMessage,
+		},
+	}, nil
+}

pkg/server/server.go

@@ -20,6 +20,7 @@ import (
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/server/analyze"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/server/config"
+	"github.com/k8sgpt-ai/k8sgpt/pkg/server/query"
 	"log"
 	"net"
 	"net/http"
@@ -50,6 +51,7 @@ type Config struct {
 	Output         string
 	ConfigHandler  *config.Handler
 	AnalyzeHandler *analyze.Handler
+	QueryHandler   *query.Handler
 	Logger         *zap.Logger
 	metricsServer  *http.Server
 	listener       net.Listener
@@ -98,6 +100,7 @@ func (s *Config) Serve() error {
 
 	s.ConfigHandler = &config.Handler{}
 	s.AnalyzeHandler = &analyze.Handler{}
+	s.QueryHandler = &query.Handler{}
 	s.listener = lis
 	s.Logger.Info(fmt.Sprintf("binding api to %s", s.Port))
 	grpcServerUnaryInterceptor := grpc.UnaryInterceptor(LogInterceptor(s.Logger))
@@ -105,7 +108,7 @@ func (s *Config) Serve() error {
 	reflection.Register(grpcServer)
 	rpc.RegisterServerConfigServiceServer(grpcServer, s.ConfigHandler)
 	rpc.RegisterServerAnalyzerServiceServer(grpcServer, s.AnalyzeHandler)
-
+	rpc.RegisterServerQueryServiceServer(grpcServer, s.QueryHandler)
 	if s.EnableHttp {
 		s.Logger.Info("enabling rest/http api")
 		gwmux := runtime.NewServeMux()