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