|
| 1 | +import type { ChatCompletionRequest, ChatMessage, LLMock } from "@copilotkit/aimock" |
| 2 | + |
| 3 | +const TASKS = ["A", "B", "C"] as const |
| 4 | +const ROUNDS = 10 |
| 5 | + |
| 6 | +const MODE_SEQUENCES: Record<(typeof TASKS)[number], string[]> = { |
| 7 | + A: ["ask", "debug", "architect", "orchestrator", "code", "ask", "debug", "architect", "orchestrator", "code"], |
| 8 | + B: ["debug", "architect", "orchestrator", "code", "ask", "debug", "architect", "orchestrator", "code", "ask"], |
| 9 | + C: ["architect", "orchestrator", "code", "ask", "debug", "architect", "orchestrator", "code", "ask", "debug"], |
| 10 | +} |
| 11 | + |
| 12 | +const markerFor = (taskName: (typeof TASKS)[number]) => `FOLLOWUP_MODE_ISOLATION_${taskName}` |
| 13 | +const answerFor = (taskName: (typeof TASKS)[number], round: number) => `${taskName} follow-up round ${round}` |
| 14 | +const callIdFor = (taskName: (typeof TASKS)[number], round: number) => |
| 15 | + `call_followup_mode_${taskName.toLowerCase()}_${String(round).padStart(2, "0")}` |
| 16 | + |
| 17 | +const lastToolResultContains = (req: ChatCompletionRequest, toolCallId: string, expected: string[]) => { |
| 18 | + const messages = Array.isArray(req?.messages) ? req.messages : [] |
| 19 | + const toolMessage = messages.filter((message: ChatMessage) => message?.role === "tool").at(-1) |
| 20 | + const content = toolMessage?.content |
| 21 | + |
| 22 | + return ( |
| 23 | + toolMessage?.tool_call_id === toolCallId && |
| 24 | + typeof content === "string" && |
| 25 | + expected.every((text) => content.includes(text)) |
| 26 | + ) |
| 27 | +} |
| 28 | + |
| 29 | +const followupToolCall = (taskName: (typeof TASKS)[number], round: number) => ({ |
| 30 | + name: "ask_followup_question", |
| 31 | + arguments: JSON.stringify({ |
| 32 | + question: `Task ${taskName}: choose mode for round ${round}`, |
| 33 | + follow_up: [ |
| 34 | + { |
| 35 | + text: answerFor(taskName, round), |
| 36 | + mode: MODE_SEQUENCES[taskName][round - 1], |
| 37 | + }, |
| 38 | + ], |
| 39 | + }), |
| 40 | + id: callIdFor(taskName, round), |
| 41 | +}) |
| 42 | + |
| 43 | +export const getFollowupModeIsolationPlan = () => |
| 44 | + TASKS.map((taskName) => ({ |
| 45 | + taskName, |
| 46 | + marker: markerFor(taskName), |
| 47 | + rounds: MODE_SEQUENCES[taskName].map((mode, index) => ({ |
| 48 | + round: index + 1, |
| 49 | + answer: answerFor(taskName, index + 1), |
| 50 | + mode, |
| 51 | + })), |
| 52 | + })) |
| 53 | + |
| 54 | +export function addViewStateFixtures(mock: InstanceType<typeof LLMock>) { |
| 55 | + for (const taskName of TASKS) { |
| 56 | + mock.addFixture({ |
| 57 | + match: { |
| 58 | + userMessage: markerFor(taskName), |
| 59 | + }, |
| 60 | + response: { |
| 61 | + toolCalls: [followupToolCall(taskName, 1)], |
| 62 | + }, |
| 63 | + }) |
| 64 | + |
| 65 | + for (let round = 1; round < ROUNDS; round++) { |
| 66 | + mock.addFixture({ |
| 67 | + match: { |
| 68 | + predicate: (req) => |
| 69 | + lastToolResultContains(req, callIdFor(taskName, round), [answerFor(taskName, round)]), |
| 70 | + }, |
| 71 | + response: { |
| 72 | + toolCalls: [followupToolCall(taskName, round + 1)], |
| 73 | + }, |
| 74 | + }) |
| 75 | + } |
| 76 | + |
| 77 | + mock.addFixture({ |
| 78 | + match: { |
| 79 | + predicate: (req) => |
| 80 | + lastToolResultContains(req, callIdFor(taskName, ROUNDS), [answerFor(taskName, ROUNDS)]), |
| 81 | + }, |
| 82 | + response: { |
| 83 | + toolCalls: [ |
| 84 | + { |
| 85 | + name: "attempt_completion", |
| 86 | + arguments: JSON.stringify({ |
| 87 | + result: `Task ${taskName} completed ${ROUNDS} follow-up mode switches.`, |
| 88 | + }), |
| 89 | + id: `call_followup_mode_${taskName.toLowerCase()}_complete`, |
| 90 | + }, |
| 91 | + ], |
| 92 | + }, |
| 93 | + }) |
| 94 | + } |
| 95 | +} |
0 commit comments