Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
32 changes: 30 additions & 2 deletions pkg/inference/models/api.go
Original file line number Diff line number Diff line change
Expand Up @@ -82,12 +82,38 @@ func ToOpenAI(m types.Model) (*OpenAIModel, error) {
id = tags[0]
}

return &OpenAIModel{
model := &OpenAIModel{
ID: id,
Object: "model",
Created: created,
OwnedBy: "docker",
}, nil
}

config, err := m.Config()
if err != nil {
return nil, fmt.Errorf("get config: %w", err)
}

if config != nil {
model.DMR = &DMRMetadata{
ContextWindow: config.GetContextSize(),
Architecture: config.GetArchitecture(),
Parameters: config.GetParameters(),
Quantization: config.GetQuantization(),
Size: config.GetSize(),
}
}

return model, nil
}

// DMRMetadata contains Docker Model Runner-specific metadata about a model.
type DMRMetadata struct {
ContextWindow *int32 `json:"context_window,omitempty"`
Architecture string `json:"architecture,omitempty"`
Parameters string `json:"parameters,omitempty"`
Quantization string `json:"quantization,omitempty"`
Size string `json:"size,omitempty"`
}

// OpenAIModel represents a locally stored model using OpenAI conventions.
Expand All @@ -100,6 +126,8 @@ type OpenAIModel struct {
Created int64 `json:"created"`
// OwnedBy is the model owner. At the moment, it is always "docker".
OwnedBy string `json:"owned_by"`
// DMR contains Docker Model Runner-specific metadata.
DMR *DMRMetadata `json:"dmr,omitempty"`
}

// OpenAIModelList represents a list of models using OpenAI conventions.
Expand Down
166 changes: 166 additions & 0 deletions pkg/inference/models/api_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,172 @@
}
}

// mockModel implements types.Model for testing ToOpenAI.
type mockModel struct {
id string
tags []string
config types.ModelConfig
desc types.Descriptor
}

func (m *mockModel) ID() (string, error) { return m.id, nil }

Check failure on line 323 in pkg/inference/models/api_test.go

View workflow job for this annotation

GitHub Actions / lint (linux)

File is not properly formatted (gci)
func (m *mockModel) Tags() []string { return m.tags }
func (m *mockModel) Config() (types.ModelConfig, error) { return m.config, nil }
func (m *mockModel) Descriptor() (types.Descriptor, error) { return m.desc, nil }
func (m *mockModel) GGUFPaths() ([]string, error) { return nil, nil }
func (m *mockModel) SafetensorsPaths() ([]string, error) { return nil, nil }
func (m *mockModel) DDUFPaths() ([]string, error) { return nil, nil }
func (m *mockModel) ConfigArchivePath() (string, error) { return "", nil }
func (m *mockModel) MMPROJPath() (string, error) { return "", nil }
func (m *mockModel) ChatTemplatePath() (string, error) { return "", nil }

func TestToOpenAIWithFullConfig(t *testing.T) {
m := &mockModel{
id: "sha256:abc123",
tags: []string{"ai/smollm2:latest"},
config: &types.Config{
Format: "gguf",
Quantization: "Q4_K_M",
Parameters: "1.7B",
Architecture: "llama",
Size: "1.7B",
ContextSize: int32Ptr(8192),
},
desc: types.Descriptor{},
}

result, err := ToOpenAI(m)
require.NoError(t, err)

assert.Equal(t, "ai/smollm2:latest", result.ID)
assert.Equal(t, "model", result.Object)
assert.Equal(t, "docker", result.OwnedBy)

require.NotNil(t, result.DMR)
require.NotNil(t, result.DMR.ContextWindow)
assert.Equal(t, int32(8192), *result.DMR.ContextWindow)
assert.Equal(t, "llama", result.DMR.Architecture)
assert.Equal(t, "1.7B", result.DMR.Parameters)
assert.Equal(t, "Q4_K_M", result.DMR.Quantization)
assert.Equal(t, "1.7B", result.DMR.Size)
}

func TestToOpenAIWithNilConfig(t *testing.T) {
m := &mockModel{
id: "sha256:abc123",
tags: []string{"ai/model:latest"},
desc: types.Descriptor{},
}

result, err := ToOpenAI(m)
require.NoError(t, err)

assert.Equal(t, "ai/model:latest", result.ID)
assert.Equal(t, "model", result.Object)
assert.Equal(t, "docker", result.OwnedBy)
assert.Nil(t, result.DMR)
}

func TestToOpenAIWithoutTags(t *testing.T) {
m := &mockModel{
id: "sha256:abc123",
desc: types.Descriptor{},
config: &types.Config{
Architecture: "mistral",
},
}

result, err := ToOpenAI(m)
require.NoError(t, err)

assert.Equal(t, "sha256:abc123", result.ID)
require.NotNil(t, result.DMR)
assert.Equal(t, "mistral", result.DMR.Architecture)
}

func TestToOpenAIDMROmittedWhenNilConfig(t *testing.T) {
m := &mockModel{
id: "sha256:abc123",
tags: []string{"ai/model:latest"},
desc: types.Descriptor{},
}

result, err := ToOpenAI(m)
require.NoError(t, err)

data, err := json.Marshal(result)
require.NoError(t, err)

var raw map[string]interface{}
err = json.Unmarshal(data, &raw)
require.NoError(t, err)

_, hasDMR := raw["dmr"]
assert.False(t, hasDMR, "dmr field should be omitted when config is nil")
}

func TestToOpenAIContextWindowOmittedWhenNil(t *testing.T) {
m := &mockModel{
id: "sha256:abc123",
tags: []string{"ai/model:latest"},
desc: types.Descriptor{},
config: &types.Config{
Architecture: "llama",
},
}

result, err := ToOpenAI(m)
require.NoError(t, err)

require.NotNil(t, result.DMR)
assert.Nil(t, result.DMR.ContextWindow)

data, err := json.Marshal(result)
require.NoError(t, err)

var raw map[string]interface{}
err = json.Unmarshal(data, &raw)
require.NoError(t, err)

dmr, ok := raw["dmr"].(map[string]interface{})
require.True(t, ok)
_, hasCtxWindow := dmr["context_window"]
assert.False(t, hasCtxWindow, "context_window should be omitted when nil")
}

func TestToOpenAIList(t *testing.T) {
models := []types.Model{
&mockModel{
id: "sha256:aaa",
tags: []string{"ai/model1:latest"},
desc: types.Descriptor{},
config: &types.Config{
Architecture: "llama",
Parameters: "7B",
},
},
&mockModel{
id: "sha256:bbb",
tags: []string{"ai/model2:latest"},
desc: types.Descriptor{},
},
}

result, err := ToOpenAIList(models)
require.NoError(t, err)

assert.Equal(t, "list", result.Object)
require.Len(t, result.Data, 2)

assert.Equal(t, "ai/model1:latest", result.Data[0].ID)
require.NotNil(t, result.Data[0].DMR)
assert.Equal(t, "llama", result.Data[0].DMR.Architecture)
assert.Equal(t, "7B", result.Data[0].DMR.Parameters)

assert.Equal(t, "ai/model2:latest", result.Data[1].ID)
assert.Nil(t, result.Data[1].DMR)
}

// Helper function to create int32 pointers
func int32Ptr(i int32) *int32 {
return &i
Expand Down
Loading