108 lines
2.8 KiB
Go
108 lines
2.8 KiB
Go
package common
|
|
|
|
import (
|
|
"strconv"
|
|
|
|
"github.com/tidwall/gjson"
|
|
"github.com/tidwall/sjson"
|
|
)
|
|
|
|
func GeminiTokenCountJSON(count int64) []byte {
|
|
out := make([]byte, 0, 96)
|
|
out = append(out, `{"totalTokens":`...)
|
|
out = strconv.AppendInt(out, count, 10)
|
|
out = append(out, `,"promptTokensDetails":[{"modality":"TEXT","tokenCount":`...)
|
|
out = strconv.AppendInt(out, count, 10)
|
|
out = append(out, `}]}`...)
|
|
return out
|
|
}
|
|
|
|
func ClaudeInputTokensJSON(count int64) []byte {
|
|
out := make([]byte, 0, 32)
|
|
out = append(out, `{"input_tokens":`...)
|
|
out = strconv.AppendInt(out, count, 10)
|
|
out = append(out, '}')
|
|
return out
|
|
}
|
|
|
|
// NewRawArrayItems creates a raw item slice sized for the expected input.
|
|
func NewRawArrayItems(capacity int64) [][]byte {
|
|
if capacity <= 0 {
|
|
return nil
|
|
}
|
|
return make([][]byte, 0, int(capacity))
|
|
}
|
|
|
|
func JoinRawArray(items [][]byte) []byte {
|
|
if len(items) == 0 {
|
|
return []byte("[]")
|
|
}
|
|
size := len(items) + 1
|
|
for _, item := range items {
|
|
size += len(item)
|
|
}
|
|
out := make([]byte, 0, size)
|
|
out = append(out, '[')
|
|
for i, item := range items {
|
|
if i > 0 {
|
|
out = append(out, ',')
|
|
}
|
|
out = append(out, item...)
|
|
}
|
|
return append(out, ']')
|
|
}
|
|
|
|
// SetRawArrayItems replaces an empty JSON array at path with raw items.
|
|
// The single-item path avoids allocating an intermediate joined array.
|
|
func SetRawArrayItems(data []byte, path string, items [][]byte) []byte {
|
|
if len(items) == 0 {
|
|
return data
|
|
}
|
|
if len(items) == 1 {
|
|
array := gjson.GetBytes(data, path)
|
|
if array.Raw == "[]" && array.Index >= 0 && array.Index+len(array.Raw) <= len(data) {
|
|
out := make([]byte, 0, len(data)+len(items[0]))
|
|
out = append(out, data[:array.Index]...)
|
|
out = append(out, '[')
|
|
out = append(out, items[0]...)
|
|
out = append(out, ']')
|
|
return append(out, data[array.Index+len(array.Raw):]...)
|
|
}
|
|
}
|
|
data, _ = sjson.SetRawBytes(data, path, JoinRawArray(items))
|
|
return data
|
|
}
|
|
|
|
func SSEEventData(event string, payload []byte) []byte {
|
|
out := make([]byte, 0, len(event)+len(payload)+14)
|
|
out = append(out, "event: "...)
|
|
out = append(out, event...)
|
|
out = append(out, '\n')
|
|
out = append(out, "data: "...)
|
|
out = append(out, payload...)
|
|
return out
|
|
}
|
|
|
|
func AppendSSEEventString(out []byte, event, payload string, trailingNewlines int) []byte {
|
|
out = append(out, "event: "...)
|
|
out = append(out, event...)
|
|
out = append(out, '\n')
|
|
out = append(out, "data: "...)
|
|
out = append(out, payload...)
|
|
for i := 0; i < trailingNewlines; i++ {
|
|
out = append(out, '\n')
|
|
}
|
|
return out
|
|
}
|
|
|
|
func AppendSSEEventBytes(out []byte, event string, payload []byte, trailingNewlines int) []byte {
|
|
out = append(out, "event: "...)
|
|
out = append(out, event...)
|
|
out = append(out, '\n')
|
|
out = append(out, "data: "...)
|
|
out = append(out, payload...)
|
|
for i := 0; i < trailingNewlines; i++ {
|
|
out = append(out, '\n')
|
|
}
|
|
return out
|
|
}
|