diff --git a/packages/editor-ui/src/components/ButtonParameter/ButtonParameter.vue b/packages/editor-ui/src/components/ButtonParameter/ButtonParameter.vue
index a321f3cc70..a28b68d3dd 100644
--- a/packages/editor-ui/src/components/ButtonParameter/ButtonParameter.vue
+++ b/packages/editor-ui/src/components/ButtonParameter/ButtonParameter.vue
@@ -105,6 +105,7 @@ async function onSubmit() {
 				const updateInformation = await generateCodeForAiTransform(
 					prompt.value,
 					getPath(target as string),
+					5,
 				);
 				if (!updateInformation) return;
 
diff --git a/packages/editor-ui/src/components/ButtonParameter/utils.test.ts b/packages/editor-ui/src/components/ButtonParameter/utils.test.ts
new file mode 100644
index 0000000000..df7e13d477
--- /dev/null
+++ b/packages/editor-ui/src/components/ButtonParameter/utils.test.ts
@@ -0,0 +1,88 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { generateCodeForAiTransform } from './utils';
+import { createPinia, setActivePinia } from 'pinia';
+import { generateCodeForPrompt } from '@/api/ai';
+
+vi.mock('./utils', async () => {
+	const actual = await vi.importActual('./utils');
+	return {
+		...actual,
+		getSchemas: vi.fn(() => ({
+			parentNodesSchemas: { test: 'parentSchema' },
+			inputSchema: { test: 'inputSchema' },
+		})),
+	};
+});
+
+vi.mock('@/stores/root.store', () => ({
+	useRootStore: () => ({
+		pushRef: 'mockRootPushRef',
+		restApiContext: {},
+	}),
+}));
+
+vi.mock('@/stores/ndv.store', () => ({
+	useNDVStore: () => ({
+		pushRef: 'mockNdvPushRef',
+	}),
+}));
+
+vi.mock('@/stores/settings.store', () => ({
+	useSettingsStore: vi.fn(() => ({ settings: {}, isAskAiEnabled: true })),
+}));
+
+vi.mock('prettier', () => ({
+	format: vi.fn(async (code) => await Promise.resolve(`formatted-${code}`)),
+}));
+
+vi.mock('@/api/ai', () => ({
+	generateCodeForPrompt: vi.fn(),
+}));
+
+describe('generateCodeForAiTransform - Retry Tests', () => {
+	beforeEach(() => {
+		vi.clearAllMocks();
+		const pinia = createPinia();
+		setActivePinia(pinia);
+	});
+
+	it('should retry and succeed on the second attempt', async () => {
+		const mockGeneratedCode = 'const example = "retry success";';
+
+		vi.mocked(generateCodeForPrompt)
+			.mockRejectedValueOnce(new Error('First attempt failed'))
+			.mockResolvedValueOnce({ code: mockGeneratedCode });
+
+		const result = await generateCodeForAiTransform('test prompt', 'test/path', 2);
+
+		expect(result).toEqual({
+			name: 'test/path',
+			value: 'formatted-const example = "retry success";',
+		});
+		expect(generateCodeForPrompt).toHaveBeenCalledTimes(2);
+	});
+
+	it('should exhaust retries and throw an error', async () => {
+		vi.mocked(generateCodeForPrompt).mockRejectedValue(new Error('All attempts failed'));
+
+		await expect(generateCodeForAiTransform('test prompt', 'test/path', 3)).rejects.toThrow(
+			'All attempts failed',
+		);
+
+		expect(generateCodeForPrompt).toHaveBeenCalledTimes(3);
+	});
+
+	it('should succeed on the first attempt without retries', async () => {
+		const mockGeneratedCode = 'const example = "no retries needed";';
+		vi.mocked(generateCodeForPrompt).mockResolvedValue({ code: mockGeneratedCode });
+
+		const result = await generateCodeForAiTransform('test prompt', 'test/path');
+
+		expect(result).toEqual({
+			name: 'test/path',
+			value: 'formatted-const example = "no retries needed";',
+		});
+		expect(generateCodeForPrompt).toHaveBeenCalledTimes(1);
+	});
+});
diff --git a/packages/editor-ui/src/components/ButtonParameter/utils.ts b/packages/editor-ui/src/components/ButtonParameter/utils.ts
index 80893f8667..b95846975f 100644
--- a/packages/editor-ui/src/components/ButtonParameter/utils.ts
+++ b/packages/editor-ui/src/components/ButtonParameter/utils.ts
@@ -4,10 +4,10 @@ import { useWorkflowsStore } from '@/stores/workflows.store';
 import { useNDVStore } from '@/stores/ndv.store';
 import { useDataSchema } from '@/composables/useDataSchema';
 import { executionDataToJson } from '@/utils/nodeTypesUtils';
-import { generateCodeForPrompt } from '../../api/ai';
-import { useRootStore } from '../../stores/root.store';
-import { type AskAiRequest } from '../../types/assistant.types';
-import { useSettingsStore } from '../../stores/settings.store';
+import { generateCodeForPrompt } from '@/api/ai';
+import { useRootStore } from '@/stores/root.store';
+import { type AskAiRequest } from '@/types/assistant.types';
+import { useSettingsStore } from '@/stores/settings.store';
 import { format } from 'prettier';
 import jsParser from 'prettier/plugins/babel';
 import * as estree from 'prettier/plugins/estree';
@@ -43,7 +43,7 @@ export function getSchemas() {
 
 			return {
 				nodeName: node?.name || '',
-				schema: getSchemaForExecutionData(executionDataToJson(inputData), true),
+				schema: getSchemaForExecutionData(executionDataToJson(inputData), false),
 			};
 		})
 		.filter((node) => node.schema?.value.length > 0);
@@ -57,7 +57,7 @@ export function getSchemas() {
 	};
 }
 
-export async function generateCodeForAiTransform(prompt: string, path: string) {
+export async function generateCodeForAiTransform(prompt: string, path: string, retries = 1) {
 	const schemas = getSchemas();
 
 	const payload: AskAiRequest.RequestPayload = {
@@ -74,7 +74,20 @@ export async function generateCodeForAiTransform(prompt: string, path: string) {
 	let value;
 	if (useSettingsStore().isAskAiEnabled) {
 		const { restApiContext } = useRootStore();
-		const { code } = await generateCodeForPrompt(restApiContext, payload);
+
+		let code = '';
+
+		while (retries > 0) {
+			try {
+				const { code: generatedCode } = await generateCodeForPrompt(restApiContext, payload);
+				code = generatedCode;
+				break;
+			} catch (e) {
+				retries--;
+				if (!retries) throw e;
+			}
+		}
+
 		value = code;
 	} else {
 		throw new ApplicationError('AI code generation is not enabled');
diff --git a/packages/editor-ui/src/components/NodeExecuteButton.vue b/packages/editor-ui/src/components/NodeExecuteButton.vue
index f801f0701c..b46c6b0a43 100644
--- a/packages/editor-ui/src/components/NodeExecuteButton.vue
+++ b/packages/editor-ui/src/components/NodeExecuteButton.vue
@@ -287,6 +287,7 @@ async function onClick() {
 			const updateInformation = await generateCodeForAiTransform(
 				prompt,
 				`parameters.${AI_TRANSFORM_JS_CODE}`,
+				5,
 			);
 			if (!updateInformation) return;