@@ -14,6 +14,7 @@ import (
1414 "github.com/kagent-dev/kagent/go/adk/pkg/telemetry"
1515 "github.com/openai/openai-go/v3"
1616 "github.com/openai/openai-go/v3/packages/param"
17+ "github.com/openai/openai-go/v3/packages/respjson"
1718 "github.com/openai/openai-go/v3/shared"
1819 "github.com/openai/openai-go/v3/shared/constant"
1920 "google.golang.org/adk/model"
@@ -28,6 +29,7 @@ const (
2829 openAIFinishLength = "length"
2930 openAIFinishContentFilter = "content_filter"
3031 openAIToolTypeFunction = "function"
32+ openAIExtraContentKey = "extra_content"
3133)
3234
3335// openAIFinishReasonToGenai maps OpenAI finish_reason to genai.FinishReason.
@@ -42,6 +44,82 @@ func openAIFinishReasonToGenai(reason string) genai.FinishReason {
4244 }
4345}
4446
47+ type openAIThoughtSignatureExtra struct {
48+ Google struct {
49+ ThoughtSignature string `json:"thought_signature"`
50+ } `json:"google"`
51+ }
52+
53+ func extractThoughtSignatureFromRaw (raw string ) []byte {
54+ if raw == "" {
55+ return nil
56+ }
57+
58+ var extra openAIThoughtSignatureExtra
59+ if err := json .Unmarshal ([]byte (raw ), & extra ); err != nil {
60+ return nil
61+ }
62+ if extra .Google .ThoughtSignature == "" {
63+ return nil
64+ }
65+
66+ decoded , err := base64 .StdEncoding .DecodeString (extra .Google .ThoughtSignature )
67+ if err != nil {
68+ return nil
69+ }
70+ return decoded
71+ }
72+
73+ func extractThoughtSignatureFromExtraFields (extraFields map [string ]respjson.Field ) []byte {
74+ if len (extraFields ) == 0 {
75+ return nil
76+ }
77+ field , ok := extraFields [openAIExtraContentKey ]
78+ if ! ok {
79+ return nil
80+ }
81+ return extractThoughtSignatureFromRaw (field .Raw ())
82+ }
83+
84+ func openAIExtraContentForThoughtSignature (thoughtSignature []byte ) map [string ]any {
85+ if len (thoughtSignature ) == 0 {
86+ return nil
87+ }
88+
89+ return map [string ]any {
90+ "google" : map [string ]any {
91+ "thought_signature" : base64 .StdEncoding .EncodeToString (thoughtSignature ),
92+ },
93+ }
94+ }
95+
96+ func thoughtSignaturesByToolCallID (contents []* genai.Content ) map [string ][]byte {
97+ thoughtSignatures := make (map [string ][]byte )
98+ for _ , content := range contents {
99+ if content == nil || content .Parts == nil {
100+ continue
101+ }
102+ for _ , part := range content .Parts {
103+ if part == nil || part .FunctionCall == nil || len (part .ThoughtSignature ) == 0 {
104+ continue
105+ }
106+ thoughtSignatures [part .FunctionCall .ID ] = part .ThoughtSignature
107+ }
108+ }
109+ return thoughtSignatures
110+ }
111+
112+ func newFunctionCallPart (name string , args map [string ]any , id string , thoughtSignature []byte ) * genai.Part {
113+ part := genai .NewPartFromFunctionCall (name , args )
114+ if part .FunctionCall != nil {
115+ part .FunctionCall .ID = id
116+ }
117+ if len (thoughtSignature ) > 0 {
118+ part .ThoughtSignature = thoughtSignature
119+ }
120+ return part
121+ }
122+
45123// Name implements model.LLM.
46124func (m * OpenAIModel ) Name () string {
47125 return m .Config .Model
@@ -126,6 +204,7 @@ func genaiContentsToOpenAIMessages(contents []*genai.Content, config *genai.Gene
126204 systemInstruction := strings .TrimSpace (systemBuilder .String ())
127205
128206 functionResponses := make (map [string ]* genai.FunctionResponse )
207+ thoughtSignatures := thoughtSignaturesByToolCallID (contents )
129208 for _ , c := range contents {
130209 if c == nil || c .Parts == nil {
131210 continue
@@ -167,21 +246,35 @@ func genaiContentsToOpenAIMessages(contents []*genai.Content, config *genai.Gene
167246 var toolResponseMessages []openai.ChatCompletionMessageParamUnion
168247 for _ , fc := range functionCalls {
169248 argsJSON , _ := json .Marshal (fc .Args )
170- toolCalls = append (toolCalls , openai.ChatCompletionMessageToolCallUnionParam {
171- OfFunction : & openai.ChatCompletionMessageFunctionToolCallParam {
172- ID : fc .ID ,
173- Type : constant .Function (openAIToolTypeFunction ),
174- Function : openai.ChatCompletionMessageFunctionToolCallFunctionParam {
175- Name : fc .Name ,
176- Arguments : string (argsJSON ),
177- },
249+ toolCall := openai.ChatCompletionMessageFunctionToolCallParam {
250+ ID : fc .ID ,
251+ Type : constant .Function (openAIToolTypeFunction ),
252+ Function : openai.ChatCompletionMessageFunctionToolCallFunctionParam {
253+ Name : fc .Name ,
254+ Arguments : string (argsJSON ),
178255 },
256+ }
257+ if extraContent := openAIExtraContentForThoughtSignature (thoughtSignatures [fc .ID ]); extraContent != nil {
258+ toolCall .SetExtraFields (map [string ]any {openAIExtraContentKey : extraContent })
259+ }
260+ toolCalls = append (toolCalls , openai.ChatCompletionMessageToolCallUnionParam {
261+ OfFunction : & toolCall ,
179262 })
180263 contentStr := "No response available for this function call."
181264 if fr := functionResponses [fc .ID ]; fr != nil {
182265 contentStr = functionResponseContentString (fr .Response )
183266 }
184- toolResponseMessages = append (toolResponseMessages , openai .ToolMessage (contentStr , fc .ID ))
267+ toolMessage := openai.ChatCompletionToolMessageParam {
268+ Content : openai.ChatCompletionToolMessageParamContentUnion {
269+ OfString : param .NewOpt (contentStr ),
270+ },
271+ ToolCallID : fc .ID ,
272+ Role : constant .Tool ("tool" ),
273+ }
274+ if extraContent := openAIExtraContentForThoughtSignature (thoughtSignatures [fc .ID ]); extraContent != nil {
275+ toolMessage .SetExtraFields (map [string ]any {openAIExtraContentKey : extraContent })
276+ }
277+ toolResponseMessages = append (toolResponseMessages , openai.ChatCompletionMessageParamUnion {OfTool : & toolMessage })
185278 }
186279 textContent := strings .Join (textParts , "\n " )
187280 asst := openai.ChatCompletionAssistantMessageParam {
@@ -357,7 +450,7 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet
357450 for _ , tc := range delta .ToolCalls {
358451 idx := tc .Index
359452 if toolCallsAcc [idx ] == nil {
360- toolCallsAcc [idx ] = map [string ]any {"id" : "" , "name" : "" , "arguments" : "" }
453+ toolCallsAcc [idx ] = map [string ]any {"id" : "" , "name" : "" , "arguments" : "" , "thought_signature" : [] byte ( nil ) }
361454 }
362455 if tc .ID != "" {
363456 toolCallsAcc [idx ]["id" ] = tc .ID
@@ -369,6 +462,9 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet
369462 prev , _ := toolCallsAcc [idx ]["arguments" ].(string )
370463 toolCallsAcc [idx ]["arguments" ] = prev + tc .Function .Arguments
371464 }
465+ if thoughtSignature := extractThoughtSignatureFromExtraFields (tc .JSON .ExtraFields ); len (thoughtSignature ) > 0 {
466+ toolCallsAcc [idx ]["thought_signature" ] = thoughtSignature
467+ }
372468 }
373469 if choice .FinishReason != "" {
374470 finishReason = choice .FinishReason
@@ -404,8 +500,8 @@ func runStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComplet
404500 name , _ := tc ["name" ].(string )
405501 id , _ := tc ["id" ].(string )
406502 if name != "" || id != "" {
407- p := genai . NewPartFromFunctionCall ( name , args )
408- p . FunctionCall . ID = id
503+ thoughtSignature , _ := tc [ "thought_signature" ].([] byte )
504+ p := newFunctionCallPart ( name , args , id , thoughtSignature )
409505 finalParts = append (finalParts , p )
410506 }
411507 }
@@ -438,6 +534,12 @@ func runNonStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComp
438534 yield (& model.LLMResponse {ErrorCode : "API_ERROR" , ErrorMessage : "No choices in response" }, nil )
439535 return
440536 }
537+ resp := chatCompletionToLLMResponse (completion )
538+ telemetry .SetLLMResponseAttributes (ctx , resp )
539+ yield (resp , nil )
540+ }
541+
542+ func chatCompletionToLLMResponse (completion * openai.ChatCompletion ) * model.LLMResponse {
441543 choice := completion .Choices [0 ]
442544 msg := choice .Message
443545 nParts := 0
@@ -455,8 +557,13 @@ func runNonStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComp
455557 if tc .Function .Arguments != "" {
456558 _ = json .Unmarshal ([]byte (tc .Function .Arguments ), & args )
457559 }
458- p := genai .NewPartFromFunctionCall (tc .Function .Name , args )
459- p .FunctionCall .ID = tc .ID
560+ functionToolCall := tc .AsFunction ()
561+ p := newFunctionCallPart (
562+ tc .Function .Name ,
563+ args ,
564+ tc .ID ,
565+ extractThoughtSignatureFromExtraFields (functionToolCall .JSON .ExtraFields ),
566+ )
460567 parts = append (parts , p )
461568 }
462569 }
@@ -467,13 +574,11 @@ func runNonStreaming(ctx context.Context, m *OpenAIModel, params openai.ChatComp
467574 CandidatesTokenCount : int32 (completion .Usage .CompletionTokens ),
468575 }
469576 }
470- resp := & model.LLMResponse {
577+ return & model.LLMResponse {
471578 Partial : false ,
472579 TurnComplete : true ,
473580 FinishReason : openAIFinishReasonToGenai (choice .FinishReason ),
474581 UsageMetadata : usage ,
475582 Content : & genai.Content {Role : string (genai .RoleModel ), Parts : parts },
476583 }
477- telemetry .SetLLMResponseAttributes (ctx , resp )
478- yield (resp , nil )
479584}
0 commit comments