Skip to content

Commit 199c660

Browse files
committed
chore: make the event context a simplified version of the engine context
This change will allow the unmarshaling of an event because it no longer contains the engine context. Signed-off-by: Donnie Adams <[email protected]>
1 parent 8e2de4d commit 199c660

File tree

3 files changed

+48
-32
lines changed

3 files changed

+48
-32
lines changed

pkg/engine/engine.go

Lines changed: 34 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,16 @@ type CallResult struct {
5353
Result string `json:"result,omitempty"`
5454
}
5555

56+
type commonContext struct {
57+
ID string `json:"id"`
58+
Tool types.Tool `json:"tool"`
59+
InputContext []InputContext `json:"inputContext"`
60+
// IsCredential indicates that the current call is for a credential tool
61+
IsCredential bool `json:"isCredential"`
62+
}
63+
5664
type Context struct {
65+
commonContext
5766
ID string
5867
Ctx context.Context
5968
Parent *Context
@@ -70,28 +79,43 @@ type InputContext struct {
7079
Content string `json:"content,omitempty"`
7180
}
7281

82+
type BasicContext struct {
83+
commonContext `json:",inline"`
84+
ToolName string `json:"toolName,omitempty"`
85+
ParentID string `json:"parentID,omitempty"`
86+
}
87+
7388
func (c *Context) ParentID() string {
7489
if c.Parent == nil {
7590
return ""
7691
}
7792
return c.Parent.ID
7893
}
7994

95+
func (c *Context) ToBasicContext() *BasicContext {
96+
var toolName string
97+
if c.Parent != nil {
98+
for name, id := range c.Parent.Tool.ToolMapping {
99+
if id == c.Tool.ID {
100+
toolName = name
101+
break
102+
}
103+
}
104+
}
105+
106+
return &BasicContext{
107+
commonContext: c.commonContext,
108+
ParentID: c.ParentID(),
109+
ToolName: toolName,
110+
}
111+
}
112+
80113
func (c *Context) UnmarshalJSON([]byte) error {
81114
panic("this data struct is circular by design and can not be read from json")
82115
}
83116

84117
func (c *Context) MarshalJSON() ([]byte, error) {
85-
var parentID string
86-
if c.Parent != nil {
87-
parentID = c.Parent.ID
88-
}
89-
return json.Marshal(map[string]any{
90-
"id": c.ID,
91-
"parentID": parentID,
92-
"tool": c.Tool,
93-
"inputContext": c.InputContext,
94-
})
118+
return json.Marshal(c.ToBasicContext())
95119
}
96120

97121
var execID int32

pkg/monitor/display.go

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ func (d *display) Event(event runner.Event) {
194194
currentIndex = len(d.dump.Calls)
195195
currentCall = call{
196196
ID: event.CallContext.ID,
197-
ParentID: event.CallContext.ParentID(),
197+
ParentID: event.CallContext.ParentID,
198198
ToolID: event.CallContext.Tool.ID,
199199
}
200200
d.dump.Calls = append(d.dump.Calls, currentCall)
@@ -213,20 +213,12 @@ func (d *display) Event(event runner.Event) {
213213
}
214214

215215
callName := callName{
216-
prettyIDMap: d.callIDMap,
217-
call: &currentCall,
218-
prg: d.dump.Program,
219-
calls: d.dump.Calls,
220-
credential: event.CallContext.IsCredential,
221-
}
222-
223-
if event.CallContext.Parent != nil {
224-
for name, id := range event.CallContext.Parent.Tool.ToolMapping {
225-
if id == event.CallContext.Tool.ID {
226-
callName.userSpecifiedToolName = name
227-
break
228-
}
229-
}
216+
prettyIDMap: d.callIDMap,
217+
call: &currentCall,
218+
prg: d.dump.Program,
219+
calls: d.dump.Calls,
220+
credential: event.CallContext.IsCredential,
221+
userSpecifiedToolName: event.CallContext.ToolName,
230222
}
231223

232224
switch event.Type {

pkg/runner/runner.go

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ func (r *Runner) Run(ctx context.Context, prg types.Program, env []string, input
108108

109109
type Event struct {
110110
Time time.Time `json:"time,omitempty"`
111-
CallContext *engine.Context `json:"callContext,omitempty"`
111+
CallContext *engine.BasicContext `json:"callContext,omitempty"`
112112
ToolSubCalls map[string]engine.Call `json:"toolSubCalls,omitempty"`
113113
ToolResults int `json:"toolResults,omitempty"`
114114
Type EventType `json:"type,omitempty"`
@@ -177,7 +177,7 @@ func (r *Runner) call(callCtx engine.Context, monitor Monitor, env []string, inp
177177

178178
monitor.Event(Event{
179179
Time: time.Now(),
180-
CallContext: &callCtx,
180+
CallContext: callCtx.ToBasicContext(),
181181
Type: EventTypeCallStart,
182182
Content: input,
183183
})
@@ -197,7 +197,7 @@ func (r *Runner) call(callCtx engine.Context, monitor Monitor, env []string, inp
197197
progressClose()
198198
monitor.Event(Event{
199199
Time: time.Now(),
200-
CallContext: &callCtx,
200+
CallContext: callCtx.ToBasicContext(),
201201
Type: EventTypeCallFinish,
202202
Content: *result.Result,
203203
})
@@ -210,7 +210,7 @@ func (r *Runner) call(callCtx engine.Context, monitor Monitor, env []string, inp
210210

211211
monitor.Event(Event{
212212
Time: time.Now(),
213-
CallContext: &callCtx,
213+
CallContext: callCtx.ToBasicContext(),
214214
Type: EventTypeCallSubCalls,
215215
ToolSubCalls: result.Calls,
216216
})
@@ -222,7 +222,7 @@ func (r *Runner) call(callCtx engine.Context, monitor Monitor, env []string, inp
222222

223223
monitor.Event(Event{
224224
Time: time.Now(),
225-
CallContext: &callCtx,
225+
CallContext: callCtx.ToBasicContext(),
226226
Type: EventTypeCallContinue,
227227
ToolResults: len(callResults),
228228
})
@@ -245,15 +245,15 @@ func streamProgress(callCtx *engine.Context, monitor Monitor) (chan<- types.Comp
245245
if message := status.PartialResponse; message != nil {
246246
monitor.Event(Event{
247247
Time: time.Now(),
248-
CallContext: callCtx,
248+
CallContext: callCtx.ToBasicContext(),
249249
Type: EventTypeCallProgress,
250250
ChatCompletionID: status.CompletionID,
251251
Content: message.String(),
252252
})
253253
} else {
254254
monitor.Event(Event{
255255
Time: time.Now(),
256-
CallContext: callCtx,
256+
CallContext: callCtx.ToBasicContext(),
257257
Type: EventTypeChat,
258258
ChatCompletionID: status.CompletionID,
259259
ChatRequest: status.Request,

0 commit comments

Comments
 (0)