+// 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{},
+ }
+}
+
+// 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()
+}
+
+// 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)
+ }
+
+ response, err = modelReadStream(httpResponse.Body)
+ if err != nil {
+ return response, fmt.Errorf("read model response stream: %w", err)
+ }
+
+ return response, nil
+}
+
+// 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
+}
+
projects / onechat / blobdiff