From af235653caa5bdb905114a63544c001842c1ad0b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 19 Feb 2026 22:21:48 +0000 Subject: [PATCH] feat: model discovery scanning directories for GGUF files Co-Authored-By: Virgil --- discover.go | 36 +++++++++++++ discover_test.go | 129 +++++++++++++++++++++++++++++++++++++++++++++++ rocm.go | 11 ++++ 3 files changed, 176 insertions(+) create mode 100644 discover.go create mode 100644 discover_test.go diff --git a/discover.go b/discover.go new file mode 100644 index 0000000..0d298ca --- /dev/null +++ b/discover.go @@ -0,0 +1,36 @@ +package rocm + +import ( + "path/filepath" + + "forge.lthn.ai/core/go-rocm/internal/gguf" +) + +// DiscoverModels scans a directory for GGUF model files and returns +// structured information about each. Files that cannot be parsed are skipped. +func DiscoverModels(dir string) ([]ModelInfo, error) { + matches, err := filepath.Glob(filepath.Join(dir, "*.gguf")) + if err != nil { + return nil, err + } + + var models []ModelInfo + for _, path := range matches { + meta, err := gguf.ReadMetadata(path) + if err != nil { + continue + } + + models = append(models, ModelInfo{ + Path: path, + Architecture: meta.Architecture, + Name: meta.Name, + Quantisation: gguf.FileTypeName(meta.FileType), + Parameters: meta.SizeLabel, + FileSize: meta.FileSize, + ContextLen: meta.ContextLength, + }) + } + + return models, nil +} diff --git a/discover_test.go b/discover_test.go new file mode 100644 index 0000000..c4cd5dd --- /dev/null +++ b/discover_test.go @@ -0,0 +1,129 @@ +package rocm + +import ( + "encoding/binary" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeDiscoverTestGGUF creates a minimal GGUF v3 file in dir with the given +// filename and metadata KV pairs. Returns the full path to the created file. +func writeDiscoverTestGGUF(t *testing.T, dir, filename string, kvs [][2]any) string { + t.Helper() + + path := filepath.Join(dir, filename) + + f, err := os.Create(path) + require.NoError(t, err) + defer f.Close() + + // Magic: "GGUF" in little-endian + require.NoError(t, binary.Write(f, binary.LittleEndian, uint32(0x46554747))) + // Version 3 + require.NoError(t, binary.Write(f, binary.LittleEndian, uint32(3))) + // Tensor count (uint64): 0 + require.NoError(t, binary.Write(f, binary.LittleEndian, uint64(0))) + // KV count (uint64) + require.NoError(t, binary.Write(f, binary.LittleEndian, uint64(len(kvs)))) + + for _, kv := range kvs { + key := kv[0].(string) + writeDiscoverKV(t, f, key, kv[1]) + } + + return path +} + +func writeDiscoverKV(t *testing.T, f *os.File, key string, val any) { + t.Helper() + + // Key: uint64 length + bytes + require.NoError(t, binary.Write(f, binary.LittleEndian, uint64(len(key)))) + _, err := f.Write([]byte(key)) + require.NoError(t, err) + + switch v := val.(type) { + case string: + // Type: 8 (string) + require.NoError(t, binary.Write(f, binary.LittleEndian, uint32(8))) + // String value: uint64 length + bytes + require.NoError(t, binary.Write(f, binary.LittleEndian, uint64(len(v)))) + _, err := f.Write([]byte(v)) + require.NoError(t, err) + case uint32: + // Type: 4 (uint32) + require.NoError(t, binary.Write(f, binary.LittleEndian, uint32(4))) + require.NoError(t, binary.Write(f, binary.LittleEndian, v)) + default: + t.Fatalf("writeDiscoverKV: unsupported value type %T", val) + } +} + +func TestDiscoverModels(t *testing.T) { + dir := t.TempDir() + + // Create two valid GGUF model files. + writeDiscoverTestGGUF(t, dir, "gemma3-4b-q4km.gguf", [][2]any{ + {"general.architecture", "gemma3"}, + {"general.name", "Gemma 3 4B Instruct"}, + {"general.file_type", uint32(15)}, + {"general.size_label", "4B"}, + {"gemma3.context_length", uint32(32768)}, + {"gemma3.block_count", uint32(34)}, + }) + + writeDiscoverTestGGUF(t, dir, "llama-3.1-8b-q4km.gguf", [][2]any{ + {"general.architecture", "llama"}, + {"general.name", "Llama 3.1 8B Instruct"}, + {"general.file_type", uint32(15)}, + {"general.size_label", "8B"}, + {"llama.context_length", uint32(131072)}, + {"llama.block_count", uint32(32)}, + }) + + // Create a non-GGUF file that should be ignored (no .gguf extension). + require.NoError(t, os.WriteFile(filepath.Join(dir, "README.txt"), []byte("not a model"), 0644)) + + models, err := DiscoverModels(dir) + require.NoError(t, err) + require.Len(t, models, 2) + + // Sort order from Glob is lexicographic, so gemma3 comes first. + gemma := models[0] + assert.Equal(t, filepath.Join(dir, "gemma3-4b-q4km.gguf"), gemma.Path) + assert.Equal(t, "gemma3", gemma.Architecture) + assert.Equal(t, "Gemma 3 4B Instruct", gemma.Name) + assert.Equal(t, "Q4_K_M", gemma.Quantisation) + assert.Equal(t, "4B", gemma.Parameters) + assert.Equal(t, uint32(32768), gemma.ContextLen) + assert.Greater(t, gemma.FileSize, int64(0)) + + llama := models[1] + assert.Equal(t, filepath.Join(dir, "llama-3.1-8b-q4km.gguf"), llama.Path) + assert.Equal(t, "llama", llama.Architecture) + assert.Equal(t, "Llama 3.1 8B Instruct", llama.Name) + assert.Equal(t, "Q4_K_M", llama.Quantisation) + assert.Equal(t, "8B", llama.Parameters) + assert.Equal(t, uint32(131072), llama.ContextLen) + assert.Greater(t, llama.FileSize, int64(0)) +} + +func TestDiscoverModels_EmptyDir(t *testing.T) { + dir := t.TempDir() + + models, err := DiscoverModels(dir) + require.NoError(t, err) + assert.Empty(t, models) +} + +func TestDiscoverModels_NotFound(t *testing.T) { + // filepath.Glob returns nil, nil for a pattern matching no files, + // even when the directory does not exist. + models, err := DiscoverModels("/nonexistent/dir") + require.NoError(t, err) + assert.Empty(t, models) +} diff --git a/rocm.go b/rocm.go index 7bdbf66..bea7178 100644 --- a/rocm.go +++ b/rocm.go @@ -30,3 +30,14 @@ type VRAMInfo struct { Used uint64 Free uint64 } + +// ModelInfo describes a GGUF model file discovered on disk. +type ModelInfo struct { + Path string // full path to .gguf file + Architecture string // GGUF architecture (e.g. "gemma3", "llama", "qwen2") + Name string // human-readable model name from GGUF metadata + Quantisation string // quantisation level (e.g. "Q4_K_M", "Q8_0") + Parameters string // parameter size label (e.g. "1B", "8B") + FileSize int64 // file size in bytes + ContextLen uint32 // native context window length +}