git.giacomolli.uk - onechat/blob - 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 }