ScuttleBot
| 5ac549c… | lmata | 1 | package llm |
| 5ac549c… | lmata | 2 | |
| 5ac549c… | lmata | 3 | import ( |
| 5ac549c… | lmata | 4 | "fmt" |
| 5ac549c… | lmata | 5 | "regexp" |
| 5ac549c… | lmata | 6 | ) |
| 5ac549c… | lmata | 7 | |
| 5ac549c… | lmata | 8 | // ModelFilter applies allow/block regex patterns to a slice of ModelInfo. |
| 5ac549c… | lmata | 9 | // Patterns are matched against model ID. |
| 5ac549c… | lmata | 10 | type ModelFilter struct { |
| 5ac549c… | lmata | 11 | allowlist []*regexp.Regexp |
| 5ac549c… | lmata | 12 | blocklist []*regexp.Regexp |
| 5ac549c… | lmata | 13 | } |
| 5ac549c… | lmata | 14 | |
| 5ac549c… | lmata | 15 | // NewModelFilter compiles regex allow/block patterns. |
| 5ac549c… | lmata | 16 | // Returns an error if any pattern is invalid. |
| 5ac549c… | lmata | 17 | func NewModelFilter(allow, block []string) (*ModelFilter, error) { |
| 5ac549c… | lmata | 18 | f := &ModelFilter{} |
| 5ac549c… | lmata | 19 | for _, pat := range allow { |
| 5ac549c… | lmata | 20 | r, err := regexp.Compile(pat) |
| 5ac549c… | lmata | 21 | if err != nil { |
| 5ac549c… | lmata | 22 | return nil, fmt.Errorf("llm: allow pattern %q: %w", pat, err) |
| 5ac549c… | lmata | 23 | } |
| 5ac549c… | lmata | 24 | f.allowlist = append(f.allowlist, r) |
| 5ac549c… | lmata | 25 | } |
| 5ac549c… | lmata | 26 | for _, pat := range block { |
| 5ac549c… | lmata | 27 | r, err := regexp.Compile(pat) |
| 5ac549c… | lmata | 28 | if err != nil { |
| 5ac549c… | lmata | 29 | return nil, fmt.Errorf("llm: block pattern %q: %w", pat, err) |
| 5ac549c… | lmata | 30 | } |
| 5ac549c… | lmata | 31 | f.blocklist = append(f.blocklist, r) |
| 5ac549c… | lmata | 32 | } |
| 5ac549c… | lmata | 33 | return f, nil |
| 5ac549c… | lmata | 34 | } |
| 5ac549c… | lmata | 35 | |
| 5ac549c… | lmata | 36 | // Apply filters models: removes those matching any blocklist pattern, then |
| 5ac549c… | lmata | 37 | // (if allowlist is non-empty) keeps only those matching at least one allowlist pattern. |
| 5ac549c… | lmata | 38 | func (f *ModelFilter) Apply(models []ModelInfo) []ModelInfo { |
| 5ac549c… | lmata | 39 | out := make([]ModelInfo, 0, len(models)) |
| 5ac549c… | lmata | 40 | for _, m := range models { |
| 5ac549c… | lmata | 41 | if f.blocked(m.ID) { |
| 5ac549c… | lmata | 42 | continue |
| 5ac549c… | lmata | 43 | } |
| 5ac549c… | lmata | 44 | if len(f.allowlist) > 0 && !f.allowed(m.ID) { |
| 5ac549c… | lmata | 45 | continue |
| 5ac549c… | lmata | 46 | } |
| 5ac549c… | lmata | 47 | out = append(out, m) |
| 5ac549c… | lmata | 48 | } |
| 5ac549c… | lmata | 49 | return out |
| 5ac549c… | lmata | 50 | } |
| 5ac549c… | lmata | 51 | |
| 5ac549c… | lmata | 52 | func (f *ModelFilter) allowed(id string) bool { |
| 5ac549c… | lmata | 53 | for _, r := range f.allowlist { |
| 5ac549c… | lmata | 54 | if r.MatchString(id) { |
| 5ac549c… | lmata | 55 | return true |
| 5ac549c… | lmata | 56 | } |
| 5ac549c… | lmata | 57 | } |
| 5ac549c… | lmata | 58 | return false |
| 5ac549c… | lmata | 59 | } |
| 5ac549c… | lmata | 60 | |
| 5ac549c… | lmata | 61 | func (f *ModelFilter) blocked(id string) bool { |
| 5ac549c… | lmata | 62 | for _, r := range f.blocklist { |
| 5ac549c… | lmata | 63 | if r.MatchString(id) { |
| 5ac549c… | lmata | 64 | return true |
| 5ac549c… | lmata | 65 | } |
| 5ac549c… | lmata | 66 | } |
| 5ac549c… | lmata | 67 | return false |
| 5ac549c… | lmata | 68 | } |