Luis Pater
3dd5095792
- 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.
30 lines
894 B
Go
30 lines
894 B
Go
package gemini
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
)
|
|
|
|
// PassthroughGeminiResponseStream forwards Gemini responses unchanged.
|
|
func PassthroughGeminiResponseStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) []string {
|
|
if bytes.HasPrefix(rawJSON, []byte("data:")) {
|
|
rawJSON = bytes.TrimSpace(rawJSON[5:])
|
|
}
|
|
|
|
if bytes.Equal(rawJSON, []byte("[DONE]")) {
|
|
return []string{}
|
|
}
|
|
|
|
return []string{string(rawJSON)}
|
|
}
|
|
|
|
// PassthroughGeminiResponseNonStream forwards Gemini responses unchanged.
|
|
func PassthroughGeminiResponseNonStream(_ context.Context, _ string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) string {
|
|
return string(rawJSON)
|
|
}
|
|
|
|
func GeminiTokenCount(ctx context.Context, count int64) string {
|
|
return fmt.Sprintf(`{"totalTokens":%d,"promptTokensDetails":[{"modality":"TEXT","tokenCount":%d}]}`, count, count)
|
|
}
|