feat: improve prompt, add ability to abort streaming, improve message parser

This commit is contained in:
Dominic Elm
2024-07-18 11:10:12 +02:00
parent 637ad2b870
commit 012b5bae80
12 changed files with 633 additions and 160 deletions

View File

@@ -0,0 +1,220 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (0) > onActionClose 1`] = `
{
"action": {
"content": "npm install",
"type": "shell",
},
"actionId": "0",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (0) > onActionOpen 1`] = `
{
"action": {
"content": "",
"type": "shell",
},
"actionId": "0",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (0) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (0) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onActionClose 1`] = `
{
"action": {
"content": "npm install",
"type": "shell",
},
"actionId": "0",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onActionClose 2`] = `
{
"action": {
"content": "some content
",
"filePath": "index.js",
"type": "file",
},
"actionId": "1",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onActionOpen 1`] = `
{
"action": {
"content": "",
"type": "shell",
},
"actionId": "0",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onActionOpen 2`] = `
{
"action": {
"content": "",
"filePath": "index.js",
"type": "file",
},
"actionId": "1",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (0) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (0) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (1) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (1) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (2) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (2) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (3) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (3) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (4) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (4) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (5) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (5) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (6) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (6) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;

View File

@@ -1,5 +1,15 @@
import { describe, expect, it } from 'vitest';
import { StreamingMessageParser } from './message-parser';
import { describe, expect, it, vi } from 'vitest';
import { StreamingMessageParser, type ActionCallback, type ArtifactCallback } from './message-parser';
interface ExpectedResult {
output: string;
callbacks?: {
onArtifactOpen?: number;
onArtifactClose?: number;
onActionOpen?: number;
onActionClose?: number;
};
}
describe('StreamingMessageParser', () => {
it('should pass through normal text', () => {
@@ -12,75 +22,186 @@ describe('StreamingMessageParser', () => {
expect(parser.parse('test_id', 'Hello <strong>world</strong>!')).toBe('Hello <strong>world</strong>!');
});
it.each([
['Foo bar', 'Foo bar'],
['Foo bar <', 'Foo bar '],
['Foo bar <p', 'Foo bar <p'],
['Foo bar <b', 'Foo bar '],
['Foo bar <ba', 'Foo bar <ba'],
['Foo bar <bol', 'Foo bar '],
['Foo bar <bolt', 'Foo bar '],
['Foo bar <bolta', 'Foo bar <bolta'],
['Foo bar <boltA', 'Foo bar '],
['Some text before <boltArtifact>foo</boltArtifact> Some more text', 'Some text before Some more text'],
[['Some text before <boltArti', 'fact>foo</boltArtifact> Some more text'], 'Some text before Some more text'],
[['Some text before <boltArti', 'fac', 't>foo</boltArtifact> Some more text'], 'Some text before Some more text'],
[['Some text before <boltArti', 'fact>fo', 'o</boltArtifact> Some more text'], 'Some text before Some more text'],
[
['Some text before <boltArti', 'fact>fo', 'o', '<', '/boltArtifact> Some more text'],
'Some text before Some more text',
],
[
['Some text before <boltArti', 'fact>fo', 'o<', '/boltArtifact> Some more text'],
'Some text before Some more text',
],
['Before <oltArtfiact>foo</boltArtifact> After', 'Before <oltArtfiact>foo</boltArtifact> After'],
['Before <boltArtifactt>foo</boltArtifact> After', 'Before <boltArtifactt>foo</boltArtifact> After'],
['Before <boltArtifact title="Some title">foo</boltArtifact> After', 'Before After'],
[
'Before <boltArtifact title="Some title" id="artifact_1"><boltAction type="shell">npm install</boltAction></boltArtifact> After',
'Before After',
[{ type: 'shell', content: 'npm install' }],
],
[
'Before <boltArtifact title="Some title" id="artifact_1"><boltAction type="shell">npm install</boltAction><boltAction type="file" filePath="index.js">some content</boltAction></boltArtifact> After',
'Before After',
[
{ type: 'shell', content: 'npm install' },
{ type: 'file', filePath: 'index.js', content: 'some content\n' },
],
],
])('should correctly parse chunks and strip out bolt artifacts', (input, expected, expectedActions = []) => {
let actionCounter = 0;
const expectedArtifactId = 'artifact_1';
const expectedMessageId = 'message_1';
const parser = new StreamingMessageParser({
artifactElement: '',
callbacks: {
onAction: ({ artifactId, messageId, action }) => {
expect(artifactId).toBe(expectedArtifactId);
expect(messageId).toBe(expectedMessageId);
expect(action).toEqual(expectedActions[actionCounter]);
actionCounter++;
},
},
describe('no artifacts', () => {
it.each<[string | string[], ExpectedResult | string]>([
['Foo bar', 'Foo bar'],
['Foo bar <', 'Foo bar '],
['Foo bar <p', 'Foo bar <p'],
[['Foo bar <', 's', 'p', 'an>some text</span>'], 'Foo bar <span>some text</span>'],
])('should correctly parse chunks and strip out bolt artifacts (%#)', (input, expected) => {
runTest(input, expected);
});
});
let message = '';
describe('invalid or incomplete artifacts', () => {
it.each<[string | string[], ExpectedResult | string]>([
['Foo bar <b', 'Foo bar '],
['Foo bar <ba', 'Foo bar <ba'],
['Foo bar <bol', 'Foo bar '],
['Foo bar <bolt', 'Foo bar '],
['Foo bar <bolta', 'Foo bar <bolta'],
['Foo bar <boltA', 'Foo bar '],
['Foo bar <boltArtifacs></boltArtifact>', 'Foo bar <boltArtifacs></boltArtifact>'],
['Before <oltArtfiact>foo</boltArtifact> After', 'Before <oltArtfiact>foo</boltArtifact> After'],
['Before <boltArtifactt>foo</boltArtifact> After', 'Before <boltArtifactt>foo</boltArtifact> After'],
])('should correctly parse chunks and strip out bolt artifacts (%#)', (input, expected) => {
runTest(input, expected);
});
});
let result = '';
describe('valid artifacts without actions', () => {
it.each<[string | string[], ExpectedResult | string]>([
[
'Some text before <boltArtifact title="Some title" id="artifact_1">foo bar</boltArtifact> Some more text',
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
['Some text before <boltArti', 'fact', ' title="Some title" id="artifact_1">foo</boltArtifact> Some more text'],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
[
'Some text before <boltArti',
'fac',
't title="Some title" id="artifact_1"',
' ',
'>',
'foo</boltArtifact> Some more text',
],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
[
'Some text before <boltArti',
'fact',
' title="Some title" id="artifact_1"',
' >fo',
'o</boltArtifact> Some more text',
],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
[
'Some text before <boltArti',
'fact tit',
'le="Some ',
'title" id="artifact_1">fo',
'o',
'<',
'/boltArtifact> Some more text',
],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
[
'Some text before <boltArti',
'fact title="Some title" id="artif',
'act_1">fo',
'o<',
'/boltArtifact> Some more text',
],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
'Before <boltArtifact title="Some title" id="artifact_1">foo</boltArtifact> After',
{
output: 'Before After',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
])('should correctly parse chunks and strip out bolt artifacts (%#)', (input, expected) => {
runTest(input, expected);
});
});
const chunks = Array.isArray(input) ? input : input.split('');
for (const chunk of chunks) {
message += chunk;
result += parser.parse(expectedMessageId, message);
}
expect(actionCounter).toBe(expectedActions.length);
expect(result).toEqual(expected);
describe('valid artifacts with actions', () => {
it.each<[string | string[], ExpectedResult | string]>([
[
'Before <boltArtifact title="Some title" id="artifact_1"><boltAction type="shell">npm install</boltAction></boltArtifact> After',
{
output: 'Before After',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 1, onActionClose: 1 },
},
],
[
'Before <boltArtifact title="Some title" id="artifact_1"><boltAction type="shell">npm install</boltAction><boltAction type="file" filePath="index.js">some content</boltAction></boltArtifact> After',
{
output: 'Before After',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 2, onActionClose: 2 },
},
],
])('should correctly parse chunks and strip out bolt artifacts (%#)', (input, expected) => {
runTest(input, expected);
});
});
});
function runTest(input: string | string[], outputOrExpectedResult: string | ExpectedResult) {
let expected: ExpectedResult;
if (typeof outputOrExpectedResult === 'string') {
expected = { output: outputOrExpectedResult };
} else {
expected = outputOrExpectedResult;
}
const callbacks = {
onArtifactOpen: vi.fn<ArtifactCallback>((data) => {
expect(data).toMatchSnapshot('onArtifactOpen');
}),
onArtifactClose: vi.fn<ArtifactCallback>((data) => {
expect(data).toMatchSnapshot('onArtifactClose');
}),
onActionOpen: vi.fn<ActionCallback>((data) => {
expect(data).toMatchSnapshot('onActionOpen');
}),
onActionClose: vi.fn<ActionCallback>((data) => {
expect(data).toMatchSnapshot('onActionClose');
}),
};
const parser = new StreamingMessageParser({
artifactElement: '',
callbacks,
});
let message = '';
let result = '';
const chunks = Array.isArray(input) ? input : input.split('');
for (const chunk of chunks) {
message += chunk;
result += parser.parse('message_1', message);
}
for (const name in expected.callbacks) {
const callbackName = name;
expect(callbacks[callbackName as keyof typeof callbacks]).toHaveBeenCalledTimes(
expected.callbacks[callbackName as keyof typeof expected.callbacks] ?? 0,
);
}
expect(result).toEqual(expected.output);
}

View File

@@ -21,20 +21,20 @@ export interface ActionCallbackData {
action: BoltAction;
}
type ArtifactOpenCallback = (data: ArtifactCallbackData) => void;
type ArtifactCloseCallback = (data: ArtifactCallbackData) => void;
type ActionCallback = (data: ActionCallbackData) => void;
export type ArtifactCallback = (data: ArtifactCallbackData) => void;
export type ActionCallback = (data: ActionCallbackData) => void;
interface Callbacks {
onArtifactOpen?: ArtifactOpenCallback;
onArtifactClose?: ArtifactCloseCallback;
onAction?: ActionCallback;
export interface ParserCallbacks {
onArtifactOpen?: ArtifactCallback;
onArtifactClose?: ArtifactCallback;
onActionOpen?: ActionCallback;
onActionClose?: ActionCallback;
}
type ElementFactory = () => string;
interface StreamingMessageParserOptions {
callbacks?: Callbacks;
export interface StreamingMessageParserOptions {
callbacks?: ParserCallbacks;
artifactElement?: string | ElementFactory;
}
@@ -95,10 +95,17 @@ export class StreamingMessageParser {
currentAction.content = content;
this._options.callbacks?.onAction?.({
this._options.callbacks?.onActionClose?.({
artifactId: currentArtifact.id,
messageId,
actionId: String(state.actionId++),
/**
* We decrement the id because it's been incremented already
* when `onActionOpen` was emitted to make sure the ids are
* the same.
*/
actionId: String(state.actionId - 1),
action: currentAction as BoltAction,
});
@@ -117,31 +124,17 @@ export class StreamingMessageParser {
const actionEndIndex = input.indexOf('>', actionOpenIndex);
if (actionEndIndex !== -1) {
const actionTag = input.slice(actionOpenIndex, actionEndIndex + 1);
const actionType = this.#extractAttribute(actionTag, 'type') as ActionType;
const actionAttributes = {
type: actionType,
content: '',
};
if (actionType === 'file') {
const filePath = this.#extractAttribute(actionTag, 'filePath') as string;
if (!filePath) {
logger.debug('File path not specified');
}
(actionAttributes as FileAction).filePath = filePath;
} else if (actionType !== 'shell') {
logger.warn(`Unknown action type '${actionType}'`);
}
state.currentAction = actionAttributes as FileAction | ShellAction;
state.insideAction = true;
state.currentAction = this.#parseActionTag(input, actionOpenIndex, actionEndIndex);
this._options.callbacks?.onActionOpen?.({
artifactId: currentArtifact.id,
messageId,
actionId: String(state.actionId++),
action: state.currentAction as BoltAction,
});
i = actionEndIndex + 1;
} else {
break;
@@ -241,6 +234,31 @@ export class StreamingMessageParser {
this.#messages.clear();
}
#parseActionTag(input: string, actionOpenIndex: number, actionEndIndex: number) {
const actionTag = input.slice(actionOpenIndex, actionEndIndex + 1);
const actionType = this.#extractAttribute(actionTag, 'type') as ActionType;
const actionAttributes = {
type: actionType,
content: '',
};
if (actionType === 'file') {
const filePath = this.#extractAttribute(actionTag, 'filePath') as string;
if (!filePath) {
logger.debug('File path not specified');
}
(actionAttributes as FileAction).filePath = filePath;
} else if (actionType !== 'shell') {
logger.warn(`Unknown action type '${actionType}'`);
}
return actionAttributes as FileAction | ShellAction;
}
#extractAttribute(tag: string, attributeName: string): string | undefined {
const match = tag.match(new RegExp(`${attributeName}="([^"]*)"`, 'i'));
return match ? match[1] : undefined;