orchard

  1// The model client. llama.cpp behind llama-swap speaks the OpenAI chat
  2// completions shape, so this is that shape and nothing more.
  3package main
  4
  5import (
  6	"bufio"
  7	"bytes"
  8	"context"
  9	"encoding/json"
 10	"fmt"
 11	"io"
 12	"net/http"
 13	"strings"
 14	"time"
 15)
 16
 17type Role string
 18
 19const (
 20	RoleSystem    Role = "system"
 21	RoleUser      Role = "user"
 22	RoleAssistant Role = "assistant"
 23	RoleTool      Role = "tool"
 24)
 25
 26// ToolCall is one request from the model to run something.
 27type ToolCall struct {
 28	ID       string `json:"id,omitempty"`
 29	Type     string `json:"type,omitempty"`
 30	Function struct {
 31		Name      string `json:"name"`
 32		Arguments string `json:"arguments"`
 33	} `json:"function"`
 34}
 35
 36// Message is one entry in the conversation as the model sees it.
 37type Message struct {
 38	Role       Role       `json:"role"`
 39	Content    string     `json:"content"`
 40	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
 41	ToolCallID string     `json:"tool_call_id,omitempty"`
 42	Name       string     `json:"name,omitempty"`
 43}
 44
 45type LLM struct {
 46	BaseURL string
 47	Model   string
 48	Key     string
 49	http    *http.Client
 50}
 51
 52func NewLLM(base, model, key string) *LLM {
 53	if model == "" {
 54		model = "local"
 55	}
 56	return &LLM{
 57		BaseURL: strings.TrimRight(base, "/"),
 58		Model:   model,
 59		Key:     key,
 60		// Generation is slow and a long answer is normal, so this is generous.
 61		// The context deadline on the request is the real limit.
 62		http: &http.Client{Timeout: 10 * time.Minute},
 63	}
 64}
 65
 66type chatReq struct {
 67	Model           string           `json:"model"`
 68	Messages        []Message        `json:"messages"`
 69	Tools           []map[string]any `json:"tools,omitempty"`
 70	ToolChoice      string           `json:"tool_choice,omitempty"`
 71	Temperature     float64          `json:"temperature"`
 72	TopP            float64          `json:"top_p"`
 73	TopK            int              `json:"top_k"`
 74	MaxTokens       int              `json:"max_tokens"`
 75	Stream          bool             `json:"stream,omitempty"`
 76	StreamOptions   map[string]any   `json:"stream_options,omitempty"`
 77	TimingsPerToken bool             `json:"timings_per_token,omitempty"`
 78	Kwargs          map[string]any   `json:"chat_template_kwargs,omitempty"`
 79	ResponseFormat  *responseFormat  `json:"response_format,omitempty"`
 80}
 81
 82type responseFormat struct {
 83	Type       string         `json:"type"`
 84	JSONSchema *schemaWrapper `json:"json_schema,omitempty"`
 85}
 86
 87type schemaWrapper struct {
 88	Name   string          `json:"name"`
 89	Strict bool            `json:"strict"`
 90	Schema json.RawMessage `json:"schema"`
 91}
 92
 93type chatResp struct {
 94	Choices []struct {
 95		FinishReason string `json:"finish_reason"`
 96		Message      struct {
 97			Content   string     `json:"content"`
 98			Reasoning string     `json:"reasoning_content"`
 99			ToolCalls []ToolCall `json:"tool_calls"`
100		} `json:"message"`
101		Delta struct {
102			Content   string     `json:"content"`
103			Reasoning string     `json:"reasoning_content"`
104			ToolCalls []ToolCall `json:"tool_calls"`
105		} `json:"delta"`
106	} `json:"choices"`
107	Usage struct {
108		PromptTokens     int `json:"prompt_tokens"`
109		CompletionTokens int `json:"completion_tokens"`
110	} `json:"usage"`
111	// llama.cpp reports its own timings, which are the honest numbers: they
112	// exclude this client, the tools and the network.
113	Timings struct {
114		PromptN         int     `json:"prompt_n"`
115		PromptPerSecond float64 `json:"prompt_per_second"`
116		PredictedN      int     `json:"predicted_n"`
117		PredictedPerSec float64 `json:"predicted_per_second"`
118	} `json:"timings"`
119	Error any `json:"error"`
120}
121
122// Stats is what one turn cost, for the meter in the bar. Prompt is the whole
123// window the model was handed, so it is the number that says how full the
124// context is, not how long the last message was.
125type Stats struct {
126	Prompt     int     `json:"prompt_tokens"`
127	Completion int     `json:"completion_tokens"`
128	Decode     float64 `json:"decode_tps"`
129	Prefill    float64 `json:"prefill_tps"`
130}
131
132func (s *Stats) merge(o Stats) {
133	if o.Prompt > s.Prompt {
134		s.Prompt = o.Prompt
135	}
136	s.Completion += o.Completion
137	if o.Decode > 0 {
138		s.Decode = o.Decode
139	}
140	if o.Prefill > 0 {
141		s.Prefill = o.Prefill
142	}
143}
144
145// sampling is Qwen's published non-thinking recipe, which the other models here
146// are close enough to. Thinking is off because llama.cpp puts a chain of
147// thought in reasoning_content and leaves content empty, which does not look
148// like an error: a 200, well formed JSON, and nothing to show the user.
149func (l *LLM) base(msgs []Message, maxTok int) chatReq {
150	return chatReq{
151		Model: l.Model, Messages: msgs, Temperature: 0.7, TopP: 0.8, TopK: 20,
152		MaxTokens: maxTok, Kwargs: map[string]any{"enable_thinking": false},
153	}
154}
155
156// Complete asks for one turn, with tools offered. It does not stream, because a
157// turn that is going to call a tool has nothing to show yet.
158func (l *LLM) Complete(ctx context.Context, msgs []Message, schemas []map[string]any, maxTok int) (Message, error) {
159	m, _, err := l.CompleteStats(ctx, msgs, schemas, maxTok)
160	return m, err
161}
162
163// CompleteStats is Complete plus what the call cost.
164func (l *LLM) CompleteStats(ctx context.Context, msgs []Message, schemas []map[string]any, maxTok int) (Message, Stats, error) {
165	return l.complete(ctx, msgs, schemas, "auto", maxTok)
166}
167
168// CompleteRequiringTool is the same call with the model given no say in whether
169// it calls something. It is what a nudge turns into once the turn has already
170// asked politely: a model that has written a deferral will write another one,
171// and asking a third time spends the turn instead of answering it.
172func (l *LLM) CompleteRequiringTool(ctx context.Context, msgs []Message, schemas []map[string]any, maxTok int) (Message, Stats, error) {
173	return l.complete(ctx, msgs, schemas, "required", maxTok)
174}
175
176func (l *LLM) complete(ctx context.Context, msgs []Message, schemas []map[string]any, choice string, maxTok int) (Message, Stats, error) {
177	req := l.base(msgs, maxTok)
178	if len(schemas) > 0 {
179		req.Tools, req.ToolChoice = schemas, choice
180	}
181	var out chatResp
182	if err := l.post(ctx, req, &out); err != nil {
183		return Message{}, Stats{}, err
184	}
185	if len(out.Choices) == 0 {
186		return Message{}, Stats{}, fmt.Errorf("the model returned no choices")
187	}
188	st := statsOf(out)
189	c := out.Choices[0].Message
190	// Empty content beside a full chain of thought means enable_thinking did
191	// not take. Loud is better than a blank bubble.
192	if strings.TrimSpace(c.Content) == "" && len(c.ToolCalls) == 0 && strings.TrimSpace(c.Reasoning) != "" {
193		return Message{}, st, fmt.Errorf("the model returned only reasoning and no answer")
194	}
195	return Message{Role: RoleAssistant, Content: c.Content, ToolCalls: c.ToolCalls}, st, nil
196}
197
198// Structured constrains an answer to a JSON schema, which llama.cpp compiles to
199// a GBNF grammar and samples against, so a field declared as an enum cannot come
200// back as anything else. Temperature is low because these steps are decisions
201// rather than prose.
202func (l *LLM) Structured(ctx context.Context, msgs []Message, maxTok int, schema, out any) (Stats, error) {
203	raw, err := json.Marshal(schema)
204	if err != nil {
205		return Stats{}, err
206	}
207	req := l.base(msgs, maxTok)
208	req.Temperature = 0.2
209	req.ResponseFormat = &responseFormat{Type: "json_schema",
210		JSONSchema: &schemaWrapper{Name: "response", Strict: true, Schema: raw}}
211
212	var resp chatResp
213	if err := l.post(ctx, req, &resp); err != nil {
214		return Stats{}, err
215	}
216	if len(resp.Choices) == 0 {
217		return Stats{}, fmt.Errorf("the model returned no choices")
218	}
219	st := statsOf(resp)
220	text := strings.TrimSpace(resp.Choices[0].Message.Content)
221	// A model that opens with a word before the object is still constrained to
222	// emit one, so find it rather than failing the whole step.
223	if i := strings.Index(text, "{"); i > 0 {
224		text = text[i:]
225	}
226	return st, json.Unmarshal([]byte(text), out)
227}
228
229func statsOf(r chatResp) Stats {
230	s := Stats{Prompt: r.Usage.PromptTokens, Completion: r.Usage.CompletionTokens,
231		Decode: r.Timings.PredictedPerSec, Prefill: r.Timings.PromptPerSecond}
232	if s.Prompt == 0 {
233		s.Prompt = r.Timings.PromptN
234	}
235	if s.Completion == 0 {
236		s.Completion = r.Timings.PredictedN
237	}
238	return s
239}
240
241// Stream asks for the final answer and calls onDelta as text arrives. Tools are
242// deliberately not offered here: by this point the turn is answering.
243func (l *LLM) Stream(ctx context.Context, msgs []Message, maxTok int, onDelta func(string)) (string, Stats, error) {
244	req := l.base(msgs, maxTok)
245	req.Stream = true
246	// llama.cpp only reports its timings on a stream when asked.
247	req.StreamOptions = map[string]any{"include_usage": true}
248	req.TimingsPerToken = true
249	body, err := json.Marshal(req)
250	if err != nil {
251		return "", Stats{}, err
252	}
253	hr, err := http.NewRequestWithContext(ctx, "POST", l.BaseURL+"/v1/chat/completions", bytes.NewReader(body))
254	if err != nil {
255		return "", Stats{}, err
256	}
257	hr.Header.Set("Content-Type", "application/json")
258	l.sign(hr)
259	resp, err := l.http.Do(hr)
260	if err != nil {
261		return "", Stats{}, err
262	}
263	defer resp.Body.Close()
264	if resp.StatusCode >= 400 {
265		return "", Stats{}, modelError(resp.StatusCode, resp.Body)
266	}
267	var st Stats
268	var sb strings.Builder
269	sc := bufio.NewScanner(resp.Body)
270	sc.Buffer(make([]byte, 0, 64<<10), 4<<20)
271	for sc.Scan() {
272		line := strings.TrimSpace(sc.Text())
273		if !strings.HasPrefix(line, "data:") {
274			continue
275		}
276		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
277		if data == "[DONE]" {
278			break
279		}
280		var ev chatResp
281		if json.Unmarshal([]byte(data), &ev) != nil {
282			continue
283		}
284		// The last chunk carries the usage and timings and no choices, which
285		// is why this is read before the choices check rather than after.
286		if got := statsOf(ev); got.Prompt > 0 || got.Decode > 0 || got.Completion > 0 {
287			st.merge(got)
288		}
289		if len(ev.Choices) == 0 {
290			continue
291		}
292		if d := ev.Choices[0].Delta.Content; d != "" {
293			sb.WriteString(d)
294			onDelta(d)
295		}
296	}
297	if err := sc.Err(); err != nil {
298		return sb.String(), st, err
299	}
300	return sb.String(), st, nil
301}
302
303// sign puts the gateway key on a request, and the incognito header when the
304// turn is one. An empty key means talking straight to a llama-swap with no
305// gateway in front, which is what a bare development run is.
306func (l *LLM) sign(r *http.Request) {
307	if l.Key != "" {
308		r.Header.Set("Authorization", "Bearer "+l.Key)
309	}
310	if IsIncognito(r.Context()) {
311		r.Header.Set(incognitoHeader, "1")
312	}
313}
314
315// The gateway writes down every prompt and completion it forwards, so a mode
316// that writes nothing down here has to say so there as well.
317const incognitoHeader = "X-Incognito"
318
319type incognitoKey struct{}
320
321// WithIncognito marks a context as belonging to an incognito turn. It rides the
322// context rather than the client because the client is shared by every turn,
323// and this way the gate, the compaction pass and the tools all inherit it.
324func WithIncognito(ctx context.Context) context.Context {
325	return context.WithValue(ctx, incognitoKey{}, true)
326}
327
328func IsIncognito(ctx context.Context) bool {
329	on, _ := ctx.Value(incognitoKey{}).(bool)
330	return on
331}
332
333func (l *LLM) post(ctx context.Context, req chatReq, out *chatResp) error {
334	body, err := json.Marshal(req)
335	if err != nil {
336		return err
337	}
338	hr, err := http.NewRequestWithContext(ctx, "POST", l.BaseURL+"/v1/chat/completions", bytes.NewReader(body))
339	if err != nil {
340		return err
341	}
342	hr.Header.Set("Content-Type", "application/json")
343	l.sign(hr)
344	resp, err := l.http.Do(hr)
345	if err != nil {
346		return fmt.Errorf("the model is not answering: %w", err)
347	}
348	defer resp.Body.Close()
349	if resp.StatusCode >= 400 {
350		return modelError(resp.StatusCode, resp.Body)
351	}
352	return json.NewDecoder(resp.Body).Decode(out)
353}
354
355// Warm fires a one token completion so llama-swap loads the weights while the
356// user is still typing. Asking /v1/models would not do it, since that answers
357// from config and loads nothing.
358func (l *LLM) Warm(ctx context.Context) {
359	ctx, cancel := context.WithTimeout(ctx, 90*time.Second)
360	defer cancel()
361	var out chatResp
362	_ = l.post(ctx, chatReq{Model: l.Model, MaxTokens: 1,
363		Messages: []Message{{Role: RoleUser, Content: "hi"}},
364		Kwargs:   map[string]any{"enable_thinking": false}}, &out)
365}
366
367// Healthy asks whether the server is up without waking the model, because
368// /v1/models answers from llama-swap's config and loads no weights. Never point
369// this at /health, which would defeat the idle unload.
370func (l *LLM) Healthy(ctx context.Context) bool {
371	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
372	defer cancel()
373	req, err := http.NewRequestWithContext(ctx, "GET", l.BaseURL+"/v1/models", nil)
374	if err != nil {
375		return false
376	}
377	l.sign(req)
378	resp, err := l.http.Do(req)
379	if err != nil {
380		return false
381	}
382	defer resp.Body.Close()
383	return resp.StatusCode == 200
384}
385
386// modelError carries the reason back rather than the number. A bare "the model
387// answered 500" is unactionable, and the body always says whether it was the
388// context, the template or the request.
389func modelError(status int, body io.Reader) error {
390	raw, _ := io.ReadAll(io.LimitReader(body, 32<<10))
391	var e struct {
392		Error struct {
393			Message string `json:"message"`
394		} `json:"error"`
395	}
396	if json.Unmarshal(raw, &e) == nil && strings.TrimSpace(e.Error.Message) != "" {
397		return fmt.Errorf("the model refused this turn: %s", strings.TrimSpace(e.Error.Message))
398	}
399	if txt := strings.TrimSpace(string(raw)); txt != "" {
400		return fmt.Errorf("the model answered %d: %s", status, trimLine(txt, 300))
401	}
402	return fmt.Errorf("the model answered %d with no reason", status)
403}
404
405func trimLine(s string, n int) string {
406	s = strings.Join(strings.Fields(s), " ")
407	if len(s) > n {
408		return s[:n] + "..."
409	}
410	return s
411}