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