mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 13:38:14 -05:00
Patrick Devine
dfabde4539
Redirects for registry and blob transfers now validate the target scheme and resolved addresses before following, re-check DNS on each redirect, and do not follow redirects that switch an https session to plain http. The --insecure option continues to relax address checks for private registries but not scheme checks.
956 lines
30 KiB
Go
956 lines
30 KiB
Go
package server
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"os"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
gguftest "github.com/ollama/ollama/internal/testutil/gguf"
|
|
"github.com/ollama/ollama/manifest"
|
|
"github.com/ollama/ollama/template"
|
|
"github.com/ollama/ollama/types/model"
|
|
)
|
|
|
|
func TestPruneLayersSkipsRecentOrphans(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
|
|
recentDigest := "sha256:0000000000000000000000000000000000000000000000000000000000000001"
|
|
oldDigest := "sha256:0000000000000000000000000000000000000000000000000000000000000002"
|
|
|
|
for _, digest := range []string{recentDigest, oldDigest} {
|
|
p, err := manifest.BlobsPath(digest)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(p, nil, 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
oldPath, err := manifest.BlobsPath(oldDigest)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
oldTime := time.Now().Add(-layerPruneGracePeriod - time.Hour)
|
|
if err := os.Chtimes(oldPath, oldTime, oldTime); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if err := PruneLayers(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
recentPath, err := manifest.BlobsPath(recentDigest)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := os.Stat(recentPath); err != nil {
|
|
t.Fatalf("recent orphan was pruned: %v", err)
|
|
}
|
|
if _, err := os.Stat(oldPath); !os.IsNotExist(err) {
|
|
t.Fatalf("old orphan still exists: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestGenerationDefaultsFromMetadata(t *testing.T) {
|
|
file, err := os.CreateTemp(t.TempDir(), "model-*.gguf")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if err := gguftest.Write(file, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"general.sampling.top_k": uint32(40),
|
|
"general.sampling.top_p": int32(1),
|
|
"general.sampling.min_p": float32(0),
|
|
"general.sampling.typ_p": float32(0.95),
|
|
"general.sampling.temp": uint32(1),
|
|
"general.sampling.penalty_last_n": float32(64),
|
|
"general.sampling.penalty_repeat": float32(1.05),
|
|
"general.sampling.penalty_freq": uint32(0),
|
|
"general.sampling.penalty_present": int32(0),
|
|
"general.sampling.xtc_threshold": float32(0.5),
|
|
"general.sampling.mirostat_tau": float32(5),
|
|
}, nil); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := file.Close(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
md, err := extractGGUFMetadata(file.Name())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
defaults := generationDefaultsFromMetadata(md)
|
|
check := func(key string, want any) {
|
|
t.Helper()
|
|
if got := defaults[key]; got != want {
|
|
t.Fatalf("%s = %#v, want %#v", key, got, want)
|
|
}
|
|
}
|
|
|
|
check("top_k", int64(40))
|
|
check("top_p", float64(1))
|
|
check("min_p", float64(0))
|
|
check("typical_p", float64(0.95))
|
|
check("temperature", float64(1))
|
|
check("repeat_last_n", int64(64))
|
|
check("repeat_penalty", float64(1.05))
|
|
check("frequency_penalty", float64(0))
|
|
check("presence_penalty", float64(0))
|
|
if _, ok := defaults["mirostat_tau"]; ok {
|
|
t.Fatal("mirostat_tau should not be mapped to an Ollama option")
|
|
}
|
|
if _, ok := defaults["xtc_threshold"]; ok {
|
|
t.Fatal("xtc_threshold should not be mapped to an Ollama option")
|
|
}
|
|
}
|
|
|
|
func TestGetModelTemplateMetadata(t *testing.T) {
|
|
customTemplate := "CUSTOM {{ .Prompt }}"
|
|
|
|
t.Run("records chat template and Go TEMPLATE layer", func(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
t.Setenv("OLLAMA_GO_TEMPLATE", "")
|
|
|
|
_, digest := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"tokenizer.chat_template": "{{ bos_token }}{{ messages[0]['content'] }}",
|
|
}, nil)
|
|
writeTestModelManifest(t, "template-disabled", digest, customTemplate)
|
|
|
|
m, err := GetModel("template-disabled")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !m.HasChatTemplate {
|
|
t.Fatal("expected GGUF chat template to be detected")
|
|
}
|
|
if !m.HasGoTemplate {
|
|
t.Fatal("expected Go TEMPLATE layer to be detected")
|
|
}
|
|
if got := m.Template.String(); got != customTemplate {
|
|
t.Fatalf("template = %q, want %q", got, customTemplate)
|
|
}
|
|
})
|
|
|
|
t.Run("prefers chat template when Go TEMPLATE has fewer capabilities", func(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
t.Setenv("OLLAMA_GO_TEMPLATE", "")
|
|
|
|
_, digest := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"tokenizer.chat_template": "{% if tools %}{{ tools }}{% endif %}{{ messages[0]['content'] }}",
|
|
}, nil)
|
|
writeTestModelManifest(t, "chat-template-tools", digest, customTemplate)
|
|
|
|
m, err := GetModel("chat-template-tools")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !m.PreferChatTemplate {
|
|
t.Fatal("expected chat template to be preferred")
|
|
}
|
|
if got := m.CheckCapabilities(model.CapabilityTools); got != nil {
|
|
t.Fatalf("expected tools capability, got %v", got)
|
|
}
|
|
})
|
|
|
|
t.Run("prefers Qwen chat template with tools and inferred thinking", func(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
t.Setenv("OLLAMA_GO_TEMPLATE", "")
|
|
|
|
_, digest := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"tokenizer.chat_template": "{% if tools %}{{ tools }}{% endif %}{% set content = (content.split('</think>')|last) %}",
|
|
}, nil)
|
|
writeTestModelManifest(t, "chat-template-tools-thinking", digest, "{{ range .Messages }}{{ if .Thinking }}<think>{{ .Thinking }}</think>{{ end }}{{ .Content }}{{ end }}")
|
|
|
|
m, err := GetModel("chat-template-tools-thinking")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !m.PreferChatTemplate {
|
|
t.Fatal("expected chat template to be preferred")
|
|
}
|
|
if got := m.CheckCapabilities(model.CapabilityTools); got != nil {
|
|
t.Fatalf("expected tools capability, got %v", got)
|
|
}
|
|
if got := m.CheckCapabilities(model.CapabilityThinking); got != nil {
|
|
t.Fatalf("expected thinking capability, got %v", got)
|
|
}
|
|
})
|
|
|
|
t.Run("prefers chat template with stronger tool round trip", func(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
t.Setenv("OLLAMA_GO_TEMPLATE", "")
|
|
|
|
_, digest := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"tokenizer.chat_template": `{% if tools %}{{ tools }}{% endif %}
|
|
{% for message in messages %}
|
|
{% if message.tool_calls %}
|
|
{% for tool_call in message.tool_calls %}{{ tool_call.function.name }}{% endfor %}
|
|
{% endif %}
|
|
{% if message.role == 'tool' %}tool_response {{ message.content }}{% endif %}
|
|
{% endfor %}`,
|
|
}, nil)
|
|
writeTestModelManifest(t, "chat-template-tool-round-trip", digest, `{{ if .Tools }}tools{{ end }}
|
|
{{ range .Messages }}
|
|
{{ range .ToolCalls }}{{ .Function.Name }}{{ end }}
|
|
{{ end }}`)
|
|
|
|
m, err := GetModel("chat-template-tool-round-trip")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !m.PreferChatTemplate {
|
|
t.Fatal("expected chat template to be preferred")
|
|
}
|
|
if got := m.CheckCapabilities(model.CapabilityTools); got != nil {
|
|
t.Fatalf("expected tools capability, got %v", got)
|
|
}
|
|
})
|
|
|
|
t.Run("keeps Go TEMPLATE when chat template has weaker tool support", func(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
t.Setenv("OLLAMA_GO_TEMPLATE", "")
|
|
|
|
_, digest := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"tokenizer.chat_template": `{%- if tools and not available_tools -%}
|
|
{{- set available_tools = tools -}}
|
|
{%- endif -%}
|
|
{%- if available_tools -%}
|
|
{{ '<|start_of_role|>available_tools<|end_of_role|>' }}{{ available_tools | tojson }}{{ '<|end_of_text|>' }}
|
|
{%- endif -%}
|
|
{%- if thinking -%}<think></think><response></response>{%- endif -%}
|
|
{%- for message in messages -%}
|
|
{{ '<|start_of_role|>' + message['role'] + '<|end_of_role|>' + message['content'] + '<|end_of_text|>' }}
|
|
{%- endfor -%}`,
|
|
}, nil)
|
|
writeTestModelManifest(t, "chat-template-weaker-tools", digest, `{{ if .Tools }}tools{{ end }}
|
|
{{ range .Messages }}
|
|
{{ if eq .Role "tool" }}tool_response{{ else }}{{ .Role }}{{ end }}
|
|
{{ if .ToolCalls }}<|tool_call|>{{ range .ToolCalls }}{{ .Function.Name }}{{ end }}{{ else }}{{ .Content }}{{ end }}
|
|
{{ end }}`)
|
|
|
|
m, err := GetModel("chat-template-weaker-tools")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if m.PreferChatTemplate {
|
|
t.Fatal("expected Go TEMPLATE to be preferred")
|
|
}
|
|
if got := m.CheckCapabilities(model.CapabilityTools); got != nil {
|
|
t.Fatalf("expected tools capability, got %v", got)
|
|
}
|
|
if got := m.CheckCapabilities(model.CapabilityThinking); got == nil {
|
|
t.Fatal("expected thinking capability to remain unavailable on Go TEMPLATE path")
|
|
}
|
|
})
|
|
|
|
t.Run("respects explicit Go TEMPLATE enablement", func(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
t.Setenv("OLLAMA_GO_TEMPLATE", "1")
|
|
|
|
_, digest := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"tokenizer.chat_template": "{% if tools %}{{ tools }}{% endif %}{{ messages[0]['content'] }}",
|
|
}, nil)
|
|
writeTestModelManifest(t, "go-template-forced", digest, customTemplate)
|
|
|
|
m, err := GetModel("go-template-forced")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if m.PreferChatTemplate {
|
|
t.Fatal("expected explicit Go TEMPLATE setting to suppress chat_template preference")
|
|
}
|
|
if got := m.CheckCapabilities(model.CapabilityTools); got == nil {
|
|
t.Fatal("expected tools capability to be unavailable when Go TEMPLATE is explicitly enabled")
|
|
}
|
|
})
|
|
|
|
t.Run("respects explicit Go TEMPLATE disablement", func(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
t.Setenv("OLLAMA_GO_TEMPLATE", "0")
|
|
|
|
_, digest := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"tokenizer.chat_template": "{% if tools %}{{ tools }}{% endif %}{{ messages[0]['content'] }}",
|
|
}, nil)
|
|
writeTestModelManifest(t, "go-template-disabled", digest, customTemplate)
|
|
|
|
m, err := GetModel("go-template-disabled")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if m.PreferChatTemplate {
|
|
t.Fatal("expected explicit Go TEMPLATE setting to suppress chat_template preference")
|
|
}
|
|
if got := m.CheckCapabilities(model.CapabilityTools); got != nil {
|
|
t.Fatalf("expected tools capability from GGUF chat_template, got %v", got)
|
|
}
|
|
})
|
|
|
|
t.Run("records missing chat template", func(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
t.Setenv("OLLAMA_GO_TEMPLATE", "")
|
|
|
|
_, digest := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
}, nil)
|
|
writeTestModelManifest(t, "missing-chat-template", digest, customTemplate)
|
|
|
|
m, err := GetModel("missing-chat-template")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if m.HasChatTemplate {
|
|
t.Fatal("expected missing GGUF chat template")
|
|
}
|
|
if !m.HasGoTemplate {
|
|
t.Fatal("expected Go TEMPLATE layer to be detected")
|
|
}
|
|
})
|
|
}
|
|
|
|
func writeTestModelManifest(t *testing.T, name, digest, tmpl string) {
|
|
t.Helper()
|
|
|
|
modelLayer, err := manifest.NewLayerFromLayer(digest, "application/vnd.ollama.image.model", "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
templateLayer, err := manifest.NewLayer(strings.NewReader(tmpl), "application/vnd.ollama.image.template")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
layers := []manifest.Layer{modelLayer, templateLayer}
|
|
configLayer, err := createConfigLayer(model.ConfigV2{
|
|
ModelFormat: "gguf",
|
|
ModelFamily: "llama",
|
|
ModelFamilies: []string{"llama"},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := manifest.WriteManifest(model.ParseName(name), *configLayer, layers); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
// loadTestMetadata fills in what GetModel would have read from the metadata
|
|
// files, so
|
|
// hand-built models resolve capabilities the same way loaded ones do.
|
|
func loadTestMetadata(t *testing.T, m *Model) {
|
|
t.Helper()
|
|
if m.ModelPath != "" {
|
|
md, err := extractGGUFMetadata(m.ModelPath)
|
|
if err != nil {
|
|
t.Fatalf("metadata for %s: %v", m.ModelPath, err)
|
|
}
|
|
m.metadata = md
|
|
}
|
|
m.projectorMetadata = nil
|
|
for _, path := range m.ProjectorPaths {
|
|
md, err := extractGGUFMetadata(path)
|
|
if err != nil {
|
|
t.Fatalf("projector metadata for %s: %v", path, err)
|
|
}
|
|
m.projectorMetadata = append(m.projectorMetadata, md)
|
|
}
|
|
}
|
|
|
|
func TestModelCapabilities(t *testing.T) {
|
|
// Create completion model (llama architecture without vision)
|
|
completionModelPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
}, []*gguftest.Tensor{})
|
|
|
|
ggufToolTemplateModelPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"tokenizer.chat_template": `{% if tools %}<tool_call>{{ tools }}</tool_call>{% endif %}<think>{{ messages[0]['content'] }}</think>`,
|
|
}, []*gguftest.Tensor{})
|
|
|
|
// Create vision model (llama architecture with vision block count)
|
|
visionModelPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"llama.vision.block_count": uint32(1),
|
|
}, []*gguftest.Tensor{})
|
|
|
|
// Create embedding model (bert architecture with pooling type)
|
|
embeddingModelPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "bert",
|
|
"bert.pooling_type": uint32(1),
|
|
}, []*gguftest.Tensor{})
|
|
|
|
audioProjectorPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "clip",
|
|
"clip.has_audio_encoder": true,
|
|
"vision.projector_type": "pixtral",
|
|
"clip.vision.block_count": uint32(1),
|
|
}, []*gguftest.Tensor{})
|
|
|
|
nemotronOmniModelPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "nemotron_h_omni",
|
|
"nemotron_h_omni.vision.block_count": uint32(1),
|
|
"nemotron_h_omni.audio.block_count": uint32(1),
|
|
"nemotron_h_omni.embedding_length": uint32(1),
|
|
"nemotron_h_omni.attention.head_count": uint32(1),
|
|
}, []*gguftest.Tensor{})
|
|
|
|
suppressedAudioProjectorPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "clip",
|
|
"clip.has_audio_encoder": true,
|
|
"vision.projector_type": "gemma4v",
|
|
"clip.vision.block_count": uint32(1),
|
|
}, []*gguftest.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)
|
|
}
|
|
|
|
testModels := []struct {
|
|
name string
|
|
model Model
|
|
expectedCaps []model.Capability
|
|
}{
|
|
{
|
|
name: "model with image generation capability via config",
|
|
model: Model{
|
|
Config: model.ConfigV2{
|
|
Capabilities: []string{"image"},
|
|
},
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityImage},
|
|
},
|
|
{
|
|
name: "model with image and vision capability (image editing)",
|
|
model: Model{
|
|
Config: model.ConfigV2{
|
|
Capabilities: []string{"image", "vision"},
|
|
},
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityImage, model.CapabilityVision},
|
|
},
|
|
{
|
|
name: "model with completion capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion},
|
|
},
|
|
|
|
{
|
|
name: "model with completion, tools, and insert capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: toolsInsertTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityInsert},
|
|
},
|
|
{
|
|
name: "model with tools capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: toolsTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools},
|
|
},
|
|
{
|
|
name: "model with GGUF chat_template tools and thinking",
|
|
model: Model{
|
|
ModelPath: ggufToolTemplateModelPath,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityThinking},
|
|
},
|
|
{
|
|
name: "model with Go TEMPLATE ignores GGUF chat_template capabilities",
|
|
model: Model{
|
|
ModelPath: ggufToolTemplateModelPath,
|
|
Template: chatTemplate,
|
|
HasGoTemplate: true,
|
|
HasChatTemplate: true,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion},
|
|
},
|
|
{
|
|
name: "model with tools capability from config and parser",
|
|
model: Model{
|
|
Config: model.ConfigV2{
|
|
Capabilities: []string{"completion", "tools"},
|
|
Parser: "qwen3-coder",
|
|
},
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools},
|
|
},
|
|
{
|
|
name: "model with vision capability",
|
|
model: Model{
|
|
ModelPath: visionModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision},
|
|
},
|
|
{
|
|
name: "model with vision, tools, and insert capability",
|
|
model: Model{
|
|
ModelPath: visionModelPath,
|
|
Template: toolsInsertTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityTools, model.CapabilityInsert},
|
|
},
|
|
{
|
|
name: "model with embedding capability",
|
|
model: Model{
|
|
ModelPath: embeddingModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityEmbedding},
|
|
},
|
|
{
|
|
name: "model with audio projector capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
ProjectorPaths: []string{audioProjectorPath},
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityAudio},
|
|
},
|
|
{
|
|
name: "model with parser and projector capabilities without template",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
ProjectorPaths: []string{audioProjectorPath},
|
|
Config: model.ConfigV2{
|
|
Parser: "functiongemma",
|
|
},
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityAudio, model.CapabilityTools},
|
|
},
|
|
{
|
|
name: "gemma4 projector exposes audio capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
ProjectorPaths: []string{suppressedAudioProjectorPath},
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityAudio},
|
|
},
|
|
{
|
|
name: "gemma4 gguf exposes audio capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
ProjectorPaths: []string{audioProjectorPath},
|
|
Config: model.ConfigV2{
|
|
Renderer: gemma4RendererSmall,
|
|
Capabilities: []string{"audio"},
|
|
},
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityAudio, model.CapabilityCompletion, model.CapabilityVision},
|
|
},
|
|
{
|
|
name: "nemotron3 gguf suppresses audio capability",
|
|
model: Model{
|
|
ModelPath: nemotronOmniModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision},
|
|
},
|
|
{
|
|
name: "nemotron3 projector suppresses audio capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
ProjectorPaths: []string{audioProjectorPath},
|
|
Config: model.ConfigV2{
|
|
ModelFamily: "nemotron_h_omni",
|
|
},
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision},
|
|
},
|
|
{
|
|
name: "nemotron3 safetensors suppresses vision and audio but keeps thinking",
|
|
model: Model{
|
|
Config: model.ConfigV2{
|
|
ModelFormat: "safetensors",
|
|
Parser: "nemotron-3-nano",
|
|
Renderer: "nemotron-3-nano",
|
|
Capabilities: []string{"completion", "vision", "audio"},
|
|
},
|
|
Template: chatTemplate,
|
|
},
|
|
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityThinking},
|
|
},
|
|
}
|
|
|
|
// compare two slices of model.Capability regardless of order
|
|
compareCapabilities := func(a, b []model.Capability) bool {
|
|
if len(a) != len(b) {
|
|
return false
|
|
}
|
|
|
|
aCount := make(map[model.Capability]int)
|
|
for _, cap := range a {
|
|
aCount[cap]++
|
|
}
|
|
|
|
bCount := make(map[model.Capability]int)
|
|
for _, cap := range b {
|
|
bCount[cap]++
|
|
}
|
|
|
|
for cap, count := range aCount {
|
|
if bCount[cap] != count {
|
|
return false
|
|
}
|
|
}
|
|
|
|
return true
|
|
}
|
|
|
|
for _, tt := range testModels {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
loadTestMetadata(t, &tt.model)
|
|
|
|
// Test Capabilities method
|
|
caps := tt.model.Capabilities()
|
|
if !compareCapabilities(caps, tt.expectedCaps) {
|
|
t.Errorf("Expected capabilities %v, got %v", tt.expectedCaps, caps)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestModelCheckCapabilities(t *testing.T) {
|
|
// Create simple model file for tests that don't depend on GGUF content
|
|
completionModelPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
}, []*gguftest.Tensor{})
|
|
|
|
// Create vision model (llama architecture with vision block count)
|
|
visionModelPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "llama",
|
|
"llama.vision.block_count": uint32(1),
|
|
}, []*gguftest.Tensor{})
|
|
|
|
// Create embedding model (bert architecture with pooling type)
|
|
embeddingModelPath, _ := createBinFile(t, gguftest.KV{
|
|
"general.architecture": "bert",
|
|
"bert.pooling_type": uint32(1),
|
|
}, []*gguftest.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",
|
|
},
|
|
{
|
|
name: "model missing image generation capability",
|
|
model: Model{
|
|
ModelPath: completionModelPath,
|
|
Template: chatTemplate,
|
|
},
|
|
checkCaps: []model.Capability{model.CapabilityImage},
|
|
expectedErrMsg: "does not support image generation",
|
|
},
|
|
{
|
|
name: "model with image generation capability",
|
|
model: Model{
|
|
Config: model.ConfigV2{
|
|
Capabilities: []string{"image"},
|
|
},
|
|
},
|
|
checkCaps: []model.Capability{model.CapabilityImage},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
loadTestMetadata(t, &tt.model)
|
|
|
|
// 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)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestPullModelManifest(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
manifest string
|
|
}{
|
|
{
|
|
name: "pretty printed",
|
|
manifest: `{ "schemaVersion": 2, "mediaType": "application/vnd.docker.distribution.manifest.v2+json",
|
|
"config": { "digest": "sha256:abc", "mediaType": "application/vnd.docker.container.image.v1+json", "size": 50 },
|
|
"layers": [{ "digest": "sha256:t1", "mediaType": "application/vnd.ollama.image.tensor", "size": 1024, "name": "model.weight" }]
|
|
}`,
|
|
},
|
|
{
|
|
name: "non-standard field order",
|
|
manifest: `{"layers":[{"size":999,"digest":"sha256:def","mediaType":"application/vnd.ollama.image.model"}],"schemaVersion":2,"config":{"size":50,"digest":"sha256:abc","mediaType":"application/vnd.docker.container.image.v1+json"},"mediaType":"application/vnd.docker.distribution.manifest.v2+json"}`,
|
|
},
|
|
}
|
|
|
|
for _, tt := range cases {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Write([]byte(tt.manifest))
|
|
}))
|
|
defer ts.Close()
|
|
|
|
n := model.ParseName("test/model:latest")
|
|
n.ProtocolScheme = "http"
|
|
n.Host = strings.TrimPrefix(ts.URL, "http://")
|
|
|
|
mf, data, err := pullModelManifest(t.Context(), n, ®istryOptions{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// Raw bytes must be byte-for-byte identical to what the server sent
|
|
if string(data) != tt.manifest {
|
|
t.Fatalf("raw bytes differ from server response")
|
|
}
|
|
|
|
// SHA256 of returned data must match the expected registry digest
|
|
expectedDigest := fmt.Sprintf("%x", sha256.Sum256([]byte(tt.manifest)))
|
|
gotDigest := fmt.Sprintf("%x", sha256.Sum256(data))
|
|
if gotDigest != expectedDigest {
|
|
t.Fatalf("digest mismatch\ngot: %s\nwant: %s", gotDigest, expectedDigest)
|
|
}
|
|
|
|
// Parsed manifest must still be usable
|
|
if mf.SchemaVersion != 2 {
|
|
t.Fatalf("schemaVersion = %d, want 2", mf.SchemaVersion)
|
|
}
|
|
if mf.Config.Digest == "" {
|
|
t.Fatal("config digest is empty")
|
|
}
|
|
if len(mf.Layers) == 0 {
|
|
t.Fatal("expected at least one layer")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestPullModelDuplicateDigestVerifiesBlob pulls a manifest whose config and
|
|
// layer share a digest. The registry redirects blob downloads via Location to
|
|
// an "internal" path serving bytes that don't match the digest, so PullModel
|
|
// must reject the pull with errDigestMismatch.
|
|
func TestPullModelDuplicateDigestVerifiesBlob(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
|
|
const bogusDigest = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
|
|
|
backendURL := ""
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch {
|
|
case strings.Contains(r.URL.Path, "/manifests/"):
|
|
w.Header().Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
|
|
fmt.Fprintf(w, `{
|
|
"schemaVersion": 2,
|
|
"mediaType": "application/vnd.docker.distribution.manifest.v2+json",
|
|
"config": {
|
|
"mediaType": "application/vnd.ollama.image.config",
|
|
"digest": %q,
|
|
"size": 5
|
|
},
|
|
"layers": [{
|
|
"mediaType": "application/vnd.ollama.image.model",
|
|
"digest": %q,
|
|
"size": 5
|
|
}]
|
|
}`, bogusDigest, bogusDigest)
|
|
case strings.Contains(r.URL.Path, "/internal/blobs/"):
|
|
w.Write([]byte("attacker-controlled-bytes"))
|
|
case strings.Contains(r.URL.Path, "/blobs/"):
|
|
w.Header().Set("Location", backendURL+"/internal"+r.URL.Path)
|
|
w.WriteHeader(http.StatusOK)
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer ts.Close()
|
|
backendURL = ts.URL
|
|
|
|
u, err := url.Parse(ts.URL)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
n := model.ParseName(u.Host + "/test/attack")
|
|
n.ProtocolScheme = "http"
|
|
|
|
err = PullModel(t.Context(), n.String(), ®istryOptions{Insecure: true}, func(api.ProgressResponse) {})
|
|
if !errors.Is(err, errDigestMismatch) {
|
|
t.Fatalf("PullModel = %v, want errDigestMismatch (unverified blob would persist)", err)
|
|
}
|
|
}
|
|
|
|
// TestPullManifestRejectsCrossHostRedirect: a manifest GET that the registry
|
|
// redirects to a different host must be refused by default, so a malicious
|
|
// registry can't turn a pull into a request to an internal address.
|
|
// --insecure opts out for trusted registries.
|
|
func TestPullManifestRejectsCrossHostRedirect(t *testing.T) {
|
|
t.Setenv("OLLAMA_MODELS", t.TempDir())
|
|
|
|
var internalHit atomic.Bool
|
|
internal := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
internalHit.Store(true)
|
|
}))
|
|
defer internal.Close()
|
|
|
|
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Redirect(w, r, internal.URL+r.URL.Path, http.StatusFound)
|
|
}))
|
|
defer ts.Close()
|
|
|
|
requestURL, err := url.Parse(ts.URL + "/v2/test/attack/manifests/latest")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// Default policy: cross-host redirect is refused before any request
|
|
// leaves for the internal host. (regOpts nil exercises the makeRequest
|
|
// default; the insecure protocol check doesn't apply at this level.)
|
|
blockedResp, err := makeRequest(t.Context(), http.MethodGet, requestURL, nil, nil, ®istryOptions{})
|
|
// On a CheckRedirect failure the client returns the pre-redirect
|
|
// response with its body already closed; close again defensively to
|
|
// satisfy bodyclose (double close is a no-op).
|
|
if blockedResp != nil && blockedResp.Body != nil {
|
|
blockedResp.Body.Close()
|
|
}
|
|
if !errors.Is(err, errBlockedRedirect) {
|
|
t.Fatalf("makeRequest = %v, want errBlockedRedirect", err)
|
|
}
|
|
if internalHit.Load() {
|
|
t.Fatal("internal host received a request despite the blocked redirect")
|
|
}
|
|
|
|
// Insecure opts out: the cross-host redirect is followed.
|
|
resp, err := makeRequest(t.Context(), http.MethodGet, requestURL, nil, nil, ®istryOptions{Insecure: true})
|
|
if err != nil {
|
|
t.Fatalf("makeRequest with Insecure = %v, want redirect followed", err)
|
|
}
|
|
resp.Body.Close()
|
|
if !internalHit.Load() {
|
|
t.Fatal("redirect target was not reached with Insecure set")
|
|
}
|
|
}
|