X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/7594b9938b7343717e7a57aa7aede5d40d547483..871743f6ab20d68dcede1673a9764730fda38308:/model.go?ds=sidebyside

diff --git a/model.go b/model.go
index f15be22..507fbb1 100644
--- a/model.go
+++ b/model.go
@@ -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,6 +217,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 +225,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()
@@ -304,8 +325,9 @@ type ModelOutputItem struct {
 // 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
@@ -416,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
@@ -436,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) {
@@ -454,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
 			}
@@ -469,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)