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