470 lines
10 KiB
Go
470 lines
10 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
_ "embed"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"os/exec"
|
|
"os/signal"
|
|
"path"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"syscall"
|
|
)
|
|
|
|
var (
|
|
Config struct {
|
|
Port int
|
|
Debug bool
|
|
ChatBot struct {
|
|
SessionD string
|
|
PromptDelimiter string
|
|
semaphore sync.Mutex
|
|
WD string
|
|
Command string
|
|
N int
|
|
}
|
|
}
|
|
|
|
//go:embed template.d/login.html
|
|
htmlLogin []byte
|
|
//go:embed template.d/index.html
|
|
htmlIndex []byte
|
|
//go:embed template.d/chatbot.html
|
|
htmlChatBot []byte
|
|
)
|
|
|
|
func main() {
|
|
ctx, can := signal.NotifyContext(context.Background(), syscall.SIGINT)
|
|
defer can()
|
|
|
|
ctx, cleanup := contextWithCleanup(ctx)
|
|
defer cleanup()
|
|
|
|
config(ctx)
|
|
log.Printf("%+v", Config)
|
|
|
|
listenAndServe(ctx)
|
|
log.Println("done")
|
|
}
|
|
|
|
func contextWithCleanup(ctx context.Context) (context.Context, func()) {
|
|
m := map[int]func(){}
|
|
ctx = context.WithValue(ctx, "__cleanup__", m)
|
|
return ctx, func() {
|
|
defer func() {
|
|
recover()
|
|
}()
|
|
for _, v := range m {
|
|
v()
|
|
}
|
|
}
|
|
}
|
|
|
|
func contextWithCleanupFunc(ctx context.Context, foo func()) {
|
|
v := ctx.Value("__cleanup__")
|
|
if v == nil {
|
|
panic("cannot get context with cleanup func that doesnt have cleanup init")
|
|
}
|
|
m := v.(map[int]func())
|
|
m[len(m)] = foo
|
|
}
|
|
|
|
func config(ctx context.Context) {
|
|
d, err := os.MkdirTemp(os.TempDir(), "ai.*")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
contextWithCleanupFunc(ctx, func() { os.RemoveAll(d) })
|
|
|
|
flag.IntVar(&Config.Port, "p", 37070, "port to listen on")
|
|
flag.StringVar(&Config.ChatBot.SessionD, "chatbot-session-d", d, "dir to store chat bot sessions")
|
|
flag.StringVar(&Config.ChatBot.PromptDelimiter, "chatbot-rp", "> ", "prompt delimiter prefixed by NAME/YOU")
|
|
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 --repeat_penalty 1.0", "chatbot cmd prefix")
|
|
flag.IntVar(&Config.ChatBot.N, "chatbot-n", 256, "chatbot items to gen")
|
|
flag.BoolVar(&Config.Debug, "debug", false, "debug mode")
|
|
flag.Parse()
|
|
}
|
|
|
|
func listenAndServe(ctx context.Context) {
|
|
wg := &sync.WaitGroup{}
|
|
s := &http.Server{
|
|
Addr: fmt.Sprintf(":%d", Config.Port),
|
|
Handler: http.HandlerFunc(handle),
|
|
}
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
if err := s.ListenAndServe(); err != nil && ctx.Err() == nil {
|
|
panic(err)
|
|
}
|
|
}()
|
|
|
|
<-ctx.Done()
|
|
s.Close()
|
|
wg.Wait()
|
|
}
|
|
|
|
func handle(w http.ResponseWriter, r *http.Request) {
|
|
cookie, _ := ParseCookie(r)
|
|
if err := _handle(w, r); err != nil {
|
|
log.Printf("%s: %s %s: %v", cookie.MyName(), r.Method, r.URL.Path, err)
|
|
} else {
|
|
log.Printf("%s: %s %s", cookie.MyName(), r.Method, r.URL.Path)
|
|
}
|
|
}
|
|
|
|
func _handle(w http.ResponseWriter, r *http.Request) error {
|
|
first := strings.Split(strings.TrimLeft(r.URL.Path, "/"), "/")[0]
|
|
switch first {
|
|
case "login":
|
|
return handleLogin(w, r)
|
|
case "api":
|
|
return handleAPI(w, r)
|
|
default:
|
|
return handleUI(w, r)
|
|
}
|
|
}
|
|
|
|
func handleLogin(w http.ResponseWriter, r *http.Request) error {
|
|
switch r.Method {
|
|
case http.MethodGet:
|
|
w.Write(htmlLogin)
|
|
return nil
|
|
case http.MethodPost:
|
|
err := r.ParseForm()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
cookie, err := ParseCookie(r)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
cookie.Inject(w)
|
|
http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
|
|
return nil
|
|
default:
|
|
return handleNotFound(w, r)
|
|
}
|
|
}
|
|
|
|
func handleUI(w http.ResponseWriter, r *http.Request) error {
|
|
if _, err := ParseCookie(r); err != nil {
|
|
http.Redirect(w, r, "/login", http.StatusTemporaryRedirect)
|
|
return err
|
|
}
|
|
|
|
switch r.URL.Path {
|
|
case "/":
|
|
w.Write(htmlIndex)
|
|
return nil
|
|
case "/chatbot":
|
|
w.Write(htmlChatBot)
|
|
return nil
|
|
default:
|
|
return handleNotFound(w, r)
|
|
}
|
|
}
|
|
|
|
type Cookie struct {
|
|
Name string
|
|
}
|
|
|
|
func ParseCookie(r *http.Request) (Cookie, error) {
|
|
if r.URL.Path != "/login" {
|
|
return parseCookieFromCookie(r)
|
|
}
|
|
cookie := Cookie{
|
|
Name: r.PostForm.Get("Name"),
|
|
}
|
|
return cookie, cookie.Verify()
|
|
}
|
|
|
|
func parseCookieFromCookie(r *http.Request) (Cookie, error) {
|
|
cookie, err := r.Cookie("root")
|
|
if err != nil {
|
|
return Cookie{}, err
|
|
}
|
|
|
|
decoded, err := base64.URLEncoding.DecodeString(cookie.Value)
|
|
if err != nil {
|
|
return Cookie{}, err
|
|
}
|
|
|
|
var result Cookie
|
|
if err := json.Unmarshal(decoded, &result); err != nil {
|
|
return Cookie{}, err
|
|
}
|
|
|
|
return result, result.Verify()
|
|
}
|
|
|
|
func (cookie Cookie) MyName() string {
|
|
return string(bytes.Join(
|
|
regexp.MustCompile(`[a-zA-Z]`).FindAll([]byte(cookie.Name), -1),
|
|
[]byte(""),
|
|
))
|
|
}
|
|
|
|
func (cookie Cookie) Verify() error {
|
|
if cookie.MyName() == "" {
|
|
return fmt.Errorf("incomplete cookie")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (cookie Cookie) Inject(w http.ResponseWriter) {
|
|
c := &http.Cookie{
|
|
Name: "root",
|
|
Value: cookie.Serialize(),
|
|
}
|
|
http.SetCookie(w, c)
|
|
}
|
|
|
|
func (cookie Cookie) Serialize() string {
|
|
b, _ := json.Marshal(cookie)
|
|
return base64.URLEncoding.EncodeToString(b)
|
|
}
|
|
|
|
func handleAPI(w http.ResponseWriter, r *http.Request) error {
|
|
if _, err := ParseCookie(r); err != nil {
|
|
http.Redirect(w, r, "/login", http.StatusTemporaryRedirect)
|
|
return err
|
|
}
|
|
|
|
switch r.URL.Path {
|
|
case "/api/v0/chatbot":
|
|
return handleAPIChatBot(w, r)
|
|
default:
|
|
return handleNotFound(w, r)
|
|
}
|
|
}
|
|
|
|
func handleNotFound(w http.ResponseWriter, r *http.Request) error {
|
|
http.NotFound(w, r)
|
|
return fmt.Errorf("not found: %s %s", r.Method, r.URL.Path)
|
|
}
|
|
|
|
func handleAPIChatBot(w http.ResponseWriter, r *http.Request) error {
|
|
err := r.ParseForm()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
switch r.Method {
|
|
case http.MethodPost:
|
|
return handleAPIChatBotPost(w, r)
|
|
case http.MethodPut:
|
|
return handleAPIChatBotPut(w, r)
|
|
default:
|
|
return handleNotFound(w, r)
|
|
}
|
|
}
|
|
|
|
func handleAPIChatBotPost(w http.ResponseWriter, r *http.Request) error {
|
|
cookie, _ := ParseCookie(r)
|
|
sessionD := path.Join(Config.ChatBot.SessionD, cookie.MyName())
|
|
if _, err := os.Stat(path.Join(sessionD)); err == nil {
|
|
if err := os.RemoveAll(path.Join(sessionD)); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if err := os.MkdirAll(path.Join(sessionD, "chat.d"), os.ModePerm); err != nil {
|
|
return err
|
|
}
|
|
|
|
prompt := r.PostForm.Get("Prompt")
|
|
if len(prompt) == 0 {
|
|
return errors.New("no prompt")
|
|
}
|
|
if err := os.WriteFile(path.Join(sessionD, "prompt.txt"), []byte(prompt), os.ModePerm); err != nil {
|
|
return err
|
|
}
|
|
|
|
return handleAPIChatBotPut(w, r)
|
|
}
|
|
|
|
func handleAPIChatBotPut(w http.ResponseWriter, r *http.Request) error {
|
|
message := r.PostForm.Get("Message")
|
|
if len(message) == 0 {
|
|
return errors.New("empty Message")
|
|
}
|
|
|
|
cookie, _ := ParseCookie(r)
|
|
sessionD := path.Join(Config.ChatBot.SessionD, cookie.MyName())
|
|
|
|
promptF := path.Join(sessionD, "prompt.txt")
|
|
inputF := path.Join(sessionD, "input.txt")
|
|
cacheF := path.Join(sessionD, "cache.bin")
|
|
reversePrompt := cookie.MyName()
|
|
if len(reversePrompt) > 8 {
|
|
reversePrompt = reversePrompt[:8]
|
|
}
|
|
reversePrompt = reversePrompt + Config.ChatBot.PromptDelimiter
|
|
forwardPrompt := "YOU" + Config.ChatBot.PromptDelimiter
|
|
if err := copyFile(inputF, promptF); err != nil {
|
|
return err
|
|
}
|
|
if err := chatBotGenerateInitCacheF(r.Context(), cacheF, inputF); err != nil {
|
|
return err
|
|
}
|
|
if err := appendFile(inputF, reversePrompt+message+"\n"+forwardPrompt); err != nil {
|
|
return err
|
|
}
|
|
|
|
justNew, err := chatBotGenerateAndFillInputF(r.Context(), cacheF, inputF, reversePrompt)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := os.Rename(inputF, promptF); err != nil {
|
|
return err
|
|
}
|
|
w.Write(bytes.TrimSuffix(justNew, []byte(reversePrompt)))
|
|
|
|
return nil
|
|
}
|
|
|
|
func chatBotGenerateInitCacheF(ctx context.Context, cacheF, inputF string) error {
|
|
if _, err := os.Stat(cacheF); !os.IsNotExist(err) {
|
|
return nil
|
|
}
|
|
commands := strings.Fields(Config.ChatBot.Command)
|
|
commands = append(commands,
|
|
"--batch-size", "8",
|
|
"--prompt-cache", cacheF,
|
|
"-f", inputF,
|
|
"--n_predict", "1",
|
|
)
|
|
command := exec.CommandContext(
|
|
ctx,
|
|
commands[0],
|
|
commands[1:]...,
|
|
)
|
|
command.Dir = Config.ChatBot.WD
|
|
|
|
Config.ChatBot.semaphore.Lock()
|
|
defer Config.ChatBot.semaphore.Unlock()
|
|
|
|
if b, err := command.CombinedOutput(); err != nil {
|
|
return fmt.Errorf("error generating cache with '%s': %w: %s", command.String(), err, b)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func chatBotGenerateAndFillInputF(ctx context.Context, cacheF, inputF, reversePrompt string) ([]byte, error) {
|
|
commands := strings.Fields(Config.ChatBot.Command)
|
|
commands = append(commands,
|
|
"-f", inputF,
|
|
"--prompt-cache-all",
|
|
"--prompt-cache", cacheF,
|
|
"-n", strconv.Itoa(Config.ChatBot.N),
|
|
)
|
|
if len(reversePrompt) > 0 {
|
|
commands = append(commands,
|
|
"--reverse-prompt", reversePrompt,
|
|
)
|
|
}
|
|
command := exec.CommandContext(
|
|
ctx,
|
|
commands[0],
|
|
commands[1:]...,
|
|
)
|
|
command.Dir = Config.ChatBot.WD
|
|
command.Stderr = log.Writer()
|
|
|
|
stdout, err := command.StdoutPipe()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
buff := bytes.NewBuffer(nil)
|
|
go func() {
|
|
stdout.Read(make([]byte, 1)) //1 BOS byte
|
|
io.Copy(buff, stdout)
|
|
}()
|
|
|
|
if err := func() error {
|
|
Config.ChatBot.semaphore.Lock()
|
|
defer Config.ChatBot.semaphore.Unlock()
|
|
|
|
return command.Run()
|
|
}(); err != nil {
|
|
return nil, err
|
|
|
|
}
|
|
|
|
oldAndNew := buff.Bytes()
|
|
log.Printf("generated: [%s]", oldAndNew)
|
|
priorContent, err := os.ReadFile(inputF)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
justNew := oldAndNew[len(priorContent):]
|
|
idx := bytes.LastIndex(
|
|
append(priorContent, justNew...),
|
|
[]byte(reversePrompt),
|
|
)
|
|
log.Printf("found lastindex(%d priorContent + %d newContent, %s) = %v", len(priorContent), len(justNew), reversePrompt, idx)
|
|
if idx+len(reversePrompt) > len(priorContent) {
|
|
justNew = justNew[:idx+len(reversePrompt)-len(priorContent)]
|
|
}
|
|
if err := _appendFile(inputF, string(justNew)); err != nil {
|
|
return nil, err
|
|
}
|
|
log.Printf("newly generated: [%s]", justNew)
|
|
|
|
if !bytes.HasSuffix(append(priorContent, justNew...), []byte(reversePrompt)) {
|
|
more, err := chatBotGenerateAndFillInputF(ctx, cacheF, inputF, reversePrompt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
justNew = append(justNew, more...)
|
|
}
|
|
return justNew, nil
|
|
}
|
|
|
|
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 {
|
|
return _appendFile(toF, "\n"+msg+"\n")
|
|
}
|
|
|
|
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()
|
|
|
|
if _, err := f.WriteString(msg); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|