AI Plugin MVP implementation - first iteration

cmd/executor/ai-face/main.go

@@ -8,10 +8,11 @@ import (
 	"fmt"
 	"io"
 	"net/http"
-
-	"github.com/kubeshop/botkube-cloud-plugins/internal/auth"
+	"strings"
 
 	"github.com/hashicorp/go-plugin"
+
+	"github.com/kubeshop/botkube-cloud-plugins/internal/auth"
 	aibrain "github.com/kubeshop/botkube-cloud-plugins/internal/source/ai-brain"
 	"github.com/kubeshop/botkube/pkg/api"
 	"github.com/kubeshop/botkube/pkg/api/executor"
@@ -77,7 +78,7 @@ func (e *AIFace) Execute(_ context.Context, in executor.ExecuteInput) (executor.
 	aiBrainWebhookURL := fmt.Sprintf("%s/%s", in.Context.IncomingWebhook.BaseSourceURL, cfg.AIBrainSourceName)
 
 	body, err := json.Marshal(aibrain.Payload{
-		Prompt:    in.Command,
+		Prompt:    strings.TrimPrefix(in.Command, pluginName),
 		MessageID: in.Context.Message.ParentActivityID,
 	})
 	if err != nil {

cmd/source/ai-brain/main.go

@@ -1,98 +1,23 @@
 package main
 
 import (
-	"context"
 	_ "embed"
-	"fmt"
-	"sync"
 
 	"github.com/kubeshop/botkube-cloud-plugins/internal/auth"
 
 	"github.com/hashicorp/go-plugin"
+
 	aibrain "github.com/kubeshop/botkube-cloud-plugins/internal/source/ai-brain"
-	"github.com/kubeshop/botkube/pkg/api"
 	"github.com/kubeshop/botkube/pkg/api/source"
-	"github.com/kubeshop/botkube/pkg/loggerx"
-	"github.com/sirupsen/logrus"
 )
 
 // version is set via ldflags by GoReleaser.
 var version = "dev"
 
-const (
-	pluginName  = "ai-brain"
-	description = "Calls AI engine with incoming webhook prompts and streams the response."
-)
-
-// AI implements Botkube source plugin.
-type AI struct {
-	incomingPrompts sync.Map
-}
-
-// Metadata returns details about plugin.
-func (*AI) Metadata(context.Context) (api.MetadataOutput, error) {
-	return api.MetadataOutput{
-		Version:     version,
-		Description: description,
-		Recommended: true,
-		JSONSchema: api.JSONSchema{
-			Value: aibrain.ConfigJSONSchema,
-		},
-		ExternalRequest: api.ExternalRequestMetadata{
-			Payload: api.ExternalRequestPayload{
-				JSONSchema: api.JSONSchema{
-					Value: aibrain.IncomingWebhookJSONSchema,
-				},
-			},
-		},
-	}, nil
-}
-
-// Stream implements Botkube source plugin.
-func (a *AI) Stream(_ context.Context, in source.StreamInput) (source.StreamOutput, error) {
-	cfg, err := aibrain.MergeConfigs(in.Configs)
-	if err != nil {
-		return source.StreamOutput{}, fmt.Errorf("while merging configuration: %w", err)
-	}
-
-	log := loggerx.New(cfg.Log)
-	out := source.StreamOutput{
-		Event: make(chan source.Event),
-	}
-	go a.processPrompts(in.Context.SourceName, out.Event, log)
-
-	log.Infof("Setup successful for source configuration %q", in.Context.SourceName)
-	return out, nil
-}
-
-func (a *AI) processPrompts(sourceName string, event chan<- source.Event, log logrus.FieldLogger) {
-	a.incomingPrompts.Store(sourceName, aibrain.NewProcessor(log, event))
-}
-
-// HandleExternalRequest handles incoming payload and returns an event based on it.
-func (a *AI) HandleExternalRequest(_ context.Context, in source.ExternalRequestInput) (source.ExternalRequestOutput, error) {
-	brain, ok := a.incomingPrompts.Load(in.Context.SourceName)
-	if !ok {
-		return source.ExternalRequestOutput{}, fmt.Errorf("source %q not found", in.Context.SourceName)
-	}
-	quickResponse, err := brain.(*aibrain.Processor).Process(in.Payload)
-	if err != nil {
-		return source.ExternalRequestOutput{}, fmt.Errorf("while processing payload: %w", err)
-	}
-
-	return source.ExternalRequestOutput{
-		Event: source.Event{
-			Message: quickResponse,
-		},
-	}, nil
-}
-
 func main() {
 	source.Serve(map[string]plugin.Plugin{
-		pluginName: &source.Plugin{
-			Source: auth.NewProtectedSource(&AI{
-				incomingPrompts: sync.Map{},
-			}),
+		aibrain.PluginName: &source.Plugin{
+			Source: auth.NewProtectedSource(aibrain.NewSource(version)),
 		},
 	})
 }

go.mod

@@ -28,9 +28,9 @@ require (
 	github.com/muesli/reflow v0.3.0
 	github.com/olekukonko/tablewriter v0.0.5
 	github.com/prometheus/client_golang v1.16.0
+	github.com/sashabaranov/go-openai v1.19.4
 	github.com/sirupsen/logrus v1.9.3
 	github.com/slack-go/slack v0.12.2
-	github.com/sourcegraph/conc v0.3.0
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.8.4
 	golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1
@@ -159,7 +159,6 @@ require (
 	go.opentelemetry.io/otel/metric v1.21.0 // indirect
 	go.opentelemetry.io/otel/trace v1.21.0 // indirect
 	go.starlark.net v0.0.0-20230525235612-a134d8f9ddca // indirect
-	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/arch v0.3.0 // indirect
 	golang.org/x/crypto v0.17.0 // indirect
 	golang.org/x/net v0.19.0 // indirect

go.sum

@@ -815,6 +815,8 @@ github.com/ryanuber/columnize v2.1.0+incompatible/go.mod h1:sm1tb6uqfes/u+d4ooFo
 github.com/ryanuber/go-glob v1.0.0/go.mod h1:807d1WSdnB0XRJzKNil9Om6lcp/3a0v4qIHxIXzX/Yc=
 github.com/sanity-io/litter v1.5.5 h1:iE+sBxPBzoK6uaEP5Lt3fHNgpKcHXc/A2HGETy0uJQo=
 github.com/sanity-io/litter v1.5.5/go.mod h1:9gzJgR2i4ZpjZHsKvUXIRQVk7P+yM3e+jAF7bU2UI5U=
+github.com/sashabaranov/go-openai v1.19.4 h1:GbaDiqvgYCabyqzuIbcEeT6/ZX1nVfur+++oTBfOgks=
+github.com/sashabaranov/go-openai v1.19.4/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/sclevine/spec v1.2.0/go.mod h1:W4J29eT/Kzv7/b9IWLB055Z+qvVC9vt0Arko24q7p+U=
 github.com/sclevine/spec v1.4.0/go.mod h1:LvpgJaFyvQzRvc1kaDs0bulYwzC70PbiYjC4QnFHkOM=
 github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc=
@@ -830,8 +832,6 @@ github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
 github.com/slack-go/slack v0.12.2 h1:x3OppyMyGIbbiyFhsBmpf9pwkUzMhthJMRNmNlA4LaQ=
 github.com/slack-go/slack v0.12.2/go.mod h1:hlGi5oXA+Gt+yWTPP0plCdRKmjsDxecdHxYQdlMQKOw=
-github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
-github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
 github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
 github.com/spf13/cast v1.3.1/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
 github.com/spf13/cast v1.5.0 h1:rj3WzYc11XZaIZMPKmwP96zkFEnnAmV8s6XbB2aY32w=

hack/openai/main.go

@@ -0,0 +1,41 @@
+package main
+
+import (
+	"context"
+	"log"
+	"os"
+	"time"
+
+	openai "github.com/sashabaranov/go-openai"
+)
+
+const (
+	assistantID = "asst_eMM9QaWLi6cajHE4PdG1yU53" // Botkube
+)
+
+func main() {
+	client := openai.NewClient(os.Getenv("OPENAI_API_KEY"))
+
+	// Update assistant with latest tools.
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	_, 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 := `You are an experienced DevOps engineer. You have deep
+	understand how to operate a kubernetes cluster and troubleshoot running
+	workloads in kubernetes. You have access to tools which can help you. Keep
+	your answers short and on the subject. Do not get involved in unrelated
+	topics.`
+
+	return c.ModifyAssistant(ctx, id, openai.AssistantRequest{
+		Model:        openai.GPT4TurboPreview,
+		Instructions: &instructions,
+		Tools:        openAITools,
+	})
+}

hack/openai/tools.go

@@ -0,0 +1,97 @@
+package main
+
+import (
+	"github.com/sashabaranov/go-openai"
+	"github.com/sashabaranov/go-openai/jsonschema"
+)
+
+var openAITools = []openai.AssistantTool{
+	{
+		Type: openai.AssistantToolType(openai.ToolTypeFunction),
+		Function: &openai.FunctionDefinition{
+			Name:        "kubectlGetPods",
+			Description: "Useful for getting one pod or all pods in a kubernetes namespace.",
+			Parameters: jsonschema.Definition{
+				Type: jsonschema.Object,
+				Properties: map[string]jsonschema.Definition{
+					"namespace": {
+						Type:        jsonschema.String,
+						Description: "Kubernetes namespace, e.g. kube-system",
+					},
+					"pod_name": {
+						Type:        jsonschema.String,
+						Description: "Kubernetes pod name, e.g. botkube-6c6fd8b4d6-f559q",
+					},
+				},
+				Required: []string{"namespace"},
+			},
+		},
+	},
+	{
+		Type: openai.AssistantToolType(openai.ToolTypeFunction),
+		Function: &openai.FunctionDefinition{
+			Name:        "kubectlGetSecrets",
+			Description: "Useful for getting one secret or all secrets in a kubernetes namespace.",
+			Parameters: jsonschema.Definition{
+				Type: jsonschema.Object,
+				Properties: map[string]jsonschema.Definition{
+					"namespace": {
+						Type:        jsonschema.String,
+						Description: "Kubernetes namespace, e.g. kube-system",
+					},
+					"secret_name": {
+						Type:        jsonschema.String,
+						Description: "Kubernetes secret name, e.g. api-key",
+					},
+				},
+				Required: []string{"namespace"},
+			},
+		},
+	},
+	{
+		Type: openai.AssistantToolType(openai.ToolTypeFunction),
+		Function: &openai.FunctionDefinition{
+			Name:        "kubectlDescribePod",
+			Description: "Useful for describing a pod in a kubernetes namespace.",
+			Parameters: jsonschema.Definition{
+				Type: jsonschema.Object,
+				Properties: map[string]jsonschema.Definition{
+					"namespace": {
+						Type:        jsonschema.String,
+						Description: "Kubernetes namespace, e.g. kube-system",
+					},
+					"pod_name": {
+						Type:        jsonschema.String,
+						Description: "Kubernetes pod name, e.g. botkube-6c6fd8b4d6-f559q",
+					},
+				},
+				Required: []string{"namespace", "pod_name"},
+			},
+		},
+	},
+	{
+		Type: openai.AssistantToolType(openai.ToolTypeFunction),
+		Function: &openai.FunctionDefinition{
+			Name:        "kubectlLogs",
+			Description: "Useful for getting container logs.",
+			Parameters: jsonschema.Definition{
+				Type: jsonschema.Object,
+				Properties: map[string]jsonschema.Definition{
+					"namespace": {
+						Type:        jsonschema.String,
+						Description: "Kubernetes namespace, e.g. kube-system",
+					},
+					"pod_name": {
+						Type:        jsonschema.String,
+						Description: "Kubernetes pod name, e.g. botkube-6c6fd8b4d6-f559q",
+					},
+					"container_name": {
+						Type:        jsonschema.String,
+						Description: "Pod container name, e.g. botkube-api-server.",
+					},
+				},
+				Required: []string{"namespace", "pod_name"},
+			},
+		},
+	},
+}