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)
|
||||
}
|
||||
|
||||
@@ -11,8 +11,11 @@ type ResponseStreamTransform func(ctx context.Context, model string, originalReq
|
||||
// ResponseNonStreamTransform converts non-stream responses between schemas.
|
||||
type ResponseNonStreamTransform func(ctx context.Context, model string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) string
|
||||
|
||||
type ResponseTokenCountTransform func(ctx context.Context, count int64) string
|
||||
|
||||
// ResponseTransform groups streaming and non-streaming transforms.
|
||||
type ResponseTransform struct {
|
||||
Stream ResponseStreamTransform
|
||||
NonStream ResponseNonStreamTransform
|
||||
Stream ResponseStreamTransform
|
||||
NonStream ResponseNonStreamTransform
|
||||
TokenCount ResponseTokenCountTransform
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user