131 lines
3.5 KiB
Go
131 lines
3.5 KiB
Go
package server
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/ollama/ollama/fs/ggml"
|
|
"github.com/ollama/ollama/template"
|
|
"github.com/ollama/ollama/types/model"
|
|
)
|
|
|
|
func TestModelCheckCapabilities(t *testing.T) {
|
|
// Create simple model file for tests that don't depend on GGUF content
|
|
completionModelPath, _ := createBinFile(t, ggml.KV{
|
|
"general.architecture": "llama",
|
|
}, []*ggml.Tensor{})
|
|
|
|
// Create vision model (llama architecture with vision block count)
|
|
visionModelPath, _ := createBinFile(t, ggml.KV{
|
|
"general.architecture": "llama",
|
|
"llama.vision.block_count": uint32(1),
|
|
}, []*ggml.Tensor{})
|
|
|
|
// Create embedding model (bert architecture with pooling type)
|
|
embeddingModelPath, _ := createBinFile(t, ggml.KV{
|
|
"general.architecture": "bert",
|
|
"bert.pooling_type": uint32(1),
|
|
}, []*ggml.Tensor{})
|
|
|
|
toolsInsertTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}{{ if .suffix }}{{ .suffix }}{{ end }}")
|
|
if err != nil {
|
|
t.Fatalf("Failed to parse template: %v", err)
|
|
}
|
|
|
|
chatTemplate, err := template.Parse("{{ .prompt }}")
|
|
if err != nil {
|
|
t.Fatalf("Failed to parse template: %v", err)
|
|
}
|
|
|
|
toolsTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}")
|
|
if err != nil {
|
|
t.Fatalf("Failed to parse template: %v", err)
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
model Model
|
|
checkCaps []model.Capability
|
|
expectedErrMsg string
|
|
}{
|
|
{
|
|
name: "completion model without tools capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
checkCaps: []model.Capability{model.CapabilityTools},
|
|
expectedErrMsg: "does not support tools",
|
|
},
|
|
{
|
|
name: "model with all needed capabilities",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: toolsInsertTemplate,
|
|
},
|
|
checkCaps: []model.Capability{model.CapabilityTools, model.CapabilityInsert},
|
|
},
|
|
{
|
|
name: "model missing insert capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: toolsTemplate,
|
|
},
|
|
checkCaps: []model.Capability{model.CapabilityInsert},
|
|
expectedErrMsg: "does not support insert",
|
|
},
|
|
{
|
|
name: "model missing vision capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: toolsTemplate,
|
|
},
|
|
checkCaps: []model.Capability{model.CapabilityVision},
|
|
expectedErrMsg: "does not support vision",
|
|
},
|
|
{
|
|
name: "model with vision capability",
|
|
model: Model{
|
|
ModelPath: visionModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
checkCaps: []model.Capability{model.CapabilityVision},
|
|
},
|
|
{
|
|
name: "model with embedding capability",
|
|
model: Model{
|
|
ModelPath: embeddingModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
checkCaps: []model.Capability{model.CapabilityEmbedding},
|
|
},
|
|
{
|
|
name: "unknown capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
checkCaps: []model.Capability{"unknown"},
|
|
expectedErrMsg: "unknown capability",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
// Test CheckCapabilities method
|
|
err := tt.model.CheckCapabilities(tt.checkCaps...)
|
|
if tt.expectedErrMsg == "" {
|
|
if err != nil {
|
|
t.Errorf("Expected no error, got: %v", err)
|
|
}
|
|
} else {
|
|
if err == nil {
|
|
t.Errorf("Expected error containing %q, got nil", tt.expectedErrMsg)
|
|
} else if !strings.Contains(err.Error(), tt.expectedErrMsg) {
|
|
t.Errorf("Expected error containing %q, got: %v", tt.expectedErrMsg, err)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|