routes.go 3.83 KB
Newer Older
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1
2
3
package server

import (
Michael Yang's avatar
Michael Yang committed
4
	"embed"
Michael Yang's avatar
Michael Yang committed
5
	"encoding/json"
6
	"errors"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
7
8
	"io"
	"log"
Michael Yang's avatar
Michael Yang committed
9
	"math"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
10
11
	"net"
	"net/http"
12
	"os"
Michael Yang's avatar
Michael Yang committed
13
14
15
	"path"
	"strings"
	"text/template"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16
17

	"github.com/gin-gonic/gin"
Michael Yang's avatar
Michael Yang committed
18
	"github.com/lithammer/fuzzysearch/fuzzy"
Michael Yang's avatar
Michael Yang committed
19
	"golang.org/x/sync/errgroup"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
20

Jeffrey Morgan's avatar
Jeffrey Morgan committed
21
	"github.com/jmorganca/ollama/api"
Michael Yang's avatar
Michael Yang committed
22
	"github.com/jmorganca/ollama/llama"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
23
24
)

Michael Yang's avatar
Michael Yang committed
25
26
27
//go:embed templates/*
var templatesFS embed.FS
var templates = template.Must(template.ParseFS(templatesFS, "templates/*.prompt"))
Michael Yang's avatar
Michael Yang committed
28

29
30
31
32
33
34
35
36
37
func cacheDir() string {
	home, err := os.UserHomeDir()
	if err != nil {
		panic(err)
	}

	return path.Join(home, ".ollama")
}

Bruce MacDonald's avatar
Bruce MacDonald committed
38
func generate(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
39
40
	req := api.GenerateRequest{
		Options: api.DefaultOptions(),
41
42
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
43
	if err := c.ShouldBindJSON(&req); err != nil {
44
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Bruce MacDonald's avatar
Bruce MacDonald committed
45
46
		return
	}
47

Bruce MacDonald's avatar
Bruce MacDonald committed
48
	if remoteModel, _ := getRemote(req.Model); remoteModel != nil {
Michael Yang's avatar
Michael Yang committed
49
50
		req.Model = remoteModel.FullName()
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
51
52
	if _, err := os.Stat(req.Model); err != nil {
		if !errors.Is(err, os.ErrNotExist) {
53
			c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Bruce MacDonald's avatar
Bruce MacDonald committed
54
55
56
57
			return
		}
		req.Model = path.Join(cacheDir(), "models", req.Model+".bin")
	}
Michael Yang's avatar
Michael Yang committed
58

Michael Yang's avatar
Michael Yang committed
59
	llm, err := llama.New(req.Model, req.Options)
Bruce MacDonald's avatar
Bruce MacDonald committed
60
	if err != nil {
Michael Yang's avatar
Michael Yang committed
61
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
Bruce MacDonald's avatar
Bruce MacDonald committed
62
63
		return
	}
Michael Yang's avatar
Michael Yang committed
64
	defer llm.Close()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
65

Michael Yang's avatar
Michael Yang committed
66
67
68
69
70
	templateNames := make([]string, 0, len(templates.Templates()))
	for _, template := range templates.Templates() {
		templateNames = append(templateNames, template.Name())
	}

Michael Yang's avatar
Michael Yang committed
71
	match, _ := matchRankOne(path.Base(req.Model), templateNames)
Michael Yang's avatar
Michael Yang committed
72
73
74
	if template := templates.Lookup(match); template != nil {
		var sb strings.Builder
		if err := template.Execute(&sb, req); err != nil {
75
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
76
77
78
79
80
81
			return
		}

		req.Prompt = sb.String()
	}

Michael Yang's avatar
Michael Yang committed
82
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
83
84
	g, _ := errgroup.WithContext(c.Request.Context())
	g.Go(func() error {
Bruce MacDonald's avatar
Bruce MacDonald committed
85
		defer close(ch)
Michael Yang's avatar
Michael Yang committed
86
		return llm.Predict(req.Prompt, func(s string) {
Michael Yang's avatar
Michael Yang committed
87
			ch <- api.GenerateResponse{Response: s}
Michael Yang's avatar
Michael Yang committed
88
89
		})
	})
Jeffrey Morgan's avatar
Jeffrey Morgan committed
90

Michael Yang's avatar
Michael Yang committed
91
	g.Go(func() error {
Michael Yang's avatar
Michael Yang committed
92
93
94
		stream(c, ch)
		return nil
	})
Michael Yang's avatar
Michael Yang committed
95

Michael Yang's avatar
Michael Yang committed
96
97
98
99
100
	if err := g.Wait(); err != nil && !errors.Is(err, io.EOF) {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
}
Michael Yang's avatar
Michael Yang committed
101

Michael Yang's avatar
Michael Yang committed
102
103
104
105
106
107
108
109
110
func pull(c *gin.Context) {
	var req api.PullRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	remote, err := getRemote(req.Model)
	if err != nil {
Michael Yang's avatar
Michael Yang committed
111
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
112
113
		return
	}
Michael Yang's avatar
Michael Yang committed
114

Michael Yang's avatar
Michael Yang committed
115
116
117
118
119
120
121
122
123
124
	ch := make(chan any)
	g, _ := errgroup.WithContext(c.Request.Context())
	g.Go(func() error {
		defer close(ch)
		return saveModel(remote, func(total, completed int64) {
			ch <- api.PullProgress{
				Total:     total,
				Completed: completed,
				Percent:   float64(total) / float64(completed) * 100,
			}
Michael Yang's avatar
Michael Yang committed
125
		})
Michael Yang's avatar
Michael Yang committed
126
	})
Michael Yang's avatar
Michael Yang committed
127

Michael Yang's avatar
Michael Yang committed
128
129
	g.Go(func() error {
		stream(c, ch)
Michael Yang's avatar
Michael Yang committed
130
		return nil
Jeffrey Morgan's avatar
Jeffrey Morgan committed
131
	})
Michael Yang's avatar
Michael Yang committed
132
133
134
135
136

	if err := g.Wait(); err != nil && !errors.Is(err, io.EOF) {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
137
138
139
140
141
}

func Serve(ln net.Listener) error {
	r := gin.Default()

142
143
144
145
	r.GET("/", func(c *gin.Context) {
		c.String(http.StatusOK, "Ollama is running")
	})

Michael Yang's avatar
Michael Yang committed
146
	r.POST("api/pull", pull)
Bruce MacDonald's avatar
Bruce MacDonald committed
147
	r.POST("/api/generate", generate)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
148
149
150
151
152
153
154
155

	log.Printf("Listening on %s", ln.Addr())
	s := &http.Server{
		Handler: r,
	}

	return s.Serve(ln)
}
Michael Yang's avatar
Michael Yang committed
156
157

func matchRankOne(source string, targets []string) (bestMatch string, bestRank int) {
Michael Yang's avatar
Michael Yang committed
158
	bestRank = math.MaxInt
Michael Yang's avatar
Michael Yang committed
159
	for _, target := range targets {
Michael Yang's avatar
Michael Yang committed
160
		if rank := fuzzy.LevenshteinDistance(source, target); bestRank > rank {
Michael Yang's avatar
Michael Yang committed
161
162
163
164
165
166
167
			bestRank = rank
			bestMatch = target
		}
	}

	return
}
Michael Yang's avatar
Michael Yang committed
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188

func stream(c *gin.Context, ch chan any) {
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		}

		bts, err := json.Marshal(val)
		if err != nil {
			return false
		}

		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
			return false
		}

		return true
	})
}