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