X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/02cdaff8be27f37f44db8892bb12554f2d702323..0e57dcfc96cd95268d5dd00af7ce0594721d3ded:/model.go?ds=sidebyside diff --git a/model.go b/model.go index bd1258b..507fbb1 100644 --- a/model.go +++ b/model.go @@ -45,7 +45,7 @@ var errModelCallLimit = errors.New("call limit reached") // 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 == endpointAPI { + if config.Endpoint == configEndpointAPI { key := os.Getenv("OPENAI_API_KEY") if key == "" { return "", nil, errors.New("endpoint api needs OPENAI_API_KEY") @@ -54,7 +54,7 @@ func modelEndpoint(config Config, dir string) (string, ModelAuthorizer, error) { return modelAPIURL, ModelAPIKey(key), nil } - return modelChatGPTURL, newChatGPTAuth(filepath.Join(dir, tokenFileName), authIssuer), nil + return modelChatGPTURL, newChatGPTAuth(filepath.Join(dir, chatgptTokenFileName), chatgptIssuer), nil } // newModelClient returns a client for the Responses endpoint at url. It records @@ -98,8 +98,10 @@ func modelUserMessage(texts ...string) any { // 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) @@ -131,6 +133,7 @@ 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() } @@ -138,9 +141,11 @@ func (k ModelAPIKey) Authorize(ctx context.Context, header http.Header) error { // 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. -// 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. +// 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 @@ -166,6 +171,16 @@ 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) { @@ -173,6 +188,11 @@ func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelRe return ModelResponse{}, errModelCallLimit } + cacheKey := c.cacheKey + if request.Agent != "" { + cacheKey += "-" + request.Agent + } + body, err := json.Marshal(ModelRequestBody{ Model: request.Model, Instructions: request.Instructions, @@ -184,7 +204,7 @@ func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelRe Store: false, Stream: true, Include: []string{"reasoning.encrypted_content"}, - PromptCacheKey: c.cacheKey, + PromptCacheKey: cacheKey, }) if err != nil { return ModelResponse{}, fmt.Errorf("encode request: %w", err) @@ -197,14 +217,15 @@ func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelRe 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) + _, chatgpt := c.auth.(*ChatGPTAuth) if chatgpt { - httpRequest.Header.Set("session_id", c.cacheKey) + httpRequest.Header.Set("session_id", cacheKey) } start := time.Now() @@ -218,6 +239,7 @@ func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelRe } record.Turn = request.Turn + record.Agent = request.Agent response, err := c.send(httpRequest, &record) if err != nil { record.Error = err.Error() @@ -302,16 +324,19 @@ type ModelOutputItem struct { // 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. Node and Turn go only to the usage file, never to the -// endpoint. +// 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 + Node *NodeKey Turn string + Agent string } // ModelResponse is the result of a model call. Output holds the output items as @@ -342,14 +367,16 @@ type ModelReasoning struct { // 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. Quota holds the -// x-codex-* headers of the response, without the turn state. +// 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"` @@ -411,6 +438,7 @@ type ModelWireError struct { // 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 @@ -431,8 +459,11 @@ func modelQuotaHeaders(header http.Header) map[string]string { // 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) { @@ -449,6 +480,7 @@ func modelReadStream(stream io.Reader) (ModelResponse, error) { } else if data.Len() > 0 { done, eventErr := modelReadEvent(&response, data.Bytes()) data.Reset() + if eventErr != nil || done { return response, eventErr } @@ -464,6 +496,7 @@ func modelReadStream(stream io.Reader) (ModelResponse, error) { // 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)