|
| 1 | +import { LLMock } from "@copilotkit/aimock" |
| 2 | +import type { ChatCompletionRequest } from "@copilotkit/aimock" |
| 3 | + |
| 4 | +import { toolResultContains } from "./tool-result" |
| 5 | + |
| 6 | +const SUBTASK_PARENT_MARKER = "SUBTASK_PARENT_CANCELLATION_SMOKE" |
| 7 | +const SUBTASK_CHILD_MARKER = "SUBTASK_CHILD_CALCULATOR_SMOKE" |
| 8 | + |
| 9 | +const SUBTASK_CHILD_PROMPT = `${SUBTASK_CHILD_MARKER}: Ask the user exactly this follow-up question: What is the square root of 81? After the user answers, complete with only the answer.` |
| 10 | +export const SUBTASK_PARENT_PROMPT = `${SUBTASK_PARENT_MARKER}: Use the new_task tool exactly once. Create an ask-mode subtask with this exact message: "${SUBTASK_CHILD_PROMPT}" Do not answer directly.` |
| 11 | +export const SUBTASK_CHILD_FOLLOWUP_ANSWER = "9" |
| 12 | + |
| 13 | +const requestContains = (req: ChatCompletionRequest, expected: string[]) => { |
| 14 | + const rawRequest = JSON.stringify(req) |
| 15 | + return expected.every((text) => rawRequest.includes(text)) |
| 16 | +} |
| 17 | + |
| 18 | +const completionAfterAnswer = (followupId: string, completionId: string) => ({ |
| 19 | + match: { |
| 20 | + predicate: (req: ChatCompletionRequest) => |
| 21 | + // Preferred: structured tool-result message carries the followup answer. |
| 22 | + toolResultContains(req, followupId, [SUBTASK_CHILD_FOLLOWUP_ANSWER]) || |
| 23 | + // Fallback 1: answer present alongside the tool-call ID but not in a role:tool message. |
| 24 | + requestContains(req, [followupId, SUBTASK_CHILD_FOLLOWUP_ANSWER]) || |
| 25 | + // Fallback 2: answer arrives as a bare user message after task resume (no tool-call ID context). |
| 26 | + requestContains(req, [ |
| 27 | + SUBTASK_CHILD_MARKER, |
| 28 | + `<user_message>\\n${SUBTASK_CHILD_FOLLOWUP_ANSWER}\\n</user_message>`, |
| 29 | + ]), |
| 30 | + }, |
| 31 | + response: { |
| 32 | + toolCalls: [ |
| 33 | + { |
| 34 | + name: "attempt_completion", |
| 35 | + arguments: JSON.stringify({ result: "9" }), |
| 36 | + id: completionId, |
| 37 | + }, |
| 38 | + ], |
| 39 | + }, |
| 40 | +}) |
| 41 | + |
| 42 | +export function addSubtaskFixtures(mock: InstanceType<typeof LLMock>) { |
| 43 | + mock.addFixture({ |
| 44 | + match: { |
| 45 | + userMessage: new RegExp(SUBTASK_PARENT_MARKER), |
| 46 | + }, |
| 47 | + response: { |
| 48 | + toolCalls: [ |
| 49 | + { |
| 50 | + name: "new_task", |
| 51 | + arguments: JSON.stringify({ |
| 52 | + mode: "ask", |
| 53 | + message: SUBTASK_CHILD_PROMPT, |
| 54 | + }), |
| 55 | + id: "call_subtasks_parent_new_task_001", |
| 56 | + }, |
| 57 | + ], |
| 58 | + }, |
| 59 | + }) |
| 60 | + |
| 61 | + mock.addFixture({ |
| 62 | + match: { |
| 63 | + userMessage: new RegExp(SUBTASK_CHILD_MARKER), |
| 64 | + }, |
| 65 | + response: { |
| 66 | + toolCalls: [ |
| 67 | + { |
| 68 | + name: "ask_followup_question", |
| 69 | + arguments: JSON.stringify({ |
| 70 | + question: "What is the square root of 81?", |
| 71 | + follow_up: [{ text: SUBTASK_CHILD_FOLLOWUP_ANSWER }], |
| 72 | + }), |
| 73 | + id: "call_subtasks_child_followup_001", |
| 74 | + }, |
| 75 | + ], |
| 76 | + }, |
| 77 | + }) |
| 78 | + |
| 79 | + mock.addFixture(completionAfterAnswer("call_subtasks_child_followup_001", "call_subtasks_child_completion_002")) |
| 80 | + |
| 81 | + mock.addFixture({ |
| 82 | + match: { |
| 83 | + toolCallId: "call_subtasks_parent_new_task_001", |
| 84 | + }, |
| 85 | + response: { |
| 86 | + toolCalls: [ |
| 87 | + { |
| 88 | + name: "attempt_completion", |
| 89 | + arguments: JSON.stringify({ result: "Parent task resumed" }), |
| 90 | + id: "call_subtasks_parent_completion_003", |
| 91 | + }, |
| 92 | + ], |
| 93 | + }, |
| 94 | + }) |
| 95 | +} |
0 commit comments