Add --depth flag for Kitfile generation

Add flag `--depth` to commands that generate Kitfiles. The provided
depth controls how far down a directory structure kit will navigate
before starting to group directories into layers -- for example, with
depth=1, files in immediate subdirectories of the context will be
processed individually and added to the Kitfile as appropriate.

Signed-off-by: Angel Misevski <amisevsk@gmail.com>

Author: amisevsk

pkg/cmd/kitimport/cmd.go

@@ -79,6 +79,7 @@ type importOptions struct {
 	kitfilePath  string
 	downloadTool string
 	concurrency  int
+	depth        int
 	modelKitRef  *registry.Reference
 }
 
@@ -100,6 +101,7 @@ func ImportCommand() *cobra.Command {
 	cmd.Flags().StringVarP(&opts.kitfilePath, "file", "f", "", "Path to Kitfile to use for packing (use '-' to read from standard input)")
 	cmd.Flags().StringVar(&opts.downloadTool, "tool", "", "Tool to use for downloading files: options are 'git' and 'hf' (default: detect based on repository)")
 	cmd.Flags().IntVar(&opts.concurrency, "concurrency", 5, "Maximum number of simultaneous downloads (for huggingface)")
+	cmd.Flags().IntVar(&opts.depth, "depth", 0, "Maximum directory depth to process when generating a Kitfile. Setting to -1 processes all files individually")
 	cmd.Flags().SortFlags = false
 	return cmd
 }

pkg/cmd/kitimport/gitimport.go

@@ -73,7 +73,7 @@ func importUsingGit(ctx context.Context, opts *importOptions) error {
 		if err != nil {
 			return fmt.Errorf("error processing directory: %w", err)
 		}
-		kf, err := generateKitfile(dirContents, opts.repo, tmpDir)
+		kf, err := generateKitfile(dirContents, opts.repo, tmpDir, opts.depth)
 		if err != nil {
 			return err
 		}

pkg/cmd/kitimport/hfimport.go

@@ -70,7 +70,7 @@ func importUsingHF(ctx context.Context, opts *importOptions) error {
 		}
 		kitfile = kf
 	} else {
-		kf, err := generateKitfile(dirListing, repo, tmpDir)
+		kf, err := generateKitfile(dirListing, repo, tmpDir, opts.depth)
 		if err != nil {
 			return err
 		}

pkg/cmd/kitimport/util.go

@@ -42,7 +42,7 @@ import (
 
 var ErrNoEditorFound = errors.New("no editor found")
 
-func generateKitfile(dirContents *kfgen.DirectoryListing, repo string, outDir string) (*artifact.KitFile, error) {
+func generateKitfile(dirContents *kfgen.DirectoryListing, repo string, outDir string, depth int) (*artifact.KitFile, error) {
 	// Fill fields in package so that they're not empty in `kit list` later.
 	sections := strings.Split(repo, "/")
 	var modelPackage *artifact.Package
@@ -52,7 +52,7 @@ func generateKitfile(dirContents *kfgen.DirectoryListing, repo string, outDir st
 			Authors: []string{sections[len(sections)-2]},
 		}
 	}
-	kitfile, err := kfgen.GenerateKitfile(dirContents, modelPackage)
+	kitfile, err := kfgen.GenerateKitfile(dirContents, modelPackage, depth)
 	if err != nil {
 		return nil, fmt.Errorf("failed to generate Kitfile: %w", err)
 	}

pkg/cmd/kitinit/cmd.go

@@ -80,6 +80,7 @@ type initOptions struct {
 	repoRef             string
 	token               string
 	outputPath          string
+	depth               int
 	// Computed fields (remote only)
 	repo     string
 	repoType hf.RepositoryType
@@ -105,6 +106,7 @@ func InitCommand() *cobra.Command {
 	cmd.Flags().StringVar(&opts.repoRef, "ref", "main", "Branch or tag for remote repository (requires --remote)")
 	cmd.Flags().StringVar(&opts.token, "token", "", "Auth token for remote repository (requires --remote)")
 	cmd.Flags().StringVarP(&opts.outputPath, "output", "o", "", "Output path for generated Kitfile ('-' writes to stdout; default: Kitfile in directory for local, stdout for remote)")
+	cmd.Flags().IntVar(&opts.depth, "depth", 0, "Maximum directory depth to process when generating a Kitfile. Setting to -1 processes all files individually")
 	cmd.Flags().SortFlags = false
 	return cmd
 }
@@ -146,7 +148,7 @@ func runCommand(opts *initOptions) func(*cobra.Command, []string) error {
 func runInit(dirContents *kfgen.DirectoryListing, opts *initOptions) error {
 	modelPackage := buildPackageFromRepo(opts.repo, opts.modelkitName, opts.modelkitDescription, opts.modelkitAuthor)
 
-	kitfile, err := kfgen.GenerateKitfile(dirContents, modelPackage)
+	kitfile, err := kfgen.GenerateKitfile(dirContents, modelPackage, opts.depth)
 	if err != nil {
 		return output.Fatalf("Error generating Kitfile: %s", err)
 	}

pkg/lib/kitfile/generate/generate.go

@@ -84,8 +84,10 @@ var promptFilePatterns = []string{
 
 // Generate a basic Kitfile by looking at the contents of a directory. Parameter
 // packageOpt can be used to define metadata for the Kitfile (i.e. the package
-// section), which is left empty if the parameter is nil.
-func GenerateKitfile(dir *DirectoryListing, packageOpt *artifact.Package) (*artifact.KitFile, error) {
+// section), which is left empty if the parameter is nil. depth controls how many
+// levels of subdirectories are processed file-by-file; at depth=0 (the default),
+// subdirectories of the root dir are analyzed as whole units.
+func GenerateKitfile(dir *DirectoryListing, packageOpt *artifact.Package, depth int) (*artifact.KitFile, error) {
 	output.Logf(output.LogLevelTrace, "Generating Kitfile in %s", dir.Path)
 	kitfile := &artifact.KitFile{
 		ManifestVersion: "1.0.0",
@@ -113,15 +115,63 @@ func GenerateKitfile(dir *DirectoryListing, packageOpt *artifact.Package) (*arti
 		return kitfile, nil
 	}
 
-	// We can make sure all files are included by including a layer with path '.'
-	// However, we only want to do this if it is necessary
-	includeCatchallSection := false
-	// Dirs we don't know how to handle automatically.
-	var unprocessedDirPaths []string
-	// Metadata files; we want these to be either model parts (if there is a model)
-	// or datasets
-	var modelFiles, metadataFiles []FileListing
-	var detectedLicenseType string
+	modelFiles, metadataFiles, unknownFiles, detectedLicenseType := processSubdirFiles(kitfile, *dir, depth)
+
+	if len(unknownFiles) > 5 {
+		output.Logf(output.LogLevelTrace, "Unrecognized files found in %s; adding catch-all code layer", dir.Path)
+		kitfile.Code = append(kitfile.Code, artifact.Code{Path: "."})
+	} else if len(unknownFiles) > 0 {
+		output.Logf(output.LogLevelTrace, "Unrecognized files found in %s; adding as code layers")
+		for _, f := range unknownFiles {
+			kitfile.Code = append(kitfile.Code, artifact.Code{Path: f.Path})
+		}
+	}
+
+	if len(modelFiles) > 0 {
+		if err := addModelToKitfile(kitfile, modelFiles); err != nil {
+			return nil, fmt.Errorf("failed to add model to Kitfile: %w", err)
+		}
+		output.Logf(output.LogLevelTrace, "Adding metadata files as model parts")
+		for _, metadataFile := range metadataFiles {
+			kitfile.Model.Parts = append(kitfile.Model.Parts, artifact.ModelPart{Path: metadataFile.Path})
+		}
+	} else {
+		output.Logf(output.LogLevelTrace, "No model detected; adding metadata files as dataset layers")
+		for _, metadataFile := range metadataFiles {
+			kitfile.DataSets = append(kitfile.DataSets, artifact.DataSet{Path: metadataFile.Path})
+		}
+	}
+
+	// If we detected a license, try to attach it to the Kitfile section that makes sense
+	if kitfile.Model != nil && detectedLicenseType != "" {
+		kitfile.Model.License = detectedLicenseType
+	} else if len(kitfile.DataSets) == 1 {
+		kitfile.DataSets[0].License = detectedLicenseType
+	} else if len(kitfile.Code) == 1 {
+		kitfile.Code[0].License = detectedLicenseType
+	} else {
+		output.Logf(output.LogLevelTrace, "Unsure what license applies to, adding to Kitfile package")
+		kitfile.Package.License = detectedLicenseType
+	}
+
+	applySkillMetadataToPackage(kitfile, *dir)
+
+	return kitfile, nil
+}
+
+// processSubdirFiles classifies each file in dir individually and adds typed layers to
+// kitfile. depth controls how many additional subdirectory levels are processed
+// file-by-file before switching to addDirToKitfile (whole-directory analysis). Returns
+// model and metadata files for the caller to pair, any files whose type could not be
+// determined (so the caller can decide how to handle them), and any detected license
+// identifier. Unclassifiable directories are added directly to kitfile.Code.
+func processSubdirFiles(kitfile *artifact.KitFile, dir DirectoryListing, depth int) (modelFiles []FileListing, metadataFiles []FileListing, unknownFiles []FileListing, detectedLicenseType string) {
+	if found, _ := dirContainsSkillMD(dir); found {
+		output.Logf(output.LogLevelTrace, "Directory %s contains SKILL.md; treating as skill", dir.Path)
+		prompt, _ := buildPromptFromSkill(dir)
+		kitfile.Prompts = append(kitfile.Prompts, prompt)
+		return nil, nil, nil, ""
+	}
 
 	output.Logf(output.LogLevelTrace, "Reading directory contents")
 	for _, file := range dir.Files {
@@ -136,14 +186,14 @@ func GenerateKitfile(dir *DirectoryListing, packageOpt *artifact.Package) (*arti
 		if strings.HasPrefix(strings.ToLower(file.Name), "readme") {
 			output.Logf(output.LogLevelTrace, "Found readme file '%s'", file.Name)
 			kitfile.Docs = append(kitfile.Docs, artifact.Docs{
-				Path:        file.Name,
+				Path:        file.Path,
 				Description: "Readme file",
 			})
 			continue
 		} else if strings.HasPrefix(strings.ToLower(file.Name), "license") {
 			output.Logf(output.LogLevelTrace, "Found license file '%s'", file.Name)
 			kitfile.Docs = append(kitfile.Docs, artifact.Docs{
-				Path:        file.Name,
+				Path:        file.Path,
 				Description: "License file",
 			})
 			licenseType, err := detectLicense(file.Path)
@@ -173,98 +223,64 @@ func GenerateKitfile(dir *DirectoryListing, packageOpt *artifact.Package) (*arti
 			kitfile.DataSets = append(kitfile.DataSets, artifact.DataSet{Path: file.Path})
 		case fileTypePrompt:
 			kitfile.Prompts = append(kitfile.Prompts, artifact.Prompt{Path: file.Path})
+		case fileTypeCode:
+			kitfile.Code = append(kitfile.Code, artifact.Code{Path: file.Path})
 		default:
-			output.Logf(output.LogLevelTrace, "File %s is either code or unknown type. Will be added as a catch-all section", file.Path)
-			// File is either code or unknown; we'll have to include it in a catch-all section
-			includeCatchallSection = true
+			output.Logf(output.LogLevelTrace, "File %s is unknown type; will be included in code section", file.Path)
+			unknownFiles = append(unknownFiles, file)
 		}
 	}
 
-	for _, subDir := range dir.Subdirs {
-		dirModelFiles, err := addDirToKitfile(kitfile, subDir)
-		if err != nil {
-			output.Logf(output.LogLevelTrace, "Failed to determine type for directory %s: %s", subDir.Path, err)
-			unprocessedDirPaths = append(unprocessedDirPaths, subDir.Path)
-		}
-		modelFiles = append(modelFiles, dirModelFiles...)
-		continue
-	}
-
-	if len(modelFiles) > 0 {
-		if err := addModelToKitfile(kitfile, modelFiles); err != nil {
-			return nil, fmt.Errorf("failed to add model to Kitfile: %w", err)
-		}
-		output.Logf(output.LogLevelTrace, "Adding metadata files as model parts")
-		for _, metadataFile := range metadataFiles {
-			kitfile.Model.Parts = append(kitfile.Model.Parts, artifact.ModelPart{Path: metadataFile.Path})
+	if depth == 0 {
+		for _, subDir := range dir.Subdirs {
+			dirModelFiles := addDirToKitfile(kitfile, subDir)
+			modelFiles = append(modelFiles, dirModelFiles...)
 		}
 	} else {
-		output.Logf(output.LogLevelTrace, "No model detected; adding metadata files as dataset layers")
-		for _, metadataFile := range metadataFiles {
-			kitfile.DataSets = append(kitfile.DataSets, artifact.DataSet{Path: metadataFile.Path})
-		}
-	}
-
-	// Decide how to handle remaining paths. Either package them in one large code layer with basePath
-	// or as separate layers for each directory.
-	output.Logf(output.LogLevelTrace, "Unable to process %d paths in %s", len(unprocessedDirPaths), dir.Path)
-	if includeCatchallSection || len(unprocessedDirPaths) > 5 {
-		output.Logf(output.LogLevelTrace, "Adding catch-all code layer to include files in %s", dir.Path)
-		// Overwrite any code layers we added before; this is cleaner than e.g. having a layer for '.' and a layer for 'src'
-		kitfile.Code = []artifact.Code{{Path: "."}}
-	} else {
-		for _, path := range unprocessedDirPaths {
-			kitfile.Code = append(kitfile.Code, artifact.Code{Path: path})
+		for _, subDir := range dir.Subdirs {
+			// Ignore licenses in subdirectories -- it's unclear where to assign them
+			subModelFiles, subMetaFiles, subUnknownFiles, _ := processSubdirFiles(kitfile, subDir, depth-1)
+			modelFiles = append(modelFiles, subModelFiles...)
+			metadataFiles = append(metadataFiles, subMetaFiles...)
+			for _, f := range subUnknownFiles {
+				output.Logf(output.LogLevelTrace, "File %s is unknown type; adding as code layer", f.Path)
+				kitfile.Code = append(kitfile.Code, artifact.Code{Path: f.Path})
+			}
 		}
 	}
-
-	// If we detected a license, try to attach it to the Kitfile section that makes sense
-	if kitfile.Model != nil && detectedLicenseType != "" {
-		kitfile.Model.License = detectedLicenseType
-	} else if len(kitfile.DataSets) == 1 {
-		kitfile.DataSets[0].License = detectedLicenseType
-	} else if len(kitfile.Code) == 1 {
-		kitfile.Code[0].License = detectedLicenseType
-	} else {
-		output.Logf(output.LogLevelTrace, "Unsure what license applies to, adding to Kitfile package")
-		kitfile.Package.License = detectedLicenseType
-	}
-
-	applySkillMetadataToPackage(kitfile, *dir)
-
-	return kitfile, nil
+	return modelFiles, metadataFiles, unknownFiles, detectedLicenseType
 }
 
-func addDirToKitfile(kitfile *artifact.KitFile, dir DirectoryListing) (modelFiles []FileListing, err error) {
+func addDirToKitfile(kitfile *artifact.KitFile, dir DirectoryListing) (modelFiles []FileListing) {
 	if found, _ := dirContainsSkillMD(dir); found {
 		output.Logf(output.LogLevelTrace, "Directory %s contains SKILL.md; treating as skill", dir.Path)
 		prompt, _ := buildPromptFromSkill(dir)
 		kitfile.Prompts = append(kitfile.Prompts, prompt)
-		return nil, nil
+		return nil
 	}
 
 	switch dir.Name {
 	case "docs":
 		output.Logf(output.LogLevelTrace, "Directory %s interpreted as documentation", dir.Name)
 		kitfile.Docs = append(kitfile.Docs, artifact.Docs{
-			Path: dir.Path,
+			Path: withTrailingSlash(dir.Path),
 		})
-		return nil, nil
+		return nil
 	case "src", "pkg", "lib", "build":
 		output.Logf(output.LogLevelTrace, "Directory %s interpreted as code", dir.Name)
 		kitfile.Code = append(kitfile.Code, artifact.Code{
-			Path: dir.Path,
+			Path: withTrailingSlash(dir.Path),
 		})
-		return nil, nil
+		return nil
 	}
 
 	// Sort entries in the directory to try and figure out what it contains. We'll reuse the
 	// fact that the fileTypes are enumerated using iota (and so are ints) to index correctly.
 	// Avoid using maps here since they iterate in a random order.
 	directoryContents := [int(fileTypeUnknown) + 1][]string{}
 	for _, subdir := range dir.Subdirs {
-		// We can, in the future, recurse deeper into the directory tree here. For now, treat secondary dirs as unknowns
-		directoryContents[int(fileTypeUnknown)] = append(directoryContents[int(fileTypeUnknown)], subdir.Path)
+		// We can, in the future, recurse deeper into the directory tree here. For now, treat additional dirs as unknowns
+		directoryContents[int(fileTypeUnknown)] = append(directoryContents[int(fileTypeUnknown)], filepath.Join(withTrailingSlash(dir.Path), subdir.Path))
 	}
 
 	var metadataFiles []FileListing
@@ -292,7 +308,9 @@ func addDirToKitfile(kitfile *artifact.KitFile, dir DirectoryListing) (modelFile
 		}
 	}
 	if directoryHasMixedContents {
-		return modelFiles, fmt.Errorf("mixed content in directory; unable to determine type")
+		output.Logf(output.LogLevelTrace, "Mixed contents in directory %s, adding as code layer", dir.Path)
+		kitfile.Code = append(kitfile.Code, artifact.Code{Path: withTrailingSlash(dir.Path)})
+		return modelFiles
 	}
 	switch overallFiletype {
 	case fileTypeModel:
@@ -301,21 +319,22 @@ func addDirToKitfile(kitfile *artifact.KitFile, dir DirectoryListing) (modelFile
 		modelFiles = append(modelFiles, metadataFiles...)
 	case fileTypeDataset:
 		output.Logf(output.LogLevelTrace, "Interpreting directory %s as a dataset directory", dir.Path)
-		kitfile.DataSets = append(kitfile.DataSets, artifact.DataSet{Path: dir.Path})
+		kitfile.DataSets = append(kitfile.DataSets, artifact.DataSet{Path: withTrailingSlash(dir.Path)})
 	case fileTypeDocs:
 		output.Logf(output.LogLevelTrace, "Interpreting directory %s as a docs directory", dir.Path)
-		kitfile.Docs = append(kitfile.Docs, artifact.Docs{Path: dir.Path})
+		kitfile.Docs = append(kitfile.Docs, artifact.Docs{Path: withTrailingSlash(dir.Path)})
 	case fileTypePrompt:
 		output.Logf(output.LogLevelTrace, "Interpreting directory %s as a prompts directory", dir.Path)
-		kitfile.Prompts = append(kitfile.Prompts, artifact.Prompt{Path: dir.Path})
+		kitfile.Prompts = append(kitfile.Prompts, artifact.Prompt{Path: withTrailingSlash(dir.Path)})
+	case fileTypeCode:
+		output.Logf(output.LogLevelTrace, "Interpreting directory %s as code directory", dir.Path)
+		kitfile.Code = append(kitfile.Code, artifact.Code{Path: withTrailingSlash(dir.Path)})
 	default:
-		output.Logf(output.LogLevelTrace, "Could not determine type for directory %s", dir.Path)
-		// If it's overall code, metadata, or unknown, just return it as unprocessed and let it be added as a Code section
-		// later
-		return modelFiles, fmt.Errorf("directory should be handled as Code")
+		output.Logf(output.LogLevelTrace, "Could not determine type for directory %s. Adding as code layer", dir.Path)
+		kitfile.Code = append(kitfile.Code, artifact.Code{Path: withTrailingSlash(dir.Path)})
 	}
 
-	return modelFiles, nil
+	return modelFiles
 }
 
 func determineFileType(filename string) fileType {
@@ -344,7 +363,6 @@ func determineFileType(filename string) fileType {
 		return fileTypeDataset
 	}
 	return fileTypeUnknown
-
 }
 
 func addModelToKitfile(kitfile *artifact.KitFile, files []FileListing) error {
@@ -440,3 +458,10 @@ func anyPattern(query string, patterns []string) bool {
 	}
 	return false
 }
+
+func withTrailingSlash(p string) string {
+	if p == "." || strings.HasSuffix(p, "/") {
+		return p
+	}
+	return p + "/"
+}

pkg/lib/kitfile/generate/skill.go

@@ -49,7 +49,7 @@ func dirContainsSkillMD(dir DirectoryListing) (bool, string) {
 
 func buildPromptFromSkill(dir DirectoryListing) (artifact.Prompt, *skill.SkillFrontmatter) {
 	prompt := artifact.Prompt{
-		Path: dir.Path,
+		Path: withTrailingSlash(dir.Path),
 	}
 
 	found, skillPath := dirContainsSkillMD(dir)