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