diff --git a/web/src/pages/agent/form-sheet/use-form-config-map.tsx b/web/src/pages/agent/form-sheet/use-form-config-map.tsx
index 13c1e95f9..de09956ba 100644
--- a/web/src/pages/agent/form-sheet/use-form-config-map.tsx
+++ b/web/src/pages/agent/form-sheet/use-form-config-map.tsx
@@ -1,3 +1,4 @@
+import { CodeTemplateStrMap, ProgrammingLanguage } from '@/constants/agent';
 import { useTranslation } from 'react-i18next';
 import { z } from 'zod';
 import { Operator } from '../constant';
@@ -133,13 +134,17 @@ export function useFormConfigMap() {
     },
     [Operator.Code]: {
       component: CodeForm,
-      defaultValues: { arguments: [] },
+      defaultValues: {
+        lang: ProgrammingLanguage.Python,
+        script: CodeTemplateStrMap[ProgrammingLanguage.Python],
+        arguments: [],
+      },
       schema: z.object({
+        lang: z.string(),
         script: z.string(),
         arguments: z.array(
           z.object({ name: z.string(), component_id: z.string() }),
         ),
-        lang: z.string(),
       }),
     },
     [Operator.Baidu]: {
diff --git a/web/src/pages/agent/form/code-form/index.tsx b/web/src/pages/agent/form/code-form/index.tsx
index c5282a6b1..4db2a96a5 100644
--- a/web/src/pages/agent/form/code-form/index.tsx
+++ b/web/src/pages/agent/form/code-form/index.tsx
@@ -1,6 +1,7 @@
 import Editor, { loader } from '@monaco-editor/react';
 import { INextOperatorForm } from '../../interface';
 
+import { FormContainer } from '@/components/form-container';
 import {
   Form,
   FormControl,
@@ -9,11 +10,17 @@ import {
   FormLabel,
   FormMessage,
 } from '@/components/ui/form';
+import { Input } from '@/components/ui/input';
 import { RAGFlowSelect } from '@/components/ui/select';
 import { CodeTemplateStrMap, ProgrammingLanguage } from '@/constants/agent';
 import { ICodeForm } from '@/interfaces/database/flow';
 import { useEffect } from 'react';
-import { DynamicInputVariable } from './next-variable';
+import { useTranslation } from 'react-i18next';
+import {
+  DynamicInputVariable,
+  TypeOptions,
+  VariableTitle,
+} from './next-variable';
 
 loader.config({ paths: { vs: '/vs' } });
 
@@ -24,6 +31,7 @@ const options = [
 
 const CodeForm = ({ form, node }: INextOperatorForm) => {
   const formData = node?.data.form as ICodeForm;
+  const { t } = useTranslation();
 
   useEffect(() => {
     // TODO: Direct operation zustand is more elegant
@@ -35,42 +43,100 @@ const CodeForm = ({ form, node }: INextOperatorForm) => {
 
   return (
     <Form {...form}>
-      <DynamicInputVariable node={node}></DynamicInputVariable>
-      <FormField
-        control={form.control}
-        name="script"
-        render={({ field }) => (
-          <FormItem>
-            <FormLabel>
+      <form
+        className="p-5 space-y-5"
+        onSubmit={(e) => {
+          e.preventDefault();
+        }}
+      >
+        <DynamicInputVariable
+          node={node}
+          title={t('flow.input')}
+        ></DynamicInputVariable>
+        <FormField
+          control={form.control}
+          name="script"
+          render={({ field }) => (
+            <FormItem>
+              <FormLabel className="flex items-center justify-between">
+                Code
+                <FormField
+                  control={form.control}
+                  name="lang"
+                  render={({ field }) => (
+                    <FormItem>
+                      <FormControl>
+                        <RAGFlowSelect {...field} options={options} />
+                      </FormControl>
+                      <FormMessage />
+                    </FormItem>
+                  )}
+                />
+              </FormLabel>
+              <FormControl>
+                <Editor
+                  height={300}
+                  theme="vs-dark"
+                  language={formData.lang}
+                  options={{
+                    minimap: { enabled: false },
+                    automaticLayout: true,
+                  }}
+                  {...field}
+                />
+              </FormControl>
+              <FormMessage />
+            </FormItem>
+          )}
+        />
+
+        {formData.lang === ProgrammingLanguage.Python ? (
+          <DynamicInputVariable
+            node={node}
+            title={'Return Values'}
+            name={'return'}
+          ></DynamicInputVariable>
+        ) : (
+          <div>
+            <VariableTitle title={'Return Values'}></VariableTitle>
+            <FormContainer className="space-y-5">
               <FormField
                 control={form.control}
-                name="lang"
+                name={'return.name'}
                 render={({ field }) => (
                   <FormItem>
+                    <FormLabel>Name</FormLabel>
                     <FormControl>
-                      <RAGFlowSelect {...field} options={options} />
+                      <Input
+                        {...field}
+                        placeholder={t('common.pleaseInput')}
+                      ></Input>
                     </FormControl>
                     <FormMessage />
                   </FormItem>
                 )}
               />
-            </FormLabel>
-            <FormControl>
-              <Editor
-                height={600}
-                theme="vs-dark"
-                language={formData.lang}
-                options={{
-                  minimap: { enabled: false },
-                  automaticLayout: true,
-                }}
-                {...field}
+              <FormField
+                control={form.control}
+                name={`return.component_id`}
+                render={({ field }) => (
+                  <FormItem className="flex-1">
+                    <FormLabel>Type</FormLabel>
+                    <FormControl>
+                      <RAGFlowSelect
+                        placeholder={t('common.pleaseSelect')}
+                        options={TypeOptions}
+                        {...field}
+                      ></RAGFlowSelect>
+                    </FormControl>
+                    <FormMessage />
+                  </FormItem>
+                )}
               />
-            </FormControl>
-            <FormMessage />
-          </FormItem>
+            </FormContainer>
+          </div>
         )}
-      />
+      </form>
     </Form>
   );
 };
diff --git a/web/src/pages/agent/form/code-form/next-variable.tsx b/web/src/pages/agent/form/code-form/next-variable.tsx
index 90b2f8099..153257ad8 100644
--- a/web/src/pages/agent/form/code-form/next-variable.tsx
+++ b/web/src/pages/agent/form/code-form/next-variable.tsx
@@ -1,23 +1,19 @@
 'use client';
 
-import { SideDown } from '@/assets/icon/Icon';
+import { FormContainer } from '@/components/form-container';
 import { Button } from '@/components/ui/button';
-import {
-  Collapsible,
-  CollapsibleContent,
-  CollapsibleTrigger,
-} from '@/components/ui/collapsible';
 import {
   FormControl,
-  FormDescription,
   FormField,
   FormItem,
   FormMessage,
 } from '@/components/ui/form';
 import { Input } from '@/components/ui/input';
 import { RAGFlowSelect } from '@/components/ui/select';
+import { Separator } from '@/components/ui/separator';
 import { RAGFlowNodeType } from '@/interfaces/database/flow';
-import { Plus, Trash2 } from 'lucide-react';
+import { Plus, X } from 'lucide-react';
+import { ReactNode } from 'react';
 import { useFieldArray, useFormContext } from 'react-hook-form';
 import { useTranslation } from 'react-i18next';
 import { useBuildComponentIdSelectOptions } from '../../hooks/use-get-begin-query';
@@ -27,6 +23,15 @@ interface IProps {
   name?: string;
 }
 
+export const TypeOptions = [
+  'String',
+  'Number',
+  'Boolean',
+  'Array[String]',
+  'Array[Number]',
+  'Object',
+].map((x) => ({ label: x, value: x }));
+
 export function DynamicVariableForm({ node, name = 'arguments' }: IProps) {
   const { t } = useTranslation();
   const form = useFormContext();
@@ -42,17 +47,16 @@ export function DynamicVariableForm({ node, name = 'arguments' }: IProps) {
   );
 
   return (
-    <div>
+    <div className="space-y-5">
       {fields.map((field, index) => {
         const typeField = `${name}.${index}.name`;
         return (
-          <div key={field.id} className="flex items-center gap-1">
+          <div key={field.id} className="flex items-center gap-2">
             <FormField
               control={form.control}
               name={typeField}
               render={({ field }) => (
                 <FormItem className="w-2/5">
-                  <FormDescription />
                   <FormControl>
                     <Input
                       {...field}
@@ -63,16 +67,18 @@ export function DynamicVariableForm({ node, name = 'arguments' }: IProps) {
                 </FormItem>
               )}
             />
+            <Separator className="w-3 text-text-sub-title" />
             <FormField
               control={form.control}
               name={`${name}.${index}.component_id`}
               render={({ field }) => (
                 <FormItem className="flex-1">
-                  <FormDescription />
                   <FormControl>
                     <RAGFlowSelect
                       placeholder={t('common.pleaseSelect')}
-                      options={valueOptions}
+                      options={
+                        name === 'arguments' ? valueOptions : TypeOptions
+                      }
                       {...field}
                     ></RAGFlowSelect>
                   </FormControl>
@@ -80,18 +86,16 @@ export function DynamicVariableForm({ node, name = 'arguments' }: IProps) {
                 </FormItem>
               )}
             />
-            <Trash2
-              className="cursor-pointer mx-3 size-4 text-colors-text-functional-danger"
-              onClick={() => remove(index)}
-            />
+            <Button variant={'ghost'} onClick={() => remove(index)}>
+              <X className="text-text-sub-title-invert " />
+            </Button>
           </div>
         );
       })}
       <Button
         onClick={() => append({ name: '', component_id: undefined })}
-        className="mt-4"
+        className="mt-4 border-dashed w-full"
         variant={'outline'}
-        size={'sm'}
       >
         <Plus />
         {t('flow.addVariable')}
@@ -100,22 +104,21 @@ export function DynamicVariableForm({ node, name = 'arguments' }: IProps) {
   );
 }
 
-export function DynamicInputVariable({ node }: IProps) {
-  const { t } = useTranslation();
+export function VariableTitle({ title }: { title: ReactNode }) {
+  return <div className="font-medium text-text-title pb-2">{title}</div>;
+}
 
+export function DynamicInputVariable({
+  node,
+  name,
+  title,
+}: IProps & { title: ReactNode }) {
   return (
-    <Collapsible defaultOpen className="group/collapsible">
-      <CollapsibleTrigger className="flex justify-between w-full pb-2">
-        <span className="font-bold text-2xl text-colors-text-neutral-strong">
-          {t('flow.input')}
-        </span>
-        <Button variant={'icon'} size={'icon'}>
-          <SideDown />
-        </Button>
-      </CollapsibleTrigger>
-      <CollapsibleContent>
-        <DynamicVariableForm node={node}></DynamicVariableForm>
-      </CollapsibleContent>
-    </Collapsible>
+    <section>
+      <VariableTitle title={title}></VariableTitle>
+      <FormContainer>
+        <DynamicVariableForm node={node} name={name}></DynamicVariableForm>
+      </FormContainer>
+    </section>
   );
 }