flow summary shows drops/forwards and increase flow limit

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

Author: huntergregory

ai/pkg/analysis/flows/types.go

@@ -20,11 +20,37 @@ type Connection struct {
 type FlowSummary map[string]*Connection
 
 func (fs FlowSummary) FormatForLM() string {
+	// FIXME hacky right now
+	forwards := fs.connStrings(flowpb.Verdict_FORWARDED)
+	drops := fs.connStrings(flowpb.Verdict_DROPPED)
+	other := fs.connStrings(flowpb.Verdict_VERDICT_UNKNOWN)
+
+	return fmt.Sprintf("SUCCESSFUL CONNECTIONS:\n%s\n\nDROPPED CONNECTIONS:\n%s\n\nOTHER CONNECTIONS:\n%s", forwards, drops, other)
+}
+
+func (fs FlowSummary) connStrings(verdict flowpb.Verdict) string {
 	connStrings := make([]string, 0, len(fs))
 	for _, conn := range fs {
+		match := false
+		for _, f := range conn.Flows {
+			// FIXME hacky right now
+			if f.GetVerdict() == verdict || (verdict == flowpb.Verdict_VERDICT_UNKNOWN && f.GetVerdict() != flowpb.Verdict_FORWARDED && f.GetVerdict() != flowpb.Verdict_DROPPED) {
+				match = true
+				break
+			}
+		}
+
+		if !match {
+			continue
+		}
+
 		connString := fmt.Sprintf("Connection: %s -> %s, Number of Flows: %d", conn.Pod1, conn.Pod2, len(conn.Flows))
 		connStrings = append(connStrings, connString)
 	}
 
+	if len(connStrings) == 0 {
+		return "none"
+	}
+
 	return strings.Join(connStrings, "\n")
 }

ai/pkg/retrieval/flows/retriever.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/microsoft/retina/ai/pkg/retrieval/flows/client"
+	"github.com/microsoft/retina/ai/pkg/util"
 
 	flowpb "github.com/cilium/cilium/api/v1/flow"
 	observerpb "github.com/cilium/cilium/api/v1/observer"
@@ -88,7 +89,7 @@ func (r *Retriever) Observe(ctx context.Context, maxFlows int) ([]*flowpb.Flow,
 
 func flowsRequest() *observerpb.GetFlowsRequest {
 	return &observerpb.GetFlowsRequest{
-		Number:    200,
+		Number:    util.MaxFlowsFromHubbleRelay,
 		Follow:    false,
 		Whitelist: []*flowpb.FlowFilter{},
 		Blacklist: nil,

ai/pkg/scenarios/flows/handler.go

@@ -8,6 +8,7 @@ import (
 	flowanalysis "github.com/microsoft/retina/ai/pkg/analysis/flows"
 	"github.com/microsoft/retina/ai/pkg/lm"
 	flowretrieval "github.com/microsoft/retina/ai/pkg/retrieval/flows"
+	"github.com/microsoft/retina/ai/pkg/util"
 
 	"github.com/sirupsen/logrus"
 	"k8s.io/client-go/kubernetes"
@@ -58,7 +59,7 @@ func (h *Handler) Handle(ctx context.Context, question string, chat lm.ChatHisto
 		return "", fmt.Errorf("error initializing flow retriever: %w", err)
 	}
 
-	flows, err := h.r.Observe(ctx, 50)
+	flows, err := h.r.Observe(ctx, util.MaxFlowsToAnalyze)
 	if err != nil {
 		return "", fmt.Errorf("error observing flows: %w", err)
 	}