feat: dns scenario and other improvements

Signed-off-by: Hunter Gregory <[email protected]>

ai/README.md

@@ -4,7 +4,7 @@
 
 - Change into this *ai/* folder.
 - `go mod tidy ; go mod vendor`
-- Set your `kubeconfigPath` in *main.go*.
+- 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.

ai/main.go

@@ -1,47 +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"
 )
 
-const kubeconfigPath = "/home/hunter/.kube/config"
+// 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
 
-// const kubeconfigPath = "C:\\Users\\hgregory\\.kube\\config"
+	// 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
-	config, err := clientcmd.BuildConfigFromFlags("", kubeconfigPath)
+	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(config)
+	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)
-	// model := lm.NewEchoModel()
-	// log.Info("initialized echo model")
-	model, err := lm.NewAzureOpenAI()
-	if err != nil {
-		log.WithError(err).Fatal("failed to create Azure OpenAI 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)
 	}
-	log.Info("initialized Azure OpenAI model")
 
-	bot := chat.NewBot(log, config, clientset, model)
-	if err := bot.Loop(); err != nil {
-		log.WithError(err).Fatal("error running chat loop")
+	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

@@ -7,41 +7,67 @@ import (
 	"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/microsoft/retina/ai/pkg/scenarios/dns"
-	"github.com/microsoft/retina/ai/pkg/scenarios/drops"
 
 	"github.com/sirupsen/logrus"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
 )
 
-var (
-	definitions = []*scenarios.Definition{
-		drops.Definition,
-		dns.Definition,
-	}
-)
-
 type Bot struct {
-	log       logrus.FieldLogger
-	config    *rest.Config
-	clientset *kubernetes.Clientset
-	model     lm.Model
+	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) *Bot {
-	return &Bot{
-		log:       log.WithField("component", "chat"),
-		config:    config,
-		clientset: clientset,
-		model:     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
-	flowRetriever := flowretrieval.NewRetriever(b.log, b.config, b.clientset)
 
 	for {
 		// TODO get user input
@@ -53,39 +79,26 @@ func (b *Bot) Loop() error {
 			return fmt.Errorf("error selecting scenario: %w", err)
 		}
 
-		// cfg.FlowRetriever.UseFile()
-
-		cfg := &scenarios.Config{
-			Log:           b.log,
-			Config:        b.config,
-			Clientset:     b.clientset,
-			Model:         b.model,
-			FlowRetriever: flowRetriever,
-		}
-
-		ctx := context.TODO()
-		response, err := definition.Handle(ctx, cfg, params, question, history)
+		newHistory, err := b.HandleScenario(question, history, definition, params)
 		if err != nil {
 			return fmt.Errorf("error handling scenario: %w", err)
 		}
 
-		fmt.Println(response)
+		fmt.Println(newHistory[len(newHistory)-1].Assistant)
 
-		// TODO keep chat loop going
-		break
+		history = newHistory
 	}
-
-	return nil
 }
 
+// FIXME fix prompts
 func (b *Bot) selectScenario(question string, history lm.ChatHistory) (*scenarios.Definition, map[string]string, error) {
-	// TODO use chat interface
-	// FIXME hard-coding the scenario and params for now
-	d := definitions[0]
-	params := map[string]string{
-		scenarios.Namespace1.Name: "default",
-		scenarios.Namespace2.Name: "default",
+	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)
 	}
 
-	return d, params, nil
+	// 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()
+}

ai/pkg/parse/flows/parser.go

@@ -34,6 +34,10 @@ func (p *Parser) Parse(flows []*flowpb.Flow) {
 }
 
 func (p *Parser) addFlow(f *flowpb.Flow) error {
+	if f == nil {
+		return nil
+	}
+
 	src := f.GetSource()
 	dst := f.GetDestination()
 	if src == nil || dst == nil {

ai/pkg/retrieval/flows/client/client.go

@@ -31,7 +31,7 @@ func New() (*Client, error) {
 	tlsDialOption := grpc.WithTransportCredentials(insecure.NewCredentials())
 
 	// FIXME make address part of a config
-	addr := ":5555"
+	addr := ":5557"
 	connection, err := grpc.NewClient(addr, tlsDialOption, connectDialOption)
 	if err != nil {
 		return nil, fmt.Errorf("failed to dial %s: %w", addr, err)

ai/pkg/retrieval/flows/retriever.go

@@ -83,13 +83,13 @@ func (r *Retriever) Observe(ctx context.Context, req *observerpb.GetFlowsRequest
 	defer portForwardCancel()
 
 	// FIXME make ports part of a config
-	cmd := exec.CommandContext(portForwardCtx, "kubectl", "port-forward", "-n", "kube-system", "svc/hubble-relay", "5555:80")
+	cmd := exec.CommandContext(portForwardCtx, "kubectl", "port-forward", "-n", "kube-system", "svc/hubble-relay", "5557:80")
 	if err := cmd.Start(); err != nil {
 		return nil, fmt.Errorf("failed to start port-forward. %v", err)
 	}
 
 	// observe flows
-	observeCtx, observeCancel := context.WithTimeout(ctx, 30*time.Second)
+	observeCtx, observeCancel := context.WithTimeout(ctx, 15*time.Second)
 	defer observeCancel()
 
 	maxFlows := req.Number
@@ -120,18 +120,25 @@ func (r *Retriever) observeFlowsGRPC(ctx context.Context, req *observerpb.GetFlo
 	}
 
 	r.flows = make([]*flowpb.Flow, 0)
+	var errReceiving error
 	for {
 		select {
 		case <-ctx.Done():
 			r.log.Info("context cancelled")
 			return r.flows, nil
 		default:
+			if errReceiving != nil {
+				// error receiving and context not done
+				// TODO handle error instead of returning error
+				return nil, fmt.Errorf("failed to receive flow. %v", err)
+			}
+
 			r.log.WithField("flowCount", len(r.flows)).Debug("processing flow")
 
 			getFlowResponse, err := stream.Recv()
 			if err != nil {
-				// TODO handle error instead of returning error
-				return nil, fmt.Errorf("failed to receive flow. %v", err)
+				errReceiving = err
+				continue
 			}
 
 			f := getFlowResponse.GetFlow()
@@ -150,7 +157,7 @@ func (r *Retriever) observeFlowsGRPC(ctx context.Context, req *observerpb.GetFlo
 // handleFlow logic is inspired by a snippet from Hubble UI
 // https://github.com/cilium/hubble-ui/blob/a06e19ba65299c63a58034a360aeedde9266ec01/backend/internal/flow_stream/flow_stream.go#L360-L395
 func (r *Retriever) handleFlow(f *flowpb.Flow) {
-	if f.GetL4() == nil || f.GetSource() == nil || f.GetDestination() == nil {
+	if (f.GetL7() == nil && f.GetL4() == nil) || f.GetSource() == nil || f.GetDestination() == nil {
 		return
 	}
 

ai/pkg/scenarios/common.go

@@ -27,7 +27,7 @@ var (
 		Name:        "namespace2",
 		DataType:    "string",
 		Description: "Namespace 2",
-		Optional:    false,
+		Optional:    true,
 		Regex:       k8sNameRegex,
 	}