X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/7594b9938b7343717e7a57aa7aede5d40d547483..0cb73c3dde36d0078f3c1aa4200f2edcb5ceed38:/model.go

diff --git a/model.go b/model.go
index f15be22..c443a12 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,107 +73,6 @@ type ModelClient struct {
 	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
-}
-
-// 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
-	}
-
-	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: c.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", c.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 {
@@ -264,33 +80,6 @@ type ModelUsageLog struct {
 	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"`
@@ -300,12 +89,20 @@ type ModelOutputItem struct {
 	} `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 only to the
-// usage file, never to the endpoint.
+// 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
@@ -412,10 +209,176 @@ type ModelWireError struct {
 	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
+}
+
+// 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)
+}
+
+// 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
+}
+
 // 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
@@ -436,8 +399,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) {
@@ -454,6 +420,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
 			}
@@ -469,6 +436,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)
@@ -517,3 +485,70 @@ func modelReadEvent(response *ModelResponse, data []byte) (bool, error) {
 
 	return false, nil
 }
+
+// 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
+}
+
+// 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
+}
+
+// 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,
+		}
+	}
+
+	return modelUserParts(content)
+}
+
+// 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,
+	}
+}