feat: publish sanitized custom.97 source

relaykit/dto/gemini.go

@@ -2,7 +2,9 @@ package dto
 
 import (
 	"encoding/json"
+	"maps"
 	"net/http"
+	"slices"
 	"strings"
 
 	kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil"
@@ -74,13 +76,47 @@ func (r *GeminiChatRequest) GetTokenCountMeta() *types.TokenCountMeta {
 	}
 
 	var inputTexts []string
+	if r.SystemInstructions != nil {
+		for _, part := range r.SystemInstructions.Parts {
+			if part.Text != "" {
+				inputTexts = append(inputTexts, part.Text)
+			}
+		}
+	}
 	for _, content := range r.Contents {
 		for _, part := range content.Parts {
 			if part.Text != "" {
 				inputTexts = append(inputTexts, part.Text)
 			}
-			if source := part.InlineData.ToFileSource(); source != nil {
-				mimeType := part.InlineData.MimeType
+			// Function calls and responses replayed in agent loops are part of the prompt.
+			if part.FunctionCall != nil {
+				args, _ := kitutil.Marshal(part.FunctionCall.Arguments)
+				inputTexts = append(inputTexts, part.FunctionCall.FunctionName, string(args))
+			}
+			mediaParts := []GeminiPart{part}
+			if part.FunctionResponse != nil {
+				inputTexts = append(inputTexts, part.FunctionResponse.Name)
+				// Count tool output as the text the model reads, not JSON-escaped.
+				for _, key := range slices.Sorted(maps.Keys(part.FunctionResponse.Response)) {
+					if text, ok := part.FunctionResponse.Response[key].(string); ok {
+						inputTexts = append(inputTexts, text)
+						continue
+					}
+					value, _ := kitutil.Marshal(part.FunctionResponse.Response[key])
+					inputTexts = append(inputTexts, string(value))
+				}
+				// Multimodal function responses carry media such as agent screenshots.
+				var responseParts []GeminiPart
+				if kitutil.Unmarshal(part.FunctionResponse.Parts, &responseParts) == nil {
+					mediaParts = append(mediaParts, responseParts...)
+				}
+			}
+			for _, media := range mediaParts {
+				source := media.InlineData.ToFileSource()
+				if source == nil {
+					continue
+				}
+				mimeType := media.InlineData.MimeType
 				var fileType types.FileType
 				if strings.HasPrefix(mimeType, "image/") {
 					fileType = types.FileTypeImage
@@ -91,13 +127,26 @@ func (r *GeminiChatRequest) GetTokenCountMeta() *types.TokenCountMeta {
 				} else {
 					fileType = types.FileTypeFile
 				}
+				// Image cost depends on the media resolution; a part's own level
+				// overrides the request-wide one.
+				detail := string(r.GenerationConfig.MediaResolution)
+				var partResolution struct {
+					Level string `json:"level"`
+				}
+				if kitutil.Unmarshal(media.MediaResolution, &partResolution) == nil && partResolution.Level != "" {
+					detail = partResolution.Level
+				}
 				files = append(files, &types.FileMeta{
 					FileType: fileType,
 					Source:   source,
+					Detail:   detail,
 				})
 			}
 		}
 	}
+	if len(r.Tools) > 0 {
+		inputTexts = append(inputTexts, string(r.Tools))
+	}
 
 	inputText := strings.Join(inputTexts, "\n")
 	return &types.TokenCountMeta{
@@ -467,6 +516,7 @@ type GeminiChatCandidate struct {
 
 type GeminiGroundingMetadata struct {
 	WebSearchQueries             []string        `json:"webSearchQueries,omitempty"`
+	ImageSearchQueries           []string        `json:"imageSearchQueries,omitempty"`
 	RetrievalQueries             []string        `json:"retrievalQueries,omitempty"`
 	GroundingChunks              json.RawMessage `json:"groundingChunks,omitempty"`
 	GroundingSupports            json.RawMessage `json:"groundingSupports,omitempty"`