From a1ad8a2b44305e790f59baaec08b5d74d9f3ab71 Mon Sep 17 00:00:00 2001 From: bel Date: Sat, 17 Jun 2023 10:15:33 -0600 Subject: [PATCH] mvp ish --- vicuna-tools.d/main.go | 85 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 83 insertions(+), 2 deletions(-) diff --git a/vicuna-tools.d/main.go b/vicuna-tools.d/main.go index d04bf0c..830d863 100644 --- a/vicuna-tools.d/main.go +++ b/vicuna-tools.d/main.go @@ -1,6 +1,7 @@ package main import ( + "bytes" "context" _ "embed" "encoding/base64" @@ -8,9 +9,11 @@ import ( "errors" "flag" "fmt" + "io" "log" "net/http" "os" + "os/exec" "os/signal" "path" "strings" @@ -25,6 +28,8 @@ var ( ChatBot struct { SessionD string semaphore sync.Mutex + WD string + Command string } } @@ -80,7 +85,9 @@ func config(ctx context.Context) { contextWithCleanupFunc(ctx, func() { os.RemoveAll(d) }) flag.IntVar(&Config.Port, "p", 37070, "port to listen on") - flag.StringVar(&Config.ChatBot.SessionD, "chatbot-d", d, "dir to store chat bot sessions") + flag.StringVar(&Config.ChatBot.SessionD, "chatbot-session-d", d, "dir to store chat bot sessions") + flag.StringVar(&Config.ChatBot.WD, "chatbot-working-d", "./llama.cpp", "working directory for chatbot") + flag.StringVar(&Config.ChatBot.Command, "chatbot-cmd", "./main -m ./models/ggml-vic7b-uncensored-q5_1.bin -n 256 --repeat_penalty 1.0", "chatbot cmd prefix") flag.BoolVar(&Config.Debug, "debug", false, "debug mode") flag.Parse() } @@ -279,8 +286,82 @@ func handleAPIChatBotPut(w http.ResponseWriter, r *http.Request) error { return errors.New("empty Message") } + cookie, _ := ParseCookie(r) + sessionD := path.Join(Config.ChatBot.SessionD, cookie.Name) + + promptF := path.Join(sessionD, "prompt.txt") + inputF := path.Join(sessionD, "input.txt") + if err := copyFile(inputF, promptF); err != nil { + return err + } + if err := appendFile(inputF, message); err != nil { + return err + } + + commands := strings.Fields(Config.ChatBot.Command) + commands = append(commands, "-f", inputF) + command := exec.CommandContext( + r.Context(), + commands[0], + commands[1:]..., + ) + command.Dir = Config.ChatBot.WD + + stdout, err := command.StdoutPipe() + if err != nil { + return err + } + buff := bytes.NewBuffer(nil) + go func() { + io.Copy(buff, stdout) + }() + + stderr, err := command.StderrPipe() + if err != nil { + return err + } + go func() { + io.Copy(os.Stderr, stderr) + }() + Config.ChatBot.semaphore.Lock() defer Config.ChatBot.semaphore.Unlock() - return errors.New("not impl") + if err := command.Run(); err != nil { + return err + } + + log.Printf("OUTPUT: %s", buff.Bytes()) + return errors.New("not impl: move input to prompt") +} + +func copyFile(toF, fromF string) error { + from, err := os.Open(fromF) + if err != nil { + return err + } + defer from.Close() + + to, err := os.Create(toF) + if err != nil { + return err + } + defer to.Close() + + io.Copy(to, from) + return nil +} + +func appendFile(toF, msg string) error { + f, err := os.OpenFile(toF, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600) + if err != nil { + return err + } + defer f.Close() + + f.Write([]byte("\n")) + if _, err := f.WriteString(msg); err != nil { + return err + } + return nil }