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

import (
Michael Yang's avatar
Michael Yang committed
4
	"encoding/json"
5
	"errors"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
6
7
8
9
	"io"
	"log"
	"net"
	"net/http"
10
	"os"
Michael Yang's avatar
Michael Yang committed
11
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
12
	"strings"
13
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
14

Michael Yang's avatar
Michael Yang committed
15
	"dario.cat/mergo"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16
17
	"github.com/gin-gonic/gin"

Jeffrey Morgan's avatar
Jeffrey Morgan committed
18
	"github.com/jmorganca/ollama/api"
Michael Yang's avatar
Michael Yang committed
19
	"github.com/jmorganca/ollama/llama"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
20
21
)

22
func GenerateHandler(c *gin.Context) {
23
24
	start := time.Now()

Michael Yang's avatar
Michael Yang committed
25
	var req api.GenerateRequest
Bruce MacDonald's avatar
Bruce MacDonald committed
26
	if err := c.ShouldBindJSON(&req); err != nil {
27
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Bruce MacDonald's avatar
Bruce MacDonald committed
28
29
		return
	}
30

31
32
33
34
	model, err := GetModel(req.Model)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
Bruce MacDonald's avatar
Bruce MacDonald committed
35
	}
Michael Yang's avatar
Michael Yang committed
36

Michael Yang's avatar
Michael Yang committed
37
38
39
40
41
42
43
44
45
46
47
	opts := api.DefaultOptions()
	if err := mergo.Merge(&opts, model.Options, mergo.WithOverride); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	if err := mergo.Merge(&opts, req.Options, mergo.WithOverride); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

48
	prompt, err := model.Prompt(req)
49
50
51
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
Michael Yang's avatar
Michael Yang committed
52
53
	}

Michael Yang's avatar
Michael Yang committed
54
	llm, err := llama.New(model.ModelPath, opts)
Michael Yang's avatar
Michael Yang committed
55
56
57
58
59
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	defer llm.Close()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
60

Michael Yang's avatar
Michael Yang committed
61
62
63
	ch := make(chan any)
	go func() {
		defer close(ch)
Michael Yang's avatar
Michael Yang committed
64
		fn := func(r api.GenerateResponse) {
Michael Yang's avatar
Michael Yang committed
65
66
67
68
69
70
71
			r.Model = req.Model
			r.CreatedAt = time.Now().UTC()
			if r.Done {
				r.TotalDuration = time.Since(start)
			}

			ch <- r
Michael Yang's avatar
Michael Yang committed
72
73
74
75
76
		}

		if err := llm.Predict(req.Context, prompt, fn); err != nil {
			ch <- gin.H{"error": err.Error()}
		}
Michael Yang's avatar
Michael Yang committed
77
	}()
Michael Yang's avatar
Michael Yang committed
78

Michael Yang's avatar
Michael Yang committed
79
	streamResponse(c, ch)
Michael Yang's avatar
Michael Yang committed
80
}
Michael Yang's avatar
Michael Yang committed
81

82
func PullModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
83
84
85
86
87
88
	var req api.PullRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

89
90
91
	ch := make(chan any)
	go func() {
		defer close(ch)
92
93
		fn := func(r api.ProgressResponse) {
			ch <- r
94
		}
95

96
		if err := PullModel(req.Name, req.Username, req.Password, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
97
			ch <- gin.H{"error": err.Error()}
98
99
100
101
102
103
		}
	}()

	streamResponse(c, ch)
}

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

111
112
113
	ch := make(chan any)
	go func() {
		defer close(ch)
114
115
		fn := func(r api.ProgressResponse) {
			ch <- r
116
		}
117

118
		if err := PushModel(req.Name, req.Username, req.Password, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
119
			ch <- gin.H{"error": err.Error()}
120
121
122
123
124
125
		}
	}()

	streamResponse(c, ch)
}

126
func CreateModelHandler(c *gin.Context) {
127
128
129
	var req api.CreateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
Michael Yang's avatar
Michael Yang committed
130
		return
131
132
	}

Michael Yang's avatar
Michael Yang committed
133
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
134
135
	go func() {
		defer close(ch)
136
137
138
		fn := func(status string) {
			ch <- api.CreateProgress{
				Status: status,
Michael Yang's avatar
Michael Yang committed
139
			}
140
141
		}

142
		if err := CreateModel(req.Name, req.Path, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
143
			ch <- gin.H{"error": err.Error()}
144
		}
Michael Yang's avatar
Michael Yang committed
145
	}()
Michael Yang's avatar
Michael Yang committed
146

Michael Yang's avatar
Michael Yang committed
147
	streamResponse(c, ch)
Bruce MacDonald's avatar
Bruce MacDonald committed
148
149
}

150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
func DeleteModelHandler(c *gin.Context) {
	var req api.DeleteRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	ch := make(chan any)
	go func() {
		defer close(ch)
		fn := func(r api.ProgressResponse) {
			ch <- r
		}

		if err := DeleteModel(req.Name, fn); err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
	}()

	streamResponse(c, ch)
}

func ListModelsHandler(c *gin.Context) {
Patrick Devine's avatar
Patrick Devine committed
174
175
176
177
178
179
180
181
	var models []api.ListResponseModel
	fp, err := GetManifestPath()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	err = filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
		if err != nil {
182
183
184
185
			if errors.Is(err, os.ErrNotExist) {
				log.Printf("manifest file does not exist: %s", fp)
				return nil
			}
Patrick Devine's avatar
Patrick Devine committed
186
187
188
189
190
			return err
		}
		if !info.IsDir() {
			fi, err := os.Stat(path)
			if err != nil {
191
192
				log.Printf("skipping file: %s", fp)
				return nil
Patrick Devine's avatar
Patrick Devine committed
193
194
195
196
197
198
199
200
201
202
			}
			path := path[len(fp)+1:]
			slashIndex := strings.LastIndex(path, "/")
			if slashIndex == -1 {
				return nil
			}
			tag := path[:slashIndex] + ":" + path[slashIndex+1:]
			mp := ParseModelPath(tag)
			manifest, err := GetManifest(mp)
			if err != nil {
203
204
				log.Printf("skipping file: %s", fp)
				return nil
Patrick Devine's avatar
Patrick Devine committed
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
			}
			model := api.ListResponseModel{
				Name:       mp.GetShortTagname(),
				Size:       manifest.GetTotalSize(),
				ModifiedAt: fi.ModTime(),
			}
			models = append(models, model)
		}
		return nil
	})
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	c.JSON(http.StatusOK, api.ListResponse{models})
}

Bruce MacDonald's avatar
Bruce MacDonald committed
223
224
225
func Serve(ln net.Listener) error {
	r := gin.Default()

226
227
228
229
	r.GET("/", func(c *gin.Context) {
		c.String(http.StatusOK, "Ollama is running")
	})

230
231
232
233
234
235
	r.POST("/api/pull", PullModelHandler)
	r.POST("/api/generate", GenerateHandler)
	r.POST("/api/create", CreateModelHandler)
	r.POST("/api/push", PushModelHandler)
	r.GET("/api/tags", ListModelsHandler)
	r.DELETE("/api/delete", DeleteModelHandler)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
236
237
238
239
240
241
242
243

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

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

Michael Yang's avatar
Michael Yang committed
245
func streamResponse(c *gin.Context, ch chan any) {
Michael Yang's avatar
Michael Yang committed
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
	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
	})
}