Skip to content

Commit 6c3b312

Browse files
committed
Fix chain
1 parent d55fbbb commit 6c3b312

File tree

4 files changed

+121
-44
lines changed

4 files changed

+121
-44
lines changed

core/services/chatservice/chatservice.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,7 @@ func (s *service) Chat(ctx context.Context, req ChatRequest) (string, int, int,
129129
switch task.ID {
130130
case "append_user_message":
131131
task.Hook.Args["subject_id"] = req.SubjectID
132+
case "append_user_message_2":
132133
case "persist_messages":
133134
task.Hook.Args["subject_id"] = req.SubjectID
134135
}

core/taskengine/taskexec.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,9 @@ func (exe *SimpleExec) TaskExec(taskCtx context.Context, startingTime time.Time,
204204
var taskErr error
205205
var output any = input
206206
var outputType DataType = dataType
207-
207+
if currentTask.Type == Noop {
208+
return output, outputType, "noop", nil
209+
}
208210
// Unified prompt extraction function
209211
getPrompt := func() (string, error) {
210212
switch outputType {

core/taskengine/tasktype.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ const (
3434
// ParseTransition will attempt to parse a transition command from the input and strip the transition prefix if it exists.
3535
ParseTransition TaskType = "parse_transition"
3636

37+
Noop TaskType = "noop"
38+
3739
// Hook indicates this task should execute an external action rather than calling the LLM.
3840
Hook TaskType = "hook"
3941
)

core/tasksrecipes/chatchain.go

Lines changed: 115 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -112,47 +112,13 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
112112
ID: "chat_chain",
113113
Description: "Standard chat processing pipeline with hooks",
114114
Tasks: []taskengine.ChainTask{
115-
{
116-
ID: "append_user_message",
117-
Description: "Append user message to chat history",
118-
Type: taskengine.Hook,
119-
Hook: &taskengine.HookCall{
120-
Type: "append_user_message",
121-
Args: map[string]string{
122-
"subject_id": req.SubjectID,
123-
},
124-
},
125-
Transition: taskengine.TaskTransition{
126-
Branches: []taskengine.TransitionBranch{
127-
{
128-
Operator: "default", Goto: "inject_context",
129-
AlertOnMatch: "Test Alert",
130-
},
131-
},
132-
},
133-
},
134-
// {
135-
// ID: "inject_context",
136-
// Description: "Add context to the conversation",
137-
// Type: taskengine.RawString,
138-
// PromptTemplate: "Evaluate if input contains question that requires knowleage if yes output /search <search query> to retrieve context. Input: {{.input}}",
139-
// InputVar: "input",
140-
// Transition: taskengine.TaskTransition{
141-
// Branches: []taskengine.TransitionBranch{
142-
// {
143-
// Operator: "default",
144-
// Goto: "mux_input",
145-
// },
146-
// },
147-
// },
148-
// },
149115
{
150116
ID: "mux_input",
151117
Description: "Check for commands like /echo",
152118
Type: taskengine.ParseTransition,
153119
Transition: taskengine.TaskTransition{
154120
Branches: []taskengine.TransitionBranch{
155-
{Operator: "default", Goto: "moderate"},
121+
{Operator: taskengine.OpDefault, Goto: "moderate"},
156122
{
157123
Operator: "equals",
158124
When: "echo",
@@ -171,6 +137,60 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
171137
},
172138
},
173139
},
140+
{
141+
ID: "moderate",
142+
Description: "Moderate the input",
143+
Type: taskengine.ParseNumber,
144+
PromptTemplate: "Classify the input as safe (0) or spam (10) respond with an numeric value between 0 and 10. Input: {{.input}}",
145+
InputVar: "input",
146+
Transition: taskengine.TaskTransition{
147+
Branches: []taskengine.TransitionBranch{
148+
{
149+
Operator: taskengine.OpGreaterThan,
150+
When: "4",
151+
Goto: "reject_request",
152+
},
153+
{
154+
Operator: "default",
155+
Goto: "do_we_need_context",
156+
},
157+
},
158+
OnFailure: "request_failed",
159+
},
160+
},
161+
{
162+
ID: "do_we_need_context",
163+
Description: "Add context to the conversation",
164+
Type: taskengine.RawString,
165+
PromptTemplate: "Classify the necessity of searching documents with (0) not likely and (10) highly likely to solve the users request with an numeric value between 0 and 10. Input {{.input}}",
166+
InputVar: "input",
167+
Transition: taskengine.TaskTransition{
168+
Branches: []taskengine.TransitionBranch{
169+
{
170+
Operator: taskengine.OpDefault,
171+
Goto: "swap_to_input",
172+
},
173+
{
174+
Operator: taskengine.OpGreaterThan,
175+
When: "4",
176+
Goto: "search_knowledge",
177+
},
178+
},
179+
},
180+
},
181+
{
182+
ID: "swap_to_input",
183+
Type: taskengine.Noop,
184+
InputVar: "input",
185+
Transition: taskengine.TaskTransition{
186+
Branches: []taskengine.TransitionBranch{
187+
{
188+
Operator: taskengine.OpDefault,
189+
Goto: "append_user_message",
190+
},
191+
},
192+
},
193+
},
174194
{
175195
ID: "echo_message",
176196
Description: "Echo the message",
@@ -180,7 +200,7 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
180200
},
181201
Transition: taskengine.TaskTransition{
182202
Branches: []taskengine.TransitionBranch{
183-
{Operator: "default", Goto: "persist_messages"},
203+
{Operator: taskengine.OpDefault, Goto: "append_user_message_2"},
184204
},
185205
},
186206
},
@@ -193,7 +213,21 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
193213
},
194214
Transition: taskengine.TaskTransition{
195215
Branches: []taskengine.TransitionBranch{
196-
{Operator: "default", Goto: "persist_messages"},
216+
{Operator: taskengine.OpDefault, Goto: "append_context_to_input"},
217+
},
218+
},
219+
},
220+
{
221+
ID: "append_context_to_input",
222+
Description: "Added context to the user input",
223+
Type: taskengine.Noop,
224+
PromptTemplate: "{{.input}} \n\n Context: {{.search_knowledge}}",
225+
Transition: taskengine.TaskTransition{
226+
Branches: []taskengine.TransitionBranch{
227+
{
228+
Operator: taskengine.OpDefault,
229+
Goto: "append_user_message",
230+
},
197231
},
198232
},
199233
},
@@ -209,7 +243,7 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
209243
},
210244
Transition: taskengine.TaskTransition{
211245
Branches: []taskengine.TransitionBranch{
212-
{Operator: "default", Goto: "persist_messages"},
246+
{Operator: taskengine.OpDefault, Goto: "append_user_message_2"},
213247
},
214248
},
215249
},
@@ -222,7 +256,7 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
222256
InputVar: "input",
223257
Transition: taskengine.TaskTransition{
224258
Branches: []taskengine.TransitionBranch{
225-
{Operator: "default", Goto: "raise_error"},
259+
{Operator: taskengine.OpDefault, Goto: "raise_error"},
226260
},
227261
},
228262
},
@@ -234,7 +268,7 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
234268
InputVar: "input",
235269
Transition: taskengine.TaskTransition{
236270
Branches: []taskengine.TransitionBranch{
237-
{Operator: "default", Goto: "raise_error"},
271+
{Operator: taskengine.OpDefault, Goto: "raise_error"},
238272
},
239273
},
240274
},
@@ -246,7 +280,26 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
246280
InputVar: "input",
247281
Transition: taskengine.TaskTransition{
248282
Branches: []taskengine.TransitionBranch{
249-
{Operator: "default", Goto: taskengine.TermEnd},
283+
{Operator: taskengine.OpDefault, Goto: taskengine.TermEnd},
284+
},
285+
},
286+
},
287+
{
288+
ID: "append_user_message",
289+
Description: "Append user message to chat history",
290+
Type: taskengine.Hook,
291+
Hook: &taskengine.HookCall{
292+
Type: "append_user_message",
293+
Args: map[string]string{
294+
"subject_id": req.SubjectID,
295+
},
296+
},
297+
Transition: taskengine.TaskTransition{
298+
Branches: []taskengine.TransitionBranch{
299+
{
300+
Operator: taskengine.OpDefault, Goto: "execute_model_on_messages",
301+
AlertOnMatch: "Test Alert",
302+
},
250303
},
251304
},
252305
},
@@ -264,7 +317,26 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
264317
// InputVar: "append_user_message",
265318
Transition: taskengine.TaskTransition{
266319
Branches: []taskengine.TransitionBranch{
267-
{Operator: "default", Goto: "persist_messages"},
320+
{Operator: taskengine.OpDefault, Goto: "persist_messages"},
321+
},
322+
},
323+
},
324+
{
325+
ID: "append_user_message_2",
326+
Description: "Append user message to chat history",
327+
Type: taskengine.Hook,
328+
Hook: &taskengine.HookCall{
329+
Type: "append_user_message",
330+
Args: map[string]string{
331+
"subject_id": req.SubjectID,
332+
},
333+
},
334+
Transition: taskengine.TaskTransition{
335+
Branches: []taskengine.TransitionBranch{
336+
{
337+
Operator: taskengine.OpDefault, Goto: "persist_messages",
338+
AlertOnMatch: "Test Alert",
339+
},
268340
},
269341
},
270342
},
@@ -280,7 +352,7 @@ func BuildChatChain(req BuildChatChainReq) *taskengine.ChainDefinition {
280352
},
281353
Transition: taskengine.TaskTransition{
282354
Branches: []taskengine.TransitionBranch{
283-
{Operator: "default", Goto: taskengine.TermEnd},
355+
{Operator: taskengine.OpDefault, Goto: taskengine.TermEnd},
284356
},
285357
},
286358
},

0 commit comments

Comments
 (0)