CLIProxyAPI/internal/provider/gemini-web/prompt.go

package geminiwebapi

import (
	"fmt"
	"math"
	"regexp"
	"strings"
	"unicode/utf8"

	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
	"github.com/tidwall/gjson"
)

var (
	reThink     = regexp.MustCompile(`(?s)^\s*<think>.*?</think>\s*`)
	reXMLAnyTag = regexp.MustCompile(`(?s)<\s*[^>]+>`)
)

// NormalizeRole converts a role to a standard format (lowercase, 'model' -> 'assistant').
func NormalizeRole(role string) string {
	r := strings.ToLower(role)
	if r == "model" {
		return "assistant"
	}
	return r
}

// NeedRoleTags checks if a list of messages requires role tags.
func NeedRoleTags(msgs []RoleText) bool {
	for _, m := range msgs {
		if strings.ToLower(m.Role) != "user" {
			return true
		}
	}
	return false
}

// AddRoleTag wraps content with a role tag.
func AddRoleTag(role, content string, unclose bool) string {
	if role == "" {
		role = "user"
	}
	if unclose {
		return "<|im_start|>" + role + "\n" + content
	}
	return "<|im_start|>" + role + "\n" + content + "\n<|im_end|>"
}

// BuildPrompt constructs the final prompt from a list of messages.
func BuildPrompt(msgs []RoleText, tagged bool, appendAssistant bool) string {
	if len(msgs) == 0 {
		if tagged && appendAssistant {
			return AddRoleTag("assistant", "", true)
		}
		return ""
	}
	if !tagged {
		var sb strings.Builder
		for i, m := range msgs {
			if i > 0 {
				sb.WriteString("\n")
			}
			sb.WriteString(m.Text)
		}
		return sb.String()
	}
	var sb strings.Builder
	for _, m := range msgs {
		sb.WriteString(AddRoleTag(m.Role, m.Text, false))
		sb.WriteString("\n")
	}
	if appendAssistant {
		sb.WriteString(AddRoleTag("assistant", "", true))
	}
	return strings.TrimSpace(sb.String())
}

// RemoveThinkTags strips <think>...</think> blocks from a string.
func RemoveThinkTags(s string) string {
	return strings.TrimSpace(reThink.ReplaceAllString(s, ""))
}

// SanitizeAssistantMessages removes think tags from assistant messages.
func SanitizeAssistantMessages(msgs []RoleText) []RoleText {
	out := make([]RoleText, 0, len(msgs))
	for _, m := range msgs {
		if strings.ToLower(m.Role) == "assistant" {
			out = append(out, RoleText{Role: m.Role, Text: RemoveThinkTags(m.Text)})
		} else {
			out = append(out, m)
		}
	}
	return out
}

// AppendXMLWrapHintIfNeeded appends an XML wrap hint to messages containing XML-like blocks.
func AppendXMLWrapHintIfNeeded(msgs []RoleText, disable bool) []RoleText {
	if disable {
		return msgs
	}
	const xmlWrapHint = "\nFor any xml block, e.g. tool call, always wrap it with: \n`````xml\n...\n`````\n"
	out := make([]RoleText, 0, len(msgs))
	for _, m := range msgs {
		t := m.Text
		if reXMLAnyTag.MatchString(t) {
			t = t + xmlWrapHint
		}
		out = append(out, RoleText{Role: m.Role, Text: t})
	}
	return out
}

// EstimateTotalTokensFromRawJSON estimates token count by summing text parts.
func EstimateTotalTokensFromRawJSON(rawJSON []byte) int {
	totalChars := 0
	contents := gjson.GetBytes(rawJSON, "contents")
	if contents.Exists() {
		contents.ForEach(func(_, content gjson.Result) bool {
			content.Get("parts").ForEach(func(_, part gjson.Result) bool {
				if t := part.Get("text"); t.Exists() {
					totalChars += utf8.RuneCountInString(t.String())
				}
				return true
			})
			return true
		})
	}
	if totalChars <= 0 {
		return 0
	}
	return int(math.Ceil(float64(totalChars) / 4.0))
}

// Request chunking helpers ------------------------------------------------

const continuationHint = "\n(More messages to come, please reply with just 'ok.')"

func ChunkByRunes(s string, size int) []string {
	if size <= 0 {
		return []string{s}
	}
	chunks := make([]string, 0, (len(s)/size)+1)
	var buf strings.Builder
	count := 0
	for _, r := range s {
		buf.WriteRune(r)
		count++
		if count >= size {
			chunks = append(chunks, buf.String())
			buf.Reset()
			count = 0
		}
	}
	if buf.Len() > 0 {
		chunks = append(chunks, buf.String())
	}
	if len(chunks) == 0 {
		return []string{""}
	}
	return chunks
}

func MaxCharsPerRequest(cfg *config.Config) int {
	// Read max characters per request from config with a conservative default.
	if cfg != nil {
		if v := cfg.GeminiWeb.MaxCharsPerRequest; v > 0 {
			return v
		}
	}
	return 1_000_000
}

func SendWithSplit(chat *ChatSession, text string, files []string, cfg *config.Config) (ModelOutput, error) {
	// Validate chat session
	if chat == nil {
		return ModelOutput{}, fmt.Errorf("nil chat session")
	}

	// Resolve maxChars characters per request
	maxChars := MaxCharsPerRequest(cfg)
	if maxChars <= 0 {
		maxChars = 1_000_000
	}

	// If within limit, send directly
	if utf8.RuneCountInString(text) <= maxChars {
		return chat.SendMessage(text, files)
	}

	// Decide whether to use continuation hint (enabled by default)
	useHint := true
	if cfg != nil && cfg.GeminiWeb.DisableContinuationHint {
		useHint = false
	}

	// Compute chunk size in runes. If the hint does not fit, disable it for this request.
	hintLen := 0
	if useHint {
		hintLen = utf8.RuneCountInString(continuationHint)
	}
	chunkSize := maxChars - hintLen
	if chunkSize <= 0 {
		// maxChars is too small to accommodate the hint; fall back to no-hint splitting
		useHint = false
		chunkSize = maxChars
	}

	// Split into rune-safe chunks
	chunks := ChunkByRunes(text, chunkSize)
	if len(chunks) == 0 {
		chunks = []string{""}
	}

	// Send all but the last chunk without files, optionally appending hint
	for i := 0; i < len(chunks)-1; i++ {
		part := chunks[i]
		if useHint {
			part += continuationHint
		}
		if _, err := chat.SendMessage(part, nil); err != nil {
			return ModelOutput{}, err
		}
	}

	// Send final chunk with files and return the actual output
	return chat.SendMessage(chunks[len(chunks)-1], files)
}