mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-19 04:40:52 +08:00
feat(translators): add token counting support for Claude and Gemini responses
- Implemented `TokenCount` transform method across translators to calculate token usage. - Integrated token counting logic into executor pipelines for Claude, Gemini, and CLI translators. - Added corresponding API endpoints and handlers (`/messages/count_tokens`) for token usage retrieval. - Enhanced translation registry to support `TokenCount` functionality alongside existing response types.
This commit is contained in:
@@ -91,6 +91,19 @@ func (r *Registry) TranslateNonStream(ctx context.Context, from, to Format, mode
|
||||
return string(rawJSON)
|
||||
}
|
||||
|
||||
// TranslateNonStream applies the registered non-stream response translator.
|
||||
func (r *Registry) TranslateTokenCount(ctx context.Context, from, to Format, count int64, rawJSON []byte) string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
if byTarget, ok := r.responses[to]; ok {
|
||||
if fn, isOk := byTarget[from]; isOk && fn.TokenCount != nil {
|
||||
return fn.TokenCount(ctx, count)
|
||||
}
|
||||
}
|
||||
return string(rawJSON)
|
||||
}
|
||||
|
||||
var defaultRegistry = NewRegistry()
|
||||
|
||||
// Default exposes the package-level registry for shared use.
|
||||
@@ -122,3 +135,8 @@ func TranslateStream(ctx context.Context, from, to Format, model string, origina
|
||||
func TranslateNonStream(ctx context.Context, from, to Format, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) string {
|
||||
return defaultRegistry.TranslateNonStream(ctx, from, to, model, originalRequestRawJSON, requestRawJSON, rawJSON, param)
|
||||
}
|
||||
|
||||
// TranslateTokenCount is a helper on the default registry.
|
||||
func TranslateTokenCount(ctx context.Context, from, to Format, count int64, rawJSON []byte) string {
|
||||
return defaultRegistry.TranslateTokenCount(ctx, from, to, count, rawJSON)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user