X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/7cc0014244124aea59672904de65a979a8e5fb2c..80e348724041a694002fc38f57e8868b9b9a2be9:/model.go diff --git a/model.go b/model.go index f87240a..074ff60 100644 --- a/model.go +++ b/model.go @@ -40,85 +40,6 @@ const ( // 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 { @@ -128,19 +49,15 @@ type ModelAuthorizer interface { // 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. -// 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 @@ -156,6 +73,163 @@ type ModelClient struct { calls atomic.Int64 } +// 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 +} + +// 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"` +} + +// ModelInputPart is a text or image part in a user message or tool result. +type ModelInputPart struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + ImageURL string `json:"image_url,omitempty"` +} + +// 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"` +} + +// 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 @@ -166,6 +240,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 +257,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 +273,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,6 +286,7 @@ 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) @@ -204,7 +294,7 @@ func (c *ModelClient) Create(ctx context.Context, request ModelRequest) (ModelRe _, chatgpt := c.auth.(*ChatGPTAuth) if chatgpt { - httpRequest.Header.Set("session_id", c.cacheKey) + httpRequest.Header.Set("session_id", cacheKey) } start := time.Now() @@ -218,6 +308,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() @@ -256,13 +347,6 @@ func (c *ModelClient) send(httpRequest *http.Request, record *ModelUsageRecord) 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) @@ -290,127 +374,78 @@ func (u *ModelUsageLog) Write(record ModelUsageRecord) (err error) { 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"` -} +// 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") + } -// 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. -type ModelRequest struct { - Model string - Effort string - Instructions string - Tools []ModelTool - Input []any - Node *nodeKey - Turn string -} + return modelAPIURL, ModelAPIKey(key), nil + } -// 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 + return modelChatGPTURL, newChatGPTAuth(filepath.Join(dir, chatgptTokenFileName), chatgptIssuer), nil } -// 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"` -} +// modelUserMessage returns a user input item with one text part for each text. +func modelUserMessage(texts ...string) any { + content := make([]ModelInputPart, len(texts)) + for n, text := range texts { + content[n] = ModelInputPart{ + Type: "input_text", + Text: text, + } + } -type ModelReasoning struct { - Effort string `json:"effort"` + return modelUserParts(content) } -// 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. -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"` - 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"` +// modelUserParts keeps text and image parts in one user message. +func modelUserParts(content []ModelInputPart) any { + return struct { + Role string `json:"role"` + Content []ModelInputPart `json:"content"` + }{ + Role: "user", + Content: content, + } } -// 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"` -} +// modelOutputText returns the text of the message items of output. +func modelOutputText(output []json.RawMessage) (string, error) { + var text strings.Builder -// 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 -} + for _, raw := range output { + var item ModelOutputItem -// 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"` -} + err := json.Unmarshal(raw, &item) + if err != nil { + return "", fmt.Errorf("decode output item: %w", err) + } -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"` -} + if item.Type != "message" { + continue + } -type ModelWireUsage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - InputTokensDetails struct { - CachedTokens int `json:"cached_tokens"` - } `json:"input_tokens_details"` -} + for _, part := range item.Content { + if part.Type == "output_text" { + text.WriteString(part.Text) + } + } + } -type ModelWireError struct { - Code string `json:"code"` - Message string `json:"message"` + return text.String(), nil } // 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 @@ -431,8 +466,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 +487,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 +503,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)