pkg/aflow: allow to specify model per-flow

We may want to use a weaker model for some workflows.
Allow to use different models for different workflows.

1 files changed, 26 insertions, 23 deletions

diff --git a/tools/syz-aflow/aflow.go b/tools/syz-aflow/aflow.go
index b9b59e90e..d915c2061 100644
--- a/tools/syz-aflow/aflow.go
+++ b/tools/syz-aflow/aflow.go
@@ -33,7 +33,7 @@ func main() {
 		flagFlow        = flag.String("workflow", "", "workflow to execute")
 		flagInput       = flag.String("input", "", "input json file with workflow arguments")
 		flagWorkdir     = flag.String("workdir", "", "directory for kernel checkout, kernel builds, etc")
-		flagModel       = flag.String("model", aflow.DefaultModel, "use this LLM model")
+		flagModel       = flag.String("model", "", "use this LLM model, if empty use the workflow default model")
 		flagCacheSize   = flag.String("cache-size", "10GB", "max cache size (e.g. 100MB, 5GB, 1TB)")
 		flagDownloadBug = flag.String("download-bug", "", "extid of a bug to download from the dashboard"+
 			" and save into -input file")
@@ -73,33 +73,14 @@ func main() {
 	}
 }
 
-func parseSize(s string) (uint64, error) {
-	var size uint64
-	var suffix string
-	if _, err := fmt.Sscanf(s, "%d%s", &size, &suffix); err != nil {
-		return 0, fmt.Errorf("failed to parse cache size %q: %w", s, err)
-	}
-	switch suffix {
-	case "KB":
-		size <<= 10
-	case "MB":
-		size <<= 20
-	case "GB":
-		size <<= 30
-	case "TB":
-		size <<= 40
-	case "":
-	default:
-		return 0, fmt.Errorf("unknown size suffix %q", suffix)
-	}
-	return size, nil
-}
-
 func run(ctx context.Context, model, flowName, inputFile, workdir string, cacheSize uint64) error {
 	flow := aflow.Flows[flowName]
 	if flow == nil {
 		return fmt.Errorf("workflow %q is not found", flowName)
 	}
+	if model == "" {
+		model = flow.Model
+	}
 	inputData, err := os.ReadFile(inputFile)
 	if err != nil {
 		return fmt.Errorf("failed to open -input file: %w", err)
@@ -194,3 +175,25 @@ func getAccessToken() (string, error) {
 
 	return token.AccessToken, nil
 }
+
+func parseSize(s string) (uint64, error) {
+	var size uint64
+	var suffix string
+	if _, err := fmt.Sscanf(s, "%d%s", &size, &suffix); err != nil {
+		return 0, fmt.Errorf("failed to parse cache size %q: %w", s, err)
+	}
+	switch suffix {
+	case "KB":
+		size <<= 10
+	case "MB":
+		size <<= 20
+	case "GB":
+		size <<= 30
+	case "TB":
+		size <<= 40
+	case "":
+	default:
+		return 0, fmt.Errorf("unknown size suffix %q", suffix)
+	}
+	return size, nil
+}