Clean-up

kubeshop · May 6, 2024 · 103ef23 · 103ef23
1 parent b028395
commit 103ef23
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 25 deletions.
diff --git a/hack/openai/main.go b/hack/openai/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"flag"
 	"log"
 	"os"
 	"time"
@@ -11,13 +12,24 @@ import (
 )
 
 const (
-	//assistantID = "asst_eMM9QaWLi6cajHE4PdG1yU53" // Botkube
-	assistantID = "asst_ejVrAgjhhvCw6jGFYq5JyBqj" // Botkube
+	prodAssistantID = "asst_eMM9QaWLi6cajHE4PdG1yU53"
+	devAssistantID  = "asst_ejVrAgjhhvCw6jGFYq5JyBqj"
 )
 
 func main() {
 	client := openai.NewClient(os.Getenv("OPENAI_API_KEY"))
 
+	env := flag.String("env", "dev", "Environment to update. Allow values: 'dev' or 'prod'.")
+	flag.Parse()
+
+	var assistantID string
+	switch *env {
+	case "dev":
+		assistantID = devAssistantID
+	case "prod":
+		assistantID = prodAssistantID
+
+	}
 	// Update assistant with latest tools.
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()

diff --git a/internal/source/ai-brain/assistant.go b/internal/source/ai-brain/assistant.go
@@ -9,8 +9,6 @@ import (
 	"os"
 	"time"
 
-	"github.com/sanity-io/litter"
-
 	"github.com/kubeshop/botkube-cloud-plugins/internal/otelx"
 	"github.com/kubeshop/botkube-cloud-plugins/internal/remote"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
@@ -66,7 +64,7 @@ func newAssistant(cfg *Config, log logrus.FieldLogger, out chan source.Event, ku
 	tracer := otel.Tracer(serviceName)
 
 	kcRunner := NewKubectlRunner(kubeConfigPath, tracer)
-	bkRunner, err := NewBotkubeRunner(kcRunner, tracer)
+	bkRunner, err := NewBotkubeRunner(tracer)
 	if err != nil {
 		return nil, fmt.Errorf("while creating Botkube runner: %w", err)
 	}
@@ -92,7 +90,6 @@ func newAssistant(cfg *Config, log logrus.FieldLogger, out chan source.Event, ku
 			"kubectlTopNodes":                     kcRunner.TopNodes,
 			"kubectlLogs":                         kcRunner.Logs,
 			"botkubeGetStartupAgentConfiguration": bkRunner.GetStartupAgentConfiguration,
-			"botkubeGetCloudAgentConfiguration":   bkRunner.GetCloudAgentConfiguration, // TODO: do we need that?
 			"botkubeGetAgentStatus":               bkRunner.GetAgentStatus,
 		},
 	}, nil
@@ -202,7 +199,6 @@ func (i *assistant) handleThread(ctx context.Context, p *Payload) (err error) {
 
 		switch run.Status {
 		case openai.RunStatusCancelling, openai.RunStatusFailed:
-			litter.Dump(run)
 			return true, fmt.Errorf("got unexpected status: %s", run.Status)
 
 		case openai.RunStatusExpired:

diff --git a/internal/source/ai-brain/botkube_tools.go b/internal/source/ai-brain/botkube_tools.go
@@ -27,16 +27,17 @@ type BotkubeRunner struct {
 	httpCli          *http.Client
 	rawStartupConfig string
 	startupConfig    *ConfigWithDetails
-	logGetter        LogGetter
 }
 
-// LogGetter provides ability to get logs.
-type LogGetter interface {
-	Logs(ctx context.Context, rawArgs []byte, _ *Payload) (string, error)
+// ConfigWithDetails represents Botkube configuration with additional details.
+type ConfigWithDetails struct {
+	*config.Config           `yaml:",inline"`
+	LoaderValidationWarnings string `yaml:"loaderValidationWarnings"`
+	IncomingRequestPrompt    string `yaml:"incomingRequestPrompt,omitempty"`
 }
 
 // NewBotkubeRunner creates new runner instance.
-func NewBotkubeRunner(logGetter LogGetter, tracer trace.Tracer) (*BotkubeRunner, error) {
+func NewBotkubeRunner(tracer trace.Tracer) (*BotkubeRunner, error) {
 	cfg, ok := remote.GetConfig()
 	if !ok {
 		return nil, fmt.Errorf("cannot get Botkube cloud related environment variables")
@@ -46,7 +47,6 @@ func NewBotkubeRunner(logGetter LogGetter, tracer trace.Tracer) (*BotkubeRunner,
 		tracer:    tracer,
 		deployCli: remote.NewDeploymentClient(cfg),
 		httpCli:   httpx.NewHTTPClient(),
-		logGetter: logGetter,
 	}
 
 	err := r.initStartupConfig()
@@ -56,18 +56,6 @@ func NewBotkubeRunner(logGetter LogGetter, tracer trace.Tracer) (*BotkubeRunner,
 	return r, nil
 }
 
-type ConfigWithDetails struct {
-	*config.Config           `yaml:",inline"`
-	LoaderValidationWarnings string `yaml:"loaderValidationWarnings"`
-	IncomingRequestPrompt    string `yaml:"incomingRequestPrompt,omitempty"`
-}
-
-// TODO: to remove?
-func (r *BotkubeRunner) GetCloudAgentConfiguration(ctx context.Context, _ []byte, _ *Payload) (string, error) {
-	rawCfg, _, err := r.fetchConfig(ctx)
-	return rawCfg, err
-}
-
 // GetStartupAgentConfiguration returns Botkube startup configuration.
 func (r *BotkubeRunner) GetStartupAgentConfiguration(ctx context.Context, _ []byte, p *Payload) (string, error) {
 	_, span := r.tracer.Start(ctx, "aibrain.BotkubeRunner.GetStartupAgentConfiguration")
@@ -91,6 +79,7 @@ func (r *BotkubeRunner) GetStartupAgentConfiguration(ctx context.Context, _ []by
 	return string(raw), nil
 }
 
+// GetAgentStatus returns Botkube Agent health status.
 func (r *BotkubeRunner) GetAgentStatus(ctx context.Context, _ []byte, _ *Payload) (string, error) {
 	_, span := r.tracer.Start(ctx, "aibrain.BotkubeRunner.GetAgentStatus")
 	defer span.End()