Skip to content

Commit 9387009

Browse files
EItanyaEItanya
authored
bump-a2a-0.2.0 (#567)
* bump-a2a-0.2.0 move util and fix client fixed some issues Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * more updates Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * get rid of internal pkg Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * finally works Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * Should be getting closer on A2A Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * streaming works: Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * works Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * use JSON Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * fixed issue with tools Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * fixed e2e test issue via tools Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * fixed the unit tests Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> * PR comments Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io> --------- Signed-off-by: Eitan Yarmush <eitan.yarmush@solo.io>
1 parent 5036e22 commit 9387009

File tree

32 files changed

+873
-712
lines changed

32 files changed

+873
-712
lines changed

‎.github/workflows/ci.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ jobs:
8484
export GW_IP=$(kubectl get svc -n kagent kagent -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
8585
export KAGENT_A2A_URL="http://${GW_IP}:8083/api/a2a"
8686
cd go
87-
go test -v -timeout 30s -run ^TestInvokeAPI$ github.com/kagent-dev/kagent/go/test/e2e
87+
go test -v -run ^TestInvokeAPI$ github.com/kagent-dev/kagent/go/test/e2e
8888
8989
9090
go-unit-tests:

‎go/autogen/client/fake/client.go‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package fake
22

33
import (
44
"context"
5+
"encoding/json"
56
"fmt"
67
"sync"
78

@@ -109,14 +110,11 @@ func (m *InMemoryAutogenClient) GetTeamByID(teamID int, userID string) (*autogen
109110
}
110111

111112
func (m *InMemoryAutogenClient) InvokeTask(req *autogen_client.InvokeTaskRequest) (*autogen_client.InvokeTaskResult, error) {
112-
// For in-memory implementation, return a basic result
113+
// For in-memory implementation, return a basic result with properly formatted TextMessage
113114
return &autogen_client.InvokeTaskResult{
114115
TaskResult: autogen_client.TaskResult{
115-
Messages: []autogen_client.TaskMessageMap{
116-
{
117-
"role": "assistant",
118-
"content": fmt.Sprintf("Task completed: %s", req.Task),
119-
},
116+
Messages: []json.RawMessage{
117+
json.RawMessage(fmt.Sprintf(`{"type": "TextMessage", "content": "Task completed: %s", "source": "assistant"}`, req.Task)),
120118
},
121119
},
122120
}, nil
@@ -144,11 +142,8 @@ func (m *InMemoryAutogenClient) InvokeSession(sessionID int, userID string, requ
144142

145143
return &autogen_client.TeamResult{
146144
TaskResult: autogen_client.TaskResult{
147-
Messages: []autogen_client.TaskMessageMap{
148-
{
149-
"role": "assistant",
150-
"content": fmt.Sprintf("Session task completed: %s", request.Task),
151-
},
145+
Messages: []json.RawMessage{
146+
json.RawMessage(fmt.Sprintf(`{"type": "TextMessage", "content": "Session task completed: %s", "source": "assistant"}`, request.Task)),
152147
},
153148
},
154149
}, nil
@@ -370,7 +365,7 @@ func (m *InMemoryAutogenClient) InvokeSessionStream(sessionID int, userID string
370365
defer close(ch)
371366
ch <- &autogen_client.SseEvent{
372367
Event: "message",
373-
Data: []byte(fmt.Sprintf("Session stream task completed: %s", request.Task)),
368+
Data: []byte(fmt.Sprintf(`{"type": "TextMessage", "content": "Session stream task completed: %s", "source": "assistant"}`, request.Task)),
374369
}
375370
}()
376371

@@ -383,7 +378,7 @@ func (m *InMemoryAutogenClient) InvokeTaskStream(req *autogen_client.InvokeTaskR
383378
defer close(ch)
384379
ch <- &autogen_client.SseEvent{
385380
Event: "message",
386-
Data: []byte(fmt.Sprintf("Task stream completed: %s", req.Task)),
381+
Data: []byte(fmt.Sprintf(`{"type": "TextMessage", "content": "Task stream completed: %s", "source": "assistant"}`, req.Task)),
387382
}
388383
}()
389384

‎go/autogen/client/messages.go‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
package client
2+
3+
import (
4+
"encoding/json"
5+
"fmt"
6+
)
7+
8+
type Event interface {
9+
GetType() string
10+
}
11+
12+
type BaseEvent struct {
13+
Type string `json:"type"`
14+
}
15+
16+
func (e *BaseEvent) GetType() string {
17+
return e.Type
18+
}
19+
20+
type BaseChatMessage struct {
21+
BaseEvent
22+
Source string `json:"source"`
23+
Metadata map[string]string `json:"metadata"`
24+
ModelsUsage *ModelsUsage `json:"models_usage"`
25+
}
26+
27+
type TextMessage struct {
28+
BaseChatMessage
29+
Content string `json:"content"`
30+
}
31+
32+
type ModelClientStreamingChunkEvent struct {
33+
BaseChatMessage
34+
Content string `json:"content"`
35+
}
36+
type FunctionCall struct {
37+
ID string `json:"id"`
38+
Arguments string `json:"arguments"`
39+
Name string `json:"name"`
40+
}
41+
type ToolCallRequestEvent struct {
42+
BaseChatMessage
43+
Content []FunctionCall `json:"content"`
44+
}
45+
46+
type FunctionExecutionResult struct {
47+
CallID string `json:"call_id"`
48+
Content string `json:"content"`
49+
}
50+
51+
type ToolCallExecutionEvent struct {
52+
BaseChatMessage
53+
Content []FunctionExecutionResult `json:"content"`
54+
}
55+
56+
type MemoryQueryEvent struct {
57+
BaseChatMessage
58+
Content []map[string]interface{} `json:"content"`
59+
}
60+
61+
type ToolCallSummaryMessage struct {
62+
BaseChatMessage
63+
ToolCalls []FunctionCall `json:"tool_calls"`
64+
Results []FunctionExecutionResult `json:"results"`
65+
}
66+
67+
const (
68+
TextMessageLabel = "TextMessage"
69+
ToolCallRequestEventLabel = "ToolCallRequestEvent"
70+
ToolCallExecutionEventLabel = "ToolCallExecutionEvent"
71+
StopMessageLabel = "StopMessage"
72+
HandoffMessageLabel = "HandoffMessage"
73+
ModelClientStreamingChunkEventLabel = "ModelClientStreamingChunkEvent"
74+
LLMCallEventMessageLabel = "LLMCallEventMessage"
75+
MemoryQueryEventLabel = "MemoryQueryEvent"
76+
ToolCallSummaryMessageLabel = "ToolCallSummaryMessage"
77+
)
78+
79+
func ParseEvent(event []byte) (Event, error) {
80+
var baseEvent BaseEvent
81+
if err := json.Unmarshal(event, &baseEvent); err != nil {
82+
return nil, err
83+
}
84+
85+
switch baseEvent.Type {
86+
case TextMessageLabel:
87+
var textMessage TextMessage
88+
if err := json.Unmarshal(event, &textMessage); err != nil {
89+
return nil, err
90+
}
91+
return &textMessage, nil
92+
case ModelClientStreamingChunkEventLabel:
93+
var modelClientStreamingChunkEvent ModelClientStreamingChunkEvent
94+
if err := json.Unmarshal(event, &modelClientStreamingChunkEvent); err != nil {
95+
return nil, err
96+
}
97+
return &modelClientStreamingChunkEvent, nil
98+
case ToolCallRequestEventLabel:
99+
var toolCallRequestEvent ToolCallRequestEvent
100+
if err := json.Unmarshal(event, &toolCallRequestEvent); err != nil {
101+
return nil, err
102+
}
103+
return &toolCallRequestEvent, nil
104+
case ToolCallExecutionEventLabel:
105+
var toolCallExecutionEvent ToolCallExecutionEvent
106+
if err := json.Unmarshal(event, &toolCallExecutionEvent); err != nil {
107+
return nil, err
108+
}
109+
return &toolCallExecutionEvent, nil
110+
case MemoryQueryEventLabel:
111+
var memoryQueryEvent MemoryQueryEvent
112+
if err := json.Unmarshal(event, &memoryQueryEvent); err != nil {
113+
return nil, err
114+
}
115+
return &memoryQueryEvent, nil
116+
case ToolCallSummaryMessageLabel:
117+
var ToolCallSummaryMessage ToolCallSummaryMessage
118+
if err := json.Unmarshal(event, &ToolCallSummaryMessage); err != nil {
119+
return nil, err
120+
}
121+
return &ToolCallSummaryMessage, nil
122+
default:
123+
return nil, fmt.Errorf("unknown event type: %s", baseEvent.Type)
124+
}
125+
}
126+
127+
func GetLastStringMessage(events []Event) string {
128+
for i := len(events) - 1; i >= 0; i-- {
129+
if _, ok := events[i].(*TextMessage); ok {
130+
return events[i].(*TextMessage).Content
131+
}
132+
}
133+
return ""
134+
}

‎go/autogen/client/types.go‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package client
33
import (
44
"bufio"
55
"bytes"
6+
"encoding/json"
67
"errors"
78
"fmt"
89
"io"
@@ -70,6 +71,13 @@ func (m *ModelsUsage) String() string {
7071
return fmt.Sprintf("Prompt Tokens: %d, Completion Tokens: %d", m.PromptTokens, m.CompletionTokens)
7172
}
7273

74+
func (m *ModelsUsage) ToMap() map[string]interface{} {
75+
return map[string]interface{}{
76+
"prompt_tokens": m.PromptTokens,
77+
"completion_tokens": m.CompletionTokens,
78+
}
79+
}
80+
7381
type TaskMessageMap map[string]interface{}
7482

7583
type RunMessage struct {
@@ -121,8 +129,10 @@ type TeamResult struct {
121129
}
122130

123131
type TaskResult struct {
124-
Messages []TaskMessageMap `json:"messages"`
125-
StopReason string `json:"stop_reason"`
132+
// These are all of type Event, but we don't want to unmarshal them here
133+
// because we want to handle them in the caller
134+
Messages []json.RawMessage `json:"messages"`
135+
StopReason string `json:"stop_reason"`
126136
}
127137

128138
// APIResponse is the common response wrapper for all API responses
@@ -174,7 +184,7 @@ var (
174184

175185
func streamSseResponse(r io.ReadCloser) chan *SseEvent {
176186
scanner := bufio.NewScanner(r)
177-
ch := make(chan *SseEvent)
187+
ch := make(chan *SseEvent, 10)
178188
go func() {
179189
defer close(ch)
180190
defer r.Close()

‎go/cli/cmd/kagent/main.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,7 @@ func main() {
128128
a2aCmd.Flags().StringVarP(&a2aCfg.AgentName, "agent-name", "a", "", "Agent Name")
129129
a2aCmd.Flags().StringVarP(&a2aCfg.Task, "task", "t", "", "Task")
130130
a2aCmd.Flags().DurationVarP(&a2aCfg.Timeout, "timeout", "T", 300*time.Second, "Timeout")
131+
a2aCmd.Flags().BoolVarP(&a2aCfg.Stream, "stream", "S", false, "Stream the response")
131132

132133
getCmd := &cobra.Command{
133134
Use: "get",
@@ -210,7 +211,7 @@ func main() {
210211

211212
getCmd.AddCommand(getSessionCmd, getRunCmd, getAgentCmd, getToolCmd)
212213

213-
rootCmd.AddCommand(installCmd, uninstallCmd, invokeCmd, bugReportCmd, versionCmd, dashboardCmd, getCmd)
214+
rootCmd.AddCommand(installCmd, uninstallCmd, invokeCmd, bugReportCmd, versionCmd, dashboardCmd, getCmd, a2aCmd)
214215

215216
// Initialize config
216217
if err := config.Init(); err != nil {

0 commit comments

Comments
 (0)