Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
SSH key fingerprint:
SHA256:GBzT2DXvAuGV9XIV5W3WrzVpjU54FThmxHXdbz95J24
1802 changed files with 503346 additions and 2 deletions
387
backend/internal/runtime/executor/helps/claude_input_tokens.go
Normal file
387
backend/internal/runtime/executor/helps/claude_input_tokens.go
Normal file
|
|
@ -0,0 +1,387 @@
|
|||
package helps
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
"github.com/tiktoken-go/tokenizer"
|
||||
|
||||
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
|
||||
)
|
||||
|
||||
var (
|
||||
claudeInputTokenizerOnce sync.Once
|
||||
claudeInputTokenizerCodec tokenizer.Codec
|
||||
claudeInputTokenizerErr error
|
||||
)
|
||||
|
||||
// ClaudeInputTokenState tracks the one-time input token update for a translated Claude stream.
|
||||
type ClaudeInputTokenState struct {
|
||||
upstreamFormat sdktranslator.Format
|
||||
responseFormat sdktranslator.Format
|
||||
originalRequest []byte
|
||||
codec tokenizer.Codec
|
||||
handled bool
|
||||
}
|
||||
|
||||
// NewClaudeInputTokenState creates request-scoped state for translated Claude input token usage.
|
||||
func NewClaudeInputTokenState(sourceFormat, upstreamFormat, responseFormat sdktranslator.Format, originalRequest []byte) *ClaudeInputTokenState {
|
||||
enabled := sourceFormat == sdktranslator.FormatClaude &&
|
||||
upstreamFormat != sdktranslator.FormatClaude &&
|
||||
responseFormat == sdktranslator.FormatClaude
|
||||
return &ClaudeInputTokenState{
|
||||
upstreamFormat: upstreamFormat,
|
||||
responseFormat: responseFormat,
|
||||
originalRequest: originalRequest,
|
||||
handled: !enabled,
|
||||
}
|
||||
}
|
||||
|
||||
// TranslateStreamWithClaudeInputTokens translates a stream chunk and estimates Claude message_start input usage once.
|
||||
func TranslateStreamWithClaudeInputTokens(
|
||||
ctx context.Context,
|
||||
upstreamFormat, responseFormat sdktranslator.Format,
|
||||
model string,
|
||||
originalRequestRawJSON, requestRawJSON, rawJSON []byte,
|
||||
param *any,
|
||||
state *ClaudeInputTokenState,
|
||||
) [][]byte {
|
||||
chunks := sdktranslator.TranslateStream(
|
||||
ctx,
|
||||
upstreamFormat,
|
||||
responseFormat,
|
||||
model,
|
||||
originalRequestRawJSON,
|
||||
requestRawJSON,
|
||||
rawJSON,
|
||||
param,
|
||||
)
|
||||
if responseFormat == sdktranslator.FormatOpenAIResponse {
|
||||
for i, chunk := range chunks {
|
||||
chunks[i] = EnsureResponsesUsageDetails(chunk)
|
||||
}
|
||||
}
|
||||
if state == nil {
|
||||
return chunks
|
||||
}
|
||||
return state.apply(ctx, chunks)
|
||||
}
|
||||
|
||||
func claudeInputTokenizer() (tokenizer.Codec, error) {
|
||||
claudeInputTokenizerOnce.Do(func() {
|
||||
claudeInputTokenizerCodec, claudeInputTokenizerErr = tokenizer.Get(tokenizer.O200kBase)
|
||||
})
|
||||
return claudeInputTokenizerCodec, claudeInputTokenizerErr
|
||||
}
|
||||
|
||||
// CountClaudeInputTokens estimates tokens for a Claude request with the O200kBase tokenizer.
|
||||
func CountClaudeInputTokens(payload []byte) (int64, error) {
|
||||
enc, err := claudeInputTokenizer()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("initialize O200kBase tokenizer: %w", err)
|
||||
}
|
||||
count, err := countClaudeInputTokens(enc, payload)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("count Claude input tokens: %w", err)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func countClaudeInputTokens(enc tokenizer.Codec, payload []byte) (int64, error) {
|
||||
if enc == nil {
|
||||
return 0, fmt.Errorf("encoder is nil")
|
||||
}
|
||||
segments, err := collectClaudeInputTokenSegments(payload)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(segments) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
count, err := enc.Count(strings.Join(segments, "\n"))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return int64(count), nil
|
||||
}
|
||||
|
||||
func collectClaudeInputTokenSegments(payload []byte) ([]string, error) {
|
||||
if len(bytes.TrimSpace(payload)) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if !gjson.ValidBytes(payload) {
|
||||
return nil, fmt.Errorf("invalid Claude request JSON")
|
||||
}
|
||||
|
||||
root := gjson.ParseBytes(payload)
|
||||
segments := make([]string, 0, 32)
|
||||
collectClaudeSystemTokenSegments(root.Get("system"), &segments)
|
||||
collectClaudeMessageTokenSegments(root.Get("messages"), &segments)
|
||||
collectClaudeToolTokenSegments(root.Get("tools"), &segments)
|
||||
collectClaudeToolChoiceTokenSegments(root.Get("tool_choice"), &segments)
|
||||
return segments, nil
|
||||
}
|
||||
|
||||
func collectClaudeSystemTokenSegments(system gjson.Result, segments *[]string) {
|
||||
if system.Type == gjson.String {
|
||||
appendClaudeTokenString(segments, system.String())
|
||||
return
|
||||
}
|
||||
if !system.IsArray() {
|
||||
return
|
||||
}
|
||||
system.ForEach(func(_, part gjson.Result) bool {
|
||||
if part.Type == gjson.String {
|
||||
appendClaudeTokenString(segments, part.String())
|
||||
} else if part.Get("type").String() == "text" {
|
||||
appendClaudeTokenString(segments, part.Get("text").String())
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
func collectClaudeMessageTokenSegments(messages gjson.Result, segments *[]string) {
|
||||
if !messages.IsArray() {
|
||||
return
|
||||
}
|
||||
messages.ForEach(func(_, message gjson.Result) bool {
|
||||
appendClaudeTokenString(segments, message.Get("role").String())
|
||||
collectClaudeContentTokenSegments(message.Get("content"), segments)
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
func collectClaudeContentTokenSegments(content gjson.Result, segments *[]string) {
|
||||
if !content.Exists() {
|
||||
return
|
||||
}
|
||||
if content.Type == gjson.String {
|
||||
appendClaudeTokenString(segments, content.String())
|
||||
return
|
||||
}
|
||||
if content.IsArray() {
|
||||
content.ForEach(func(_, part gjson.Result) bool {
|
||||
collectClaudeContentTokenSegments(part, segments)
|
||||
return true
|
||||
})
|
||||
return
|
||||
}
|
||||
if !content.IsObject() {
|
||||
return
|
||||
}
|
||||
|
||||
switch content.Get("type").String() {
|
||||
case "text":
|
||||
appendClaudeTokenString(segments, content.Get("text").String())
|
||||
case "thinking":
|
||||
appendClaudeTokenString(segments, content.Get("thinking").String())
|
||||
case "document":
|
||||
collectClaudeDocumentTokenSegments(content, segments)
|
||||
case "tool_use", "server_tool_use", "mcp_tool_use":
|
||||
appendClaudeTokenString(segments, content.Get("id").String())
|
||||
appendClaudeTokenString(segments, content.Get("name").String())
|
||||
appendClaudeTokenJSON(segments, content.Get("input"))
|
||||
case "tool_result", "mcp_tool_result", "web_search_tool_result", "web_fetch_tool_result", "code_execution_tool_result", "bash_code_execution_tool_result", "text_editor_code_execution_tool_result":
|
||||
appendClaudeTokenString(segments, content.Get("tool_use_id").String())
|
||||
appendClaudeTokenString(segments, content.Get("tool_call_id").String())
|
||||
collectClaudeContentTokenSegments(content.Get("content"), segments)
|
||||
case "web_search_result", "search_result":
|
||||
if source := content.Get("source"); source.Type == gjson.String {
|
||||
appendClaudeTokenString(segments, source.String())
|
||||
}
|
||||
appendClaudeTokenString(segments, content.Get("title").String())
|
||||
appendClaudeTokenString(segments, content.Get("url").String())
|
||||
appendClaudeTokenString(segments, content.Get("page_age").String())
|
||||
collectClaudeContentTokenSegments(content.Get("content"), segments)
|
||||
case "web_fetch_result":
|
||||
appendClaudeTokenString(segments, content.Get("url").String())
|
||||
appendClaudeTokenString(segments, content.Get("retrieved_at").String())
|
||||
collectClaudeContentTokenSegments(content.Get("content"), segments)
|
||||
case "code_execution_result", "bash_code_execution_result", "text_editor_code_execution_result":
|
||||
appendClaudeTokenString(segments, content.Get("stdout").String())
|
||||
appendClaudeTokenString(segments, content.Get("stderr").String())
|
||||
appendClaudeTokenString(segments, content.Get("return_code").String())
|
||||
collectClaudeContentTokenSegments(content.Get("content"), segments)
|
||||
collectClaudeContentTokenSegments(content.Get("output"), segments)
|
||||
case "tool_reference":
|
||||
appendClaudeTokenString(segments, content.Get("tool_name").String())
|
||||
case "image", "input_audio", "audio", "video", "redacted_thinking":
|
||||
return
|
||||
case "":
|
||||
appendClaudeTokenJSON(segments, content)
|
||||
default:
|
||||
appendClaudeTokenString(segments, content.Get("text").String())
|
||||
}
|
||||
}
|
||||
|
||||
func collectClaudeDocumentTokenSegments(document gjson.Result, segments *[]string) {
|
||||
source := document.Get("source")
|
||||
if source.Get("type").String() != "text" {
|
||||
return
|
||||
}
|
||||
appendClaudeTokenString(segments, document.Get("title").String())
|
||||
appendClaudeTokenString(segments, document.Get("context").String())
|
||||
appendClaudeTokenString(segments, source.Get("data").String())
|
||||
appendClaudeTokenString(segments, source.Get("content").String())
|
||||
}
|
||||
|
||||
func collectClaudeToolTokenSegments(tools gjson.Result, segments *[]string) {
|
||||
if !tools.IsArray() {
|
||||
return
|
||||
}
|
||||
tools.ForEach(func(_, tool gjson.Result) bool {
|
||||
appendClaudeTokenString(segments, tool.Get("type").String())
|
||||
appendClaudeTokenString(segments, tool.Get("name").String())
|
||||
appendClaudeTokenString(segments, tool.Get("description").String())
|
||||
appendClaudeTokenJSON(segments, tool.Get("input_schema"))
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
func collectClaudeToolChoiceTokenSegments(toolChoice gjson.Result, segments *[]string) {
|
||||
if !toolChoice.Exists() {
|
||||
return
|
||||
}
|
||||
if toolChoice.Type == gjson.String {
|
||||
appendClaudeTokenString(segments, toolChoice.String())
|
||||
return
|
||||
}
|
||||
appendClaudeTokenString(segments, toolChoice.Get("type").String())
|
||||
appendClaudeTokenString(segments, toolChoice.Get("name").String())
|
||||
}
|
||||
|
||||
func appendClaudeTokenString(segments *[]string, value string) {
|
||||
if segments == nil {
|
||||
return
|
||||
}
|
||||
if trimmed := strings.TrimSpace(value); trimmed != "" {
|
||||
*segments = append(*segments, trimmed)
|
||||
}
|
||||
}
|
||||
|
||||
func appendClaudeTokenJSON(segments *[]string, value gjson.Result) {
|
||||
if !value.Exists() {
|
||||
return
|
||||
}
|
||||
if value.Type == gjson.String {
|
||||
appendClaudeTokenString(segments, value.String())
|
||||
return
|
||||
}
|
||||
raw := strings.TrimSpace(value.Raw)
|
||||
if raw == "" {
|
||||
return
|
||||
}
|
||||
var compact bytes.Buffer
|
||||
if err := json.Compact(&compact, []byte(raw)); err == nil {
|
||||
appendClaudeTokenString(segments, compact.String())
|
||||
return
|
||||
}
|
||||
appendClaudeTokenString(segments, raw)
|
||||
}
|
||||
|
||||
func (state *ClaudeInputTokenState) apply(ctx context.Context, chunks [][]byte) [][]byte {
|
||||
if state == nil || state.handled {
|
||||
return chunks
|
||||
}
|
||||
for i := range chunks {
|
||||
updated, found := state.applyChunk(ctx, chunks[i])
|
||||
if !found {
|
||||
continue
|
||||
}
|
||||
state.handled = true
|
||||
chunks[i] = updated
|
||||
break
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
func (state *ClaudeInputTokenState) applyChunk(ctx context.Context, chunk []byte) ([]byte, bool) {
|
||||
for lineStart := 0; lineStart < len(chunk); {
|
||||
lineEnd := bytes.IndexByte(chunk[lineStart:], '\n')
|
||||
if lineEnd < 0 {
|
||||
lineEnd = len(chunk)
|
||||
} else {
|
||||
lineEnd += lineStart
|
||||
}
|
||||
|
||||
contentEnd := lineEnd
|
||||
if contentEnd > lineStart && chunk[contentEnd-1] == '\r' {
|
||||
contentEnd--
|
||||
}
|
||||
line := chunk[lineStart:contentEnd]
|
||||
trimmedLeft := bytes.TrimLeft(line, " \t")
|
||||
if bytes.HasPrefix(trimmedLeft, []byte("data:")) {
|
||||
payloadOffset := len(line) - len(trimmedLeft) + len("data:")
|
||||
for payloadOffset < len(line) && (line[payloadOffset] == ' ' || line[payloadOffset] == '\t') {
|
||||
payloadOffset++
|
||||
}
|
||||
payloadEnd := len(line)
|
||||
for payloadEnd > payloadOffset && (line[payloadEnd-1] == ' ' || line[payloadEnd-1] == '\t') {
|
||||
payloadEnd--
|
||||
}
|
||||
payload := line[payloadOffset:payloadEnd]
|
||||
if gjson.GetBytes(payload, "type").String() == "message_start" {
|
||||
inputTokens := gjson.GetBytes(payload, "message.usage.input_tokens")
|
||||
if inputTokens.Exists() && inputTokens.Int() != 0 {
|
||||
return chunk, true
|
||||
}
|
||||
count, err := state.estimate()
|
||||
if err != nil {
|
||||
state.logEstimateError(ctx, err)
|
||||
return chunk, true
|
||||
}
|
||||
if count == 0 {
|
||||
return chunk, true
|
||||
}
|
||||
updatedPayload, errSet := sjson.SetBytes(payload, "message.usage.input_tokens", count)
|
||||
if errSet != nil {
|
||||
state.logEstimateError(ctx, fmt.Errorf("set message_start usage: %w", errSet))
|
||||
return chunk, true
|
||||
}
|
||||
payloadStart := lineStart + payloadOffset
|
||||
payloadStop := lineStart + payloadEnd
|
||||
updated := make([]byte, 0, len(chunk)+len(updatedPayload)-len(payload))
|
||||
updated = append(updated, chunk[:payloadStart]...)
|
||||
updated = append(updated, updatedPayload...)
|
||||
updated = append(updated, chunk[payloadStop:]...)
|
||||
return updated, true
|
||||
}
|
||||
}
|
||||
|
||||
if lineEnd == len(chunk) {
|
||||
break
|
||||
}
|
||||
lineStart = lineEnd + 1
|
||||
}
|
||||
return chunk, false
|
||||
}
|
||||
|
||||
func (state *ClaudeInputTokenState) estimate() (int64, error) {
|
||||
enc := state.codec
|
||||
if enc == nil {
|
||||
var err error
|
||||
enc, err = claudeInputTokenizer()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("initialize O200kBase tokenizer: %w", err)
|
||||
}
|
||||
}
|
||||
count, err := countClaudeInputTokens(enc, state.originalRequest)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("count Claude input tokens: %w", err)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func (state *ClaudeInputTokenState) logEstimateError(ctx context.Context, err error) {
|
||||
LogWithRequestID(ctx).WithFields(log.Fields{
|
||||
"upstream_format": state.upstreamFormat.String(),
|
||||
"response_format": state.responseFormat.String(),
|
||||
}).WithError(err).Warn("failed to estimate Claude input tokens")
|
||||
}
|
||||
Loading…
Reference in a new issue