mvp ish

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
 }