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