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