feat: Add thought_signature support for Gemini tool calls, including a compatibility function for old sessions and refactoring OpenAI types.

This commit is contained in:
viettranx committed 2026-02-27 18:46:01 +07:00
1 parent f3d2baa0bc
commit 0844969819
4 files changed
+196 -79

No files matched your search

+32 -79
View File
@@ -142,6 +142,9 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun
acc.Name = tc.Function.Name
}
acc.rawArgs += tc.Function.Arguments
if tc.Function.ThoughtSignature != "" {
acc.thoughtSig = tc.Function.ThoughtSignature
}
}
if chunk.Choices[0].FinishReason != "" {
@@ -167,6 +170,9 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun
args := make(map[string]interface{})
_ = json.Unmarshal([]byte(acc.rawArgs), &args)
acc.Arguments = args
if acc.thoughtSig != "" {
acc.Metadata = map[string]string{"thought_signature": acc.thoughtSig}
}
result.ToolCalls = append(result.ToolCalls, acc.ToolCall)
}
@@ -182,12 +188,20 @@ func (p *OpenAIProvider) ChatStream(ctx context.Context, req ChatRequest, onChun
}
func (p *OpenAIProvider) buildRequestBody(model string, req ChatRequest, stream bool) map[string]interface{} {
// Gemini 2.5+: collapse old tool_call cycles missing thought_signature into plain text.
// Old session messages stored before the thought_signature fix don't have it,
// and Gemini rejects them with HTTP 400. Collapse preserves context as text.
inputMessages := req.Messages
if strings.Contains(strings.ToLower(p.name), "gemini") {
inputMessages = collapseToolCallsWithoutSig(inputMessages)
}
// Convert messages to proper OpenAI wire format.
// This is necessary because our internal Message/ToolCall structs don't match
// the OpenAI API format (tool_calls need type+function wrapper, arguments as JSON string).
// Also omits empty content on assistant messages with tool_calls (Gemini compatibility).
msgs := make([]map[string]interface{}, 0, len(req.Messages))
for _, m := range req.Messages {
msgs := make([]map[string]interface{}, 0, len(inputMessages))
for _, m := range inputMessages {
msg := map[string]interface{}{
"role": m.Role,
}
@@ -221,13 +235,17 @@ func (p *OpenAIProvider) buildRequestBody(model string, req ChatRequest, stream
toolCalls := make([]map[string]interface{}, len(m.ToolCalls))
for i, tc := range m.ToolCalls {
argsJSON, _ := json.Marshal(tc.Arguments)
fn := map[string]interface{}{
"name": tc.Name,
"arguments": string(argsJSON),
}
if sig := tc.Metadata["thought_signature"]; sig != "" {
fn["thought_signature"] = sig
}
toolCalls[i] = map[string]interface{}{
"id": tc.ID,
"type": "function",
"function": map[string]interface{}{
"name": tc.Name,
"arguments": string(argsJSON),
},
"id": tc.ID,
"type": "function",
"function": fn,
}
}
msg["tool_calls"] = toolCalls
@@ -312,11 +330,15 @@ func (p *OpenAIProvider) parseResponse(resp *openAIResponse) *ChatResponse {
for _, tc := range msg.ToolCalls {
args := make(map[string]interface{})
_ = json.Unmarshal([]byte(tc.Function.Arguments), &args)
result.ToolCalls = append(result.ToolCalls, ToolCall{
call := ToolCall{
ID: tc.ID,
Name: tc.Function.Name,
Arguments: args,
})
}
if tc.Function.ThoughtSignature != "" {
call.Metadata = map[string]string{"thought_signature": tc.Function.ThoughtSignature}
}
result.ToolCalls = append(result.ToolCalls, call)
}
if len(result.ToolCalls) > 0 {
@@ -338,72 +360,3 @@ func (p *OpenAIProvider) parseResponse(resp *openAIResponse) *ChatResponse {
return result
}
// OpenAI API response types (internal)
type openAIResponse struct {
Choices []openAIChoice `json:"choices"`
Usage *openAIUsage `json:"usage,omitempty"`
}
type openAIChoice struct {
Message openAIMessage `json:"message"`
FinishReason string `json:"finish_reason"`
}
type openAIMessage struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []openAIToolCall `json:"tool_calls,omitempty"`
}
type openAIToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function openAIFunctionCall `json:"function"`
}
type openAIFunctionCall struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
}
type openAIUsage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
PromptTokensDetails *openAIPromptDetails `json:"prompt_tokens_details,omitempty"`
}
type openAIPromptDetails struct {
CachedTokens int `json:"cached_tokens"`
}
// Streaming types
type openAIStreamChunk struct {
Choices []openAIStreamChoice `json:"choices"`
Usage *openAIUsage `json:"usage,omitempty"`
}
type openAIStreamChoice struct {
Delta openAIStreamDelta `json:"delta"`
FinishReason string `json:"finish_reason,omitempty"`
}
type openAIStreamDelta struct {
Content string `json:"content,omitempty"`
ToolCalls []openAIStreamToolCall `json:"tool_calls,omitempty"`
}
type openAIStreamToolCall struct {
Index int `json:"index"`
ID string `json:"id,omitempty"`
Function openAIFunctionCall `json:"function"`
}
// rawArgs is a temporary field for accumulating streamed arguments.
// We extend ToolCall with it during streaming only.
type toolCallAccumulator struct {
ToolCall
rawArgs string
}
+90
View File
@@ -0,0 +1,90 @@
package providers
import (
"encoding/json"
"fmt"
"strings"
)
// collapseToolCallsWithoutSig rewrites assistant messages whose tool_calls lack
// thought_signature (required by Gemini 2.5+) into plain text. This handles old
// session history stored before the thought_signature capture fix.
// Tool call cycles (assistant + tool results) are collapsed into:
// - assistant: original content + "[Called tool_name({args})]"
// - user: "[tool_name result]: content"
func collapseToolCallsWithoutSig(msgs []Message) []Message {
// Collect tool_call IDs that need collapsing.
collapseIDs := make(map[string]bool)
for _, m := range msgs {
if m.Role != "assistant" || len(m.ToolCalls) == 0 {
continue
}
for _, tc := range m.ToolCalls {
if tc.Metadata["thought_signature"] == "" {
// If any tool_call in the message is missing sig, collapse all of them.
for _, tc2 := range m.ToolCalls {
collapseIDs[tc2.ID] = true
}
break
}
}
}
if len(collapseIDs) == 0 {
return msgs
}
result := make([]Message, 0, len(msgs))
for i := 0; i < len(msgs); i++ {
m := msgs[i]
// Collapse assistant with tool_calls → plain text assistant
if m.Role == "assistant" && len(m.ToolCalls) > 0 && collapseIDs[m.ToolCalls[0].ID] {
var sb strings.Builder
if m.Content != "" {
sb.WriteString(m.Content)
sb.WriteString("\n\n")
}
for _, tc := range m.ToolCalls {
argsJSON, _ := json.Marshal(tc.Arguments)
fmt.Fprintf(&sb, "[Called %s(%s)]\n", tc.Name, string(argsJSON))
}
result = append(result, Message{
Role: "assistant",
Content: strings.TrimSpace(sb.String()),
})
// Collect consecutive tool results → one user message
var toolBuf strings.Builder
for i+1 < len(msgs) && msgs[i+1].Role == "tool" && collapseIDs[msgs[i+1].ToolCallID] {
i++
toolMsg := msgs[i]
toolName := toolMsg.ToolCallID
for _, tc := range m.ToolCalls {
if tc.ID == toolMsg.ToolCallID {
toolName = tc.Name
break
}
}
if toolBuf.Len() > 0 {
toolBuf.WriteString("\n\n")
}
fmt.Fprintf(&toolBuf, "[%s result]: %s", toolName, toolMsg.Content)
}
if toolBuf.Len() > 0 {
result = append(result, Message{
Role: "user",
Content: toolBuf.String(),
})
}
continue
}
// Skip orphaned tool results whose assistant was already collapsed
if m.Role == "tool" && collapseIDs[m.ToolCallID] {
continue
}
result = append(result, m)
}
return result
}
+73
View File
@@ -0,0 +1,73 @@
package providers
// OpenAI API response types (internal)
type openAIResponse struct {
Choices []openAIChoice `json:"choices"`
Usage *openAIUsage `json:"usage,omitempty"`
}
type openAIChoice struct {
Message openAIMessage `json:"message"`
FinishReason string `json:"finish_reason"`
}
type openAIMessage struct {
Role string `json:"role"`
Content string `json:"content"`
ToolCalls []openAIToolCall `json:"tool_calls,omitempty"`
}
type openAIToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function openAIFunctionCall `json:"function"`
}
type openAIFunctionCall struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
ThoughtSignature string `json:"thought_signature,omitempty"` // Gemini 2.5/3: must echo back
}
type openAIUsage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
PromptTokensDetails *openAIPromptDetails `json:"prompt_tokens_details,omitempty"`
}
type openAIPromptDetails struct {
CachedTokens int `json:"cached_tokens"`
}
// Streaming types
type openAIStreamChunk struct {
Choices []openAIStreamChoice `json:"choices"`
Usage *openAIUsage `json:"usage,omitempty"`
}
type openAIStreamChoice struct {
Delta openAIStreamDelta `json:"delta"`
FinishReason string `json:"finish_reason,omitempty"`
}
type openAIStreamDelta struct {
Content string `json:"content,omitempty"`
ToolCalls []openAIStreamToolCall `json:"tool_calls,omitempty"`
}
type openAIStreamToolCall struct {
Index int `json:"index"`
ID string `json:"id,omitempty"`
Function openAIFunctionCall `json:"function"`
}
// toolCallAccumulator extends ToolCall with temporary fields for accumulating
// streamed arguments and thought_signature during SSE streaming.
type toolCallAccumulator struct {
ToolCall
rawArgs string
thoughtSig string
}
+1
View File
@@ -62,6 +62,7 @@ type ToolCall struct {
ID string `json:"id"`
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments"`
Metadata map[string]string `json:"metadata,omitempty"` // provider-specific (e.g. Gemini thought_signature)
}
// ToolDefinition describes a tool available to the LLM.