Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 86 additions & 6 deletions catalog/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -320,33 +320,113 @@ catalogs:
enabled: true
# Required: List of model identifiers to include
# Format: "organization/model-name" or "username/model-name"
# Supports wildcard patterns: "organization/*" or "organization/prefix*"
includedModels:
- "meta-llama/Llama-3.1-8B-Instruct"
- "ibm-granite/granite-4.0-h-small"
- "microsoft/phi-2"

- "microsoft/phi-3*" # All models starting with "phi-3"

# Optional: Exclude specific models or patterns
# Supports exact matches or patterns ending with "*"
excludedModels:
- "some-org/unwanted-model"
- "another-org/test-*" # Excludes all models starting with "test-"

# Optional: Configure a custom environment variable name for the API key
# Defaults to "HF_API_KEY" if not specified
properties:
apiKeyEnvVar: "MY_CUSTOM_API_KEY_VAR"
```

#### Organization-Restricted Sources

You can restrict a source to only fetch models from a specific organization using the `allowedOrganization` property. This automatically prefixes all model patterns with the organization name:

```yaml
catalogs:
- name: "Meta LLaMA Models"
id: "meta-llama-models"
type: "hf"
enabled: true
properties:
allowedOrganization: "meta-llama"
apiKeyEnvVar: "HF_API_KEY"
includedModels:
# These patterns are automatically prefixed with "meta-llama/"
- "*" # Expands to: meta-llama/*
- "Llama-3*" # Expands to: meta-llama/Llama-3*
- "CodeLlama-*" # Expands to: meta-llama/CodeLlama-*
excludedModels:
- "*-4bit" # Excludes: meta-llama/*-4bit
- "*-GGUF" # Excludes: meta-llama/*-GGUF
```

**Benefits of organization-restricted sources:**
- **Simplified configuration**: No need to repeat organization name in every pattern
- **Security**: Prevents accidental inclusion of models from other organizations
- **Convenience**: Use `"*"` to get all models from an organization
- **Performance**: Optimized API calls when fetching from a single organization

#### Model Filtering

Both `includedModels` and `excludedModels` are top-level properties (not nested under `properties`):

- **`includedModels`** (required): List of model identifiers to fetch from Hugging Face. Format: `"organization/model-name"` or `"username/model-name"`
- **`includedModels`** (required): List of model identifiers to fetch from Hugging Face
- **`excludedModels`** (optional): List of models or patterns to exclude from the results

The `excludedModels` property supports:
#### Supported Pattern Types

**Exact Model Names:**
```yaml
includedModels:
- "meta-llama/Llama-3.1-8B-Instruct" # Specific model
- "microsoft/phi-2" # Specific model
```

**Wildcard Patterns:**

In `includedModels`, wildcards can match model names by a prefix.

```yaml
includedModels:
- "microsoft/phi-*" # All models starting with "phi-"
- "meta-llama/Llama-3*" # All models starting with "Llama-3"
- "huggingface/*" # All models from huggingface organization
```

**Organization-Only Patterns (with `allowedOrganization`):**
```yaml
properties:
allowedOrganization: "meta-llama"
includedModels:
- "*" # All models from meta-llama organization
- "Llama-3*" # All meta-llama models starting with "Llama-3"
- "CodeLlama-*" # All meta-llama models starting with "CodeLlama-"
```

#### Pattern Validation

**Valid patterns:**
- `"org/model"` - Exact model name
- `"org/prefix*"` - Models starting with prefix
- `"org/*"` - All models from organization
- `"*"` - All models (only when using `allowedOrganization`)

**Invalid patterns (will be rejected):**
- `"*"` - Global wildcard (without `allowedOrganization`)
- `"*/*"` - Global organization wildcard
- `"org*"` - Wildcard in organization name
- `"org/"` - Empty model name
- `"*prefix*"` - Multiple wildcards
Comment on lines +409 to +420
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Doesn't mention if it's valid to have wildcard at beginning or middle: *something or some*thing. Maybe it's implied well enough. Just noting in case we want to add now or in a follow-up

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good call out. I pushed a commit that attempts to clarify this. includedModels only allows a prefix or an exact name, but excludedModels supports wildcards anywhere.


#### Exclusion Patterns

The `excludedModels` property supports prefixes like `includedModels` and also suffixes and mid-name wildcards:
- **Exact matches**: `"meta-llama/Llama-3.1-8B-Instruct"` - excludes this specific model
- **Pattern matching**: `"test-*"` - excludes all models starting with "test-"
- **Pattern matching**:
- `"*-draft"` - excludes all models ending with "-draft"
- `"Llama-3.*-Instruct"` - excludes all Llama 3.x models ending with "-Instruct"
- **Organization patterns**: `"test-org/*"` - excludes all models from test-org

## Development

Expand Down
132 changes: 118 additions & 14 deletions catalog/internal/catalog/hf_catalog.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ const (
apiKeyEnvVarKey = "apiKeyEnvVar"
maxModelsKey = "maxModels"
syncIntervalKey = "syncInterval"
allowedOrgKey = "allowedOrganization"

// defaultMaxModels is the default limit for models fetched PER PATTERN.
// This limit is applied independently to each pattern in includedModels
Expand Down Expand Up @@ -360,13 +361,72 @@ func (p *hfModelProvider) Models(ctx context.Context) (<-chan ModelProviderRecor
return ch, nil
}

// expandModelNames takes a list of model identifiers (which may include wildcards)
// and returns a list of concrete model names by expanding any wildcard patterns.
// Uses the same logic as FetchModelNamesForPreview.
func (p *hfModelProvider) expandModelNames(ctx context.Context, modelIdentifiers []string) ([]string, error) {
var allNames []string
var failedPatterns []string
var wildcardPatterns []string

for _, pattern := range modelIdentifiers {
select {
case <-ctx.Done():
return nil, ctx.Err()
default:
}

patternType, org, searchPrefix := parseModelPattern(pattern)

switch patternType {
case PatternInvalid:
return nil, fmt.Errorf("wildcard pattern %q is not supported - Hugging Face requires a specific organization (e.g., 'ibm-granite/*' or 'meta-llama/Llama-2-*')", pattern)

case PatternOrgAll, PatternOrgPrefix:
wildcardPatterns = append(wildcardPatterns, pattern)
glog.Infof("Expanding wildcard pattern: %s (org=%s, prefix=%s)", pattern, org, searchPrefix)
models, err := p.listModelsByAuthor(ctx, org, searchPrefix)
if err != nil {
failedPatterns = append(failedPatterns, pattern)
glog.Warningf("Failed to expand wildcard pattern %s: %v", pattern, err)
continue
}
allNames = append(allNames, models...)

case PatternExact:
// Direct model name - no expansion needed
allNames = append(allNames, pattern)
}
}

// Check error conditions for wildcard pattern failures
if len(wildcardPatterns) > 0 && len(allNames) == 0 {
// All wildcard patterns failed AND no results from exact patterns - this is an error
if len(failedPatterns) > 0 {
return nil, fmt.Errorf("no models found: %v", failedPatterns)
} else {
return nil, fmt.Errorf("no models found")
}
} else if len(failedPatterns) > 0 {
// Some patterns failed but we have results - log warning and continue with partial results
glog.Warningf("Some wildcard patterns failed to expand and were skipped: %v", failedPatterns)
}

return allNames, nil
}

func (p *hfModelProvider) getModelsFromHF(ctx context.Context) ([]ModelProviderRecord, error) {
var records []ModelProviderRecord
// First expand any wildcard patterns to concrete model names
expandedModels, err := p.expandModelNames(ctx, p.includedModels)
if err != nil {
return nil, fmt.Errorf("failed to expand model patterns: %w", err)
}

var records []ModelProviderRecord
currentTime := time.Now().UnixMilli()
lastSyncedStr := strconv.FormatInt(currentTime, 10)

for _, modelName := range p.includedModels {
for _, modelName := range expandedModels {
// Skip if excluded - check before fetching to avoid unnecessary API calls
if !p.filter.Allows(modelName) {
glog.V(2).Infof("Skipping excluded model: %s", modelName)
Expand Down Expand Up @@ -726,6 +786,9 @@ func newHFModelProvider(ctx context.Context, source *Source, reldir string) (<-c
p.baseURL = strings.TrimSuffix(url, "/")
}

allowedOrg, _ := source.Properties[allowedOrgKey].(string)
restrictToOrg(allowedOrg, &source.IncludedModels, &source.ExcludedModels)

// Parse sync interval (optional, defaults to 24 hours)
// This can be configured as a duration string (e.g., "1s", "10s", "1m", "24h").
// For testing, a shorter interval can be used to speed up tests.
Expand Down Expand Up @@ -805,6 +868,13 @@ func NewHFPreviewProvider(config *PreviewConfig) (*hfModelProvider, error) {
p.baseURL = strings.TrimSuffix(url, "/")
}

allowedOrg, _ := config.Properties[allowedOrgKey].(string)
restrictToOrg(allowedOrg, &config.IncludedModels, &config.ExcludedModels)

if len(config.IncludedModels) == 0 {
return nil, fmt.Errorf("includedModels is required for HuggingFace source preview (specifies which models to fetch from HuggingFace)")
}

// Parse maxModels limit (optional, defaults to 500)
// This limit is applied PER PATTERN (e.g., each "org/*" pattern gets its own limit)
// to prevent overloading the Hugging Face API and respect rate limiting.
Expand Down Expand Up @@ -868,19 +938,27 @@ func parseModelPattern(pattern string) (PatternType, string, string) {
return PatternOrgAll, org, ""
}

parts := strings.SplitN(pattern, "/", 2)

org := parts[0]
// Ensure org is not empty or a wildcard
if org == "" || strings.Contains(org, "*") {
return PatternInvalid, "", ""
}

var model string
if len(parts) == 2 {
model = parts[1]
if model == "" {
return PatternInvalid, "", ""
}
}

// Check if it has a wildcard after org/prefix
if strings.Contains(pattern, "/") && strings.HasSuffix(pattern, "*") {
parts := strings.SplitN(pattern, "/", 2)
if len(parts) == 2 {
org := parts[0]
// Ensure org is not empty or a wildcard
if org == "" || org == "*" {
return PatternInvalid, "", ""
}
prefix := strings.TrimSuffix(parts[1], "*")
if prefix != "" {
return PatternOrgPrefix, org, prefix
}
if strings.HasSuffix(model, "*") {
prefix := strings.TrimSuffix(model, "*")
if prefix != "" {
return PatternOrgPrefix, org, prefix
}
}

Expand Down Expand Up @@ -1085,3 +1163,29 @@ func (p *hfModelProvider) FetchModelNamesForPreview(ctx context.Context, modelId

return names, nil
}

// restrictToOrg prefixes included and excluded model lists with an
// organization name for convenience and to prevent any other organization from
// being retrieved.
func restrictToOrg(org string, included *[]string, excluded *[]string) {
if org == "" {
// No op
return
}

prefix := org + "/"

if included == nil || len(*included) == 0 {
*included = []string{prefix + "*"}
} else {
for i := range *included {
(*included)[i] = prefix + (*included)[i]
}
}

if excluded != nil {
for i := range *excluded {
(*excluded)[i] = prefix + (*excluded)[i]
}
}
}
Loading
Loading