diff --git a/lib/chat.ts b/lib/chat.ts index 880e949..a069fad 100644 --- a/lib/chat.ts +++ b/lib/chat.ts @@ -4,7 +4,7 @@ export type Pdf = { name: string; data: string }; // base64, appl export type Message = { role: Role; content: string; html?: string; images?: Image[]; pdfs?: Pdf[] }; export type Provider = 'openai' | 'anthropic' | 'google'; -const FOLLOWUPS_BLOCK_RE = /(?:\r?\n|\s)*([\s\S]*?)<\/followups>\s*$/i; +const FOLLOWUPS_BLOCK_RE = /(?:\r?\n|\s)*(?:([\s\S]*?)<\/followups>|((?:[\s\S]*?<\/followup>\s*)+))\s*$/i; const FOLLOWUP_RE = /([\s\S]*?)<\/followup>/gi; export function getProvider(model: string): Provider { @@ -17,7 +17,8 @@ export function splitMessageFollowups(content: string): { content: string; follo const match = content.match(FOLLOWUPS_BLOCK_RE); if (!match) return { content, followups: [] }; - const followups = Array.from(match[1].matchAll(FOLLOWUP_RE)) + const followupsContent = match[1] ?? match[2] ?? ''; + const followups = Array.from(followupsContent.matchAll(FOLLOWUP_RE)) .map(([, followup]) => followup.replace(/\s+/g, ' ').trim()) .filter(Boolean); if (followups.length === 0) return { content, followups: [] }; diff --git a/tests/unit.test.ts b/tests/unit.test.ts index d155bbf..2911460 100644 --- a/tests/unit.test.ts +++ b/tests/unit.test.ts @@ -212,6 +212,14 @@ test('splitMessageFollowups: strips the trailing followups block and returns fol assert.deepEqual(parsed.followups, ['First follow-up.', 'Second follow-up.']); }); +test('splitMessageFollowups: treats trailing bare followup tags like a followups wrapper', () => { + const parsed = splitMessageFollowups( + 'Main answer.\n\nFirst follow-up.\nSecond follow-up.', + ); + assert.equal(parsed.content, 'Main answer.'); + assert.deepEqual(parsed.followups, ['First follow-up.', 'Second follow-up.']); +}); + test('splitMessageFollowups: ignores followups tags that are not at the end of the message', () => { const content = 'Example\nStill part of the visible message.'; const parsed = splitMessageFollowups(content);