orchard

  1// The gateway itself: an OpenAI shaped endpoint in front of llama-swap that
  2// checks an API key, forwards the request, and writes down what was asked and
  3// what came back.
  4//
  5// It is a proxy and not a client library on purpose. Every caller here already
  6// speaks the OpenAI chat completions shape, so putting this in the path costs
  7// them a base url and a header rather than a rewrite, and anything that shape
  8// supports keeps working without this file knowing about it.
  9package main
 10
 11import (
 12	"bufio"
 13	"bytes"
 14	"encoding/json"
 15	"fmt"
 16	"io"
 17	"log/slog"
 18	"net/http"
 19	"strings"
 20	"time"
 21)
 22
 23// The whole request body is held so it can be both forwarded and logged, so
 24// this is a real ceiling rather than a formality. A 64k window of text is a
 25// long way under it.
 26const maxRequestBytes = 32 << 20
 27
 28// incognitoHeader is a caller saying this turn is not to be written down. Chat
 29// and search set it when the user is in incognito, and it is taken at its word
 30// because every caller here holds a key that was handed out by hand.
 31const incognitoHeader = "X-Incognito"
 32
 33type upstreamReq struct {
 34	Model    string            `json:"model"`
 35	Stream   bool              `json:"stream"`
 36	Messages []json.RawMessage `json:"messages"`
 37	Tools    []json.RawMessage `json:"tools,omitempty"`
 38}
 39
 40// authenticate resolves the bearer token on a request. Sessions do not work
 41// here: the callers are other containers with no browser and no cookie.
 42func (s *site) authenticate(r *http.Request) (Key, bool) {
 43	h := r.Header.Get("Authorization")
 44	if !strings.HasPrefix(h, "Bearer ") {
 45		// llama.cpp's own clients send the key as an api key header, so both
 46		// spellings are accepted rather than making every caller special.
 47		if v := r.Header.Get("X-Api-Key"); v != "" {
 48			return s.store.Authenticate(strings.TrimSpace(v))
 49		}
 50		return Key{}, false
 51	}
 52	return s.store.Authenticate(strings.TrimSpace(strings.TrimPrefix(h, "Bearer ")))
 53}
 54
 55func (s *site) requireKey(next func(http.ResponseWriter, *http.Request, Key)) http.HandlerFunc {
 56	return func(w http.ResponseWriter, r *http.Request) {
 57		k, ok := s.authenticate(r)
 58		if !ok {
 59			w.Header().Set("Content-Type", "application/json; charset=utf-8")
 60			w.Header().Set("Cache-Control", "no-store")
 61			w.WriteHeader(http.StatusUnauthorized)
 62			_, _ = w.Write([]byte(`{"error":{"message":"a valid api key is required","type":"invalid_request_error"}}`))
 63			return
 64		}
 65		s.store.TouchKey(k.ID)
 66		next(w, r, k)
 67	}
 68}
 69
 70// completions forwards one chat completion and logs it. Both shapes go through
 71// here, since a streamed answer has to be reassembled to be written down and a
 72// caller that streams is exactly the one whose output is worth keeping.
 73func (s *site) completions(w http.ResponseWriter, r *http.Request, k Key) {
 74	started := time.Now()
 75	body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxRequestBytes))
 76	if err != nil {
 77		http.Error(w, "request too large", http.StatusRequestEntityTooLarge)
 78		return
 79	}
 80	var req upstreamReq
 81	_ = json.Unmarshal(body, &req)
 82
 83	// A row with its text blanked would still say who asked something and when,
 84	// so incognito writes no row at all and the prompt is never copied out of
 85	// the body it arrived in.
 86	keep := r.Header.Get(incognitoHeader) != "1"
 87
 88	call := Call{KeyID: k.ID, Caller: k.Name, Model: req.Model}
 89	if keep {
 90		call.Messages = string(mustJSON(req.Messages))
 91		if len(req.Tools) > 0 {
 92			call.Tools = fmt.Sprintf("%d offered", len(req.Tools))
 93		}
 94	}
 95	defer func() {
 96		if !keep {
 97			return
 98		}
 99		call.MS = time.Since(started).Milliseconds()
100		s.store.LogCall(call)
101	}()
102
103	up, err := http.NewRequestWithContext(r.Context(), http.MethodPost, s.upstream+r.URL.Path, bytes.NewReader(body))
104	if err != nil {
105		call.Err = err.Error()
106		http.Error(w, "bad gateway", http.StatusBadGateway)
107		return
108	}
109	up.Header.Set("Content-Type", "application/json")
110
111	resp, err := s.client.Do(up)
112	if err != nil {
113		call.Err = err.Error()
114		call.Status = http.StatusBadGateway
115		slog.Error("upstream refused", "err", err, "caller", k.Name)
116		http.Error(w, "the model server is not answering", http.StatusBadGateway)
117		return
118	}
119	defer resp.Body.Close()
120	call.Status = resp.StatusCode
121
122	for h, vs := range resp.Header {
123		for _, v := range vs {
124			w.Header().Add(h, v)
125		}
126	}
127	w.Header().Set("Cache-Control", "no-store")
128	w.WriteHeader(resp.StatusCode)
129
130	// An upstream refusal carries its reason in the body, and that body is the
131	// one thing worth having when a turn fails. Without this the log says 500
132	// and nothing else, which is what made a template error look like a size
133	// limit for an afternoon.
134	if resp.StatusCode >= 400 {
135		body, _ := io.ReadAll(io.LimitReader(resp.Body, 64<<10))
136		_, _ = w.Write(body)
137		call.Err = upstreamReason(body)
138		slog.Error("upstream refused a call", "status", resp.StatusCode,
139			"caller", k.Name, "reason", call.Err)
140		return
141	}
142
143	if req.Stream {
144		call.Completion, call.PromptTok, call.OutputTok, call.DecodeTPS = s.pipeStream(w, resp.Body)
145		return
146	}
147	out, err := io.ReadAll(io.LimitReader(resp.Body, maxRequestBytes))
148	if err != nil {
149		call.Err = err.Error()
150		return
151	}
152	_, _ = w.Write(out)
153	call.Completion, call.PromptTok, call.OutputTok, call.DecodeTPS = summarise(out)
154}
155
156// pipeStream copies the event stream straight through while reading the deltas
157// out of it. The bytes the caller receives are the bytes upstream sent, since
158// re-serialising them would put this service in the position of having to keep
159// up with a format it only wants to record.
160func (s *site) pipeStream(w http.ResponseWriter, body io.Reader) (string, int, int, float64) {
161	rc := http.NewResponseController(w)
162	var text strings.Builder
163	var promptTok, outTok int
164	var tps float64
165
166	sc := bufio.NewScanner(body)
167	sc.Buffer(make([]byte, 0, 64<<10), 8<<20)
168	for sc.Scan() {
169		line := sc.Text()
170		fmt.Fprintln(w, line)
171		// A blank line ends an event, and flushing there rather than per line
172		// keeps a half written event from reaching the caller.
173		if line == "" {
174			_ = rc.Flush()
175			continue
176		}
177		data, ok := strings.CutPrefix(line, "data: ")
178		if !ok || data == "[DONE]" {
179			continue
180		}
181		var chunk struct {
182			Choices []struct {
183				Delta struct {
184					Content string `json:"content"`
185				} `json:"delta"`
186			} `json:"choices"`
187			Usage struct {
188				PromptTokens     int `json:"prompt_tokens"`
189				CompletionTokens int `json:"completion_tokens"`
190			} `json:"usage"`
191			Timings struct {
192				PredictedPerSec float64 `json:"predicted_per_second"`
193			} `json:"timings"`
194		}
195		if json.Unmarshal([]byte(data), &chunk) != nil {
196			continue
197		}
198		for _, c := range chunk.Choices {
199			text.WriteString(c.Delta.Content)
200		}
201		if chunk.Usage.PromptTokens > 0 {
202			promptTok = chunk.Usage.PromptTokens
203		}
204		if chunk.Usage.CompletionTokens > 0 {
205			outTok = chunk.Usage.CompletionTokens
206		}
207		if chunk.Timings.PredictedPerSec > 0 {
208			tps = chunk.Timings.PredictedPerSec
209		}
210	}
211	_ = rc.Flush()
212	return text.String(), promptTok, outTok, tps
213}
214
215func summarise(out []byte) (string, int, int, float64) {
216	var resp struct {
217		Choices []struct {
218			Message struct {
219				Content   string            `json:"content"`
220				ToolCalls []json.RawMessage `json:"tool_calls"`
221			} `json:"message"`
222		} `json:"choices"`
223		Usage struct {
224			PromptTokens     int `json:"prompt_tokens"`
225			CompletionTokens int `json:"completion_tokens"`
226		} `json:"usage"`
227		Timings struct {
228			PredictedPerSec float64 `json:"predicted_per_second"`
229		} `json:"timings"`
230	}
231	if json.Unmarshal(out, &resp) != nil || len(resp.Choices) == 0 {
232		return "", 0, 0, 0
233	}
234	text := resp.Choices[0].Message.Content
235	// A turn that only called tools has no content, and logging it as empty
236	// loses the thing that actually happened.
237	if text == "" && len(resp.Choices[0].Message.ToolCalls) > 0 {
238		text = "[tool calls: " + string(mustJSON(resp.Choices[0].Message.ToolCalls)) + "]"
239	}
240	return text, resp.Usage.PromptTokens, resp.Usage.CompletionTokens, resp.Timings.PredictedPerSec
241}
242
243// passthrough carries the endpoints that are not a completion, like the model
244// list. They are not logged, because there is no prompt in them and a health
245// check every thirty seconds would bury the log that matters.
246func (s *site) passthrough(w http.ResponseWriter, r *http.Request, _ Key) {
247	up, err := http.NewRequestWithContext(r.Context(), r.Method, s.upstream+r.URL.Path, r.Body)
248	if err != nil {
249		http.Error(w, "bad gateway", http.StatusBadGateway)
250		return
251	}
252	up.Header.Set("Content-Type", r.Header.Get("Content-Type"))
253	resp, err := s.client.Do(up)
254	if err != nil {
255		http.Error(w, "the model server is not answering", http.StatusBadGateway)
256		return
257	}
258	defer resp.Body.Close()
259	for h, vs := range resp.Header {
260		for _, v := range vs {
261			w.Header().Add(h, v)
262		}
263	}
264	w.WriteHeader(resp.StatusCode)
265	_, _ = io.Copy(w, io.LimitReader(resp.Body, maxRequestBytes))
266}
267
268func mustJSON(v any) []byte {
269	b, err := json.Marshal(v)
270	if err != nil {
271		return []byte("[]")
272	}
273	return b
274}
275
276// upstreamReason pulls the message out of llama.cpp's error shape and falls
277// back to the raw body, since a body that does not parse is still the evidence.
278func upstreamReason(body []byte) string {
279	var e struct {
280		Error struct {
281			Message string `json:"message"`
282			Type    string `json:"type"`
283		} `json:"error"`
284	}
285	if json.Unmarshal(body, &e) == nil && e.Error.Message != "" {
286		if e.Error.Type != "" {
287			return e.Error.Type + ": " + e.Error.Message
288		}
289		return e.Error.Message
290	}
291	return strings.TrimSpace(string(body))
292}