wip go engine

Co-authored-by: Patrick Devine <pdevine@sonic.net>

server/routes.go

+package server
+
+import (
+	"fmt"
+	"io"
+	"log"
+	"net"
+	"net/http"
+	"runtime"
+
+	"github.com/gin-gonic/gin"
+	llama "github.com/go-skynet/go-llama.cpp"
+
+	"github.com/ollama/ollama/api"
+)
+
+func Serve(ln net.Listener) error {
+	r := gin.Default()
+
+	var l *llama.LLama
+
+	gpulayers := 1
+	tokens := 512
+	threads := runtime.NumCPU()
+	model := "/Users/pdevine/.cache/gpt4all/GPT4All-13B-snoozy.ggmlv3.q4_0.bin"
+
+	r.POST("/api/load", func(c *gin.Context) {
+		var err error
+		l, err = llama.New(model, llama.EnableF16Memory, llama.SetContext(128), llama.EnableEmbeddings, llama.SetGPULayers(gpulayers))
+		if err != nil {
+			fmt.Println("Loading the model failed:", err.Error())
+		}
+	})
+
+	r.POST("/api/unload", func(c *gin.Context) {
+	})
+
+	r.POST("/api/generate", func(c *gin.Context) {
+		var req api.GenerateRequest
+		if err := c.ShouldBindJSON(&req); err != nil {
+			c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
+			return
+		}
+
+		ch := make(chan string)
+
+		go func() {
+			defer close(ch)
+			_, err := l.Predict(req.Prompt, llama.Debug, llama.SetTokenCallback(func(token string) bool {
+				ch <- token
+                        	return true
+                	}), llama.SetTokens(tokens), llama.SetThreads(threads), llama.SetTopK(90), llama.SetTopP(0.86), llama.SetStopWords("llama"))
+                	if err != nil {
+				panic(err)
+			}
+                }()
+
+		c.Stream(func(w io.Writer) bool {
+			tok, ok := <-ch
+			if !ok {
+				return false
+			}
+			c.SSEvent("token", tok)
+			return true
+		})
+
+/*
+                embeds, err := l.Embeddings(text)
+                if err != nil {
+                        fmt.Printf("Embeddings: error %s \n", err.Error())
+                }
+*/
+		
+	})
+
+	log.Printf("Listening on %s", ln.Addr())
+	s := &http.Server{
+		Handler: r,
+	}
+
+	return s.Serve(ln)
+}

signature/signature.go

+package signature
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"encoding/base64"
+	"encoding/hex"
+	"fmt"
+	"strings"
+
+	"golang.org/x/crypto/ssh"
+)
+
+type SignatureData struct {
+	Method string
+	Path   string
+	Data   []byte
+}
+
+func GetBytesToSign(s SignatureData) []byte {
+	// contentHash = base64(hex(sha256(s.Data)))
+	hash := sha256.Sum256(s.Data)
+	hashHex := make([]byte, hex.EncodedLen(len(hash)))
+	hex.Encode(hashHex, hash[:])
+	contentHash := base64.StdEncoding.EncodeToString(hashHex)
+
+	// bytesToSign e.g.: "GET,http://localhost,OTdkZjM1O...
+	bytesToSign := []byte(strings.Join([]string{s.Method, s.Path, contentHash}, ","))
+
+	return bytesToSign
+}
+
+// SignData takes a SignatureData object and signs it with a raw private key
+func SignAuthData(s SignatureData, rawKey []byte) (string, error) {
+	bytesToSign := GetBytesToSign(s)
+
+	// TODO replace this w/ a non-SSH based private key
+	privateKey, err := ssh.ParseRawPrivateKey(rawKey)
+	if err != nil {
+		return "", err
+	}
+
+	signer, err := ssh.NewSignerFromKey(privateKey)
+	if err != nil {
+		return "", err
+	}
+
+	// get the pubkey, but remove the type
+	pubKey := ssh.MarshalAuthorizedKey(signer.PublicKey())
+	parts := bytes.Split(pubKey, []byte(" "))
+	if len(parts) < 2 {
+		return "", fmt.Errorf("malformed private key")
+	}
+
+	signedData, err := signer.Sign(nil, bytesToSign)
+	if err != nil {
+		return "", err
+	}
+
+	// signature is <pubkey>:<signature>
+	sig := fmt.Sprintf("%s:%s", bytes.TrimSpace(parts[1]), base64.StdEncoding.EncodeToString(signedData.Blob))
+	return sig, nil
+}