orchard

  1// llm.bythewood.me, the model gateway.
  2//
  3// One model on one card, in front of every service that wants one. Before this,
  4// chat and search each carried their own llama-swap and their own copy of the
  5// weights, which on an 8GB card means whichever one you used last evicted the
  6// other. This holds the model, hands out API keys, and writes down every prompt
  7// and every completion that passes through it, apart from the calls a caller
  8// marks incognito, which are forwarded and never recorded.
  9//
 10// The web half is behind auth.bythewood.me like the other dashboards. The API
 11// half is behind a key, because the callers are containers with no browser.
 12package main
 13
 14import (
 15	"encoding/json"
 16	"flag"
 17	"fmt"
 18	"html/template"
 19	"log/slog"
 20	"net/http"
 21	"os"
 22	"strconv"
 23	"strings"
 24	"time"
 25
 26	"llm.bythewood.me/web"
 27)
 28
 29const (
 30	listenAddr = ":8000"
 31	sourceURL  = "https://github.com/overshard/orchard/tree/main/sites/llm.bythewood.me"
 32	selfSource = "llm"
 33)
 34
 35// csp is the same shape as every other site here. 'unsafe-inline' is for the
 36// analytics collector snippet, which is an inline script by design, and for the
 37// handful of inline style attributes the templates set.
 38func csp() string {
 39	return strings.Join([]string{
 40		"default-src 'self'",
 41		"script-src 'self' 'unsafe-inline' https://analytics.bythewood.me",
 42		"style-src 'self' 'unsafe-inline'",
 43		"img-src 'self' data:",
 44		"font-src 'self'",
 45		"connect-src 'self' https://analytics.bythewood.me",
 46		"base-uri 'self'",
 47		"form-action 'self'",
 48		"frame-ancestors 'none'",
 49	}, "; ")
 50}
 51
 52type site struct {
 53	auth     *web.Authenticator
 54	store    *Store
 55	tpl      *template.Template
 56	client   *http.Client
 57	upstream string
 58	model    string
 59	retain   time.Duration
 60}
 61
 62func env(k, def string) string {
 63	if v := os.Getenv(k); v != "" {
 64		return v
 65	}
 66	return def
 67}
 68
 69func main() {
 70	var (
 71		addr     = flag.String("addr", env("LLM_ADDR", listenAddr), "listen address")
 72		upstream = flag.String("upstream", env("LLM_UPSTREAM", "http://swap:8091"), "llama-swap base url")
 73		model    = flag.String("model", env("LLM_MODEL", "local"), "the model name callers ask for")
 74		dbPath   = flag.String("db", env("LLM_DB", "data/llm.db"), "database path")
 75		retain   = flag.Duration("retain", 90*24*time.Hour, "how long a logged call is kept")
 76		check    = flag.Bool("healthcheck", false, "probe this process and exit")
 77		newKey   = flag.String("newkey", "", "mint an api key with this name, print it, and exit")
 78	)
 79	flag.Parse()
 80
 81	if *check {
 82		if err := web.HealthCheck("http://127.0.0.1"+*addr+"/healthz", 3*time.Second); err != nil {
 83			os.Exit(1)
 84		}
 85		return
 86	}
 87
 88	web.SetupLogging()
 89	web.ShipLogs(selfSource, web.HTTPSink())
 90
 91	store, err := OpenStore(*dbPath)
 92	if err != nil {
 93		slog.Error("opening the database", "err", err)
 94		os.Exit(1)
 95	}
 96	defer store.Close()
 97
 98	// The first key cannot come from the web UI, because that is behind auth
 99	// and auth is reached over the tunnel this gateway is meant to be feeding.
100	// Same shape as auth's own recovery codes: a flag, printed once.
101	if *newKey != "" {
102		secret, _, err := store.NewKey(*newKey)
103		if err != nil {
104			fmt.Fprintln(os.Stderr, err)
105			os.Exit(1)
106		}
107		fmt.Println(secret)
108		store.Checkpoint()
109		return
110	}
111
112	s := &site{
113		auth:     web.NewAuthenticator(),
114		store:    store,
115		upstream: strings.TrimSuffix(*upstream, "/"),
116		model:    *model,
117		retain:   *retain,
118		// No timeout on the client. A cold model load plus a long generation
119		// runs well past any value worth picking, and the caller's own context
120		// is what bounds this instead.
121		client: &http.Client{},
122	}
123	if err := s.loadTemplates(); err != nil {
124		slog.Error("templates", "err", err)
125		os.Exit(1)
126	}
127	go s.prune()
128
129	mux := http.NewServeMux()
130
131	// The dashboard. Behind auth like every other one here.
132	// The root is the one public page: what this is and a way in. A signed in
133	// visitor gets the dashboard here rather than the pitch.
134	mux.HandleFunc("GET /{$}", s.root)
135	mux.HandleFunc("GET /calls", s.auth.RequireAuth(s.callsPage))
136	mux.Handle("GET /static/", s.static())
137	mux.HandleFunc("GET /login", func(w http.ResponseWriter, r *http.Request) {
138		http.Redirect(w, r, web.LoginURL(r), http.StatusSeeOther)
139	})
140	mux.HandleFunc("POST /keys", s.auth.RequireAuth(s.keyCreate))
141	mux.HandleFunc("POST /keys/{id}/revoke", s.auth.RequireAuth(s.keyRevoke))
142	mux.HandleFunc("POST /keys/{id}/delete", s.auth.RequireAuth(s.keyDelete))
143
144	// The gateway. A key, not a cookie, because the callers are containers.
145	mux.HandleFunc("POST /v1/chat/completions", s.requireKey(s.completions))
146	mux.HandleFunc("POST /v1/completions", s.requireKey(s.completions))
147	mux.HandleFunc("GET /v1/models", s.requireKey(s.passthrough))
148	mux.HandleFunc("POST /v1/embeddings", s.requireKey(s.passthrough))
149
150	// Unkeyed, and it says nothing but whether this process is up. Asking the
151	// upstream here would wake the weights every thirty seconds and defeat the
152	// idle unload this whole service exists to make possible.
153	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
154		w.Header().Set("Cache-Control", "no-store")
155		fmt.Fprintln(w, "ok")
156	})
157
158	slog.Info("llm listening", "addr", *addr, "upstream", s.upstream, "model", s.model)
159	web.Serve(*addr, web.Chain(mux, web.Recovered, web.Logged, web.SecurityHeaders(csp())))
160}
161
162// prune trims the call log on a timer. The prompts are the whole of what
163// anybody asked this estate, so they age out rather than accumulating until the
164// volume fills.
165func (s *site) prune() {
166	for {
167		if n, err := s.store.Prune(s.retain); err != nil {
168			slog.Error("pruning the call log", "err", err)
169		} else if n > 0 {
170			slog.Info("pruned the call log", "rows", n, "keep", s.retain.String())
171		}
172		time.Sleep(6 * time.Hour)
173	}
174}
175
176func (s *site) loadTemplates() error {
177	// Reading assets off disk means their hashes change while the process is
178	// up, so the reload path drops the cached versions with the templates.
179	if Reloaded {
180		resetAssetVersions()
181	}
182	t, err := template.New("").Funcs(template.FuncMap{
183		"when":  fmtWhen,
184		"short": short,
185		"comma": comma,
186		"asset": assetURL,
187	}).ParseFS(assets(), "templates/*.html")
188	if err != nil {
189		return err
190	}
191	s.tpl = t
192	return nil
193}
194
195func (s *site) static() http.Handler {
196	fs := http.FileServerFS(assets())
197	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
198		if Reloaded {
199			w.Header().Set("Cache-Control", "no-store")
200		} else {
201			w.Header().Set("Cache-Control", "public, max-age=31536000, immutable")
202		}
203		fs.ServeHTTP(w, r)
204	})
205}
206
207func (s *site) render(w http.ResponseWriter, name string, data map[string]any) {
208	if Reloaded {
209		if err := s.loadTemplates(); err != nil {
210			http.Error(w, err.Error(), 500)
211			return
212		}
213	}
214	data["Source"] = sourceURL
215	data["Model"] = s.model
216	w.Header().Set("Content-Type", "text/html; charset=utf-8")
217	w.Header().Set("Cache-Control", "no-store")
218	if err := s.tpl.ExecuteTemplate(w, name, data); err != nil {
219		slog.Error("render", "name", name, "err", err)
220	}
221}
222
223func (s *site) root(w http.ResponseWriter, r *http.Request) {
224	if s.auth.Authenticated(r) {
225		s.overview(w, r)
226		return
227	}
228	s.render(w, "landing.html", map[string]any{"Title": "Overview"})
229}
230
231func (s *site) overview(w http.ResponseWriter, r *http.Request) {
232	keys, err := s.store.Keys()
233	if err != nil {
234		http.Error(w, err.Error(), 500)
235		return
236	}
237	usage, _ := s.store.Usage(time.Now().Add(-24 * time.Hour))
238	recent, _ := s.store.Calls("", 12)
239	nKeys, nCalls := s.store.Counts()
240
241	// A key is shown once, on the redirect straight after it is made, because
242	// nothing here can produce it again.
243	fresh := r.URL.Query().Get("key")
244
245	s.render(w, "overview.html", map[string]any{
246		"Keys": keys, "Usage": usage, "Recent": recent,
247		"Fresh": fresh, "NKeys": nKeys, "NCalls": nCalls, "Title": "Overview",
248		"Upstream": s.upstream, "Retain": s.retain,
249	})
250}
251
252func (s *site) callsPage(w http.ResponseWriter, r *http.Request) {
253	caller := r.URL.Query().Get("caller")
254	limit := 100
255	if n, err := strconv.Atoi(r.URL.Query().Get("limit")); err == nil && n > 0 && n <= 500 {
256		limit = n
257	}
258	calls, err := s.store.Calls(caller, limit)
259	if err != nil {
260		http.Error(w, err.Error(), 500)
261		return
262	}
263	s.render(w, "calls.html", map[string]any{"Calls": calls, "Caller": caller, "Limit": limit, "Title": "Calls"})
264}
265
266func (s *site) keyCreate(w http.ResponseWriter, r *http.Request) {
267	secret, _, err := s.store.NewKey(r.FormValue("name"))
268	if err != nil {
269		http.Error(w, err.Error(), 400)
270		return
271	}
272	// Carried in the query rather than rendered here so a refresh of the
273	// resulting page does not mint a second key.
274	http.Redirect(w, r, "/?key="+secret, http.StatusSeeOther)
275}
276
277func (s *site) keyRevoke(w http.ResponseWriter, r *http.Request) {
278	id, _ := strconv.ParseInt(r.PathValue("id"), 10, 64)
279	if err := s.store.Revoke(id); err != nil {
280		http.Error(w, err.Error(), 500)
281		return
282	}
283	http.Redirect(w, r, "/", http.StatusSeeOther)
284}
285
286func (s *site) keyDelete(w http.ResponseWriter, r *http.Request) {
287	id, _ := strconv.ParseInt(r.PathValue("id"), 10, 64)
288	if err := s.store.DeleteKey(id); err != nil {
289		http.Error(w, err.Error(), 500)
290		return
291	}
292	http.Redirect(w, r, "/", http.StatusSeeOther)
293}
294
295func fmtWhen(t time.Time) string {
296	if t.IsZero() || t.Unix() <= 0 {
297		return "never"
298	}
299	d := time.Since(t)
300	switch {
301	case d < time.Minute:
302		return "just now"
303	case d < time.Hour:
304		return fmt.Sprintf("%dm ago", int(d.Minutes()))
305	case d < 24*time.Hour:
306		return fmt.Sprintf("%dh ago", int(d.Hours()))
307	}
308	return t.Format("2 Jan")
309}
310
311// short renders the messages array as something readable in a table cell. The
312// whole of it is still in the database and on the call's own row.
313func short(s string, n int) string {
314	var msgs []struct {
315		Role    string `json:"role"`
316		Content string `json:"content"`
317	}
318	text := s
319	if json.Unmarshal([]byte(s), &msgs) == nil && len(msgs) > 0 {
320		parts := make([]string, 0, len(msgs))
321		for _, m := range msgs {
322			parts = append(parts, m.Role+": "+m.Content)
323		}
324		text = strings.Join(parts, " | ")
325	}
326	text = strings.Join(strings.Fields(text), " ")
327	if len(text) <= n {
328		return text
329	}
330	return text[:n] + "..."
331}
332
333func comma(n int) string {
334	s := strconv.Itoa(n)
335	if len(s) <= 3 {
336		return s
337	}
338	var out []byte
339	for i, c := range []byte(s) {
340		if i > 0 && (len(s)-i)%3 == 0 {
341			out = append(out, ',')
342		}
343		out = append(out, c)
344	}
345	return string(out)
346}