orchard

  1package main
  2
  3import (
  4	"bytes"
  5	"context"
  6	"encoding/json"
  7	"fmt"
  8	"net/http"
  9	"strings"
 10	"sync"
 11	"time"
 12)
 13
 14// LLM talks to llama-server's OpenAI compatible endpoint.
 15//
 16// The model runs in a separate container with the GPU attached, so this is
 17// always a network call even in development.
 18type LLM struct {
 19	BaseURL string
 20	Model   string
 21	Key     string
 22	client  *http.Client
 23
 24	// served is the model string llama.cpp answers with, which is the real
 25	// repository and quant rather than the "local" alias this asks for. It is
 26	// read off a response instead of configured, so a stamp on a logged answer
 27	// says what actually produced it even when the config changed underneath.
 28	mu     sync.Mutex
 29	served string
 30}
 31
 32func NewLLM(baseURL, key string) *LLM {
 33	return &LLM{
 34		BaseURL: strings.TrimRight(baseURL, "/"),
 35		Model:   "local",
 36		Key:     key,
 37		client:  &http.Client{Timeout: 4 * time.Minute},
 38	}
 39}
 40
 41// sign puts the gateway key on a request, and the incognito header when the
 42// question is one. An empty key means talking straight to a llama-swap with no
 43// gateway in front, which is what a bare development run is.
 44func (l *LLM) sign(r *http.Request) {
 45	if l.Key != "" {
 46		r.Header.Set("Authorization", "Bearer "+l.Key)
 47	}
 48	if IsIncognito(r.Context()) {
 49		r.Header.Set(incognitoHeader, "1")
 50	}
 51}
 52
 53// The gateway writes down every prompt and completion it forwards, so a
 54// question that leaves no history row here has to say so there as well.
 55const incognitoHeader = "X-Incognito"
 56
 57type incognitoKey struct{}
 58
 59// WithIncognito marks a context as belonging to an incognito question. It rides
 60// the context rather than the client because the client is shared by every
 61// question, and every step of the pipeline already carries the context.
 62func WithIncognito(ctx context.Context) context.Context {
 63	return context.WithValue(ctx, incognitoKey{}, true)
 64}
 65
 66func IsIncognito(ctx context.Context) bool {
 67	on, _ := ctx.Value(incognitoKey{}).(bool)
 68	return on
 69}
 70
 71type chatMessage struct {
 72	Role    string `json:"role"`
 73	Content string `json:"content"`
 74}
 75
 76type chatRequest struct {
 77	Model           string          `json:"model"`
 78	Messages        []chatMessage   `json:"messages"`
 79	Temperature     float64         `json:"temperature"`
 80	TopP            float64         `json:"top_p"`
 81	TopK            int             `json:"top_k"`
 82	MinP            float64         `json:"min_p"`
 83	PresencePenalty float64         `json:"presence_penalty"`
 84	MaxTokens       int             `json:"max_tokens"`
 85	Stream          bool            `json:"stream"`
 86	ResponseFormat  *responseFormat `json:"response_format,omitempty"`
 87	TemplateKwargs  map[string]any  `json:"chat_template_kwargs,omitempty"`
 88}
 89
 90// These numbers came off Qwen's model card and were kept when the model behind
 91// this became Ornith 1.5 9B, since they are the same recipe its publisher gives
 92// and the reasoning below is about the shape of each step rather than about any
 93// one model. This pipeline needs two of them rather than one setting for
 94// everything.
 95//
 96// A step handing over a JSON schema wants the likeliest token inside the
 97// grammar, since the schema is doing the deciding and creativity there is only
 98// a way to pick the wrong enum. Synthesis is the one step writing prose, and
 99// there the model card's own non-thinking numbers apply. The presence penalty
100// matters most: Qwen names it as the fix for the model repeating itself, which
101// is exactly the failure this site keeps hitting, a closing paragraph that says
102// the bullet list again in weaker words.
103//
104// Before this everything ran at temperature 0.2 with llama.cpp's defaults for
105// the rest, so prose was sampled almost greedily with nothing discouraging
106// repetition.
107// The min_p on the prose set is the one number here that is about the quant
108// rather than the model. Quantization damages the tail of the distribution
109// first, since that is where the least of the model's confidence lives, and a
110// floor cuts exactly that tail. Qwen say 0.0 because they are describing the
111// full precision weights. Unsloth say 0.01 on a quantized one.
112//
113// The constrained set leaves it at zero and takes no penalty at all. The
114// grammar is already refusing every token that would not parse, and a presence
115// penalty on JSON pushes against the braces and quotes that have to repeat for
116// the output to be valid.
117var (
118	exact = sampling{Temperature: 0.2, TopP: 0.8, TopK: 20}
119	prose = sampling{Temperature: 0.7, TopP: 0.8, TopK: 20, MinP: 0.01, PresencePenalty: 1.5}
120)
121
122type sampling struct {
123	Temperature     float64
124	TopP            float64
125	TopK            int
126	MinP            float64
127	PresencePenalty float64
128}
129
130type responseFormat struct {
131	Type       string         `json:"type"`
132	JSONSchema *schemaWrapper `json:"json_schema,omitempty"`
133}
134
135type schemaWrapper struct {
136	Name   string          `json:"name"`
137	Strict bool            `json:"strict"`
138	Schema json.RawMessage `json:"schema"`
139}
140
141type chatResponse struct {
142	Model   string `json:"model"`
143	Choices []struct {
144		Message struct {
145			Content   string `json:"content"`
146			Reasoning string `json:"reasoning_content"`
147		} `json:"message"`
148	} `json:"choices"`
149	Error *struct {
150		Message string `json:"message"`
151	} `json:"error"`
152}
153
154// Complete runs a free-form completion. Used only for the synthesis step.
155func (l *LLM) Complete(ctx context.Context, system, user string, maxTokens int) (string, error) {
156	return l.call(ctx, system, user, maxTokens, nil, prose)
157}
158
159// Structured constrains the model to a JSON schema. llama.cpp turns the schema
160// into a GBNF grammar and constrains sampling to it, so the model cannot emit a
161// citation ID outside the enum it was given, and cannot emit a refusal either.
162func (l *LLM) Structured(ctx context.Context, system, user string, maxTokens int, schema any, out any) error {
163	raw, err := json.Marshal(schema)
164	if err != nil {
165		return err
166	}
167	format := &responseFormat{
168		Type:       "json_schema",
169		JSONSchema: &schemaWrapper{Name: "response", Strict: true, Schema: raw},
170	}
171	text, err := l.call(ctx, system, user, maxTokens, format, exact)
172	if err != nil {
173		return err
174	}
175	text = strings.TrimSpace(text)
176	if i := strings.Index(text, "{"); i > 0 {
177		text = text[i:]
178	}
179	return json.Unmarshal([]byte(text), out)
180}
181
182func (l *LLM) call(ctx context.Context, system, user string, maxTokens int, format *responseFormat, s sampling) (string, error) {
183	body, err := json.Marshal(chatRequest{
184		Model:           l.Model,
185		Temperature:     s.Temperature,
186		TopP:            s.TopP,
187		TopK:            s.TopK,
188		MinP:            s.MinP,
189		PresencePenalty: s.PresencePenalty,
190		MaxTokens:       maxTokens,
191		ResponseFormat:  format,
192		// The model behind this is a thinking one and llama.cpp puts the chain
193		// of thought in reasoning_content, leaving content empty until the
194		// budget runs out. Every step here is either schema constrained or
195		// wants prose directly, so thinking only burns tokens.
196		TemplateKwargs: map[string]any{"enable_thinking": false},
197		Messages: []chatMessage{
198			{Role: "system", Content: system},
199			{Role: "user", Content: user},
200		},
201	})
202	if err != nil {
203		return "", err
204	}
205	req, err := http.NewRequestWithContext(ctx, "POST", l.BaseURL+"/v1/chat/completions", bytes.NewReader(body))
206	if err != nil {
207		return "", err
208	}
209	req.Header.Set("Content-Type", "application/json")
210	l.sign(req)
211
212	resp, err := l.client.Do(req)
213	if err != nil {
214		return "", err
215	}
216	defer resp.Body.Close()
217
218	var out chatResponse
219	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
220		return "", err
221	}
222	if out.Error != nil {
223		return "", fmt.Errorf("llm: %s", out.Error.Message)
224	}
225	if out.Model != "" {
226		l.mu.Lock()
227		l.served = out.Model
228		l.mu.Unlock()
229	}
230	if len(out.Choices) == 0 {
231		return "", fmt.Errorf("llm: no choices")
232	}
233	msg := out.Choices[0].Message
234	if msg.Content == "" && msg.Reasoning != "" {
235		return "", fmt.Errorf("llm: answered with reasoning only, thinking is not disabled")
236	}
237	return msg.Content, nil
238}
239
240// Warm fires a one token completion so llama-swap loads the model while the
241// search and the fetches are still in flight. The cold start then happens
242// inside time that was already being spent.
243func (l *LLM) Warm(ctx context.Context) {
244	ctx, cancel := context.WithTimeout(ctx, 90*time.Second)
245	defer cancel()
246	l.call(ctx, "", "hi", 1, nil, exact)
247}
248
249// Served is the model that answered last, empty until one has.
250func (l *LLM) Served() string {
251	l.mu.Lock()
252	defer l.mu.Unlock()
253	return l.served
254}
255
256// Healthy asks whether the model server is up, without waking the model.
257//
258// /v1/models answers from llama-swap's config and loads nothing. Asking /health
259// would risk pulling the weights back onto the card every time somebody opens
260// the page, which would quietly defeat the idle unload.
261func (l *LLM) Healthy(ctx context.Context) bool {
262	ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
263	defer cancel()
264	req, err := http.NewRequestWithContext(ctx, "GET", l.BaseURL+"/v1/models", nil)
265	if err != nil {
266		return false
267	}
268	resp, err := l.client.Do(req)
269	if err != nil {
270		return false
271	}
272	defer resp.Body.Close()
273	return resp.StatusCode == http.StatusOK
274}