ollama/server/routes.go

202 lines
4 KiB
Go
Raw Normal View History

package server
import (
2023-07-06 18:33:29 +00:00
"embed"
2023-07-06 17:40:11 +00:00
"encoding/json"
"errors"
"io"
"log"
2023-07-07 00:03:00 +00:00
"math"
"net"
"net/http"
"os"
2023-07-06 17:40:11 +00:00
"path"
"strings"
"text/template"
2023-07-13 01:18:06 +00:00
"time"
"github.com/gin-gonic/gin"
2023-07-06 17:40:11 +00:00
"github.com/lithammer/fuzzysearch/fuzzy"
2023-07-03 20:32:48 +00:00
"github.com/jmorganca/ollama/api"
2023-07-06 17:40:11 +00:00
"github.com/jmorganca/ollama/llama"
)
2023-07-06 18:33:29 +00:00
//go:embed templates/*
var templatesFS embed.FS
var templates = template.Must(template.ParseFS(templatesFS, "templates/*.prompt"))
2023-07-06 17:40:11 +00:00
func cacheDir() string {
home, err := os.UserHomeDir()
if err != nil {
panic(err)
}
return path.Join(home, ".ollama")
}
2023-07-05 19:37:33 +00:00
func generate(c *gin.Context) {
2023-07-13 01:18:06 +00:00
start := time.Now()
2023-07-07 22:29:17 +00:00
req := api.GenerateRequest{
Options: api.DefaultOptions(),
}
2023-07-05 19:37:33 +00:00
if err := c.ShouldBindJSON(&req); err != nil {
2023-07-07 21:04:43 +00:00
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
2023-07-05 19:37:33 +00:00
return
}
if remoteModel, _ := getRemote(req.Model); remoteModel != nil {
2023-07-06 22:43:04 +00:00
req.Model = remoteModel.FullName()
}
if _, err := os.Stat(req.Model); err != nil {
if !errors.Is(err, os.ErrNotExist) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
req.Model = path.Join(cacheDir(), "models", req.Model+".bin")
}
2023-07-06 22:43:04 +00:00
2023-07-11 21:57:17 +00:00
ch := make(chan any)
go stream(c, ch)
2023-07-06 17:40:11 +00:00
templateNames := make([]string, 0, len(templates.Templates()))
for _, template := range templates.Templates() {
templateNames = append(templateNames, template.Name())
}
2023-07-07 00:03:00 +00:00
match, _ := matchRankOne(path.Base(req.Model), templateNames)
2023-07-06 17:40:11 +00:00
if template := templates.Lookup(match); template != nil {
var sb strings.Builder
if err := template.Execute(&sb, req); err != nil {
2023-07-07 21:04:43 +00:00
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
2023-07-06 17:40:11 +00:00
return
}
req.Prompt = sb.String()
}
2023-07-11 21:57:17 +00:00
llm, err := llama.New(req.Model, req.Options)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
defer llm.Close()
2023-07-04 04:47:00 +00:00
2023-07-13 01:18:06 +00:00
fn := func(r api.GenerateResponse) {
r.Model = req.Model
r.CreatedAt = time.Now().UTC()
if r.Done {
r.TotalDuration = time.Since(start)
}
ch <- r
2023-07-11 21:57:17 +00:00
}
2023-07-06 17:40:11 +00:00
2023-07-11 21:57:17 +00:00
if err := llm.Predict(req.Prompt, fn); err != nil {
2023-07-11 18:54:22 +00:00
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
2023-07-11 21:57:17 +00:00
2023-07-11 18:54:22 +00:00
}
2023-07-06 17:40:11 +00:00
2023-07-11 18:54:22 +00:00
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 {
2023-07-11 20:05:51 +00:00
c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
2023-07-11 18:54:22 +00:00
return
}
2023-07-06 17:40:11 +00:00
2023-07-13 02:07:15 +00:00
// check if completed file exists
fi, err := os.Stat(remote.FullName())
switch {
case errors.Is(err, os.ErrNotExist):
// noop, file doesn't exist so create it
case err != nil:
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
default:
c.JSON(http.StatusOK, api.PullProgress{
Total: fi.Size(),
Completed: fi.Size(),
Percent: 100,
})
return
}
2023-07-11 18:54:22 +00:00
ch := make(chan any)
2023-07-11 21:57:17 +00:00
go stream(c, ch)
2023-07-06 17:40:11 +00:00
2023-07-11 21:57:17 +00:00
fn := func(total, completed int64) {
ch <- api.PullProgress{
Total: total,
Completed: completed,
2023-07-12 16:55:07 +00:00
Percent: float64(completed) / float64(total) * 100,
2023-07-11 21:57:17 +00:00
}
}
2023-07-07 22:29:17 +00:00
2023-07-11 21:57:17 +00:00
if err := saveModel(remote, fn); err != nil {
2023-07-07 22:29:17 +00:00
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
2023-07-05 19:37:33 +00:00
}
func Serve(ln net.Listener) error {
r := gin.Default()
2023-07-08 03:46:15 +00:00
r.GET("/", func(c *gin.Context) {
c.String(http.StatusOK, "Ollama is running")
})
2023-07-13 00:19:03 +00:00
r.POST("/api/pull", pull)
2023-07-05 19:37:33 +00:00
r.POST("/api/generate", generate)
log.Printf("Listening on %s", ln.Addr())
s := &http.Server{
Handler: r,
}
return s.Serve(ln)
}
2023-07-06 17:40:11 +00:00
func matchRankOne(source string, targets []string) (bestMatch string, bestRank int) {
2023-07-07 00:03:00 +00:00
bestRank = math.MaxInt
2023-07-06 17:40:11 +00:00
for _, target := range targets {
2023-07-07 00:03:00 +00:00
if rank := fuzzy.LevenshteinDistance(source, target); bestRank > rank {
2023-07-06 17:40:11 +00:00
bestRank = rank
bestMatch = target
}
}
return
}
2023-07-11 18:54:22 +00:00
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
})
}