mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-02 20:40:52 +08:00
631 lines
24 KiB
Go
631 lines
24 KiB
Go
// Package gemini provides response translation functionality for Claude Code to Gemini API compatibility.
|
|
// This package handles the conversion of Claude Code API responses into Gemini-compatible
|
|
// JSON format, transforming streaming events and non-streaming responses into the format
|
|
// expected by Gemini API clients. It supports both streaming and non-streaming modes,
|
|
// handling text content, tool calls, and usage metadata appropriately.
|
|
package gemini
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/tidwall/gjson"
|
|
"github.com/tidwall/sjson"
|
|
)
|
|
|
|
var (
|
|
dataTag = []byte("data:")
|
|
)
|
|
|
|
// ConvertAnthropicResponseToGeminiParams holds parameters for response conversion
|
|
// It also carries minimal streaming state across calls to assemble tool_use input_json_delta.
|
|
// This structure maintains state information needed for proper conversion of streaming responses
|
|
// from Claude Code format to Gemini format, particularly for handling tool calls that span
|
|
// multiple streaming events.
|
|
type ConvertAnthropicResponseToGeminiParams struct {
|
|
Model string
|
|
CreatedAt int64
|
|
ResponseID string
|
|
LastStorageOutput string
|
|
IsStreaming bool
|
|
|
|
// Streaming state for tool_use assembly
|
|
// Keyed by content_block index from Claude SSE events
|
|
ToolUseNames map[int]string // function/tool name per block index
|
|
ToolUseArgs map[int]*strings.Builder // accumulates partial_json across deltas
|
|
}
|
|
|
|
// ConvertClaudeResponseToGemini converts Claude Code streaming response format to Gemini format.
|
|
// This function processes various Claude Code event types and transforms them into Gemini-compatible JSON responses.
|
|
// It handles text content, tool calls, reasoning content, and usage metadata, outputting responses that match
|
|
// the Gemini API format. The function supports incremental updates for streaming responses and maintains
|
|
// state information to properly assemble multi-part tool calls.
|
|
//
|
|
// Parameters:
|
|
// - ctx: The context for the request, used for cancellation and timeout handling
|
|
// - modelName: The name of the model being used for the response
|
|
// - rawJSON: The raw JSON response from the Claude Code API
|
|
// - param: A pointer to a parameter object for maintaining state between calls
|
|
//
|
|
// Returns:
|
|
// - []string: A slice of strings, each containing a Gemini-compatible JSON response
|
|
func ConvertClaudeResponseToGemini(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) []string {
|
|
if *param == nil {
|
|
*param = &ConvertAnthropicResponseToGeminiParams{
|
|
Model: modelName,
|
|
CreatedAt: 0,
|
|
ResponseID: "",
|
|
}
|
|
}
|
|
|
|
if !bytes.HasPrefix(rawJSON, dataTag) {
|
|
return []string{}
|
|
}
|
|
rawJSON = bytes.TrimSpace(rawJSON[5:])
|
|
|
|
root := gjson.ParseBytes(rawJSON)
|
|
eventType := root.Get("type").String()
|
|
|
|
// Base Gemini response template with default values
|
|
template := `{"candidates":[{"content":{"role":"model","parts":[]}}],"usageMetadata":{"trafficType":"PROVISIONED_THROUGHPUT"},"modelVersion":"","createTime":"","responseId":""}`
|
|
|
|
// Set model version
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).Model != "" {
|
|
// Map Claude model names back to Gemini model names
|
|
template, _ = sjson.Set(template, "modelVersion", (*param).(*ConvertAnthropicResponseToGeminiParams).Model)
|
|
}
|
|
|
|
// Set response ID and creation time
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).ResponseID != "" {
|
|
template, _ = sjson.Set(template, "responseId", (*param).(*ConvertAnthropicResponseToGeminiParams).ResponseID)
|
|
}
|
|
|
|
// Set creation time to current time if not provided
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).CreatedAt == 0 {
|
|
(*param).(*ConvertAnthropicResponseToGeminiParams).CreatedAt = time.Now().Unix()
|
|
}
|
|
template, _ = sjson.Set(template, "createTime", time.Unix((*param).(*ConvertAnthropicResponseToGeminiParams).CreatedAt, 0).Format(time.RFC3339Nano))
|
|
|
|
switch eventType {
|
|
case "message_start":
|
|
// Initialize response with message metadata when a new message begins
|
|
if message := root.Get("message"); message.Exists() {
|
|
(*param).(*ConvertAnthropicResponseToGeminiParams).ResponseID = message.Get("id").String()
|
|
(*param).(*ConvertAnthropicResponseToGeminiParams).Model = message.Get("model").String()
|
|
}
|
|
return []string{}
|
|
|
|
case "content_block_start":
|
|
// Start of a content block - record tool_use name by index for functionCall assembly
|
|
if cb := root.Get("content_block"); cb.Exists() {
|
|
if cb.Get("type").String() == "tool_use" {
|
|
idx := int(root.Get("index").Int())
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseNames == nil {
|
|
(*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseNames = map[int]string{}
|
|
}
|
|
if name := cb.Get("name"); name.Exists() {
|
|
(*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseNames[idx] = name.String()
|
|
}
|
|
}
|
|
}
|
|
return []string{}
|
|
|
|
case "content_block_delta":
|
|
// Handle content delta (text, thinking, or tool use arguments)
|
|
if delta := root.Get("delta"); delta.Exists() {
|
|
deltaType := delta.Get("type").String()
|
|
|
|
switch deltaType {
|
|
case "text_delta":
|
|
// Regular text content delta for normal response text
|
|
if text := delta.Get("text"); text.Exists() && text.String() != "" {
|
|
textPart := `{"text":""}`
|
|
textPart, _ = sjson.Set(textPart, "text", text.String())
|
|
template, _ = sjson.SetRaw(template, "candidates.0.content.parts.-1", textPart)
|
|
}
|
|
case "thinking_delta":
|
|
// Thinking/reasoning content delta for models with reasoning capabilities
|
|
if text := delta.Get("thinking"); text.Exists() && text.String() != "" {
|
|
thinkingPart := `{"thought":true,"text":""}`
|
|
thinkingPart, _ = sjson.Set(thinkingPart, "text", text.String())
|
|
template, _ = sjson.SetRaw(template, "candidates.0.content.parts.-1", thinkingPart)
|
|
}
|
|
case "input_json_delta":
|
|
// Tool use input delta - accumulate partial_json by index for later assembly at content_block_stop
|
|
idx := int(root.Get("index").Int())
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseArgs == nil {
|
|
(*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseArgs = map[int]*strings.Builder{}
|
|
}
|
|
b, ok := (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseArgs[idx]
|
|
if !ok || b == nil {
|
|
bb := &strings.Builder{}
|
|
(*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseArgs[idx] = bb
|
|
b = bb
|
|
}
|
|
if pj := delta.Get("partial_json"); pj.Exists() {
|
|
b.WriteString(pj.String())
|
|
}
|
|
return []string{}
|
|
}
|
|
}
|
|
return []string{template}
|
|
|
|
case "content_block_stop":
|
|
// End of content block - finalize tool calls if any
|
|
idx := int(root.Get("index").Int())
|
|
// Claude's content_block_stop often doesn't include content_block payload (see docs/response-claude.txt)
|
|
// So we finalize using accumulated state captured during content_block_start and input_json_delta.
|
|
name := ""
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseNames != nil {
|
|
name = (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseNames[idx]
|
|
}
|
|
var argsTrim string
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseArgs != nil {
|
|
if b := (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseArgs[idx]; b != nil {
|
|
argsTrim = strings.TrimSpace(b.String())
|
|
}
|
|
}
|
|
if name != "" || argsTrim != "" {
|
|
functionCall := `{"functionCall":{"name":"","args":{}}}`
|
|
if name != "" {
|
|
functionCall, _ = sjson.Set(functionCall, "functionCall.name", name)
|
|
}
|
|
if argsTrim != "" {
|
|
functionCall, _ = sjson.SetRaw(functionCall, "functionCall.args", argsTrim)
|
|
}
|
|
template, _ = sjson.SetRaw(template, "candidates.0.content.parts.-1", functionCall)
|
|
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
|
|
(*param).(*ConvertAnthropicResponseToGeminiParams).LastStorageOutput = template
|
|
// cleanup used state for this index
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseArgs != nil {
|
|
delete((*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseArgs, idx)
|
|
}
|
|
if (*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseNames != nil {
|
|
delete((*param).(*ConvertAnthropicResponseToGeminiParams).ToolUseNames, idx)
|
|
}
|
|
return []string{template}
|
|
}
|
|
return []string{}
|
|
|
|
case "message_delta":
|
|
// Handle message-level changes (like stop reason and usage information)
|
|
if delta := root.Get("delta"); delta.Exists() {
|
|
if stopReason := delta.Get("stop_reason"); stopReason.Exists() {
|
|
switch stopReason.String() {
|
|
case "end_turn":
|
|
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
|
|
case "tool_use":
|
|
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
|
|
case "max_tokens":
|
|
template, _ = sjson.Set(template, "candidates.0.finishReason", "MAX_TOKENS")
|
|
case "stop_sequence":
|
|
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
|
|
default:
|
|
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
|
|
}
|
|
}
|
|
}
|
|
|
|
if usage := root.Get("usage"); usage.Exists() {
|
|
// Basic token counts for prompt and completion
|
|
inputTokens := usage.Get("input_tokens").Int()
|
|
outputTokens := usage.Get("output_tokens").Int()
|
|
|
|
// Set basic usage metadata according to Gemini API specification
|
|
template, _ = sjson.Set(template, "usageMetadata.promptTokenCount", inputTokens)
|
|
template, _ = sjson.Set(template, "usageMetadata.candidatesTokenCount", outputTokens)
|
|
template, _ = sjson.Set(template, "usageMetadata.totalTokenCount", inputTokens+outputTokens)
|
|
|
|
// Add cache-related token counts if present (Claude Code API cache fields)
|
|
if cacheCreationTokens := usage.Get("cache_creation_input_tokens"); cacheCreationTokens.Exists() {
|
|
template, _ = sjson.Set(template, "usageMetadata.cachedContentTokenCount", cacheCreationTokens.Int())
|
|
}
|
|
if cacheReadTokens := usage.Get("cache_read_input_tokens"); cacheReadTokens.Exists() {
|
|
// Add cache read tokens to cached content count
|
|
existingCacheTokens := usage.Get("cache_creation_input_tokens").Int()
|
|
totalCacheTokens := existingCacheTokens + cacheReadTokens.Int()
|
|
template, _ = sjson.Set(template, "usageMetadata.cachedContentTokenCount", totalCacheTokens)
|
|
}
|
|
|
|
// Add thinking tokens if present (for models with reasoning capabilities)
|
|
if thinkingTokens := usage.Get("thinking_tokens"); thinkingTokens.Exists() {
|
|
template, _ = sjson.Set(template, "usageMetadata.thoughtsTokenCount", thinkingTokens.Int())
|
|
}
|
|
|
|
// Set traffic type (required by Gemini API)
|
|
template, _ = sjson.Set(template, "usageMetadata.trafficType", "PROVISIONED_THROUGHPUT")
|
|
}
|
|
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
|
|
|
|
return []string{template}
|
|
case "message_stop":
|
|
// Final message with usage information - no additional output needed
|
|
return []string{}
|
|
case "error":
|
|
// Handle error responses and convert to Gemini error format
|
|
errorMsg := root.Get("error.message").String()
|
|
if errorMsg == "" {
|
|
errorMsg = "Unknown error occurred"
|
|
}
|
|
|
|
// Create error response in Gemini format
|
|
errorResponse := `{"error":{"code":400,"message":"","status":"INVALID_ARGUMENT"}}`
|
|
errorResponse, _ = sjson.Set(errorResponse, "error.message", errorMsg)
|
|
return []string{errorResponse}
|
|
|
|
default:
|
|
// Unknown event type, return empty response
|
|
return []string{}
|
|
}
|
|
}
|
|
|
|
// convertArrayToJSON converts []interface{} to JSON array string
|
|
func convertArrayToJSON(arr []interface{}) string {
|
|
result := "[]"
|
|
for _, item := range arr {
|
|
switch itemData := item.(type) {
|
|
case map[string]interface{}:
|
|
itemJSON := convertMapToJSON(itemData)
|
|
result, _ = sjson.SetRaw(result, "-1", itemJSON)
|
|
case string:
|
|
result, _ = sjson.Set(result, "-1", itemData)
|
|
case bool:
|
|
result, _ = sjson.Set(result, "-1", itemData)
|
|
case float64, int, int64:
|
|
result, _ = sjson.Set(result, "-1", itemData)
|
|
default:
|
|
result, _ = sjson.Set(result, "-1", itemData)
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
// convertMapToJSON converts map[string]interface{} to JSON object string
|
|
func convertMapToJSON(m map[string]interface{}) string {
|
|
result := "{}"
|
|
for key, value := range m {
|
|
switch val := value.(type) {
|
|
case map[string]interface{}:
|
|
nestedJSON := convertMapToJSON(val)
|
|
result, _ = sjson.SetRaw(result, key, nestedJSON)
|
|
case []interface{}:
|
|
arrayJSON := convertArrayToJSON(val)
|
|
result, _ = sjson.SetRaw(result, key, arrayJSON)
|
|
case string:
|
|
result, _ = sjson.Set(result, key, val)
|
|
case bool:
|
|
result, _ = sjson.Set(result, key, val)
|
|
case float64, int, int64:
|
|
result, _ = sjson.Set(result, key, val)
|
|
default:
|
|
result, _ = sjson.Set(result, key, val)
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
// ConvertClaudeResponseToGeminiNonStream converts a non-streaming Claude Code response to a non-streaming Gemini response.
|
|
// This function processes the complete Claude Code response and transforms it into a single Gemini-compatible
|
|
// JSON response. It handles message content, tool calls, reasoning content, and usage metadata, combining all
|
|
// the information into a single response that matches the Gemini API format.
|
|
//
|
|
// Parameters:
|
|
// - ctx: The context for the request, used for cancellation and timeout handling
|
|
// - modelName: The name of the model being used for the response
|
|
// - rawJSON: The raw JSON response from the Claude Code API
|
|
// - param: A pointer to a parameter object for the conversion (unused in current implementation)
|
|
//
|
|
// Returns:
|
|
// - string: A Gemini-compatible JSON response containing all message content and metadata
|
|
func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, _ *any) string {
|
|
// Base Gemini response template for non-streaming with default values
|
|
template := `{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}],"usageMetadata":{"trafficType":"PROVISIONED_THROUGHPUT"},"modelVersion":"","createTime":"","responseId":""}`
|
|
|
|
// Set model version
|
|
template, _ = sjson.Set(template, "modelVersion", modelName)
|
|
|
|
streamingEvents := make([][]byte, 0)
|
|
|
|
scanner := bufio.NewScanner(bytes.NewReader(rawJSON))
|
|
buffer := make([]byte, 10240*1024)
|
|
scanner.Buffer(buffer, 10240*1024)
|
|
for scanner.Scan() {
|
|
line := scanner.Bytes()
|
|
// log.Debug(string(line))
|
|
if bytes.HasPrefix(line, dataTag) {
|
|
jsonData := bytes.TrimSpace(line[5:])
|
|
streamingEvents = append(streamingEvents, jsonData)
|
|
}
|
|
}
|
|
// log.Debug("streamingEvents: ", streamingEvents)
|
|
// log.Debug("rawJSON: ", string(rawJSON))
|
|
|
|
// Initialize parameters for streaming conversion with proper state management
|
|
newParam := &ConvertAnthropicResponseToGeminiParams{
|
|
Model: modelName,
|
|
CreatedAt: 0,
|
|
ResponseID: "",
|
|
LastStorageOutput: "",
|
|
IsStreaming: false,
|
|
ToolUseNames: nil,
|
|
ToolUseArgs: nil,
|
|
}
|
|
|
|
// Process each streaming event and collect parts
|
|
var allParts []interface{}
|
|
var finalUsage map[string]interface{}
|
|
var responseID string
|
|
var createdAt int64
|
|
|
|
for _, eventData := range streamingEvents {
|
|
if len(eventData) == 0 {
|
|
continue
|
|
}
|
|
|
|
root := gjson.ParseBytes(eventData)
|
|
eventType := root.Get("type").String()
|
|
|
|
switch eventType {
|
|
case "message_start":
|
|
// Extract response metadata including ID, model, and creation time
|
|
if message := root.Get("message"); message.Exists() {
|
|
responseID = message.Get("id").String()
|
|
newParam.ResponseID = responseID
|
|
newParam.Model = message.Get("model").String()
|
|
|
|
// Set creation time to current time if not provided
|
|
createdAt = time.Now().Unix()
|
|
newParam.CreatedAt = createdAt
|
|
}
|
|
|
|
case "content_block_start":
|
|
// Prepare for content block; record tool_use name by index for later functionCall assembly
|
|
idx := int(root.Get("index").Int())
|
|
if cb := root.Get("content_block"); cb.Exists() {
|
|
if cb.Get("type").String() == "tool_use" {
|
|
if newParam.ToolUseNames == nil {
|
|
newParam.ToolUseNames = map[int]string{}
|
|
}
|
|
if name := cb.Get("name"); name.Exists() {
|
|
newParam.ToolUseNames[idx] = name.String()
|
|
}
|
|
}
|
|
}
|
|
continue
|
|
|
|
case "content_block_delta":
|
|
// Handle content delta (text, thinking, or tool input)
|
|
if delta := root.Get("delta"); delta.Exists() {
|
|
deltaType := delta.Get("type").String()
|
|
switch deltaType {
|
|
case "text_delta":
|
|
// Process regular text content
|
|
if text := delta.Get("text"); text.Exists() && text.String() != "" {
|
|
partJSON := `{"text":""}`
|
|
partJSON, _ = sjson.Set(partJSON, "text", text.String())
|
|
part := gjson.Parse(partJSON).Value().(map[string]interface{})
|
|
allParts = append(allParts, part)
|
|
}
|
|
case "thinking_delta":
|
|
// Process reasoning/thinking content
|
|
if text := delta.Get("thinking"); text.Exists() && text.String() != "" {
|
|
partJSON := `{"thought":true,"text":""}`
|
|
partJSON, _ = sjson.Set(partJSON, "text", text.String())
|
|
part := gjson.Parse(partJSON).Value().(map[string]interface{})
|
|
allParts = append(allParts, part)
|
|
}
|
|
case "input_json_delta":
|
|
// accumulate args partial_json for this index
|
|
idx := int(root.Get("index").Int())
|
|
if newParam.ToolUseArgs == nil {
|
|
newParam.ToolUseArgs = map[int]*strings.Builder{}
|
|
}
|
|
if _, ok := newParam.ToolUseArgs[idx]; !ok || newParam.ToolUseArgs[idx] == nil {
|
|
newParam.ToolUseArgs[idx] = &strings.Builder{}
|
|
}
|
|
if pj := delta.Get("partial_json"); pj.Exists() {
|
|
newParam.ToolUseArgs[idx].WriteString(pj.String())
|
|
}
|
|
}
|
|
}
|
|
|
|
case "content_block_stop":
|
|
// Handle tool use completion by assembling accumulated arguments
|
|
idx := int(root.Get("index").Int())
|
|
// Claude's content_block_stop often doesn't include content_block payload (see docs/response-claude.txt)
|
|
// So we finalize using accumulated state captured during content_block_start and input_json_delta.
|
|
name := ""
|
|
if newParam.ToolUseNames != nil {
|
|
name = newParam.ToolUseNames[idx]
|
|
}
|
|
var argsTrim string
|
|
if newParam.ToolUseArgs != nil {
|
|
if b := newParam.ToolUseArgs[idx]; b != nil {
|
|
argsTrim = strings.TrimSpace(b.String())
|
|
}
|
|
}
|
|
if name != "" || argsTrim != "" {
|
|
functionCallJSON := `{"functionCall":{"name":"","args":{}}}`
|
|
if name != "" {
|
|
functionCallJSON, _ = sjson.Set(functionCallJSON, "functionCall.name", name)
|
|
}
|
|
if argsTrim != "" {
|
|
functionCallJSON, _ = sjson.SetRaw(functionCallJSON, "functionCall.args", argsTrim)
|
|
}
|
|
// Parse back to interface{} for allParts
|
|
functionCall := gjson.Parse(functionCallJSON).Value().(map[string]interface{})
|
|
allParts = append(allParts, functionCall)
|
|
// cleanup used state for this index
|
|
if newParam.ToolUseArgs != nil {
|
|
delete(newParam.ToolUseArgs, idx)
|
|
}
|
|
if newParam.ToolUseNames != nil {
|
|
delete(newParam.ToolUseNames, idx)
|
|
}
|
|
}
|
|
|
|
case "message_delta":
|
|
// Extract final usage information using sjson for token counts and metadata
|
|
if usage := root.Get("usage"); usage.Exists() {
|
|
usageJSON := `{}`
|
|
|
|
// Basic token counts for prompt and completion
|
|
inputTokens := usage.Get("input_tokens").Int()
|
|
outputTokens := usage.Get("output_tokens").Int()
|
|
|
|
// Set basic usage metadata according to Gemini API specification
|
|
usageJSON, _ = sjson.Set(usageJSON, "promptTokenCount", inputTokens)
|
|
usageJSON, _ = sjson.Set(usageJSON, "candidatesTokenCount", outputTokens)
|
|
usageJSON, _ = sjson.Set(usageJSON, "totalTokenCount", inputTokens+outputTokens)
|
|
|
|
// Add cache-related token counts if present (Claude Code API cache fields)
|
|
if cacheCreationTokens := usage.Get("cache_creation_input_tokens"); cacheCreationTokens.Exists() {
|
|
usageJSON, _ = sjson.Set(usageJSON, "cachedContentTokenCount", cacheCreationTokens.Int())
|
|
}
|
|
if cacheReadTokens := usage.Get("cache_read_input_tokens"); cacheReadTokens.Exists() {
|
|
// Add cache read tokens to cached content count
|
|
existingCacheTokens := usage.Get("cache_creation_input_tokens").Int()
|
|
totalCacheTokens := existingCacheTokens + cacheReadTokens.Int()
|
|
usageJSON, _ = sjson.Set(usageJSON, "cachedContentTokenCount", totalCacheTokens)
|
|
}
|
|
|
|
// Add thinking tokens if present (for models with reasoning capabilities)
|
|
if thinkingTokens := usage.Get("thinking_tokens"); thinkingTokens.Exists() {
|
|
usageJSON, _ = sjson.Set(usageJSON, "thoughtsTokenCount", thinkingTokens.Int())
|
|
}
|
|
|
|
// Set traffic type (required by Gemini API)
|
|
usageJSON, _ = sjson.Set(usageJSON, "trafficType", "PROVISIONED_THROUGHPUT")
|
|
|
|
// Convert to map[string]interface{} using gjson
|
|
finalUsage = gjson.Parse(usageJSON).Value().(map[string]interface{})
|
|
}
|
|
}
|
|
}
|
|
|
|
// Set response metadata
|
|
if responseID != "" {
|
|
template, _ = sjson.Set(template, "responseId", responseID)
|
|
}
|
|
if createdAt > 0 {
|
|
template, _ = sjson.Set(template, "createTime", time.Unix(createdAt, 0).Format(time.RFC3339Nano))
|
|
}
|
|
|
|
// Consolidate consecutive text parts and thinking parts for cleaner output
|
|
consolidatedParts := consolidateParts(allParts)
|
|
|
|
// Set the consolidated parts array
|
|
if len(consolidatedParts) > 0 {
|
|
template, _ = sjson.SetRaw(template, "candidates.0.content.parts", convertToJSONString(consolidatedParts))
|
|
}
|
|
|
|
// Set usage metadata
|
|
if finalUsage != nil {
|
|
template, _ = sjson.SetRaw(template, "usageMetadata", convertToJSONString(finalUsage))
|
|
}
|
|
|
|
return template
|
|
}
|
|
|
|
func GeminiTokenCount(ctx context.Context, count int64) string {
|
|
return fmt.Sprintf(`{"totalTokens":%d,"promptTokensDetails":[{"modality":"TEXT","tokenCount":%d}]}`, count, count)
|
|
}
|
|
|
|
// consolidateParts merges consecutive text parts and thinking parts to create a cleaner response.
|
|
// This function processes the parts array to combine adjacent text elements and thinking elements
|
|
// into single consolidated parts, which results in a more readable and efficient response structure.
|
|
// Tool calls and other non-text parts are preserved as separate elements.
|
|
func consolidateParts(parts []interface{}) []interface{} {
|
|
if len(parts) == 0 {
|
|
return parts
|
|
}
|
|
|
|
var consolidated []interface{}
|
|
var currentTextPart strings.Builder
|
|
var currentThoughtPart strings.Builder
|
|
var hasText, hasThought bool
|
|
|
|
flushText := func() {
|
|
// Flush accumulated text content to the consolidated parts array
|
|
if hasText && currentTextPart.Len() > 0 {
|
|
textPartJSON := `{"text":""}`
|
|
textPartJSON, _ = sjson.Set(textPartJSON, "text", currentTextPart.String())
|
|
textPart := gjson.Parse(textPartJSON).Value().(map[string]interface{})
|
|
consolidated = append(consolidated, textPart)
|
|
currentTextPart.Reset()
|
|
hasText = false
|
|
}
|
|
}
|
|
|
|
flushThought := func() {
|
|
// Flush accumulated thinking content to the consolidated parts array
|
|
if hasThought && currentThoughtPart.Len() > 0 {
|
|
thoughtPartJSON := `{"thought":true,"text":""}`
|
|
thoughtPartJSON, _ = sjson.Set(thoughtPartJSON, "text", currentThoughtPart.String())
|
|
thoughtPart := gjson.Parse(thoughtPartJSON).Value().(map[string]interface{})
|
|
consolidated = append(consolidated, thoughtPart)
|
|
currentThoughtPart.Reset()
|
|
hasThought = false
|
|
}
|
|
}
|
|
|
|
for _, part := range parts {
|
|
partMap, ok := part.(map[string]interface{})
|
|
if !ok {
|
|
// Flush any pending parts and add this non-text part
|
|
flushText()
|
|
flushThought()
|
|
consolidated = append(consolidated, part)
|
|
continue
|
|
}
|
|
|
|
if thought, isThought := partMap["thought"]; isThought && thought == true {
|
|
// This is a thinking part - flush any pending text first
|
|
flushText() // Flush any pending text first
|
|
|
|
if text, hasTextContent := partMap["text"].(string); hasTextContent {
|
|
currentThoughtPart.WriteString(text)
|
|
hasThought = true
|
|
}
|
|
} else if text, hasTextContent := partMap["text"].(string); hasTextContent {
|
|
// This is a regular text part - flush any pending thought first
|
|
flushThought() // Flush any pending thought first
|
|
|
|
currentTextPart.WriteString(text)
|
|
hasText = true
|
|
} else {
|
|
// This is some other type of part (like function call) - flush both text and thought
|
|
flushText()
|
|
flushThought()
|
|
consolidated = append(consolidated, part)
|
|
}
|
|
}
|
|
|
|
// Flush any remaining parts
|
|
flushThought() // Flush thought first to maintain order
|
|
flushText()
|
|
|
|
return consolidated
|
|
}
|
|
|
|
// convertToJSONString converts interface{} to JSON string using sjson/gjson.
|
|
// This function provides a consistent way to serialize different data types to JSON strings
|
|
// for inclusion in the Gemini API response structure.
|
|
func convertToJSONString(v interface{}) string {
|
|
switch val := v.(type) {
|
|
case []interface{}:
|
|
return convertArrayToJSON(val)
|
|
case map[string]interface{}:
|
|
return convertMapToJSON(val)
|
|
default:
|
|
// For simple types, create a temporary JSON and extract the value
|
|
temp := `{"temp":null}`
|
|
temp, _ = sjson.Set(temp, "temp", val)
|
|
return gjson.Get(temp, "temp").Raw
|
|
}
|
|
}
|