Add Botkube tools

hack/openai/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"flag"
 	"log"
 	"os"
 	"time"
@@ -11,29 +12,49 @@ import (
 )
 
 const (
-	assistantID = "asst_eMM9QaWLi6cajHE4PdG1yU53" // 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
+	default:
+		log.Fatalf("Invalid environment: %s", *env)
+	}
+
 	// Update assistant with latest tools.
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
+	log.Printf("Updating %s assistant with ID %s", *env, assistantID)
 	_, err := updateAssistant(ctx, client, assistantID)
 	if err != nil {
 		log.Fatal(err)
 	}
 }
 
 func updateAssistant(ctx context.Context, c *openai.Client, id string) (openai.Assistant, error) {
-	instructions := heredoc.Doc(`Your role involves deeply understanding how to operate and troubleshoot Kubernetes clusters and their workloads. 
-You possess extensive expertise in Kubernetes and cloud-native networking. Take kubernetes cluster and namespace configuration, such as security policies, events, during troubleshooting.
+	instructions := heredoc.Doc(`
+Your role involves deeply understanding how to operate and troubleshoot Kubernetes clusters and their workloads. 
+You possess extensive expertise in Kubernetes and cloud-native networking. 
+Take kubernetes cluster and namespace configuration, such as security policies, events, during troubleshooting.
 Employ a Chain of Thought (CoT) process for diagnosing and resolving cluster issues. 
 Ensure your explanations are simplified for clarity to non-technical users. 
 Utilize available tools for diagnosing the specific cluster in question. Focus your responses to be concise and relevant.
-Keep responses concise, within 2000 characters. Provide extended answers only upon request.`)
+Keep responses concise, within 2000 characters. Provide extended answers only upon request.
+
+Make sure you fetch Botkube Agent configuration to answer questions about Botkube or channel configuration.
+`)
 
 	return c.ModifyAssistant(ctx, id, openai.AssistantRequest{
 		Model:        openai.GPT4Turbo,

hack/openai/tools.go

@@ -7,6 +7,47 @@ import (
 )
 
 var openAITools = []openai.AssistantTool{
+	{
+		Type: openai.AssistantToolType(openai.ToolTypeFunction),
+		Function: &openai.FunctionDefinition{
+			Name: "botkubeGetStartupAgentConfiguration",
+			Description: heredoc.Doc(`
+			Retrieve the Botkube Agent configuration in YAML format. This includes various details such as the configured communication platform, 
+			aliases, enabled plugins, and their settings like RBAC.
+
+			Use this tool to answer questions like:
+			- What sources are configured for this channel?
+			- Is the Kubernetes plugin enabled?
+			- What Kubernetes alerts are currently configured on my cluster?
+
+			Avoid using configuration names like 'botkube/kubernetes_6ulLY' in responses unless explicitly requested. 
+			When a user asks about "this channel", prompt them to provide the exact channel name.
+
+			The "loaderValidationWarnings" field lists all the warnings about the configuration. Use it to suggest fixes.
+			The "incomingRequestPrompt" field defines the communicating platform associated with a prompt. Use this information to address platform-specific questions.
+
+			If the "incomingRequestPrompt" field is not specified, prompt user provide platform name.
+
+			Terms "cloudSlack" and "cloudTeams" refer to the Slack and Teams platforms, respectively.`),
+		},
+	},
+	{
+		Type: openai.AssistantToolType(openai.ToolTypeFunction),
+		Function: &openai.FunctionDefinition{
+			Name: "botkubeGetAgentStatus",
+			Description: heredoc.Doc(`
+			Get the current Botkube Agent status in JSON format. 
+            It represents the status of each plugin and communication platform, indicating whether it's enabled, crashing, and if so, with what reason. 
+			If the "enabled" field is not specified, the plugin is disabled.
+
+			Use this tool to answer questions like:
+			- Which sources are enabled for this channel?
+			- Is the Kubernetes plugin enabled?
+			- Is the kubectl plugin disabled?
+			- What is the status of a specific plugin?
+			- Why am I not receiving Kubernetes notifications?`),
+		},
+	},
 	{
 		Type: openai.AssistantToolType(openai.ToolTypeFunction),
 		Function: &openai.FunctionDefinition{

internal/remote/deploy_client.go

@@ -26,11 +26,6 @@ func NewDeploymentClient(cfg Config) *DeploymentClient {
 	return &DeploymentClient{client: NewDefaultGqlClient(cfg)}
 }
 
-// Deployment returns deployment with Botkube configuration.
-type Deployment struct {
-	ID string
-}
-
 // IsConnectedWithCould returns whether connected to Botkube Cloud
 func (d *DeploymentClient) IsConnectedWithCould() error {
 	var query struct {
@@ -57,6 +52,28 @@ func (d *DeploymentClient) IsConnectedWithCould() error {
 	return nil
 }
 
+// GetConfigOutput returns deployment with Botkube configuration.
+type GetConfigOutput struct {
+	ResourceVersion int
+	YAMLConfig      string
+}
+
+// GetConfig retrieves deployment configuration.
+func (d *DeploymentClient) GetConfig(ctx context.Context) (GetConfigOutput, error) {
+	var query struct {
+		Deployment GetConfigOutput `graphql:"deployment(id: $id)"`
+	}
+	deployID := d.client.DeploymentID()
+	variables := map[string]interface{}{
+		"id": graphql.ID(deployID),
+	}
+	err := d.client.Client().Query(ctx, &query, variables)
+	if err != nil {
+		return GetConfigOutput{}, fmt.Errorf("while getting config with resource version for %q: %w", deployID, err)
+	}
+	return query.Deployment, nil
+}
+
 func (d *DeploymentClient) withRetries(fn func() error) error {
 	return retry.Do(
 		func() error {

internal/source/ai-brain/assistant.go

@@ -35,7 +35,7 @@ const (
 
 var temperature float32 = 0.1
 
-type tool func(ctx context.Context, args []byte) (string, error)
+type tool func(ctx context.Context, args []byte, p *Payload) (string, error)
 
 // Payload represents incoming webhook payload.
 type Payload struct {
@@ -53,7 +53,7 @@ type assistant struct {
 	tracer       trace.Tracer
 }
 
-func newAssistant(cfg *Config, log logrus.FieldLogger, out chan source.Event, kubeConfigPath string) *assistant {
+func newAssistant(cfg *Config, log logrus.FieldLogger, out chan source.Event, kubeConfigPath string) (*assistant, error) {
 	if cfg.HoneycombAPIKey != "" {
 		log.Debug("Setting up opentelemetry with honeycomb")
 		_, err := otelx.Init(cfg.HoneycombAPIKey, serviceName, cfg.HoneycombSampleRate, cfg.Version)
@@ -64,6 +64,10 @@ func newAssistant(cfg *Config, log logrus.FieldLogger, out chan source.Event, ku
 	tracer := otel.Tracer(serviceName)
 
 	kcRunner := NewKubectlRunner(kubeConfigPath, tracer)
+	bkRunner, err := NewBotkubeRunner(tracer)
+	if err != nil {
+		return nil, fmt.Errorf("while creating Botkube runner: %w", err)
+	}
 
 	config := openai.DefaultConfig("")
 	config.HTTPClient = &http.Client{
@@ -79,14 +83,16 @@ func newAssistant(cfg *Config, log logrus.FieldLogger, out chan source.Event, ku
 		assistID:     cfg.OpenAIAssistantID,
 		cache:        newCache(cacheTTL),
 		tools: map[string]tool{
-			"kubectlGetResource":      kcRunner.GetResource,
-			"kubectlDescribeResource": kcRunner.DescribeResource,
-			"kubectlGetEvents":        kcRunner.GetEvents,
-			"kubectlTopPods":          kcRunner.TopPods,
-			"kubectlTopNodes":         kcRunner.TopNodes,
-			"kubectlLogs":             kcRunner.Logs,
+			"kubectlGetResource":                  kcRunner.GetResource,
+			"kubectlDescribeResource":             kcRunner.DescribeResource,
+			"kubectlGetEvents":                    kcRunner.GetEvents,
+			"kubectlTopPods":                      kcRunner.TopPods,
+			"kubectlTopNodes":                     kcRunner.TopNodes,
+			"kubectlLogs":                         kcRunner.Logs,
+			"botkubeGetStartupAgentConfiguration": bkRunner.GetStartupAgentConfiguration,
+			"botkubeGetAgentStatus":               bkRunner.GetAgentStatus,
 		},
-	}
+	}, nil
 }
 
 func (i *assistant) handle(ctx context.Context, in source.ExternalRequestInput) (api.Message, error) {
@@ -209,7 +215,7 @@ func (i *assistant) handleThread(ctx context.Context, p *Payload) (err error) {
 			return true, nil // success
 
 		case openai.RunStatusRequiresAction:
-			if err = i.handleStatusRequiresAction(ctx, run); err != nil {
+			if err = i.handleStatusRequiresAction(ctx, run, p); err != nil {
 				toolsRetries++
 				return toolsRetries >= maxToolExecutionRetries, fmt.Errorf("while handling requires action: %w", err)
 			}
@@ -286,7 +292,7 @@ func (i *assistant) handleStatusCompleted(ctx context.Context, run openai.Run, p
 	return nil
 }
 
-func (i *assistant) handleStatusRequiresAction(ctx context.Context, run openai.Run) error {
+func (i *assistant) handleStatusRequiresAction(ctx context.Context, run openai.Run, p *Payload) error {
 	ctx, span := i.tracer.Start(ctx, "aibrain.assistant.handleStatusRequiresAction")
 	defer span.End()
 
@@ -307,7 +313,7 @@ func (i *assistant) handleStatusRequiresAction(ctx context.Context, run openai.R
 			continue
 		}
 
-		out, err := doer(ctx, []byte(t.Function.Arguments))
+		out, err := doer(ctx, []byte(t.Function.Arguments), p)
 		if err != nil {
 			span.RecordError(err)
 			span.SetStatus(codes.Error, err.Error())