feat: Refactor enrichment command to make viper enrichment file work with file strategy mode

Author: wiebe-vandendriessche

cmd/enrich.go

@@ -87,12 +87,12 @@ from Hugging Face API and README before enrichment.`,
 		}
 
 		// Load config file values if using file strategy
-		var fileValues map[string]interface{}
+		var configViper *viper.Viper
 		if strategy == "file" {
 			if enrichConfigFile == "" {
 				return fmt.Errorf("--file is required when using --strategy file")
 			}
-			fileValues, err = loadEnrichmentConfig(enrichConfigFile)
+			configViper, err = loadEnrichmentConfig(enrichConfigFile)
 			if err != nil {
 				return fmt.Errorf("failed to load config file: %w", err)
 			}
@@ -106,7 +106,7 @@ from Hugging Face API and README before enrichment.`,
 		})
 
 		// Run enrichment
-		enriched, err := e.Enrich(bom, fileValues)
+		enriched, err := e.Enrich(bom, configViper)
 		if err != nil {
 			return fmt.Errorf("enrichment failed: %w", err)
 		}
@@ -150,7 +150,7 @@ func init() {
 	enrichCmd.Flags().StringVar(&enrichSpecVersion, "spec", "", "CycloneDX spec version for output (default: same as input)")
 
 	enrichCmd.Flags().StringVar(&enrichStrategy, "strategy", "interactive", "Enrichment strategy: interactive|file")
-	enrichCmd.Flags().StringVar(&enrichConfigFile, "file", "/config/enrichment.yaml", "Path to enrichment config file (YAML)")
+	enrichCmd.Flags().StringVar(&enrichConfigFile, "file", "./config/enrichment.yaml", "Path to enrichment config file (YAML)")
 	enrichCmd.Flags().BoolVar(&enrichRequiredOnly, "required-only", false, "Only prompt for required fields")
 	enrichCmd.Flags().Float64Var(&enrichMinWeight, "min-weight", 0.0, "Only prompt for fields with weight >= this value")
 	enrichCmd.Flags().BoolVar(&enrichRefetch, "refetch", false, "Refetch model metadata from Hugging Face before enrichment")
@@ -165,7 +165,7 @@ func init() {
 }
 
 // loadEnrichmentConfig loads enrichment values from a YAML config file
-func loadEnrichmentConfig(path string) (map[string]interface{}, error) {
+func loadEnrichmentConfig(path string) (*viper.Viper, error) {
 	v := viper.New()
 	v.SetConfigFile(path)
 	v.SetConfigType("yaml")
@@ -174,5 +174,5 @@ func loadEnrichmentConfig(path string) (map[string]interface{}, error) {
 		return nil, err
 	}
 
-	return v.AllSettings(), nil
+	return v, nil
 }

config/enrichment.yaml

@@ -1,32 +1,27 @@
-# Component-level fields
-licenses: Apache-2.0
-tags:
+# AIBoMGen Enrichment Configuration
+# This file provides values to enrich an existing AIBOM when using --strategy file
+#
+# IMPORTANT: Use the full BOM path as the key to avoid ambiguity
+# Keys are case-insensitive (they are automatically lowercased during processing)
+
+BOM.metadata.component.name: "my-model-name"
+BOM.metadata.component.licenses: "Apache-2.0"
+BOM.metadata.component.manufacturer: "Google"
+BOM.metadata.component.group: "google"
+BOM.metadata.component.tags:
   - nlp
   - transformer
   - bert
-manufacturer: Google
-group: google
-
-# Model card fields
-task: text-classification
-architectureFamily: transformer
-modelArchitecture: bert-base
-
-# Hugging Face properties
-language: en
-libraryName: transformers
 
-# Model card sections
-useCases: "Sentiment analysis, text classification"
-technicalLimitations: "Limited to English text, max 512 tokens"
-ethicalConsiderations: "May reflect biases present in training data"
+BOM.metadata.component.properties.huggingface:language: "en"
+BOM.metadata.component.properties.huggingface:libraryName: "transformers"
+BOM.metadata.component.properties.huggingface:baseModel: "distilbert-base-uncased"
+BOM.metadata.component.properties.huggingface:modelCardContact: "[email protected]"
 
-# External references
-paperUrl: https://arxiv.org/abs/1810.04805
-demoUrl: https://huggingface.co/spaces/demo
+BOM.metadata.component.modelCard.modelParameters.task: "text-classification"
+BOM.metadata.component.modelCard.modelParameters.architectureFamily: "transformer"
+BOM.metadata.component.modelCard.modelParameters.modelArchitecture: "bert-base"
 
-# Environmental impact
-hardwareType: "NVIDIA V100"
-hoursUsed: "96"
-cloudProvider: "Google Cloud"
-computeRegion: "us-central1"
+BOM.metadata.component.modelCard.considerations.useCases: "Text generation, dialogue systems, creative writing"
+BOM.metadata.component.modelCard.considerations.technicalLimitations: "Limited context window of 512 tokens, may generate inconsistent or factually incorrect text"
+BOM.metadata.component.modelCard.considerations.ethicalConsiderations: "May generate biased or harmful content based on training data. Should not be used for sensitive applications without additional safeguards."

internal/enricher/enricher.go

@@ -54,7 +54,7 @@ func New(opts Options) *Enricher {
 }
 
 // Enrich enriches a BOM with additional metadata
-func (e *Enricher) Enrich(bom *cdx.BOM, fileValues map[string]interface{}) (*cdx.BOM, error) {
+func (e *Enricher) Enrich(bom *cdx.BOM, configViper interface{}) (*cdx.BOM, error) {
 	if bom == nil {
 		return nil, fmt.Errorf("nil BOM")
 	}
@@ -111,7 +111,7 @@ func (e *Enricher) Enrich(bom *cdx.BOM, fileValues map[string]interface{}) (*cdx
 
 		switch e.config.Strategy {
 		case "file":
-			value, err = e.getValueFromFile(spec, fileValues)
+			value, err = e.getValueFromFile(spec, configViper)
 		case "interactive":
 			value, err = e.getValueInteractive(spec, src, tgt.BOM)
 		default:
@@ -209,12 +209,27 @@ func (e *Enricher) refetchMetadata(modelID string) (*fetcher.ModelAPIResponse, *
 }
 
 // getValueFromFile extracts a value from the config file
-func (e *Enricher) getValueFromFile(spec metadata.FieldSpec, values map[string]interface{}) (interface{}, error) {
-	// Map field spec key to config file key
-	// e.g. "BOM.metadata.component.licenses" -> "licenses"
-	key := configKeyFromSpec(spec.Key)
+func (e *Enricher) getValueFromFile(spec metadata.FieldSpec, configViper interface{}) (interface{}, error) {
+	// If no config provided, return nil
+	if configViper == nil {
+		return nil, nil
+	}
+
+	// Type assert to viper.Viper
+	type viperGetter interface {
+		Get(key string) interface{}
+	}
 
-	if val, ok := values[key]; ok {
+	v, ok := configViper.(viperGetter)
+	if !ok {
+		return nil, fmt.Errorf("invalid config type")
+	}
+
+	// Use the full key - viper will handle the nested lookup and lowercasing
+	key := string(spec.Key)
+	val := v.Get(key)
+
+	if val != nil {
 		logf("", "loaded %s from config file", spec.Key)
 		return val, nil
 	}
@@ -349,11 +364,8 @@ func bomModelCard(bom *cdx.BOM) *cdx.MLModelCard {
 }
 
 func configKeyFromSpec(key metadata.Key) string {
-	// Convert "BOM.metadata.component.licenses" to "licenses"
-	parts := strings.Split(string(key), ".")
-	if len(parts) > 0 {
-		return parts[len(parts)-1]
-	}
+	// Use the full key to avoid ambiguity
+	// e.g. "BOM.metadata.component.properties.huggingface:baseModel" stays as-is
 	return string(key)
 }