ollama/server/model.go

636 lines
16 KiB
Go

package server
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"os"
"regexp"
"slices"
"strings"
"text/template/parse"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/fs/ggml"
"github.com/ollama/ollama/template"
"github.com/ollama/ollama/types/model"
)
var intermediateBlobs map[string]string = make(map[string]string)
type layerGGML struct {
Layer
*ggml.GGML
}
func parseFromModel(ctx context.Context, name model.Name, fn func(api.ProgressResponse)) (layers []*layerGGML, err error) {
m, err := ParseNamedManifest(name)
switch {
case errors.Is(err, os.ErrNotExist):
if err := PullModel(ctx, name.String(), &registryOptions{}, fn); err != nil {
return nil, err
}
m, err = ParseNamedManifest(name)
if err != nil {
return nil, err
}
case err != nil:
return nil, err
}
for _, layer := range m.Layers {
layer, err := NewLayerFromLayer(layer.Digest, layer.MediaType, name.DisplayShortest())
if err != nil {
return nil, err
}
switch layer.MediaType {
case "application/vnd.ollama.image.model",
"application/vnd.ollama.image.projector",
"application/vnd.ollama.image.adapter":
blobpath, err := GetBlobsPath(layer.Digest)
if err != nil {
return nil, err
}
blob, err := os.Open(blobpath)
if err != nil {
return nil, err
}
defer blob.Close()
f, _, err := ggml.Decode(blob, 0)
if err != nil {
return nil, err
}
layers = append(layers, &layerGGML{layer, f})
default:
layers = append(layers, &layerGGML{layer, nil})
}
}
return layers, nil
}
func detectChatTemplate(layers []*layerGGML) ([]*layerGGML, error) {
for _, layer := range layers {
if s := layer.GGML.KV().ChatTemplate(); s != "" {
if t, err := template.Named(s); err != nil {
slog.Debug("template detection", "error", err, "template", s)
} else {
layer, err := NewLayer(t.Reader(), "application/vnd.ollama.image.template")
if err != nil {
return nil, err
}
layer.status = fmt.Sprintf("using autodetected template %s", t.Name)
layers = append(layers, &layerGGML{layer, nil})
if t.Parameters != nil {
var b bytes.Buffer
if err := json.NewEncoder(&b).Encode(t.Parameters); err != nil {
return nil, err
}
layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
if err != nil {
return nil, err
}
layers = append(layers, &layerGGML{layer, nil})
}
}
}
}
return layers, nil
}
func detectContentType(r io.Reader) (string, error) {
var b bytes.Buffer
if _, err := io.Copy(&b, r); err != nil {
return "", err
}
if contentType := ggml.DetectContentType(b.Bytes()); contentType != "" {
return contentType, nil
}
if contentType := http.DetectContentType(b.Bytes()); contentType != "application/octet-stream" {
return contentType, nil
}
return "unknown", nil
}
func parseObjects(s string) []map[string]any {
var objs []map[string]any
for offset := 0; offset < len(s); {
var obj map[string]any
decoder := json.NewDecoder(strings.NewReader(s[offset:]))
if err := decoder.Decode(&obj); errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
break
} else if syntax := &(json.SyntaxError{}); errors.As(err, &syntax) {
// skip over any syntax errors
offset += int(syntax.Offset)
} else if unmarshalType := &(json.UnmarshalTypeError{}); errors.As(err, &unmarshalType) {
// skip over any unmarshalable types
offset += int(unmarshalType.Offset)
} else if err != nil {
return nil
} else {
offset += int(decoder.InputOffset())
objs = append(objs, obj)
}
}
return objs
}
// parseToolCalls attempts to parse a JSON string into a slice of ToolCalls.
// mxyng: this only really works if the input contains tool calls in some JSON format
func (m *Model) parseToolCalls(s string) ([]api.ToolCall, bool) {
// create a subtree from the node that ranges over .ToolCalls
tmpl := m.Template.Subtree(func(n parse.Node) bool {
if t, ok := n.(*parse.RangeNode); ok {
return slices.Contains(template.Identifiers(t.Pipe), "ToolCalls")
}
return false
})
if tmpl == nil {
slog.Debug("parseToolCalls: no ToolCalls template found")
return nil, false
}
slog.Debug("parseToolCalls: executing template with test data", "input", s)
var b bytes.Buffer
if err := tmpl.Execute(&b, map[string][]api.ToolCall{
"ToolCalls": {
{
Function: api.ToolCallFunction{
Name: "@@name@@",
Arguments: api.ToolCallFunctionArguments{
"@@argument@@": 1,
},
},
},
},
}); err != nil {
slog.Debug("parseToolCalls: template execution failed", "error", err)
return nil, false
}
slog.Debug("parseToolCalls: template executed successfully", "output", b.String())
templateObjects := parseObjects(b.String())
if len(templateObjects) == 0 {
return nil, false
}
slog.Debug("parseToolCalls: template objects", "objects", templateObjects)
// find the keys that correspond to the name and arguments fields
var name, arguments string
for k, v := range templateObjects[0] {
switch v.(type) {
case string:
name = k
case map[string]any:
arguments = k
}
}
if name == "" || arguments == "" {
return nil, false
}
responseObjects := parseObjects(s)
if len(responseObjects) == 0 {
return nil, false
}
// collect all nested objects
var collect func(any) []map[string]any
collect = func(obj any) (all []map[string]any) {
switch o := obj.(type) {
case map[string]any:
all = append(all, o)
for _, v := range o {
all = append(all, collect(v)...)
}
case []any:
for _, v := range o {
all = append(all, collect(v)...)
}
}
return all
}
var objs []map[string]any
for _, p := range responseObjects {
objs = append(objs, collect(p)...)
}
var toolCalls []api.ToolCall
for _, kv := range objs {
n, nok := kv[name].(string)
a, aok := kv[arguments].(map[string]any)
if nok && aok {
toolCalls = append(toolCalls, api.ToolCall{
Function: api.ToolCallFunction{
Name: n,
Arguments: a,
},
})
}
}
return toolCalls, len(toolCalls) > 0
}
// ToolCallFormat represents different possible formats for tool calls
type toolCallFormat struct {
// Direct format
Name string `json:"name,omitempty"`
Arguments map[string]any `json:"arguments,omitempty"`
// Command-r-plus format
ToolName string `json:"tool_name,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
// Function format
Function *struct {
Name string `json:"name"`
Arguments map[string]any `json:"arguments,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
} `json:"function,omitempty"`
// Xlam format
ToolCalls []toolCallFormat `json:"tool_calls,omitempty"`
}
func parseJSONToolCalls(obj map[string]any) ([]api.ToolCall, bool) {
// Helper to convert any to []any safely
toArray := func(v any) []any {
if arr, ok := v.([]any); ok {
return arr
}
return nil
}
// Convert a single format to a tool call
makeToolCall := func(f toolCallFormat) (api.ToolCall, bool) {
switch {
case f.Name != "" && f.Arguments != nil:
return api.ToolCall{
Function: api.ToolCallFunction{
Name: f.Name,
Arguments: f.Arguments,
},
}, true
case f.Name != "" && f.Parameters != nil: // Handle parameters field
return api.ToolCall{
Function: api.ToolCallFunction{
Name: f.Name,
Arguments: f.Parameters,
},
}, true
case f.ToolName != "" && f.Parameters != nil:
return api.ToolCall{
Function: api.ToolCallFunction{
Name: f.ToolName,
Arguments: f.Parameters,
},
}, true
case f.Function != nil && f.Function.Name != "":
args := f.Function.Arguments
if args == nil {
args = f.Function.Parameters
}
if args != nil {
return api.ToolCall{
Function: api.ToolCallFunction{
Name: f.Function.Name,
Arguments: args,
},
}, true
}
}
return api.ToolCall{}, false
}
// Try parsing as array first
if arr := toArray(obj); arr != nil {
var calls []api.ToolCall
for _, item := range arr {
if itemMap, ok := item.(map[string]any); ok {
var format toolCallFormat
data, _ := json.Marshal(itemMap)
if err := json.Unmarshal(data, &format); err == nil {
if call, ok := makeToolCall(format); ok {
calls = append(calls, call)
}
}
}
}
if len(calls) > 0 {
return calls, true
}
}
// Try parsing as single object
var format toolCallFormat
data, _ := json.Marshal(obj)
if err := json.Unmarshal(data, &format); err != nil {
return nil, false
}
// Handle xlam format (tool_calls array)
if len(format.ToolCalls) > 0 {
var calls []api.ToolCall
for _, f := range format.ToolCalls {
if call, ok := makeToolCall(f); ok {
calls = append(calls, call)
}
}
if len(calls) > 0 {
return calls, true
}
}
// Try as single tool call
if call, ok := makeToolCall(format); ok {
return []api.ToolCall{call}, true
}
return nil, false
}
func (m *Model) GetToolCallFormat(s string) (string, string, bool) {
// Try to detect the tool call format from the model's template
tmpl := m.Template.Subtree(func(n parse.Node) bool {
if t, ok := n.(*parse.RangeNode); ok {
return slices.Contains(template.Identifiers(t.Pipe), "Content")
}
return false
})
if tmpl != nil {
// Execute template with test data to see the format
var b bytes.Buffer
if err := tmpl.Execute(&b, map[string][]api.ToolCall{
"ToolCalls": {
{
Function: api.ToolCallFunction{
Name: "function_name",
Arguments: api.ToolCallFunctionArguments{
"argument1": "value1",
// "argument2": "value2",
},
},
},
},
}); err == nil {
// Look for special tokens in the template output
output := strings.TrimSpace(b.String())
slog.Debug("tool call template output", "output", output)
if strings.Contains(output, "<") {
// Extract the special token between < and >
start := strings.Index(output, "<")
end := strings.Index(output, ">")
if start >= 0 && end > start {
token := output[start : end+1]
return output, token, true
}
} else if strings.Contains(output, "[") {
// Check if it's a tool call token rather than JSON array
start := strings.Index(output, "[")
end := strings.Index(output, "]")
if start >= 0 && end > start {
token := output[start : end+1]
// Only consider it a token if it's not valid JSON
var jsonTest any
if err := json.Unmarshal([]byte(token), &jsonTest); err != nil {
return output, token, true
}
}
}
}
}
return "", "", false
}
func parsePythonFunctionCall(s string) (api.ToolCall, bool) {
re := regexp.MustCompile(`(\w+)\((.*?)\)`)
if match := re.FindStringSubmatchIndex(s); match != nil {
name := s[match[2]:match[3]]
args := s[match[4]:match[5]]
// Check if there's a < after the closing bracket
if idx := strings.Index(s[match[5]:], "<"); idx >= 0 {
// Wait for closing > by returning false
if !strings.Contains(s[match[5]+idx:], ">") {
return api.ToolCall{}, false
}
}
arguments := make(api.ToolCallFunctionArguments)
if strings.Contains(args, "=") { // Keyword args
pairs := strings.SplitSeq(args, ",")
for pair := range pairs {
pair = strings.TrimSpace(pair)
kv := strings.Split(pair, "=")
if len(kv) == 2 {
key := strings.TrimSpace(kv[0])
value := strings.TrimSpace(kv[1])
arguments[key] = value
}
}
return api.ToolCall{
Function: api.ToolCallFunction{
Name: name,
Arguments: arguments,
},
}, true
}
}
return api.ToolCall{}, false
}
func (m *Model) ParseToolCallsStream(s string, prefix *string, specialToken *string) ([]api.ToolCall, bool, bool) {
// The prefix check for for the tags shouldn't really be used and we should be consuming this from the model
// Knowing what the tool token enables quicker and more reliable parsing
// TODO: not sure how we're going to handle chatting before the tool call
// TODO: detection would be relying on the model to know what the tool token is
// fmt.Println("parsing tool calls", s)
if prefix == nil {
prefix = new(string)
*prefix = ""
}
if specialToken == nil {
specialToken = new(string)
*specialToken = ""
}
// TODO: cache this
// _, token, ok := m.GetToolCallFormat(s)
// if ok && token != "" {
// fmt.Println("token", token)
// *specialToken = token
// }
// fmt.Println("prefix", *prefix)
// fmt.Println("special token", *specialToken)
var partial bool
s = strings.TrimSpace(s)
if len(s) == 0 {
return nil, false, false
}
if specialToken != nil && len(*specialToken) > 0 {
s2 := *specialToken
if strings.HasPrefix(s, string(s2[0])) {
// fmt.Println("prefix 1 is", string(s2[0]))
partial = true
*prefix = string(s2[0])
}
}
if len(s) > 0 {
if s[0] == '[' {
s = strings.ReplaceAll(s, "\n", "")
// tool call list with no special token
if len(s) > 1 && s[1] == '{' {
// fmt.Println("prefix 2 in [{", string(s[0]))
partial = true
*specialToken = "[{"
*prefix = "[{"
} else if *specialToken == "" {
// possible tool call with special token but not in template
// split s over spaces to check for special token
if len(s) > 0 && s[len(s)-1] == ']' {
partial = true
*specialToken = s
*prefix = "["
}
}
} else if s[0] == '{' {
// fmt.Println("prefix 2 in {", string(s[0]))
partial = true
*specialToken = "{"
*prefix = "{"
} else if s[0] == '<' {
// TODO: the only issue here is that we might miss a > if the token is weird
// The </ works as that would only happen for a tool call mainly for when we create a tag
if len(s) > 1 && s[1] == '/' {
// fmt.Println("prefix3 in <", string(s[0]))
// returning a partial here is a hack to ensure that we don't send the content downstream
return nil, true, true
// TODO: jank hack to get special token right
// special token might not be set yet
} else if s[len(s)-1] == '>' {
partial = true
*specialToken = s
*prefix = "<"
} else if specialToken != nil && *specialToken == "" {
partial = true
*specialToken = "<"
*prefix = "<"
}
}
}
// fmt.Println("special token", *specialToken)
// fmt.Println("prefix", *prefix)
if !partial {
return nil, false, false
}
// Look for <function_call> tags
// fmt.Println("looking for special token", *specialToken)
start := strings.Index(s, *specialToken)
if start == -1 {
if partial {
// fmt.Println("did not find opening tag, partial match", *specialToken)
return nil, true, true
}
return nil, false, false
}
end := len(s)
// Extract content between tags
var content string
// fmt.Println("prefix before is", *prefix)
if *prefix == "[{" || *prefix == "{" {
content = s[start:end]
} else {
content = s[start+len(*specialToken) : end]
}
content = strings.TrimSpace(content)
// fmt.Println("content", content)
var toolCalls []api.ToolCall
// Try parsing as JSON first - could be single object or array
var jsonObj any
if err := json.Unmarshal([]byte(content), &jsonObj); err == nil {
// Try as single object
if obj, ok := jsonObj.(map[string]any); ok {
// fmt.Println("obj", obj)
if calls, ok := parseJSONToolCalls(obj); ok {
toolCalls = append(toolCalls, calls...)
}
}
// Try as array of objects
if arr, ok := jsonObj.([]any); ok {
for _, item := range arr {
if obj, ok := item.(map[string]any); ok {
if calls, ok := parseJSONToolCalls(obj); ok {
toolCalls = append(toolCalls, calls...)
}
}
}
}
} else {
// TODO: review this case
// Check for partial JSON before trying Python style
if strings.HasPrefix(content, "{") || strings.HasPrefix(content, "[{") {
// We have an opening brace/bracket but failed to parse - likely partial JSON
return nil, true, true
}
// Try parsing as Python function call
if toolCall, ok := parsePythonFunctionCall(content); ok {
toolCalls = append(toolCalls, toolCall)
}
}
// Only return success if we found valid tool calls and no errors
if len(toolCalls) > 0 {
// Check if any of the tool calls are malformed
for _, call := range toolCalls {
if call.Function.Name == "" || len(call.Function.Arguments) == 0 {
return nil, false, false
}
}
return toolCalls, false, true
}
// fmt.Println("no tool calls found, partial match", partial)
if partial {
return nil, true, true
}
return nil, false, false
}