Skip to content

Commit

Permalink
feat: adding a query mode for the schednex scheduler (#1257)
Browse files Browse the repository at this point in the history
Signed-off-by: AlexsJones <[email protected]>
  • Loading branch information
AlexsJones committed Oct 24, 2024
1 parent ce5ade8 commit 1dea6dc
Show file tree
Hide file tree
Showing 5 changed files with 60 additions and 4 deletions.
7 changes: 4 additions & 3 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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/go v1.5.1-20240807134501-ea98c104104d.1
buf.build/gen/go/k8sgpt-ai/k8sgpt/protocolbuffers/go v1.34.2-20240807134501-ea98c104104d.2
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-20240920204244-7a91c8620515.1
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
Expand All @@ -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
Expand Down
18 changes: 18 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
7 changes: 7 additions & 0 deletions pkg/server/query/handler.go
Original file line number Diff line number Diff line change
@@ -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
}
27 changes: 27 additions & 0 deletions pkg/server/query/query.go
Original file line number Diff line number Diff line change
@@ -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
}
5 changes: 4 additions & 1 deletion pkg/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -98,14 +100,15 @@ 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))
grpcServer := grpc.NewServer(grpcServerUnaryInterceptor)
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()
Expand Down

0 comments on commit 1dea6dc

Please sign in to comment.