Skip to content

Commit 0bf8984

Browse files
authored
Merge branch 'main' into fix/ui
2 parents 51953dc + 1bf4617 commit 0bf8984

File tree

21 files changed

+1711
-455
lines changed

21 files changed

+1711
-455
lines changed

.github/workflows/ci.yaml

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ jobs:
5656
driver-opts: network=host
5757

5858
- name: Set up Helm
59-
uses: azure/setup-helm@v4.2.0
59+
uses: azure/setup-helm@v5.0.0
6060
with:
6161
version: v3.18.0
6262

@@ -164,7 +164,7 @@ jobs:
164164
uses: actions/checkout@v6
165165

166166
- name: Set up Helm
167-
uses: azure/setup-helm@v4.2.0
167+
uses: azure/setup-helm@v5.0.0
168168
with:
169169
version: v3.18.0
170170
# Install unittest plugin
@@ -229,7 +229,7 @@ jobs:
229229
- name: Publish to Chromatic
230230
# Requires repo secret CHROMATIC_PROJECT_TOKEN. Skipped when unset, or on fork PRs (no token access).
231231
if: ${{ env.CHROMATIC_PROJECT_TOKEN != '' && (github.event.pull_request == null || github.event.pull_request.head.repo.full_name == github.repository) }}
232-
uses: chromaui/action@v16.0.0
232+
uses: chromaui/action@v16.1.0
233233
with:
234234
projectToken: ${{ env.CHROMATIC_PROJECT_TOKEN }}
235235
workingDir: ui
@@ -314,7 +314,7 @@ jobs:
314314
uses: actions/checkout@v6
315315

316316
- name: Install uv
317-
uses: astral-sh/setup-uv@v5
317+
uses: astral-sh/setup-uv@v7
318318

319319
- name: Install python
320320
run: uv python install ${{ matrix.python-version }}
@@ -338,7 +338,7 @@ jobs:
338338
uses: actions/checkout@v6
339339

340340
- name: Install uv
341-
uses: astral-sh/setup-uv@v5
341+
uses: astral-sh/setup-uv@v7
342342

343343
- name: Install python
344344
run: uv python install 3.10

.github/workflows/tag.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ jobs:
9595
- name: 'Checkout GitHub Action'
9696
uses: actions/checkout@main
9797
- name: Install uv
98-
uses: astral-sh/setup-uv@v6
98+
uses: astral-sh/setup-uv@v7
9999
- name: 'Release Python Packages'
100100
working-directory: python
101101
run: |

docker/skills-init/Dockerfile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,5 +14,5 @@ RUN CGO_ENABLED=0 go build -trimpath -ldflags="-s -w" -o /build/krane .
1414

1515
FROM alpine:3.23
1616

17-
RUN apk add --no-cache git
17+
RUN apk upgrade --no-cache && apk add --no-cache git
1818
COPY --from=krane-builder /build/krane /usr/local/bin/krane

go/adk/pkg/models/openai_adk.go

Lines changed: 122 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -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.
46124
func (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

Comments
 (0)