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