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