summaryrefslogtreecommitdiff
path: root/server.go
blob: 5a1a1c37551cba892123907146a56531506262ab (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
package main

import (
	"elefant/config"
	"encoding/json"
	"fmt"
	"net/http"
	"time"
)

type Server struct {
	// nolint
	config config.Config
}

func (srv *Server) ListenToRequests(port string) {
	// h := srv.actions
	mux := http.NewServeMux()
	server := &http.Server{
		Addr:         "localhost:" + port,
		Handler:      mux,
		ReadTimeout:  time.Second * 5,
		WriteTimeout: time.Second * 5,
	}
	mux.HandleFunc("GET /ping", pingHandler)
	mux.HandleFunc("GET /model", modelHandler)
	mux.HandleFunc("POST /completion", completionHandler)
	fmt.Println("Listening", "addr", server.Addr)
	if err := server.ListenAndServe(); err != nil {
		panic(err)
	}
}

// create server
// listen to the completion endpoint handler
func pingHandler(w http.ResponseWriter, req *http.Request) {
	if _, err := w.Write([]byte("pong")); err != nil {
		logger.Error("server ping", "error", err)
	}
}

func completionHandler(w http.ResponseWriter, req *http.Request) {
	// post request
	body := req.Body
	// get body as io.reader
	// pass it to the /completion
	go sendMsgToLLM(body)
out:
	for {
		select {
		case chunk := <-chunkChan:
			fmt.Print(chunk)
			if _, err := w.Write([]byte(chunk)); err != nil {
				logger.Warn("failed to write chunk", "value", chunk)
				continue
			}
		case <-streamDone:
			break out
		}
	}
}

func modelHandler(w http.ResponseWriter, req *http.Request) {
	llmModel := fetchModelName()
	payload, err := json.Marshal(llmModel)
	if err != nil {
		logger.Error("model handler", "error", err)
		// return err
		return
	}
	if _, err := w.Write(payload); err != nil {
		logger.Error("model handler", "error", err)
	}
}