Check test assertions with testifylint

[onechat] / model.go

package main

import (
        "bufio"
        "bytes"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "net/http"
        "os"
        "path/filepath"
        "strings"
        "sync"
        "sync/atomic"
        "time"
)

const (
        // modelChatGPTURL is the Responses endpoint of the ChatGPT backend.
        modelChatGPTURL = "https://chatgpt.com/backend-api/codex/responses"

        // modelAPIURL is the Responses endpoint of the OpenAI API.
        modelAPIURL = "https://api.openai.com/v1/responses"

        // modelQuotaPrefix starts the names of the quota headers of the ChatGPT
        // backend, in canonical form.
        modelQuotaPrefix = "X-Codex-"

        // modelTurnStateHeader is an x-codex-* header that holds an opaque
        // state of the backend, not quota.
        modelTurnStateHeader = "X-Codex-Turn-State"

        // modelErrorBytes is the most bytes of an error response that go into
        // the error.
        modelErrorBytes = 4096
)

// errModelCallLimit tells that a client sent its limit of calls.
var errModelCallLimit = errors.New("call limit reached")

// modelEndpoint returns the URL and the authorizer of the endpoint in config.
// The ChatGPT backend uses the token file in dir. The OpenAI API uses the key in
// OPENAI_API_KEY, and there is no fallback from one to the other. All clients
// share the authorizer, so one refresh runs at a time.
func modelEndpoint(config Config, dir string) (string, ModelAuthorizer, error) {
        if config.Endpoint == configEndpointAPI {
                key := os.Getenv("OPENAI_API_KEY")
                if key == "" {
                        return "", nil, errors.New("endpoint api needs OPENAI_API_KEY")
                }

                return modelAPIURL, ModelAPIKey(key), nil
        }

        return modelChatGPTURL, newChatGPTAuth(filepath.Join(dir, chatgptTokenFileName), chatgptIssuer), nil
}

// newModelClient returns a client for the Responses endpoint at url. It records
// each call in the usage file at usagePath.
func newModelClient(url string, auth ModelAuthorizer, cacheKey, usagePath string) *ModelClient {
        return &ModelClient{
                url:      url,
                auth:     auth,
                cacheKey: cacheKey,
                usage: &ModelUsageLog{
                        path: usagePath,
                },
                http: &http.Client{},
        }
}

// modelUserMessage returns a user input item with one text part for each text.
func modelUserMessage(texts ...string) any {
        type part struct {
                Type string `json:"type"`
                Text string `json:"text"`
        }

        content := make([]part, len(texts))
        for n, text := range texts {
                content[n] = part{
                        Type: "input_text",
                        Text: text,
                }
        }

        return struct {
                Role    string `json:"role"`
                Content []part `json:"content"`
        }{
                Role:    "user",
                Content: content,
        }
}

// modelOutputText returns the text of the message items of output.
func modelOutputText(output []json.RawMessage) (string, error) {
        var text strings.Builder
        for _, raw := range output {
                var item ModelOutputItem
                err := json.Unmarshal(raw, &item)
                if err != nil {
                        return "", fmt.Errorf("decode output item: %w", err)
                }

                if item.Type != "message" {
                        continue
                }

                for _, part := range item.Content {
                        if part.Type == "output_text" {
                                text.WriteString(part.Text)
                        }
                }
        }

        return text.String(), nil
}

// ModelAuthorizer adds the credentials of a model call to the headers of its
// request.
type ModelAuthorizer interface {
        Authorize(ctx context.Context, header http.Header) error
}

// ModelAPIKey authorizes calls to the OpenAI API with a key.
type ModelAPIKey string

// Authorize sets the key as the bearer token.
func (k ModelAPIKey) Authorize(ctx context.Context, header http.Header) error {
        header.Set("Authorization", "Bearer "+string(k))
        return ctx.Err()
}

// ModelClient sends model calls to a Responses endpoint for one agent. The
// ChatGPT backend and the OpenAI API take the same request body and differ only
// in the URL and the credentials. Every call sends the fixed cache key of the
// agent, so the turns and the compactions of the agent share one prompt cache.
// A call of a subagent sends the key of the agent with "-" and the name of the
// subagent, so that the calls of a subagent do not share a cache entry with the
// calls of the agent that run at the same time. The ChatGPT backend caches only
// with the key in the session_id header, so a call to it sends the key there
// and in the body. Every call is recorded in the usage file.
//
// The request never holds max_output_tokens or metadata, because the ChatGPT
// backend refuses them. It sets store to false, so the client sends the whole
// input on each call, and it asks for the encrypted reasoning, so a later call
// can send the reasoning items back.
type ModelClient struct {
        url      string
        auth     ModelAuthorizer
        cacheKey string
        usage    *ModelUsageLog
        http     *http.Client
        limit    int64
        calls    atomic.Int64
}

// LimitCalls makes the client refuse each call after the first limit calls.
func (c *ModelClient) LimitCalls(limit int64) {
        c.limit = limit
}

// LimitReached tells if the client refuses further calls.
func (c *ModelClient) LimitReached() bool {
        return c.limit > 0 && c.calls.Load() >= c.limit
}

// CallLimit returns the configured call limit, or zero for no limit.
func (c *ModelClient) CallLimit() int64 {
        return c.limit
}

// UsagePath returns the file that records the calls of this client.
func (c *ModelClient) UsagePath() string {
        return c.usage.path
}

// Create sends the call and reads its event stream to the end. Over the call
// limit it sends nothing and returns errModelCallLimit.
func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelResponse, error) {
        if c.calls.Add(1) > c.limit && c.limit > 0 {
                return ModelResponse{}, errModelCallLimit
        }

        cacheKey := c.cacheKey
        if request.Agent != "" {
                cacheKey += "-" + request.Agent
        }

        body, err := json.Marshal(ModelRequestBody{
                Model:        request.Model,
                Instructions: request.Instructions,
                Input:        request.Input,
                Tools:        request.Tools,
                Reasoning: ModelReasoning{
                        Effort: request.Effort,
                },
                Store:          false,
                Stream:         true,
                Include:        []string{"reasoning.encrypted_content"},
                PromptCacheKey: cacheKey,
        })
        if err != nil {
                return ModelResponse{}, fmt.Errorf("encode request: %w", err)
        }

        httpRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
        if err != nil {
                return ModelResponse{}, fmt.Errorf("create request: %w", err)
        }

        httpRequest.Header.Set("Content-Type", "application/json")
        httpRequest.Header.Set("Accept", "text/event-stream")
        err = c.auth.Authorize(ctx, httpRequest.Header)
        if err != nil {
                return ModelResponse{}, fmt.Errorf("authorize model call: %w", err)
        }

        _, chatgpt := c.auth.(*ChatGPTAuth)
        if chatgpt {
                httpRequest.Header.Set("session_id", cacheKey)
        }

        start := time.Now()
        record := ModelUsageRecord{
                Date:  start.Format(time.RFC3339Nano),
                Model: request.Model,
        }

        if request.Node != nil {
                record.Node = &[2]int{request.Node.level, request.Node.index}
        }

        record.Turn = request.Turn
        record.Agent = request.Agent
        response, err := c.send(httpRequest, &record)
        if err != nil {
                record.Error = err.Error()
        }

        record.Seconds = time.Since(start).Seconds()
        record.InputTokens = response.Usage.Input
        record.CachedTokens = response.Usage.Cached
        record.OutputTokens = response.Usage.Output
        record.TextBytes = len(response.Text)

        return response, errors.Join(err, c.usage.Write(record))
}

// send makes the HTTP call and reads its stream. It puts the status and the
// quota headers into record.
func (c *ModelClient) send(httpRequest *http.Request, record *ModelUsageRecord) (response ModelResponse, err error) {
        httpResponse, err := c.http.Do(httpRequest)
        if err != nil {
                return ModelResponse{}, fmt.Errorf("send model call: %w", err)
        }

        defer func() {
                err = errors.Join(err, httpResponse.Body.Close())
        }()

        record.Status = httpResponse.StatusCode
        record.Quota = modelQuotaHeaders(httpResponse.Header)
        if httpResponse.StatusCode != http.StatusOK {
                text, readErr := io.ReadAll(io.LimitReader(httpResponse.Body, modelErrorBytes))
                err = fmt.Errorf("model call: HTTP %d: %s", httpResponse.StatusCode, strings.TrimSpace(string(text)))

                return ModelResponse{}, errors.Join(err, readErr)
        }

        return modelReadStream(httpResponse.Body)
}

// ModelUsageLog appends one JSON line for each model call to usage.jsonl. Calls
// can run at once, so a mutex serializes the writes.
type ModelUsageLog struct {
        mutex sync.Mutex
        path  string
}

// Write appends record to the file with one write.
func (u *ModelUsageLog) Write(record ModelUsageRecord) (err error) {
        line, err := json.Marshal(record)
        if err != nil {
                return fmt.Errorf("encode usage: %w", err)
        }

        u.mutex.Lock()
        defer u.mutex.Unlock()

        file, err := os.OpenFile(u.path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
        if err != nil {
                return fmt.Errorf("open usage file: %w", err)
        }

        defer func() {
                err = errors.Join(err, file.Close())
        }()

        _, err = file.Write(append(line, '\n'))
        if err != nil {
                return fmt.Errorf("write usage: %w", err)
        }

        return nil
}

// ModelOutputItem holds the fields of an output item that give its text.
type ModelOutputItem struct {
        Type    string `json:"type"`
        Content []struct {
                Type string `json:"type"`
                Text string `json:"text"`
        } `json:"content"`
}

// ModelRequest is one model call. Input holds the input items in order, such as
// modelUserMessage values or output items of an earlier response. Node names
// the node that a compaction builds, or is nil. Turn names the request of a
// turn: first or next. Agent names the subagent that makes the call, or is
// empty for a call of the agent itself. Node, Turn, and Agent go to the usage
// file, never to the endpoint as fields. Agent also gives the call the cache
// key of its subagent.
type ModelRequest struct {
        Model        string
        Effort       string
        Instructions string
        Tools        []ModelTool
        Input        []any
        Node         *NodeKey
        Turn         string
        Agent        string
}

// ModelResponse is the result of a model call. Output holds the output items as
// the endpoint sent them. Text is the text of the message items.
type ModelResponse struct {
        Output []json.RawMessage
        Text   string
        Usage  ModelUsage
}

// ModelRequestBody is the JSON body of a call.
type ModelRequestBody struct {
        Model          string         `json:"model"`
        Instructions   string         `json:"instructions"`
        Input          []any          `json:"input"`
        Tools          []ModelTool    `json:"tools"`
        Reasoning      ModelReasoning `json:"reasoning"`
        Store          bool           `json:"store"`
        Stream         bool           `json:"stream"`
        Include        []string       `json:"include"`
        PromptCacheKey string         `json:"prompt_cache_key"`
}

type ModelReasoning struct {
        Effort string `json:"effort"`
}

// ModelUsageRecord is one line of usage.jsonl. Date is the start of the call,
// and Seconds its length. Node is the node [l, i] of a compaction. Status is
// the HTTP status, or 0 when no response came. TextBytes is the size of the
// text of the reply. Turn names the request of a turn. Agent names the subagent
// of the call. Quota holds the x-codex-* headers of the response, without the
// turn state.
type ModelUsageRecord struct {
        Date         string            `json:"date"`
        Seconds      float64           `json:"seconds"`
        Model        string            `json:"model"`
        Node         *[2]int           `json:"node,omitempty"`
        Turn         string            `json:"turn,omitempty"`
        Agent        string            `json:"agent,omitempty"`
        Status       int               `json:"status"`
        InputTokens  int               `json:"input_tokens"`
        CachedTokens int               `json:"cached_tokens"`
        OutputTokens int               `json:"output_tokens"`
        TextBytes    int               `json:"text_bytes"`
        Quota        map[string]string `json:"quota,omitempty"`
        Error        string            `json:"error,omitempty"`
}

// ModelTool is the definition of a function tool.
type ModelTool struct {
        Type        string `json:"type"`
        Name        string `json:"name"`
        Description string `json:"description"`
        Parameters  any    `json:"parameters"`
        Strict      bool   `json:"strict"`
}

// ModelUsage gives the token counts of a call. Cached is the part of the input
// tokens that came from the prompt cache.
type ModelUsage struct {
        Input  int
        Cached int
        Output int
}

// ModelStreamEvent holds the fields of the stream events that the client reads.
type ModelStreamEvent struct {
        Type     string              `json:"type"`
        Item     json.RawMessage     `json:"item"`
        Response *ModelEventResponse `json:"response"`
        Code     string              `json:"code"`
        Message  string              `json:"message"`
}

type ModelEventResponse struct {
        Output            []json.RawMessage `json:"output"`
        Usage             *ModelWireUsage   `json:"usage"`
        Error             *ModelWireError   `json:"error"`
        IncompleteDetails *struct {
                Reason string `json:"reason"`
        } `json:"incomplete_details"`
}

type ModelWireUsage struct {
        InputTokens        int `json:"input_tokens"`
        OutputTokens       int `json:"output_tokens"`
        InputTokensDetails struct {
                CachedTokens int `json:"cached_tokens"`
        } `json:"input_tokens_details"`
}

type ModelWireError struct {
        Code    string `json:"code"`
        Message string `json:"message"`
}

// modelQuotaHeaders returns the x-codex-* headers, with lowercase names,
// without the turn state.
func modelQuotaHeaders(header http.Header) map[string]string {
        var quota map[string]string
        for name, values := range header {
                if !strings.HasPrefix(name, modelQuotaPrefix) || name == modelTurnStateHeader || len(values) == 0 {
                        continue
                }

                if quota == nil {
                        quota = make(map[string]string)
                }

                quota[strings.ToLower(name)] = values[0]
        }

        return quota
}

// modelReadStream reads the server-sent events of a call until the response
// completes. The output items come from the output_item.done events, or from
// the completed response when the stream sent none.
func modelReadStream(stream io.Reader) (ModelResponse, error) {
        var response ModelResponse
        reader := bufio.NewReader(stream)
        var data bytes.Buffer
        for {
                line, err := reader.ReadBytes('\n')
                if err != nil && !errors.Is(err, io.EOF) {
                        return response, fmt.Errorf("read event stream: %w", err)
                }

                line = bytes.TrimRight(line, "\r\n")
                if len(line) > 0 {
                        value, found := bytes.CutPrefix(line, []byte("data:"))
                        if found {
                                data.Write(bytes.TrimPrefix(value, []byte(" ")))
                                data.WriteByte('\n')
                        }
                } else if data.Len() > 0 {
                        done, eventErr := modelReadEvent(&response, data.Bytes())
                        data.Reset()
                        if eventErr != nil || done {
                                return response, eventErr
                        }
                }

                if errors.Is(err, io.EOF) {
                        return response, errors.New("event stream ended before the response completed")
                }
        }
}

// modelReadEvent adds one stream event to response, and tells if it ends the
// stream.
func modelReadEvent(response *ModelResponse, data []byte) (bool, error) {
        var event ModelStreamEvent
        err := json.Unmarshal(data, &event)
        if err != nil {
                return true, fmt.Errorf("decode stream event: %w", err)
        }

        switch event.Type {
        case "response.output_item.done":
                response.Output = append(response.Output, event.Item)
        case "response.completed":
                if event.Response == nil {
                        return true, errors.New("completed event without a response")
                }

                if len(response.Output) == 0 {
                        response.Output = event.Response.Output
                }

                if event.Response.Usage != nil {
                        response.Usage = ModelUsage{
                                Input:  event.Response.Usage.InputTokens,
                                Cached: event.Response.Usage.InputTokensDetails.CachedTokens,
                                Output: event.Response.Usage.OutputTokens,
                        }
                }

                response.Text, err = modelOutputText(response.Output)

                return true, err
        case "response.incomplete":
                reason := "unknown"
                if event.Response != nil && event.Response.IncompleteDetails != nil {
                        reason = event.Response.IncompleteDetails.Reason
                }

                return true, fmt.Errorf("response incomplete: %s", reason)
        case "response.failed":
                message := "no details"
                if event.Response != nil && event.Response.Error != nil {
                        message = event.Response.Error.Code + ": " + event.Response.Error.Message
                }

                return true, fmt.Errorf("response failed: %s", message)
        case "error":
                return true, fmt.Errorf("stream error: %s: %s", event.Code, event.Message)
        }

        return false, nil
}