Skip to content

Commit c140925

Browse files
Pavel910Pavel910
committed
fix: improve LLM prompt and response handling
1 parent a12e471 commit c140925

4 files changed

Lines changed: 346 additions & 4 deletions

File tree

‎packages/ai-powerups/__tests__/features/WbGeneratePageContent/ComponentFilter.test.ts‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
import { describe, it, expect, beforeEach } from "vitest";
2+
import { ComponentFilter } from "~/api/features/WbGeneratePageContent/ComponentFilter.js";
3+
4+
const catalog = [
5+
{ name: "Webiny/Grid" },
6+
{ name: "Webiny/GridColumn" },
7+
{ name: "Webiny/Lexical" },
8+
{ name: "Webiny/Image" },
9+
{ name: "Webiny/Box" }
10+
];
11+
12+
describe("ComponentFilter", () => {
13+
let filter: ComponentFilter;
14+
15+
beforeEach(() => {
16+
filter = new ComponentFilter(catalog);
17+
});
18+
19+
it("should pass through all valid components unchanged", () => {
20+
const elements = [
21+
{ component: "Webiny/Lexical", inputs: { content: "hello" } },
22+
{ component: "Webiny/Image", inputs: { src: "img.png" } }
23+
];
24+
25+
expect(filter.filter(elements)).toEqual(elements);
26+
});
27+
28+
it("should remove root-level elements with invalid component names", () => {
29+
const elements = [
30+
{ component: "Webiny/Lexical", inputs: { content: "hello" } },
31+
{ component: "Webiny/FakeComponent", inputs: { content: "bad" } },
32+
{ component: "Webiny/Image", inputs: { src: "img.png" } }
33+
];
34+
35+
expect(filter.filter(elements)).toEqual([
36+
{ component: "Webiny/Lexical", inputs: { content: "hello" } },
37+
{ component: "Webiny/Image", inputs: { src: "img.png" } }
38+
]);
39+
});
40+
41+
it("should remove nested CreateElement actions with invalid components", () => {
42+
const elements = [
43+
{
44+
component: "Webiny/Box",
45+
inputs: {
46+
children: [
47+
{
48+
action: "CreateElement",
49+
params: {
50+
component: "Webiny/Lexical",
51+
inputs: { content: "valid" }
52+
}
53+
},
54+
{
55+
action: "CreateElement",
56+
params: {
57+
component: "Webiny/MadeUp",
58+
inputs: { content: "invalid" }
59+
}
60+
}
61+
]
62+
}
63+
}
64+
];
65+
66+
const result = filter.filter(elements);
67+
const children = (result[0] as any).inputs.children;
68+
expect(children).toHaveLength(1);
69+
expect(children[0].params.component).toBe("Webiny/Lexical");
70+
});
71+
72+
it("should handle deeply nested Grid > GridColumn > children", () => {
73+
const elements = [
74+
{
75+
component: "Webiny/Grid",
76+
inputs: {
77+
gridLayout: "6-6",
78+
columns: [
79+
{
80+
children: {
81+
action: "CreateElement",
82+
params: {
83+
component: "Webiny/GridColumn",
84+
inputs: {
85+
children: [
86+
{
87+
action: "CreateElement",
88+
params: {
89+
component: "Webiny/Lexical",
90+
inputs: { content: "valid" }
91+
}
92+
},
93+
{
94+
action: "CreateElement",
95+
params: {
96+
component: "Webiny/Hallucinated",
97+
inputs: { content: "bad" }
98+
}
99+
}
100+
]
101+
}
102+
}
103+
}
104+
}
105+
]
106+
}
107+
}
108+
];
109+
110+
const result = filter.filter(elements);
111+
const column = (result[0] as any).inputs.columns[0].children;
112+
const innerChildren = column.params.inputs.children;
113+
expect(innerChildren).toHaveLength(1);
114+
expect(innerChildren[0].params.component).toBe("Webiny/Lexical");
115+
});
116+
117+
it("should leave tool envelopes untouched", () => {
118+
const elements = [
119+
{
120+
component: "Webiny/Lexical",
121+
inputs: {
122+
content: {
123+
tool: "textToLexical",
124+
params: { text: "hello" }
125+
}
126+
}
127+
}
128+
];
129+
130+
expect(filter.filter(elements)).toEqual(elements);
131+
});
132+
133+
it("should return empty array when all components are invalid", () => {
134+
const elements = [
135+
{ component: "Fake/One", inputs: {} },
136+
{ component: "Fake/Two", inputs: {} }
137+
];
138+
139+
expect(filter.filter(elements)).toEqual([]);
140+
});
141+
142+
it("should return empty array for empty input", () => {
143+
expect(filter.filter([])).toEqual([]);
144+
});
145+
146+
it("should keep parent when removing children empties a slot", () => {
147+
const elements = [
148+
{
149+
component: "Webiny/Box",
150+
inputs: {
151+
children: [
152+
{
153+
action: "CreateElement",
154+
params: {
155+
component: "Webiny/NonExistent",
156+
inputs: {}
157+
}
158+
}
159+
]
160+
}
161+
}
162+
];
163+
164+
const result = filter.filter(elements);
165+
expect(result).toHaveLength(1);
166+
expect((result[0] as any).inputs.children).toEqual([]);
167+
});
168+
169+
it("should remove a single CreateElement in a slot input when invalid", () => {
170+
const elements = [
171+
{
172+
component: "Webiny/Grid",
173+
inputs: {
174+
columns: [
175+
{
176+
children: {
177+
action: "CreateElement",
178+
params: {
179+
component: "Webiny/NonExistent",
180+
inputs: {}
181+
}
182+
}
183+
}
184+
]
185+
}
186+
}
187+
];
188+
189+
const result = filter.filter(elements);
190+
const column = (result[0] as any).inputs.columns[0].children;
191+
expect(column).toBeNull();
192+
});
193+
194+
it("should preserve non-element values in inputs", () => {
195+
const elements = [
196+
{
197+
component: "Webiny/Grid",
198+
inputs: {
199+
gridLayout: "6-6",
200+
someFlag: true,
201+
count: 42
202+
}
203+
}
204+
];
205+
206+
expect(filter.filter(elements)).toEqual(elements);
207+
});
208+
});

‎packages/ai-powerups/src/api/features/WbGeneratePageContent/ComponentFilter.ts‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
/**
2+
* Validates LLM-generated page content against the component catalog.
3+
* Recursively walks the element tree (including nested CreateElement actions
4+
* inside slot inputs) and removes any element whose component name is not
5+
* in the catalog. Tool envelopes are left untouched.
6+
*/
7+
export class ComponentFilter {
8+
private readonly validNames: Set<string>;
9+
10+
constructor(components: Array<{ name: string }>) {
11+
this.validNames = new Set(components.map(c => c.name));
12+
}
13+
14+
filter(elements: unknown[]): unknown[] {
15+
return elements
16+
.filter(el => this.isValid(el))
17+
.map(el => this.processInputs(el as Record<string, unknown>));
18+
}
19+
20+
private isValid(el: unknown): boolean {
21+
if (!el || typeof el !== "object") {
22+
return true;
23+
}
24+
25+
const obj = el as Record<string, unknown>;
26+
27+
if ("component" in obj && typeof obj.component === "string") {
28+
return this.validNames.has(obj.component);
29+
}
30+
31+
if (obj.action === "CreateElement" && obj.params && typeof obj.params === "object") {
32+
const params = obj.params as Record<string, unknown>;
33+
if ("component" in params && typeof params.component === "string") {
34+
return this.validNames.has(params.component);
35+
}
36+
}
37+
38+
return true;
39+
}
40+
41+
private processInputs(el: Record<string, unknown>): Record<string, unknown> {
42+
const inputs = this.getInputs(el);
43+
if (!inputs) {
44+
return el;
45+
}
46+
47+
const cleaned: Record<string, unknown> = {};
48+
for (const [key, value] of Object.entries(inputs)) {
49+
cleaned[key] = this.processValue(value);
50+
}
51+
52+
return this.setInputs(el, cleaned);
53+
}
54+
55+
private getInputs(el: Record<string, unknown>): Record<string, unknown> | null {
56+
if ("inputs" in el && el.inputs && typeof el.inputs === "object") {
57+
return el.inputs as Record<string, unknown>;
58+
}
59+
60+
if (
61+
el.action === "CreateElement" &&
62+
el.params &&
63+
typeof el.params === "object" &&
64+
"inputs" in (el.params as Record<string, unknown>)
65+
) {
66+
const params = el.params as Record<string, unknown>;
67+
return params.inputs as Record<string, unknown>;
68+
}
69+
70+
return null;
71+
}
72+
73+
private setInputs(
74+
el: Record<string, unknown>,
75+
inputs: Record<string, unknown>
76+
): Record<string, unknown> {
77+
if ("inputs" in el) {
78+
return { ...el, inputs };
79+
}
80+
81+
if (el.action === "CreateElement" && el.params && typeof el.params === "object") {
82+
return { ...el, params: { ...el.params, inputs } };
83+
}
84+
85+
return el;
86+
}
87+
88+
private processValue(value: unknown): unknown {
89+
if (Array.isArray(value)) {
90+
return value.filter(item => this.isValid(item)).map(item => this.processValue(item));
91+
}
92+
93+
if (!value || typeof value !== "object") {
94+
return value;
95+
}
96+
97+
const obj = value as Record<string, unknown>;
98+
99+
if ("tool" in obj) {
100+
return value;
101+
}
102+
103+
if (obj.action === "CreateElement") {
104+
if (!this.isValid(obj)) {
105+
return null;
106+
}
107+
return this.processInputs(obj);
108+
}
109+
110+
if ("component" in obj && "inputs" in obj) {
111+
return this.processInputs(obj);
112+
}
113+
114+
const result: Record<string, unknown> = {};
115+
let changed = false;
116+
for (const [k, v] of Object.entries(obj)) {
117+
const processed = this.processValue(v);
118+
result[k] = processed;
119+
if (processed !== v) {
120+
changed = true;
121+
}
122+
}
123+
return changed ? result : value;
124+
}
125+
}

‎packages/ai-powerups/src/api/features/WbGeneratePageContent/WbGeneratePageContentUseCase.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import type {
1414
} from "./abstractions.js";
1515
import { buildDomainPrompt } from "./buildPrompt.js";
1616
import { LlmJsonResponse } from "./LlmJsonResponse.js";
17+
import { ComponentFilter } from "./ComponentFilter.js";
1718

1819
class WbGeneratePageContentUseCaseImpl implements WbGeneratePageContentUseCase.Interface {
1920
constructor(
@@ -60,7 +61,8 @@ class WbGeneratePageContentUseCaseImpl implements WbGeneratePageContentUseCase.I
6061
Object.assign(sdkTools, projectFileTool);
6162
}
6263

63-
const systemText = buildDomainPrompt(params.components, params.tools) + context.toString();
64+
const components = params.components as Array<{ name: string }>;
65+
const systemText = buildDomainPrompt(components, params.tools) + context.toString();
6466

6567
const system = {
6668
role: "system" as const,
@@ -91,7 +93,9 @@ class WbGeneratePageContentUseCaseImpl implements WbGeneratePageContentUseCase.I
9193
aiResult.text ||
9294
(aiResult.steps.filter(step => step.text.length > 0).pop()?.text ?? "");
9395

94-
const output = LlmJsonResponse.fromRawText(text).toString();
96+
const elements = LlmJsonResponse.fromRawText(text).toArray();
97+
const componentFilter = new ComponentFilter(components);
98+
const output = JSON.stringify(componentFilter.filter(elements));
9599

96100
const filesRead = new Set<string>();
97101
let toolCallsMade = 0;

‎packages/ai-powerups/src/api/features/WbGeneratePageContent/buildPrompt.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
export function buildDomainPrompt(components: unknown, tools: unknown): string {
1+
export function buildDomainPrompt(components: Array<{ name: string }>, tools: unknown): string {
2+
const componentNames = components.map(c => `"${c.name}"`).join(" | ");
23
return `You are a page content generator. Given a user prompt, generate structured page content using the provided component catalog and available tools.
34
45
###
@@ -72,8 +73,10 @@ ${JSON.stringify(tools, null, 2)}
7273
### Page Schema
7374
7475
\`\`\`typescript
76+
type ComponentName = ${componentNames};
77+
7578
type ElementSchema = {
76-
component: string;
79+
component: ComponentName;
7780
inputs: Record<string, unknown>;
7881
};
7982
@@ -138,5 +141,7 @@ Key rules:
138141
- Webiny/GridColumn's "children" is an array of CreateElement actions for
139142
the actual content
140143
144+
IMPORTANT: Only use components listed in the Component Catalog above. Do NOT invent component names. Any element with an unrecognized component name will be silently removed from the output.
145+
141146
You MUST return a parsable JSON object with a "page" key containing the array of elements. No extra text outside the JSON.`;
142147
}

0 commit comments

Comments
 (0)