diff --git a/src/compiler/base/nodes/for.test.ts b/src/compiler/base/nodes/for.test.ts index 9e0269d..2fa4bb4 100644 --- a/src/compiler/base/nodes/for.test.ts +++ b/src/compiler/base/nodes/for.test.ts @@ -1,7 +1,7 @@ import { Message, MessageContent, TextContent } from '$promptl/types' import { describe, expect, it } from 'vitest' -import { render } from '../../index' +import { Adapters, render } from '$promptl/index' import { removeCommonIndent } from '../../utils' async function getCompiledText( @@ -25,6 +25,14 @@ async function getCompiledText( }, '') } +function getMessageTexts(messages: Message[]): string[] { + return messages.map((m) => + (m.content as MessageContent[]) + .map((c) => (c as TextContent).text) + .join(''), + ) +} + describe('each loops', async () => { it('iterates over any iterable object', async () => { const prompt1 = `{{ for element in [1, 2, 3] }} {{element}} {{ endfor }}` @@ -78,3 +86,323 @@ describe('each loops', async () => { expect(result3).toBe('11') }) }) + +describe('nested loops', async () => { + it('renders all inner elements for each outer iteration', async () => { + const prompt = removeCommonIndent(` + {{ for category in categories }} + + Category: {{category.name}} + {{ for fruit in category.fruits }} + - {{ fruit }} + {{ endfor }} + + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + categories: [ + { name: 'andres', fruits: ['banana'] }, + { name: 'manu', fruits: ['apple', 'tomato'] }, + { name: 'paula', fruits: ['watermelon', 'banana'] }, + ], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + + expect(result.messages.length).toBe(3) + expect(texts[0]).toContain('andres') + expect(texts[0]).toContain('banana') + expect(texts[1]).toContain('manu') + expect(texts[1]).toContain('apple') + expect(texts[1]).toContain('tomato') + expect(texts[2]).toContain('paula') + expect(texts[2]).toContain('watermelon') + expect(texts[2]).toContain('banana') + }) + + it('does not skip inner elements when outer arrays have different lengths', async () => { + const prompt = removeCommonIndent(` + {{ for group in groups }} + + {{ for item in group }} + {{item}} + {{ endfor }} + + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + groups: [['a'], ['b', 'c'], ['d', 'e', 'f']], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + + expect(result.messages.length).toBe(3) + expect(texts[0]).toContain('a') + expect(texts[1]).toContain('b') + expect(texts[1]).toContain('c') + expect(texts[2]).toContain('d') + expect(texts[2]).toContain('e') + expect(texts[2]).toContain('f') + }) + + it('handles inner loop with single-element arrays followed by multi-element arrays', async () => { + const prompt = removeCommonIndent(` + {{ for group in groups }} + + {{ for item in group }} + {{item}} + {{ endfor }} + + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + groups: [['x'], ['y', 'z']], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + + expect(result.messages.length).toBe(2) + expect(texts[0]).toContain('x') + expect(texts[1]).toContain('y') + expect(texts[1]).toContain('z') + }) + + it('handles empty inner arrays correctly', async () => { + const prompt = removeCommonIndent(` + {{ for group in groups }} + + Group: + {{ for item in group }} + {{item}} + {{ endfor }} + + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + groups: [[], ['a', 'b'], []], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + + expect(result.messages.length).toBe(3) + expect(texts[1]).toContain('a') + expect(texts[1]).toContain('b') + }) + + it('renders three levels of nested loops correctly', async () => { + const prompt = removeCommonIndent(` + {{ for a in [1, 2] }} + {{ for b in [1, 2] }} + {{ for c in [1, 2] }} + {{a}}.{{b}}.{{c}} + {{ endfor }} + {{ endfor }} + {{ endfor }} + `) + + const result = await render({ + prompt, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + expect(result.messages.length).toBe(8) + expect(texts).toEqual([ + '1.1.1', + '1.1.2', + '1.2.1', + '1.2.2', + '2.1.1', + '2.1.2', + '2.2.1', + '2.2.2', + ]) + }) + + it('handles nested loops with conditionals inside', async () => { + const prompt = removeCommonIndent(` + {{ for group in groups }} + + {{ for item in group }} + {{ if item > 2 }}big:{{item}}{{ else }}small:{{item}}{{ endif }} + {{ endfor }} + + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + groups: [ + [1, 3], + [4, 2, 5], + ], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + + expect(result.messages.length).toBe(2) + expect(texts[0]).toContain('small:1') + expect(texts[0]).toContain('big:3') + expect(texts[1]).toContain('big:4') + expect(texts[1]).toContain('small:2') + expect(texts[1]).toContain('big:5') + }) + + it('handles nested loops with index variables', async () => { + const prompt = removeCommonIndent(` + {{ for group, gi in groups }} + + {{ for item, ii in group }} + {{gi}}.{{ii}}:{{item}} + {{ endfor }} + + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + groups: [['a'], ['b', 'c', 'd']], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + + expect(result.messages.length).toBe(2) + expect(texts[0]).toContain('0.0:a') + expect(texts[1]).toContain('1.0:b') + expect(texts[1]).toContain('1.1:c') + expect(texts[1]).toContain('1.2:d') + }) + + it('correctly aggregates values across nested loop iterations', async () => { + const prompt = removeCommonIndent(` + {{ for outer in [[1, 2], [3, 4], [5, 6]] }} + {{ for inner in outer }} + {{inner}} + {{ endfor }} + {{ endfor }} + `) + + const result = await render({ + prompt, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + expect(texts).toEqual(['1', '2', '3', '4', '5', '6']) + }) + + it('does not produce duplicate or missing messages with varying inner array sizes', async () => { + const prompt = removeCommonIndent(` + {{ for row in rows }} + {{ for cell in row.cells }} + {{row.id}}-{{cell}} + {{ endfor }} + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + rows: [ + { id: 'A', cells: ['1'] }, + { id: 'B', cells: ['2', '3'] }, + { id: 'C', cells: ['4'] }, + { id: 'D', cells: ['5', '6', '7'] }, + ], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + expect(texts).toEqual([ + 'A-1', + 'B-2', + 'B-3', + 'C-4', + 'D-5', + 'D-6', + 'D-7', + ]) + }) + + it('handles nested loops where first outer iteration has more inner elements than subsequent ones', async () => { + const prompt = removeCommonIndent(` + {{ for group in groups }} + {{ for item in group }} + {{item}} + {{ endfor }} + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + groups: [ + ['a', 'b', 'c'], + ['d'], + ['e', 'f'], + ], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + expect(texts).toEqual(['a', 'b', 'c', 'd', 'e', 'f']) + }) + + it('handles nested loops with else blocks in the inner loop', async () => { + const prompt = removeCommonIndent(` + {{ for group in groups }} + + {{ for item in group }} + item:{{item}} + {{ else }} + empty + {{ endfor }} + + {{ endfor }} + `) + + const result = await render({ + prompt, + parameters: { + groups: [['a'], [], ['b', 'c']], + }, + adapter: Adapters.default, + }) + + const texts = getMessageTexts(result.messages) + + expect(result.messages.length).toBe(3) + expect(texts[0]).toContain('item:a') + expect(texts[1]).toContain('empty') + expect(texts[2]).toContain('item:b') + expect(texts[2]).toContain('item:c') + }) +}) diff --git a/src/compiler/base/nodes/for.ts b/src/compiler/base/nodes/for.ts index cade16f..af452f1 100644 --- a/src/compiler/base/nodes/for.ts +++ b/src/compiler/base/nodes/for.ts @@ -1,12 +1,31 @@ import { hasContent, isIterable } from '$promptl/compiler/utils' import errors from '$promptl/error/errors' -import { ForBlock } from '$promptl/parser/interfaces' +import { ForBlock, TemplateNode } from '$promptl/parser/interfaces' import { CompileNodeContext, TemplateNodeWithStatus } from '../types' type ForNodeWithStatus = TemplateNodeWithStatus & { status: TemplateNodeWithStatus['status'] & { loopIterationIndex: number + loopInvocationCount: number + } +} + +function clearNodeStatus(node: TemplateNode): void { + const n = node as TemplateNodeWithStatus + if (n.status) { + delete n.status.completedAs + delete n.status.scopePointers + } + if (node.children) { + for (const child of node.children) { + clearNodeStatus(child) + } + } + if ('else' in node && node.else?.children) { + for (const child of node.else.children) { + clearNodeStatus(child) + } } } @@ -59,6 +78,7 @@ export async function compile({ ) } + const invocationCount = nodeWithStatus.status.loopInvocationCount ?? 0 let i = 0 for await (const element of iterableElement) { @@ -80,7 +100,7 @@ export async function compile({ isInsideMessageTag, isInsideContentTag, fullPath, - completedValue: `step_${i}`, + completedValue: `step_${invocationCount}_${i}`, }) } @@ -90,5 +110,10 @@ export async function compile({ nodeWithStatus.status = { ...nodeWithStatus.status, loopIterationIndex: 0, + loopInvocationCount: invocationCount + 1, + } + + for (const child of node.children ?? []) { + clearNodeStatus(child) } } diff --git a/src/compiler/chain.test.ts b/src/compiler/chain.test.ts index da53525..7ca63cd 100644 --- a/src/compiler/chain.test.ts +++ b/src/compiler/chain.test.ts @@ -448,6 +448,37 @@ describe('chain', async () => { }) }) + it('handles nested loops with steps and varying inner array sizes', async () => { + const prompt = removeCommonIndent(` + {{for group in groups}} + {{for item in group.items}} + {{group.name}}-{{item}} + <${TAG_NAMES.step} /> + {{endfor}} + {{endfor}} + `) + + const chain = new Chain({ + prompt, + parameters: { + groups: [ + { name: 'A', items: ['x'] }, + { name: 'B', items: ['y', 'z'] }, + { name: 'C', items: ['w'] }, + ], + }, + adapter: Adapters.default, + }) + + const { messages } = await complete({ chain }) + const userMessages = messages.filter((m) => m.role === MessageRole.user) + const userTexts = userMessages.map((m) => + m.content.map((c) => (c as TextContent).text).join(''), + ) + + expect(userTexts).toEqual(['A-x', 'B-y', 'B-z', 'C-w']) + }) + it('saves the response in a variable', async () => { const prompt = removeCommonIndent(` <${TAG_NAMES.step} raw="rawResponse" as="responseText"/>