+package main
+
+import (
+ "bufio"
+ "bytes"
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "io"
+ "net/http"
+ "os"
+ "strings"
+ "sync"
+ "time"
+)
+
+const (
+ // quotaPrefix starts the names of the quota headers of the ChatGPT
+ // backend, in canonical form.
+ quotaPrefix = "X-Codex-"
+
+ // errorBytes is the most bytes of an error response that go into the
+ // error.
+ errorBytes = 4096
+)
+
+// authorizer adds the credentials of a model call to the headers of its
+// request.
+type authorizer interface {
+ Authorize(ctx context.Context, header http.Header) error
+}
+
+// apiKey authorizes calls to the OpenAI API with a key.
+type apiKey string
+
+// Client 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. 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 Client struct {
+ url string
+ auth authorizer
+ cacheKey string
+ usage *usageLog
+ http *http.Client
+ clock func() time.Time
+}
+
+// Request is one model call. Input holds the input items in order, such as
+// userMessage values or output items of an earlier response.
+type Request struct {
+ Model string
+ Effort string
+ Instructions string
+ Tools []Tool
+ Input []any
+}
+
+// Tool is the definition of a function tool.
+type Tool struct {
+ Type string `json:"type"`
+ Name string `json:"name"`
+ Description string `json:"description"`
+ Parameters any `json:"parameters"`
+ Strict bool `json:"strict"`
+}
+
+// Response 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 Response struct {
+ Output []json.RawMessage
+ Text string
+ Usage Usage
+}
+
+// Usage gives the token counts of a call. Cached is the part of the input
+// tokens that came from the prompt cache.
+type Usage struct {
+ Input int
+ Cached int
+ Output int
+}
+
+// requestBody is the JSON body of a call.
+type requestBody struct {
+ Model string `json:"model"`
+ Instructions string `json:"instructions"`
+ Input []any `json:"input"`
+ Tools []Tool `json:"tools"`
+ Reasoning reasoning `json:"reasoning"`
+ Store bool `json:"store"`
+ Stream bool `json:"stream"`
+ Include []string `json:"include"`
+ PromptCacheKey string `json:"prompt_cache_key"`
+}
+
+type reasoning struct {
+ Effort string `json:"effort"`
+}
+
+// streamEvent holds the fields of the stream events that the client reads.
+type streamEvent struct {
+ Type string `json:"type"`
+ Item json.RawMessage `json:"item"`
+ Response *eventResponse `json:"response"`
+ Code string `json:"code"`
+ Message string `json:"message"`
+}
+
+type eventResponse struct {
+ Output []json.RawMessage `json:"output"`
+ Usage *wireUsage `json:"usage"`
+ Error *wireError `json:"error"`
+ IncompleteDetails *struct {
+ Reason string `json:"reason"`
+ } `json:"incomplete_details"`
+}
+
+type wireUsage struct {
+ InputTokens int `json:"input_tokens"`
+ OutputTokens int `json:"output_tokens"`
+ InputTokensDetails struct {
+ CachedTokens int `json:"cached_tokens"`
+ } `json:"input_tokens_details"`
+}
+
+type wireError struct {
+ Code string `json:"code"`
+ Message string `json:"message"`
+}
+
+// outputItem holds the fields of an output item that give its text.
+type outputItem struct {
+ Type string `json:"type"`
+ Content []struct {
+ Type string `json:"type"`
+ Text string `json:"text"`
+ } `json:"content"`
+}
+
+// newClient returns a client for the Responses endpoint at url. It records
+// each call in the usage file at usagePath.
+func newClient(url string, auth authorizer, cacheKey, usagePath string) *Client {
+ return &Client{
+ url: url,
+ auth: auth,
+ cacheKey: cacheKey,
+ usage: &usageLog{
+ path: usagePath,
+ },
+ http: &http.Client{},
+ clock: time.Now,
+ }
+}
+
+// Create sends the call and reads its event stream to the end.
+func (c *Client) Create(ctx context.Context, request Request) (Response, error) {
+ body, err := json.Marshal(requestBody{
+ Model: request.Model,
+ Instructions: request.Instructions,
+ Input: request.Input,
+ Tools: request.Tools,
+ Reasoning: reasoning{
+ Effort: request.Effort,
+ },
+ Store: false,
+ Stream: true,
+ Include: []string{"reasoning.encrypted_content"},
+ PromptCacheKey: c.cacheKey,
+ })
+ if err != nil {
+ return Response{}, fmt.Errorf("encode request: %w", err)
+ }
+
+ httpRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
+ if err != nil {
+ return Response{}, 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 Response{}, fmt.Errorf("authorize model call: %w", err)
+ }
+
+ record := usageRecord{
+ Date: c.clock().Format(time.RFC3339Nano),
+ Model: request.Model,
+ }
+
+ response, err := c.send(httpRequest, &record)
+ if err != nil {
+ record.Error = err.Error()
+ }
+
+ record.InputTokens = response.Usage.Input
+ record.CachedTokens = response.Usage.Cached
+ record.OutputTokens = response.Usage.Output
+
+ 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 *Client) send(httpRequest *http.Request, record *usageRecord) (response Response, err error) {
+ httpResponse, err := c.http.Do(httpRequest)
+ if err != nil {
+ return Response{}, fmt.Errorf("send model call: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, httpResponse.Body.Close())
+ }()
+
+ record.Status = httpResponse.StatusCode
+ record.Quota = quotaHeaders(httpResponse.Header)
+ if httpResponse.StatusCode != http.StatusOK {
+ text, readErr := io.ReadAll(io.LimitReader(httpResponse.Body, errorBytes))
+ err = fmt.Errorf("model call: HTTP %d: %s", httpResponse.StatusCode, strings.TrimSpace(string(text)))
+
+ return Response{}, errors.Join(err, readErr)
+ }
+
+ return readStream(httpResponse.Body)
+}
+
+// usageLog appends one JSON line for each model call to usage.jsonl. Calls can
+// run at once, so a mutex serializes the writes.
+type usageLog struct {
+ mutex sync.Mutex
+ path string
+}
+
+// usageRecord is one line of usage.jsonl. Status is the HTTP status, or 0 when
+// no response came. Quota holds the x-codex-* headers of the response.
+type usageRecord struct {
+ Date string `json:"date"`
+ Model string `json:"model"`
+ Status int `json:"status"`
+ InputTokens int `json:"input_tokens"`
+ CachedTokens int `json:"cached_tokens"`
+ OutputTokens int `json:"output_tokens"`
+ Quota map[string]string `json:"quota,omitempty"`
+ Error string `json:"error,omitempty"`
+}
+
+// Write appends record to the file with one write.
+func (u *usageLog) Write(record usageRecord) (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
+}
+
+// Authorize sets the key as the bearer token.
+func (k apiKey) Authorize(ctx context.Context, header http.Header) error {
+ header.Set("Authorization", "Bearer "+string(k))
+ return ctx.Err()
+}
+
+// userMessage returns a user input item with one text part for each text.
+func userMessage(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,
+ }
+}
+
+// quotaHeaders returns the x-codex-* headers, with lowercase names.
+func quotaHeaders(header http.Header) map[string]string {
+ var quota map[string]string
+ for name, values := range header {
+ if !strings.HasPrefix(name, quotaPrefix) || len(values) == 0 {
+ continue
+ }
+
+ if quota == nil {
+ quota = make(map[string]string)
+ }
+
+ quota[strings.ToLower(name)] = values[0]
+ }
+
+ return quota
+}
+
+// readStream 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 readStream(stream io.Reader) (Response, error) {
+ var response Response
+ 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 := readEvent(&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")
+ }
+ }
+}
+
+// readEvent adds one stream event to response, and tells if it ends the
+// stream.
+func readEvent(response *Response, data []byte) (bool, error) {
+ var event streamEvent
+ 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 = Usage{
+ Input: event.Response.Usage.InputTokens,
+ Cached: event.Response.Usage.InputTokensDetails.CachedTokens,
+ Output: event.Response.Usage.OutputTokens,
+ }
+ }
+
+ response.Text, err = outputText(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
+}
+
+// outputText returns the text of the message items of output.
+func outputText(output []json.RawMessage) (string, error) {
+ var text strings.Builder
+ for _, raw := range output {
+ var item outputItem
+ 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
+}
projects / onechat / commitdiff