]> git.giacomolli.uk - onechat/blob - responses.go
Move shared helpers to files of their own
[onechat] / responses.go
1 package main
2
3 import (
4         "bufio"
5         "bytes"
6         "context"
7         "encoding/json"
8         "errors"
9         "fmt"
10         "io"
11         "net/http"
12         "os"
13         "path/filepath"
14         "strings"
15         "sync"
16         "time"
17 )
18
19 const (
20         // chatgptURL is the Responses endpoint of the ChatGPT backend.
21         chatgptURL = "https://chatgpt.com/backend-api/codex/responses"
22
23         // apiURL is the Responses endpoint of the OpenAI API.
24         apiURL = "https://api.openai.com/v1/responses"
25
26         // quotaPrefix starts the names of the quota headers of the ChatGPT
27         // backend, in canonical form.
28         quotaPrefix = "X-Codex-"
29
30         // errorBytes is the most bytes of an error response that go into the
31         // error.
32         errorBytes = 4096
33 )
34
35 // authorizer adds the credentials of a model call to the headers of its
36 // request.
37 type authorizer interface {
38         Authorize(ctx context.Context, header http.Header) error
39 }
40
41 // apiKey authorizes calls to the OpenAI API with a key.
42 type apiKey string
43
44 // Client sends model calls to a Responses endpoint for one agent. The ChatGPT
45 // backend and the OpenAI API take the same request body and differ only in the
46 // URL and the credentials. Every call sends the fixed cache key of the agent,
47 // so the turns and the compactions of the agent share one prompt cache. Every
48 // call is recorded in the usage file.
49 //
50 // The request never holds max_output_tokens or metadata, because the ChatGPT
51 // backend refuses them. It sets store to false, so the client sends the whole
52 // input on each call, and it asks for the encrypted reasoning, so a later call
53 // can send the reasoning items back.
54 type Client struct {
55         url      string
56         auth     authorizer
57         cacheKey string
58         usage    *usageLog
59         http     *http.Client
60         clock    func() time.Time
61 }
62
63 // Request is one model call. Input holds the input items in order, such as
64 // userMessage values or output items of an earlier response.
65 type Request struct {
66         Model        string
67         Effort       string
68         Instructions string
69         Tools        []Tool
70         Input        []any
71 }
72
73 // Tool is the definition of a function tool.
74 type Tool struct {
75         Type        string `json:"type"`
76         Name        string `json:"name"`
77         Description string `json:"description"`
78         Parameters  any    `json:"parameters"`
79         Strict      bool   `json:"strict"`
80 }
81
82 // Response is the result of a model call. Output holds the output items as the
83 // endpoint sent them. Text is the text of the message items.
84 type Response struct {
85         Output []json.RawMessage
86         Text   string
87         Usage  Usage
88 }
89
90 // Usage gives the token counts of a call. Cached is the part of the input
91 // tokens that came from the prompt cache.
92 type Usage struct {
93         Input  int
94         Cached int
95         Output int
96 }
97
98 // requestBody is the JSON body of a call.
99 type requestBody struct {
100         Model          string    `json:"model"`
101         Instructions   string    `json:"instructions"`
102         Input          []any     `json:"input"`
103         Tools          []Tool    `json:"tools"`
104         Reasoning      reasoning `json:"reasoning"`
105         Store          bool      `json:"store"`
106         Stream         bool      `json:"stream"`
107         Include        []string  `json:"include"`
108         PromptCacheKey string    `json:"prompt_cache_key"`
109 }
110
111 type reasoning struct {
112         Effort string `json:"effort"`
113 }
114
115 // streamEvent holds the fields of the stream events that the client reads.
116 type streamEvent struct {
117         Type     string          `json:"type"`
118         Item     json.RawMessage `json:"item"`
119         Response *eventResponse  `json:"response"`
120         Code     string          `json:"code"`
121         Message  string          `json:"message"`
122 }
123
124 type eventResponse struct {
125         Output            []json.RawMessage `json:"output"`
126         Usage             *wireUsage        `json:"usage"`
127         Error             *wireError        `json:"error"`
128         IncompleteDetails *struct {
129                 Reason string `json:"reason"`
130         } `json:"incomplete_details"`
131 }
132
133 type wireUsage struct {
134         InputTokens        int `json:"input_tokens"`
135         OutputTokens       int `json:"output_tokens"`
136         InputTokensDetails struct {
137                 CachedTokens int `json:"cached_tokens"`
138         } `json:"input_tokens_details"`
139 }
140
141 type wireError struct {
142         Code    string `json:"code"`
143         Message string `json:"message"`
144 }
145
146 // outputItem holds the fields of an output item that give its text.
147 type outputItem struct {
148         Type    string `json:"type"`
149         Content []struct {
150                 Type string `json:"type"`
151                 Text string `json:"text"`
152         } `json:"content"`
153 }
154
155 // modelEndpoint returns the URL and the authorizer of the endpoint in config.
156 // The ChatGPT backend uses the token file in dir. The OpenAI API uses the key in
157 // OPENAI_API_KEY, and there is no fallback from one to the other. All clients
158 // share the authorizer, so one refresh runs at a time.
159 func modelEndpoint(config Config, dir string) (string, authorizer, error) {
160         if config.Endpoint == endpointAPI {
161                 key := os.Getenv("OPENAI_API_KEY")
162                 if key == "" {
163                         return "", nil, errors.New("endpoint api needs OPENAI_API_KEY")
164                 }
165
166                 return apiURL, apiKey(key), nil
167         }
168
169         return chatgptURL, newChatGPTAuth(filepath.Join(dir, tokenFileName), authIssuer), nil
170 }
171
172 // newClient returns a client for the Responses endpoint at url. It records
173 // each call in the usage file at usagePath.
174 func newClient(url string, auth authorizer, cacheKey, usagePath string) *Client {
175         return &Client{
176                 url:      url,
177                 auth:     auth,
178                 cacheKey: cacheKey,
179                 usage: &usageLog{
180                         path: usagePath,
181                 },
182                 http:  &http.Client{},
183                 clock: time.Now,
184         }
185 }
186
187 // Create sends the call and reads its event stream to the end.
188 func (c *Client) Create(ctx context.Context, request Request) (Response, error) {
189         body, err := json.Marshal(requestBody{
190                 Model:        request.Model,
191                 Instructions: request.Instructions,
192                 Input:        request.Input,
193                 Tools:        request.Tools,
194                 Reasoning: reasoning{
195                         Effort: request.Effort,
196                 },
197                 Store:          false,
198                 Stream:         true,
199                 Include:        []string{"reasoning.encrypted_content"},
200                 PromptCacheKey: c.cacheKey,
201         })
202         if err != nil {
203                 return Response{}, fmt.Errorf("encode request: %w", err)
204         }
205
206         httpRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
207         if err != nil {
208                 return Response{}, fmt.Errorf("create request: %w", err)
209         }
210
211         httpRequest.Header.Set("Content-Type", "application/json")
212         httpRequest.Header.Set("Accept", "text/event-stream")
213         err = c.auth.Authorize(ctx, httpRequest.Header)
214         if err != nil {
215                 return Response{}, fmt.Errorf("authorize model call: %w", err)
216         }
217
218         record := usageRecord{
219                 Date:  c.clock().Format(time.RFC3339Nano),
220                 Model: request.Model,
221         }
222
223         response, err := c.send(httpRequest, &record)
224         if err != nil {
225                 record.Error = err.Error()
226         }
227
228         record.InputTokens = response.Usage.Input
229         record.CachedTokens = response.Usage.Cached
230         record.OutputTokens = response.Usage.Output
231
232         return response, errors.Join(err, c.usage.Write(record))
233 }
234
235 // send makes the HTTP call and reads its stream. It puts the status and the
236 // quota headers into record.
237 func (c *Client) send(httpRequest *http.Request, record *usageRecord) (response Response, err error) {
238         httpResponse, err := c.http.Do(httpRequest)
239         if err != nil {
240                 return Response{}, fmt.Errorf("send model call: %w", err)
241         }
242
243         defer func() {
244                 err = errors.Join(err, httpResponse.Body.Close())
245         }()
246
247         record.Status = httpResponse.StatusCode
248         record.Quota = quotaHeaders(httpResponse.Header)
249         if httpResponse.StatusCode != http.StatusOK {
250                 text, readErr := io.ReadAll(io.LimitReader(httpResponse.Body, errorBytes))
251                 err = fmt.Errorf("model call: HTTP %d: %s", httpResponse.StatusCode, strings.TrimSpace(string(text)))
252
253                 return Response{}, errors.Join(err, readErr)
254         }
255
256         return readStream(httpResponse.Body)
257 }
258
259 // usageLog appends one JSON line for each model call to usage.jsonl. Calls can
260 // run at once, so a mutex serializes the writes.
261 type usageLog struct {
262         mutex sync.Mutex
263         path  string
264 }
265
266 // usageRecord is one line of usage.jsonl. Status is the HTTP status, or 0 when
267 // no response came. Quota holds the x-codex-* headers of the response.
268 type usageRecord struct {
269         Date         string            `json:"date"`
270         Model        string            `json:"model"`
271         Status       int               `json:"status"`
272         InputTokens  int               `json:"input_tokens"`
273         CachedTokens int               `json:"cached_tokens"`
274         OutputTokens int               `json:"output_tokens"`
275         Quota        map[string]string `json:"quota,omitempty"`
276         Error        string            `json:"error,omitempty"`
277 }
278
279 // Write appends record to the file with one write.
280 func (u *usageLog) Write(record usageRecord) (err error) {
281         line, err := json.Marshal(record)
282         if err != nil {
283                 return fmt.Errorf("encode usage: %w", err)
284         }
285
286         u.mutex.Lock()
287         defer u.mutex.Unlock()
288
289         file, err := os.OpenFile(u.path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
290         if err != nil {
291                 return fmt.Errorf("open usage file: %w", err)
292         }
293
294         defer func() {
295                 err = errors.Join(err, file.Close())
296         }()
297
298         _, err = file.Write(append(line, '\n'))
299         if err != nil {
300                 return fmt.Errorf("write usage: %w", err)
301         }
302
303         return nil
304 }
305
306 // Authorize sets the key as the bearer token.
307 func (k apiKey) Authorize(ctx context.Context, header http.Header) error {
308         header.Set("Authorization", "Bearer "+string(k))
309         return ctx.Err()
310 }
311
312 // userMessage returns a user input item with one text part for each text.
313 func userMessage(texts ...string) any {
314         type part struct {
315                 Type string `json:"type"`
316                 Text string `json:"text"`
317         }
318
319         content := make([]part, len(texts))
320         for n, text := range texts {
321                 content[n] = part{
322                         Type: "input_text",
323                         Text: text,
324                 }
325         }
326
327         return struct {
328                 Role    string `json:"role"`
329                 Content []part `json:"content"`
330         }{
331                 Role:    "user",
332                 Content: content,
333         }
334 }
335
336 // quotaHeaders returns the x-codex-* headers, with lowercase names.
337 func quotaHeaders(header http.Header) map[string]string {
338         var quota map[string]string
339         for name, values := range header {
340                 if !strings.HasPrefix(name, quotaPrefix) || len(values) == 0 {
341                         continue
342                 }
343
344                 if quota == nil {
345                         quota = make(map[string]string)
346                 }
347
348                 quota[strings.ToLower(name)] = values[0]
349         }
350
351         return quota
352 }
353
354 // readStream reads the server-sent events of a call until the response
355 // completes. The output items come from the output_item.done events, or from
356 // the completed response when the stream sent none.
357 func readStream(stream io.Reader) (Response, error) {
358         var response Response
359         reader := bufio.NewReader(stream)
360         var data bytes.Buffer
361         for {
362                 line, err := reader.ReadBytes('\n')
363                 if err != nil && !errors.Is(err, io.EOF) {
364                         return response, fmt.Errorf("read event stream: %w", err)
365                 }
366
367                 line = bytes.TrimRight(line, "\r\n")
368                 if len(line) > 0 {
369                         value, found := bytes.CutPrefix(line, []byte("data:"))
370                         if found {
371                                 data.Write(bytes.TrimPrefix(value, []byte(" ")))
372                                 data.WriteByte('\n')
373                         }
374                 } else if data.Len() > 0 {
375                         done, eventErr := readEvent(&response, data.Bytes())
376                         data.Reset()
377                         if eventErr != nil || done {
378                                 return response, eventErr
379                         }
380                 }
381
382                 if errors.Is(err, io.EOF) {
383                         return response, errors.New("event stream ended before the response completed")
384                 }
385         }
386 }
387
388 // readEvent adds one stream event to response, and tells if it ends the
389 // stream.
390 func readEvent(response *Response, data []byte) (bool, error) {
391         var event streamEvent
392         err := json.Unmarshal(data, &event)
393         if err != nil {
394                 return true, fmt.Errorf("decode stream event: %w", err)
395         }
396
397         switch event.Type {
398         case "response.output_item.done":
399                 response.Output = append(response.Output, event.Item)
400         case "response.completed":
401                 if event.Response == nil {
402                         return true, errors.New("completed event without a response")
403                 }
404
405                 if len(response.Output) == 0 {
406                         response.Output = event.Response.Output
407                 }
408
409                 if event.Response.Usage != nil {
410                         response.Usage = Usage{
411                                 Input:  event.Response.Usage.InputTokens,
412                                 Cached: event.Response.Usage.InputTokensDetails.CachedTokens,
413                                 Output: event.Response.Usage.OutputTokens,
414                         }
415                 }
416
417                 response.Text, err = outputText(response.Output)
418
419                 return true, err
420         case "response.incomplete":
421                 reason := "unknown"
422                 if event.Response != nil && event.Response.IncompleteDetails != nil {
423                         reason = event.Response.IncompleteDetails.Reason
424                 }
425
426                 return true, fmt.Errorf("response incomplete: %s", reason)
427         case "response.failed":
428                 message := "no details"
429                 if event.Response != nil && event.Response.Error != nil {
430                         message = event.Response.Error.Code + ": " + event.Response.Error.Message
431                 }
432
433                 return true, fmt.Errorf("response failed: %s", message)
434         case "error":
435                 return true, fmt.Errorf("stream error: %s: %s", event.Code, event.Message)
436         }
437
438         return false, nil
439 }
440
441 // outputText returns the text of the message items of output.
442 func outputText(output []json.RawMessage) (string, error) {
443         var text strings.Builder
444         for _, raw := range output {
445                 var item outputItem
446                 err := json.Unmarshal(raw, &item)
447                 if err != nil {
448                         return "", fmt.Errorf("decode output item: %w", err)
449                 }
450
451                 if item.Type != "message" {
452                         continue
453                 }
454
455                 for _, part := range item.Content {
456                         if part.Type == "output_text" {
457                                 text.WriteString(part.Text)
458                         }
459                 }
460         }
461
462         return text.String(), nil
463 }