feat: chat bot foundations with flow retrieval/analysis

.gitignore

@@ -15,7 +15,7 @@
 *.out
 
 # Dependency directories (remove the comment below to include it)
-# vendor/
+vendor/
 
 # Go workspace file
 go.work

ai/README.md

@@ -0,0 +1,17 @@
+# Retina AI
+
+## Usage
+
+- Change into this *ai/* folder.
+- `go mod tidy ; go mod vendor`
+- Modify the `defaultConfig` values in *main.go*
+- If using Azure OpenAI:
+  - Make sure you're logged into your account/subscription in your terminal.
+  - Specify environment variables for Deployment name and Endpoint URL. Get deployment from e.g. [https://oai.azure.com/portal/deployment](https://oai.azure.com/portal/deployment) and Endpoint from e.g. Deployment > Playground > Code.
+    - Linux:
+      - `read -p "Enter AOAI_COMPLETIONS_ENDPOINT: " AOAI_COMPLETIONS_ENDPOINT && export AOAI_COMPLETIONS_ENDPOINT=$AOAI_COMPLETIONS_ENDPOINT`
+      - `read -p "Enter AOAI_DEPLOYMENT_NAME: " AOAI_DEPLOYMENT_NAME && export AOAI_DEPLOYMENT_NAME=$AOAI_DEPLOYMENT_NAME`
+    - Windows:
+      - `$env:AOAI_COMPLETIONS_ENDPOINT = Read-Host 'Enter AOAI_COMPLETIONS_ENDPOINT'`
+      - `$env:AOAI_DEPLOYMENT_NAME = Read-Host 'Enter AOAI_DEPLOYMENT_NAME'`
+- `go run main.go`

ai/go.mod

@@ -0,0 +1,61 @@
+module github.com/microsoft/retina/ai
+
+go 1.22.3
+
+require (
+	github.com/Azure/azure-sdk-for-go/sdk/ai/azopenai v0.6.0
+	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.13.0
+	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0
+	github.com/cilium/cilium v1.15.7
+	github.com/cilium/hubble-ui/backend v0.0.0-20240603143312-a06e19ba6529
+	github.com/sirupsen/logrus v1.9.3
+	google.golang.org/grpc v1.65.0
+	k8s.io/client-go v0.30.3
+)
+
+require (
+	github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 // indirect
+	github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/emicklei/go-restful/v3 v3.12.1 // indirect
+	github.com/go-logr/logr v1.4.2 // indirect
+	github.com/go-openapi/jsonpointer v0.21.0 // indirect
+	github.com/go-openapi/jsonreference v0.21.0 // indirect
+	github.com/go-openapi/swag v0.23.0 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang-jwt/jwt/v5 v5.2.1 // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
+	github.com/google/gnostic-models v0.6.8 // indirect
+	github.com/google/gofuzz v1.2.0 // indirect
+	github.com/google/uuid v1.6.0 // indirect
+	github.com/imdario/mergo v0.3.16 // indirect
+	github.com/josharian/intern v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/kylelemons/godebug v1.1.0 // indirect
+	github.com/mailru/easyjson v0.7.7 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
+	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect
+	github.com/spf13/pflag v1.0.6-0.20210604193023-d5e0c0615ace // indirect
+	golang.org/x/crypto v0.25.0 // indirect
+	golang.org/x/net v0.27.0 // indirect
+	golang.org/x/oauth2 v0.20.0 // indirect
+	golang.org/x/sys v0.22.0 // indirect
+	golang.org/x/term v0.22.0 // indirect
+	golang.org/x/text v0.16.0 // indirect
+	golang.org/x/time v0.5.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20240528184218-531527333157 // indirect
+	google.golang.org/protobuf v1.34.1 // indirect
+	gopkg.in/inf.v0 v0.9.1 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+	k8s.io/api v0.30.3 // indirect
+	k8s.io/apimachinery v0.30.3 // indirect
+	k8s.io/klog/v2 v2.120.1 // indirect
+	k8s.io/kube-openapi v0.0.0-20240521193020-835d969ad83a // indirect
+	k8s.io/utils v0.0.0-20240502163921-fe8a2dddb1d0 // indirect
+	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
+	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
+	sigs.k8s.io/yaml v1.4.0 // indirect
+)

ai/main.go

@@ -0,0 +1,108 @@
+package main
+
+import (
+	"fmt"
+	"os/user"
+
+	"github.com/microsoft/retina/ai/pkg/chat"
+	"github.com/microsoft/retina/ai/pkg/lm"
+	"github.com/microsoft/retina/ai/pkg/scenarios"
+	"github.com/microsoft/retina/ai/pkg/scenarios/drops"
+
+	"github.com/sirupsen/logrus"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/tools/clientcmd"
+)
+
+// TODO incorporate this code into a CLI tool someday
+
+type config struct {
+	// currently supports "echo" or "AOAI"
+	model string
+
+	// optional. defaults to ~/.kube/config
+	kubeconfigPath string
+
+	// retrieved flows are currently written to ./flows.json
+	useFlowsFromFile bool
+
+	// eventually, the below should be optional once user input is implemented
+	question string
+	history  lm.ChatHistory
+
+	// eventually, the below should be optional once scenario selection is implemented
+	scenario   *scenarios.Definition
+	parameters map[string]string
+}
+
+var defaultConfig = &config{
+	model:            "echo", // echo or AOAI
+	useFlowsFromFile: false,
+	question:         "What's wrong with my app?",
+	history:          nil,
+	scenario:         drops.Definition, // drops.Definition or dns.Definition
+	parameters: map[string]string{
+		scenarios.Namespace1.Name: "default",
+		// scenarios.PodPrefix1.Name: "toolbox-pod",
+		// scenarios.Namespace2.Name: "default",
+		// scenarios.PodPrefix2.Name: "toolbox-pod",
+		// dns.DNSQuery.Name:         "google.com",
+		// scenarios.Nodes.Name:      "[node1,node2]",
+	},
+}
+
+func main() {
+	run(defaultConfig)
+}
+
+func run(cfg *config) {
+	log := logrus.New()
+	// log.SetLevel(logrus.DebugLevel)
+
+	log.Info("starting app...")
+
+	// retrieve configs
+	if cfg.kubeconfigPath == "" {
+		usr, err := user.Current()
+		if err != nil {
+			log.WithError(err).Fatal("failed to get current user")
+		}
+		cfg.kubeconfigPath = usr.HomeDir + "/.kube/config"
+	}
+
+	kconfig, err := clientcmd.BuildConfigFromFlags("", cfg.kubeconfigPath)
+	if err != nil {
+		log.WithError(err).Fatal("failed to get kubeconfig")
+	}
+
+	clientset, err := kubernetes.NewForConfig(kconfig)
+	if err != nil {
+		log.WithError(err).Fatal("failed to create clientset")
+	}
+	log.Info("retrieved kubeconfig and clientset")
+
+	// configure LM (language model)
+	var model lm.Model
+	switch cfg.model {
+	case "echo":
+		model = lm.NewEchoModel()
+		log.Info("initialized echo model")
+	case "AOAI":
+		model, err = lm.NewAzureOpenAI()
+		if err != nil {
+			log.WithError(err).Fatal("failed to create Azure OpenAI model")
+		}
+		log.Info("initialized Azure OpenAI model")
+	default:
+		log.Fatalf("unsupported model: %s", cfg.model)
+	}
+
+	bot := chat.NewBot(log, kconfig, clientset, model, cfg.useFlowsFromFile)
+	newHistory, err := bot.HandleScenario(cfg.question, cfg.history, cfg.scenario, cfg.parameters)
+	if err != nil {
+		log.WithError(err).Fatal("error handling scenario")
+	}
+
+	log.Info("handled scenario")
+	fmt.Println(newHistory[len(newHistory)-1].Assistant)
+}

ai/pkg/chat/chat.go

@@ -0,0 +1,104 @@
+package chat
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/microsoft/retina/ai/pkg/lm"
+	flowretrieval "github.com/microsoft/retina/ai/pkg/retrieval/flows"
+	"github.com/microsoft/retina/ai/pkg/scenarios"
+
+	"github.com/sirupsen/logrus"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
+)
+
+type Bot struct {
+	log           logrus.FieldLogger
+	config        *rest.Config
+	clientset     *kubernetes.Clientset
+	model         lm.Model
+	flowRetriever *flowretrieval.Retriever
+}
+
+// input log, config, clientset, model
+func NewBot(log logrus.FieldLogger, config *rest.Config, clientset *kubernetes.Clientset, model lm.Model, useFlowsFromFile bool) *Bot {
+	b := &Bot{
+		log:           log.WithField("component", "chat"),
+		config:        config,
+		clientset:     clientset,
+		model:         model,
+		flowRetriever: flowretrieval.NewRetriever(log, config, clientset),
+	}
+
+	if useFlowsFromFile {
+		b.flowRetriever.UseFile()
+	}
+
+	return b
+}
+
+func (b *Bot) HandleScenario(question string, history lm.ChatHistory, definition *scenarios.Definition, parameters map[string]string) (lm.ChatHistory, error) {
+	if definition == nil {
+		return history, fmt.Errorf("no scenario selected")
+	}
+
+	cfg := &scenarios.Config{
+		Log:           b.log,
+		Config:        b.config,
+		Clientset:     b.clientset,
+		Model:         b.model,
+		FlowRetriever: b.flowRetriever,
+	}
+
+	ctx := context.TODO()
+	response, err := definition.Handle(ctx, cfg, parameters, question, history)
+	if err != nil {
+		return history, fmt.Errorf("error handling scenario: %w", err)
+	}
+
+	history = append(history, lm.MessagePair{
+		User:      question,
+		Assistant: response,
+	})
+
+	return history, nil
+}
+
+// FIXME get user input and implement scenario selection
+func (b *Bot) Loop() error {
+	var history lm.ChatHistory
+
+	for {
+		// TODO get user input
+		question := "what's wrong with my app?"
+
+		// select scenario and get parameters
+		definition, params, err := b.selectScenario(question, history)
+		if err != nil {
+			return fmt.Errorf("error selecting scenario: %w", err)
+		}
+
+		newHistory, err := b.HandleScenario(question, history, definition, params)
+		if err != nil {
+			return fmt.Errorf("error handling scenario: %w", err)
+		}
+
+		fmt.Println(newHistory[len(newHistory)-1].Assistant)
+
+		history = newHistory
+	}
+}
+
+// FIXME fix prompts
+func (b *Bot) selectScenario(question string, history lm.ChatHistory) (*scenarios.Definition, map[string]string, error) {
+	ctx := context.TODO()
+	response, err := b.model.Generate(ctx, selectionSystemPrompt, nil, selectionPrompt(question, history))
+	if err != nil {
+		return nil, nil, fmt.Errorf("error generating response: %w", err)
+	}
+
+	// TODO parse response and return scenario definition and parameters
+	_ = response
+	return nil, nil, nil
+}

ai/pkg/chat/prompt.go

@@ -0,0 +1,30 @@
+package chat
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/microsoft/retina/ai/pkg/lm"
+	"github.com/microsoft/retina/ai/pkg/scenarios"
+	"github.com/microsoft/retina/ai/pkg/scenarios/dns"
+	"github.com/microsoft/retina/ai/pkg/scenarios/drops"
+)
+
+const selectionSystemPrompt = "Select a scenario"
+
+var (
+	definitions = []*scenarios.Definition{
+		drops.Definition,
+		dns.Definition,
+	}
+)
+
+func selectionPrompt(question string, history lm.ChatHistory) string {
+	// TODO include parameters etc. and reference the user chat as context
+	var sb strings.Builder
+	sb.WriteString("Select a scenario:\n")
+	for i, d := range definitions {
+		sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, d.Name))
+	}
+	return sb.String()
+}