refactor: update command flag handling to focus on --label flags

Modified the command parsing logic to only accept --label flags, treating other flags as regular arguments. Updated related unit tests to reflect changes in flag handling and ensure accurate extraction of flag values.

Signed-off-by: Zhao Chen <zhaochen.zju@gmail.com>

Author: aftersnow

pkg/modelfile/parser/parser.go

@@ -185,6 +185,7 @@ func splitCommand(line string) (string, []string, []string, error) {
 
 // extractCommandFlags parses the command flags and returns the remaining part of the line
 // and the command flags (with values only, without the flag names).
+// Only accepts --label flags, other flags are ignored and treated as arguments.
 func extractCommandFlags(line string) (string, []string, error) {
 	flags := []string{}
 	var i int
@@ -214,10 +215,17 @@ func extractCommandFlags(line string) (string, []string, error) {
 		}
 
 		if flag != "" {
-			// Extract the value part from --flag=value format
-			flagValue := extractFlagValue(flag)
-			if flagValue != "" {
-				flags = append(flags, flagValue)
+			// Only process --label flags, ignore all others
+			if strings.HasPrefix(flag, "--label") {
+				// Extract the value part from --label=value format
+				flagValue := extractFlagValue(flag)
+				if flagValue != "" {
+					flags = append(flags, flagValue)
+				}
+			} else {
+				// For non-label flags, treat them as part of the remaining arguments
+				// We need to backtrack to include this flag in the remaining content
+				return line[start:], flags, nil
 			}
 		}
 	}
@@ -231,9 +239,9 @@ func isFlag(line string, i int) bool {
 	return i+1 < len(line) && line[i] == '-' && line[i+1] == '-'
 }
 
-// extractFlagValue extracts the value from a flag string
+// extractFlagValue extracts the value from a --label flag string
 // Example: "--label=key=value" returns "key=value"
-// Example: "--untested" returns ""
+// Example: "--label=" returns ""
 func extractFlagValue(flag string) string {
 	// Remove the leading "--"
 	if strings.HasPrefix(flag, "--") {

pkg/modelfile/parser/parser_test.go

@@ -200,10 +200,10 @@ func TestSplitCommand(t *testing.T) {
 		{"MODEL foo", false, "MODEL", []string{"foo"}, []string{}},
 		{"NAME bar", false, "NAME", []string{"bar"}, []string{}},
 		{"MODEL --label=key=value /home/user/model.safetensors", false, "MODEL", []string{"/home/user/model.safetensors"}, []string{"key=value"}},
-		{"MODEL --untested --experimental=test model.safetensors", false, "MODEL", []string{"model.safetensors"}, []string{"test"}},
-		{"CONFIG --format=json config.yaml", false, "CONFIG", []string{"config.yaml"}, []string{"json"}},
-		{"MODEL --flag1=value1 --flag2=value2 model.bin", false, "MODEL", []string{"model.bin"}, []string{"value1", "value2"}},
-		{"MODEL --untested model.safetensors", false, "MODEL", []string{"model.safetensors"}, []string{}}, // flag without value
+		{"MODEL --untested --experimental=test model.safetensors", false, "MODEL", []string{"--untested", "--experimental=test", "model.safetensors"}, []string{}},
+		{"CONFIG --format=json config.yaml", false, "CONFIG", []string{"--format=json", "config.yaml"}, []string{}},
+		{"MODEL --label=flag1=value1 --label=flag2=value2 model.bin", false, "MODEL", []string{"model.bin"}, []string{"flag1=value1", "flag2=value2"}},
+		{"MODEL --untested model.safetensors", false, "MODEL", []string{"--untested", "model.safetensors"}, []string{}},
 		{"invalid", true, "", nil, nil},
 	}
 
@@ -232,30 +232,25 @@ func TestExtractFlagValue(t *testing.T) {
 		expectedValue string
 	}{
 		{
-			name:          "flag with key=value",
+			name:          "label flag with key=value",
 			flag:          "--label=key=value",
 			expectedValue: "key=value",
 		},
 		{
-			name:          "flag with simple value",
-			flag:          "--format=json",
-			expectedValue: "json",
-		},
-		{
-			name:          "flag without value",
-			flag:          "--untested",
-			expectedValue: "",
-		},
-		{
-			name:          "flag with empty value",
+			name:          "label flag with empty value",
 			flag:          "--label=",
 			expectedValue: "",
 		},
 		{
-			name:          "complex flag value",
+			name:          "complex label flag value",
 			flag:          fmt.Sprintf("--label=%s=true", modelspec.AnnotationMediaTypeUntested),
 			expectedValue: fmt.Sprintf("%s=true", modelspec.AnnotationMediaTypeUntested),
 		},
+		{
+			name:          "label flag without value",
+			flag:          "--label",
+			expectedValue: "",
+		},
 	}
 
 	assert := assert.New(t)