Skip to content

Commit dfeb60d

Browse files
author
syd
committed
feat(message): add tool_calls/tool_callbacks for easy marshal
1 parent 1b5e4a7 commit dfeb60d

3 files changed

Lines changed: 382 additions & 9 deletions

File tree

‎agents/agent.go‎

Lines changed: 232 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package agents
22

33
import (
44
"context"
5+
"encoding/json"
56
"errors"
67

78
"github.com/bububa/instructor-go"
@@ -255,17 +256,34 @@ func (a *Agent[I, O]) chat(ctx context.Context, userInput *I, response *O, llmRe
255256
}
256257
if memory := clt.Memory(); memory != nil {
257258
for _, v := range memory.List() {
259+
msg := new(components.Message)
258260
var role components.MessageRole
259261
if v.OfAssistant != nil {
260262
role = components.AssistantRole
263+
if toolCalls := v.OfAssistant.ToolCalls; len(toolCalls) > 0 {
264+
calls := make([]components.ToolCall, 0, len(toolCalls))
265+
for _, tool := range toolCalls {
266+
calls = append(calls, components.ToolCall{
267+
ID: tool.ID,
268+
Name: tool.Function.Name,
269+
Arguments: tool.Function.Arguments,
270+
})
271+
}
272+
msg.SetToolCalls(calls)
273+
}
261274
} else if v.OfUser != nil {
262275
role = components.UserRole
263276
} else if v.OfTool != nil {
264277
role = components.ToolRole
278+
call := components.ToolCallback{
279+
ID: v.OfTool.ToolCallID,
280+
Content: v.OfTool.Content.OfString.String(),
281+
}
282+
msg.SetToolCallbacks([]components.ToolCallback{call})
265283
} else if v.OfFunction != nil {
266284
role = components.FunctionRole
267285
}
268-
msg := components.NewMessage(role, nil)
286+
msg.SetRole(role)
269287
msg.SetRaw(v)
270288
a.memory.AddMessage(msg)
271289
}
@@ -306,13 +324,41 @@ func (a *Agent[I, O]) chat(ctx context.Context, userInput *I, response *O, llmRe
306324
if memory := clt.Memory(); memory != nil {
307325
for _, v := range memory.List() {
308326
var role components.MessageRole
327+
msg := new(components.Message)
309328
switch v.Role {
310329
case anthropic.RoleAssistant:
311330
role = components.AssistantRole
331+
var toolCalls []components.ToolCall
332+
for _, content := range v.Content {
333+
if toolUse := content.MessageContentToolUse; toolUse != nil {
334+
bs, _ := json.Marshal(toolUse.Input)
335+
toolCalls = append(toolCalls, components.ToolCall{
336+
ID: toolUse.ID,
337+
Name: toolUse.Name,
338+
Arguments: string(bs),
339+
})
340+
}
341+
}
342+
if len(toolCalls) > 0 {
343+
msg.SetToolCalls(toolCalls)
344+
}
312345
case anthropic.RoleUser:
313346
role = components.UserRole
347+
var calls []components.ToolCallback
348+
for _, content := range v.Content {
349+
if toolResult := content.MessageContentToolResult; toolResult != nil {
350+
calls = append(calls, components.ToolCallback{
351+
ID: *toolResult.ToolUseID,
352+
Content: *toolResult.Content[0].Text,
353+
IsError: *toolResult.IsError,
354+
})
355+
}
356+
}
357+
if len(calls) > 0 {
358+
msg.SetToolCallbacks(calls)
359+
}
314360
}
315-
msg := components.NewMessage(role, nil)
361+
msg.SetRole(role)
316362
msg.SetRaw(v)
317363
a.memory.AddMessage(msg)
318364
}
@@ -397,13 +443,42 @@ func (a *Agent[I, O]) chat(ctx context.Context, userInput *I, response *O, llmRe
397443
if memory := clt.Memory(); memory != nil {
398444
for _, v := range memory.List() {
399445
var role components.MessageRole
446+
msg := new(components.Message)
400447
switch v.Role {
401448
case geminiAPI.RoleModel:
402449
role = components.AssistantRole
450+
var tools []components.ToolCall
451+
for _, part := range v.Parts {
452+
if call := part.FunctionCall; call != nil {
453+
bs, _ := json.Marshal(call.Args)
454+
tools = append(tools, components.ToolCall{
455+
ID: call.ID,
456+
Name: call.Name,
457+
Arguments: string(bs),
458+
})
459+
}
460+
}
461+
if len(tools) > 0 {
462+
msg.SetToolCalls(tools)
463+
}
403464
case geminiAPI.RoleUser:
404465
role = components.UserRole
466+
var calls []components.ToolCallback
467+
for _, part := range v.Parts {
468+
if fn := part.FunctionResponse; fn != nil {
469+
bs, _ := json.Marshal(fn.Response)
470+
calls = append(calls, components.ToolCallback{
471+
ID: fn.ID,
472+
Name: fn.Name,
473+
Content: string(bs),
474+
})
475+
}
476+
}
477+
if len(calls) > 0 {
478+
msg.SetToolCallbacks(calls)
479+
}
405480
}
406-
msg := components.NewMessage(role, nil)
481+
msg.SetRole(role)
407482
msg.SetRaw(v)
408483
a.memory.AddMessage(msg)
409484
}
@@ -486,17 +561,34 @@ func (a *Agent[I, O]) stream(ctx context.Context, userInput *I) (<-chan instruct
486561
}
487562
if memory := clt.Memory(); memory != nil {
488563
for _, v := range memory.List() {
564+
msg := new(components.Message)
489565
var role components.MessageRole
490566
if v.OfAssistant != nil {
491567
role = components.AssistantRole
568+
if toolCalls := v.OfAssistant.ToolCalls; len(toolCalls) > 0 {
569+
calls := make([]components.ToolCall, 0, len(toolCalls))
570+
for _, tool := range toolCalls {
571+
calls = append(calls, components.ToolCall{
572+
ID: tool.ID,
573+
Name: tool.Function.Name,
574+
Arguments: tool.Function.Arguments,
575+
})
576+
}
577+
msg.SetToolCalls(calls)
578+
}
492579
} else if v.OfUser != nil {
493580
role = components.UserRole
494581
} else if v.OfTool != nil {
495582
role = components.ToolRole
583+
call := components.ToolCallback{
584+
ID: v.OfTool.ToolCallID,
585+
Content: v.OfTool.Content.OfString.String(),
586+
}
587+
msg.SetToolCallbacks([]components.ToolCallback{call})
496588
} else if v.OfFunction != nil {
497589
role = components.FunctionRole
498590
}
499-
msg := components.NewMessage(role, nil)
591+
msg.SetRole(role)
500592
msg.SetRaw(v)
501593
a.memory.AddMessage(msg)
502594
}
@@ -553,13 +645,41 @@ func (a *Agent[I, O]) stream(ctx context.Context, userInput *I) (<-chan instruct
553645
if memory := clt.Memory(); memory != nil {
554646
for _, v := range memory.List() {
555647
var role components.MessageRole
648+
msg := new(components.Message)
556649
switch v.Role {
557650
case anthropic.RoleAssistant:
558651
role = components.AssistantRole
652+
var toolCalls []components.ToolCall
653+
for _, content := range v.Content {
654+
if toolUse := content.MessageContentToolUse; toolUse != nil {
655+
bs, _ := json.Marshal(toolUse.Input)
656+
toolCalls = append(toolCalls, components.ToolCall{
657+
ID: toolUse.ID,
658+
Name: toolUse.Name,
659+
Arguments: string(bs),
660+
})
661+
}
662+
}
663+
if len(toolCalls) > 0 {
664+
msg.SetToolCalls(toolCalls)
665+
}
559666
case anthropic.RoleUser:
560667
role = components.UserRole
668+
var calls []components.ToolCallback
669+
for _, content := range v.Content {
670+
if toolResult := content.MessageContentToolResult; toolResult != nil {
671+
calls = append(calls, components.ToolCallback{
672+
ID: *toolResult.ToolUseID,
673+
Content: *toolResult.Content[0].Text,
674+
IsError: *toolResult.IsError,
675+
})
676+
}
677+
}
678+
if len(calls) > 0 {
679+
msg.SetToolCallbacks(calls)
680+
}
561681
}
562-
msg := components.NewMessage(role, nil)
682+
msg.SetRole(role)
563683
msg.SetRaw(v)
564684
a.memory.AddMessage(msg)
565685
}
@@ -669,13 +789,42 @@ func (a *Agent[I, O]) stream(ctx context.Context, userInput *I) (<-chan instruct
669789
if memory := clt.Memory(); memory != nil {
670790
for _, v := range memory.List() {
671791
var role components.MessageRole
792+
msg := new(components.Message)
672793
switch v.Role {
673794
case geminiAPI.RoleModel:
674795
role = components.AssistantRole
796+
var tools []components.ToolCall
797+
for _, part := range v.Parts {
798+
if call := part.FunctionCall; call != nil {
799+
bs, _ := json.Marshal(call.Args)
800+
tools = append(tools, components.ToolCall{
801+
ID: call.ID,
802+
Name: call.Name,
803+
Arguments: string(bs),
804+
})
805+
}
806+
}
807+
if len(tools) > 0 {
808+
msg.SetToolCalls(tools)
809+
}
675810
case geminiAPI.RoleUser:
676811
role = components.UserRole
812+
var calls []components.ToolCallback
813+
for _, part := range v.Parts {
814+
if fn := part.FunctionResponse; fn != nil {
815+
bs, _ := json.Marshal(fn.Response)
816+
calls = append(calls, components.ToolCallback{
817+
ID: fn.ID,
818+
Name: fn.Name,
819+
Content: string(bs),
820+
})
821+
}
822+
}
823+
if len(calls) > 0 {
824+
msg.SetToolCallbacks(calls)
825+
}
677826
}
678-
msg := components.NewMessage(role, nil)
827+
msg.SetRole(role)
679828
msg.SetRaw(v)
680829
a.memory.AddMessage(msg)
681830
}
@@ -779,17 +928,34 @@ func (a *Agent[I, O]) schemaStream(ctx context.Context, userInput *I) (<-chan an
779928
}
780929
if memory := clt.Memory(); memory != nil {
781930
for _, v := range memory.List() {
931+
msg := new(components.Message)
782932
var role components.MessageRole
783933
if v.OfAssistant != nil {
784934
role = components.AssistantRole
935+
if toolCalls := v.OfAssistant.ToolCalls; len(toolCalls) > 0 {
936+
calls := make([]components.ToolCall, 0, len(toolCalls))
937+
for _, tool := range toolCalls {
938+
calls = append(calls, components.ToolCall{
939+
ID: tool.ID,
940+
Name: tool.Function.Name,
941+
Arguments: tool.Function.Arguments,
942+
})
943+
}
944+
msg.SetToolCalls(calls)
945+
}
785946
} else if v.OfUser != nil {
786947
role = components.UserRole
787948
} else if v.OfTool != nil {
788949
role = components.ToolRole
950+
call := components.ToolCallback{
951+
ID: v.OfTool.ToolCallID,
952+
Content: v.OfTool.Content.OfString.String(),
953+
}
954+
msg.SetToolCallbacks([]components.ToolCallback{call})
789955
} else if v.OfFunction != nil {
790956
role = components.FunctionRole
791957
}
792-
msg := components.NewMessage(role, nil)
958+
msg.SetRole(role)
793959
msg.SetRaw(v)
794960
a.memory.AddMessage(msg)
795961
}
@@ -868,13 +1034,41 @@ func (a *Agent[I, O]) schemaStream(ctx context.Context, userInput *I) (<-chan an
8681034
if memory := clt.Memory(); memory != nil {
8691035
for _, v := range memory.List() {
8701036
var role components.MessageRole
1037+
msg := new(components.Message)
8711038
switch v.Role {
8721039
case anthropic.RoleAssistant:
8731040
role = components.AssistantRole
1041+
var toolCalls []components.ToolCall
1042+
for _, content := range v.Content {
1043+
if toolUse := content.MessageContentToolUse; toolUse != nil {
1044+
bs, _ := json.Marshal(toolUse.Input)
1045+
toolCalls = append(toolCalls, components.ToolCall{
1046+
ID: toolUse.ID,
1047+
Name: toolUse.Name,
1048+
Arguments: string(bs),
1049+
})
1050+
}
1051+
}
1052+
if len(toolCalls) > 0 {
1053+
msg.SetToolCalls(toolCalls)
1054+
}
8741055
case anthropic.RoleUser:
8751056
role = components.UserRole
1057+
var calls []components.ToolCallback
1058+
for _, content := range v.Content {
1059+
if toolResult := content.MessageContentToolResult; toolResult != nil {
1060+
calls = append(calls, components.ToolCallback{
1061+
ID: *toolResult.ToolUseID,
1062+
Content: *toolResult.Content[0].Text,
1063+
IsError: *toolResult.IsError,
1064+
})
1065+
}
1066+
}
1067+
if len(calls) > 0 {
1068+
msg.SetToolCallbacks(calls)
1069+
}
8761070
}
877-
msg := components.NewMessage(role, nil)
1071+
msg.SetRole(role)
8781072
msg.SetRaw(v)
8791073
a.memory.AddMessage(msg)
8801074
}
@@ -1006,13 +1200,42 @@ func (a *Agent[I, O]) schemaStream(ctx context.Context, userInput *I) (<-chan an
10061200
if memory := clt.Memory(); memory != nil {
10071201
for _, v := range memory.List() {
10081202
var role components.MessageRole
1203+
msg := new(components.Message)
10091204
switch v.Role {
10101205
case geminiAPI.RoleModel:
10111206
role = components.AssistantRole
1207+
var tools []components.ToolCall
1208+
for _, part := range v.Parts {
1209+
if call := part.FunctionCall; call != nil {
1210+
bs, _ := json.Marshal(call.Args)
1211+
tools = append(tools, components.ToolCall{
1212+
ID: call.ID,
1213+
Name: call.Name,
1214+
Arguments: string(bs),
1215+
})
1216+
}
1217+
}
1218+
if len(tools) > 0 {
1219+
msg.SetToolCalls(tools)
1220+
}
10121221
case geminiAPI.RoleUser:
10131222
role = components.UserRole
1223+
var calls []components.ToolCallback
1224+
for _, part := range v.Parts {
1225+
if fn := part.FunctionResponse; fn != nil {
1226+
bs, _ := json.Marshal(fn.Response)
1227+
calls = append(calls, components.ToolCallback{
1228+
ID: fn.ID,
1229+
Name: fn.Name,
1230+
Content: string(bs),
1231+
})
1232+
}
1233+
}
1234+
if len(calls) > 0 {
1235+
msg.SetToolCallbacks(calls)
1236+
}
10141237
}
1015-
msg := components.NewMessage(role, nil)
1238+
msg.SetRole(role)
10161239
msg.SetRaw(v)
10171240
a.memory.AddMessage(msg)
10181241
}

0 commit comments

Comments
 (0)