feat: add code generator (#9051)

Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com>

api/.env.example

@@ -239,6 +239,7 @@ UPLOAD_AUDIO_FILE_SIZE_LIMIT=50
 # Model Configuration
 MULTIMODAL_SEND_IMAGE_FORMAT=base64
 PROMPT_GENERATION_MAX_TOKENS=512
+CODE_GENERATION_MAX_TOKENS=1024
 
 # Mail configuration, support: resend, smtp
 MAIL_TYPE=

api/controllers/console/app/generator.py

@@ -52,4 +52,39 @@ class RuleGenerateApi(Resource):
         return rules
 
 
+class RuleCodeGenerateApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("instruction", type=str, required=True, nullable=False, location="json")
+        parser.add_argument("model_config", type=dict, required=True, nullable=False, location="json")
+        parser.add_argument("no_variable", type=bool, required=True, default=False, location="json")
+        parser.add_argument("code_language", type=str, required=False, default="javascript", location="json")
+        args = parser.parse_args()
+
+        account = current_user
+        CODE_GENERATION_MAX_TOKENS = int(os.getenv("CODE_GENERATION_MAX_TOKENS", "1024"))
+        try:
+            code_result = LLMGenerator.generate_code(
+                tenant_id=account.current_tenant_id,
+                instruction=args["instruction"],
+                model_config=args["model_config"],
+                code_language=args["code_language"],
+                max_tokens=CODE_GENERATION_MAX_TOKENS,
+            )
+        except ProviderTokenNotInitError as ex:
+            raise ProviderNotInitializeError(ex.description)
+        except QuotaExceededError:
+            raise ProviderQuotaExceededError()
+        except ModelCurrentlyNotSupportError:
+            raise ProviderModelCurrentlyNotSupportError()
+        except InvokeError as e:
+            raise CompletionRequestError(e.description)
+
+        return code_result
+
+
 api.add_resource(RuleGenerateApi, "/rule-generate")
+api.add_resource(RuleCodeGenerateApi, "/rule-code-generate")

api/core/llm_generator/llm_generator.py

@@ -8,6 +8,8 @@ from core.llm_generator.output_parser.suggested_questions_after_answer import Su
 from core.llm_generator.prompts import (
     CONVERSATION_TITLE_PROMPT,
     GENERATOR_QA_PROMPT,
+    JAVASCRIPT_CODE_GENERATOR_PROMPT_TEMPLATE,
+    PYTHON_CODE_GENERATOR_PROMPT_TEMPLATE,
     WORKFLOW_RULE_CONFIG_PROMPT_GENERATE_TEMPLATE,
 )
 from core.model_manager import ModelManager
@@ -239,6 +241,54 @@ class LLMGenerator:
 
         return rule_config
 
+    @classmethod
+    def generate_code(
+        cls,
+        tenant_id: str,
+        instruction: str,
+        model_config: dict,
+        code_language: str = "javascript",
+        max_tokens: int = 1000,
+    ) -> dict:
+        if code_language == "python":
+            prompt_template = PromptTemplateParser(PYTHON_CODE_GENERATOR_PROMPT_TEMPLATE)
+        else:
+            prompt_template = PromptTemplateParser(JAVASCRIPT_CODE_GENERATOR_PROMPT_TEMPLATE)
+
+        prompt = prompt_template.format(
+            inputs={
+                "INSTRUCTION": instruction,
+                "CODE_LANGUAGE": code_language,
+            },
+            remove_template_variables=False,
+        )
+
+        model_manager = ModelManager()
+        model_instance = model_manager.get_model_instance(
+            tenant_id=tenant_id,
+            model_type=ModelType.LLM,
+            provider=model_config.get("provider") if model_config else None,
+            model=model_config.get("name") if model_config else None,
+        )
+
+        prompt_messages = [UserPromptMessage(content=prompt)]
+        model_parameters = {"max_tokens": max_tokens, "temperature": 0.01}
+
+        try:
+            response = model_instance.invoke_llm(
+                prompt_messages=prompt_messages, model_parameters=model_parameters, stream=False
+            )
+
+            generated_code = response.message.content
+            return {"code": generated_code, "language": code_language, "error": ""}
+
+        except InvokeError as e:
+            error = str(e)
+            return {"code": "", "language": code_language, "error": f"Failed to generate code. Error: {error}"}
+        except Exception as e:
+            logging.exception(e)
+            return {"code": "", "language": code_language, "error": f"An unexpected error occurred: {str(e)}"}
+
     @classmethod
     def generate_qa_document(cls, tenant_id: str, query, document_language: str):
         prompt = GENERATOR_QA_PROMPT.format(language=document_language)

api/core/llm_generator/prompts.py

@@ -61,6 +61,73 @@ User Input: yo, 你今天咋样？
 User Input: 
 """  # noqa: E501
 
+PYTHON_CODE_GENERATOR_PROMPT_TEMPLATE = (
+    "You are an expert programmer. Generate code based on the following instructions:\n\n"
+    "Instructions: {{INSTRUCTION}}\n\n"
+    "Write the code in {{CODE_LANGUAGE}}.\n\n"
+    "Please ensure that you meet the following requirements:\n"
+    "1. Define a function named 'main'.\n"
+    "2. The 'main' function must return a dictionary (dict).\n"
+    "3. You may modify the arguments of the 'main' function, but include appropriate type hints.\n"
+    "4. The returned dictionary should contain at least one key-value pair.\n\n"
+    "5. You may ONLY use the following libraries in your code: \n"
+    "- json\n"
+    "- datetime\n"
+    "- math\n"
+    "- random\n"
+    "- re\n"
+    "- string\n"
+    "- sys\n"
+    "- time\n"
+    "- traceback\n"
+    "- uuid\n"
+    "- os\n"
+    "- base64\n"
+    "- hashlib\n"
+    "- hmac\n"
+    "- binascii\n"
+    "- collections\n"
+    "- functools\n"
+    "- operator\n"
+    "- itertools\n\n"
+    "Example:\n"
+    "def main(arg1: str, arg2: int) -> dict:\n"
+    "    return {\n"
+    '        "result": arg1 * arg2,\n'
+    "    }\n\n"
+    "IMPORTANT:\n"
+    "- Provide ONLY the code without any additional explanations, comments, or markdown formatting.\n"
+    "- DO NOT use markdown code blocks (``` or ``` python). Return the raw code directly.\n"
+    "- The code should start immediately after this instruction, without any preceding newlines or spaces.\n"
+    "- The code should be complete, functional, and follow best practices for {{CODE_LANGUAGE}}.\n\n"
+    "- Always use the format return {'result': ...} for the output.\n\n"
+    "Generated Code:\n"
+)
+JAVASCRIPT_CODE_GENERATOR_PROMPT_TEMPLATE = (
+    "You are an expert programmer. Generate code based on the following instructions:\n\n"
+    "Instructions: {{INSTRUCTION}}\n\n"
+    "Write the code in {{CODE_LANGUAGE}}.\n\n"
+    "Please ensure that you meet the following requirements:\n"
+    "1. Define a function named 'main'.\n"
+    "2. The 'main' function must return an object.\n"
+    "3. You may modify the arguments of the 'main' function, but include appropriate JSDoc annotations.\n"
+    "4. The returned object should contain at least one key-value pair.\n\n"
+    "5. The returned object should always be in the format: {result: ...}\n\n"
+    "Example:\n"
+    "function main(arg1, arg2) {\n"
+    "    return {\n"
+    "        result: arg1 * arg2\n"
+    "    };\n"
+    "}\n\n"
+    "IMPORTANT:\n"
+    "- Provide ONLY the code without any additional explanations, comments, or markdown formatting.\n"
+    "- DO NOT use markdown code blocks (``` or ``` javascript). Return the raw code directly.\n"
+    "- The code should start immediately after this instruction, without any preceding newlines or spaces.\n"
+    "- The code should be complete, functional, and follow best practices for {{CODE_LANGUAGE}}.\n\n"
+    "Generated Code:\n"
+)
+
+
 SUGGESTED_QUESTIONS_AFTER_ANSWER_INSTRUCTION_PROMPT = (
     "Please help me predict the three most likely questions that human would ask, "
     "and keeping each question under 20 characters.\n"

web/app/components/app/configuration/config/code-generator/get-code-generator-res.tsx

@@ -0,0 +1,200 @@
+import type { FC } from 'react'
+import React from 'react'
+import cn from 'classnames'
+import useBoolean from 'ahooks/lib/useBoolean'
+import { useTranslation } from 'react-i18next'
+import ConfigPrompt from '../../config-prompt'
+import { languageMap } from '../../../../workflow/nodes/_base/components/editor/code-editor/index'
+import { generateRuleCode } from '@/service/debug'
+import type { CodeGenRes } from '@/service/debug'
+import { ModelModeType } from '@/types/app'
+import type { AppType, Model } from '@/types/app'
+import Modal from '@/app/components/base/modal'
+import Button from '@/app/components/base/button'
+import { Generator } from '@/app/components/base/icons/src/vender/other'
+import Toast from '@/app/components/base/toast'
+import Loading from '@/app/components/base/loading'
+import Confirm from '@/app/components/base/confirm'
+import type { CodeLanguage } from '@/app/components/workflow/nodes/code/types'
+export type IGetCodeGeneratorResProps = {
+  mode: AppType
+  isShow: boolean
+  codeLanguages: CodeLanguage
+  onClose: () => void
+  onFinished: (res: CodeGenRes) => void
+}
+
+export const GetCodeGeneratorResModal: FC<IGetCodeGeneratorResProps> = (
+  {
+    mode,
+    isShow,
+    codeLanguages,
+    onClose,
+    onFinished,
+
+  },
+) => {
+  const { t } = useTranslation()
+  const [instruction, setInstruction] = React.useState<string>('')
+  const [isLoading, { setTrue: setLoadingTrue, setFalse: setLoadingFalse }] = useBoolean(false)
+  const [res, setRes] = React.useState<CodeGenRes | null>(null)
+  const isValid = () => {
+    if (instruction.trim() === '') {
+      Toast.notify({
+        type: 'error',
+        message: t('common.errorMsg.fieldRequired', {
+          field: t('appDebug.code.instruction'),
+        }),
+      })
+      return false
+    }
+    return true
+  }
+  const model: Model = {
+    provider: 'openai',
+    name: 'gpt-4o-mini',
+    mode: ModelModeType.chat,
+    completion_params: {
+      temperature: 0.7,
+      max_tokens: 0,
+      top_p: 0,
+      echo: false,
+      stop: [],
+      presence_penalty: 0,
+      frequency_penalty: 0,
+    },
+  }
+  const isInLLMNode = true
+  const onGenerate = async () => {
+    if (!isValid())
+      return
+    if (isLoading)
+      return
+    setLoadingTrue()
+    try {
+      const { error, ...res } = await generateRuleCode({
+        instruction,
+        model_config: model,
+        no_variable: !!isInLLMNode,
+        code_language: languageMap[codeLanguages] || 'javascript',
+      })
+      setRes(res)
+      if (error) {
+        Toast.notify({
+          type: 'error',
+          message: error,
+        })
+      }
+    }
+    finally {
+      setLoadingFalse()
+    }
+  }
+  const [showConfirmOverwrite, setShowConfirmOverwrite] = React.useState(false)
+
+  const renderLoading = (
+    <div className='w-0 grow flex flex-col items-center justify-center h-full space-y-3'>
+      <Loading />
+      <div className='text-[13px] text-gray-400'>{t('appDebug.codegen.loading')}</div>
+    </div>
+  )
+
+  return (
+    <Modal
+      isShow={isShow}
+      onClose={onClose}
+      className='!p-0 min-w-[1140px]'
+      closable
+    >
+      <div className='relative flex h-[680px] flex-wrap'>
+        <div className='w-[570px] shrink-0 p-8 h-full overflow-y-auto border-r border-gray-100'>
+          <div className='mb-8'>
+            <div className={'leading-[28px] text-lg font-bold'}>{t('appDebug.codegen.title')}</div>
+            <div className='mt-1 text-[13px] font-normal text-gray-500'>{t('appDebug.codegen.description')}</div>
+          </div>
+          <div className='mt-6'>
+            <div className='text-[0px]'>
+              <div className='mb-2 leading-5 text-sm font-medium text-gray-900'>{t('appDebug.codegen.instruction')}</div>
+              <textarea
+                className="w-full h-[200px] overflow-y-auto px-3 py-2 text-sm bg-gray-50 rounded-lg"
+                placeholder={t('appDebug.codegen.instructionPlaceholder') || ''}
+                value={instruction}
+                onChange={e => setInstruction(e.target.value)}
+              />
+            </div>
+
+            <div className='mt-5 flex justify-end'>
+              <Button
+                className='flex space-x-1'
+                variant='primary'
+                onClick={onGenerate}
+                disabled={isLoading}
+              >
+                <Generator className='w-4 h-4 text-white' />
+                <span className='text-xs font-semibold text-white'>{t('appDebug.codegen.generate')}</span>
+              </Button>
+            </div>
+          </div>
+        </div>
+        {isLoading && renderLoading}
+        {(!isLoading && res) && (
+          <div className='w-0 grow p-6 pb-0 h-full'>
+            <div className='shrink-0 mb-3 leading-[160%] text-base font-semibold text-gray-800'>{t('appDebug.codegen.resTitle')}</div>
+            <div className={cn('max-h-[555px] overflow-y-auto', !isInLLMNode && 'pb-2')}>
+              <ConfigPrompt
+                mode={mode}
+                promptTemplate={res?.code || ''}
+                promptVariables={[]}
+                readonly
+                noTitle={isInLLMNode}
+                gradientBorder
+                editorHeight={isInLLMNode ? 524 : 0}
+                noResize={isInLLMNode}
+              />
+              {!isInLLMNode && (
+                <>
+                  {res?.code && (
+                    <div className='mt-4'>
+                      <h3 className='mb-2 text-sm font-medium text-gray-900'>{t('appDebug.codegen.generatedCode')}</h3>
+                      <pre className='p-4 bg-gray-50 rounded-lg overflow-x-auto'>
+                        <code className={`language-${res.language}`}>
+                          {res.code}
+                        </code>
+                      </pre>
+                    </div>
+                  )}
+                  {res?.error && (
+                    <div className='mt-4 p-4 bg-red-50 rounded-lg'>
+                      <p className='text-sm text-red-600'>{res.error}</p>
+                    </div>
+                  )}
+                </>
+              )}
+            </div>
+
+            <div className='flex justify-end py-4 bg-white'>
+              <Button onClick={onClose}>{t('common.operation.cancel')}</Button>
+              <Button variant='primary' className='ml-2' onClick={() => {
+                setShowConfirmOverwrite(true)
+              }}>{t('appDebug.codegen.apply')}</Button>
+            </div>
+          </div>
+        )}
+      </div>
+      {showConfirmOverwrite && (
+        <Confirm
+          title={t('appDebug.codegen.overwriteConfirmTitle')}
+          content={t('appDebug.codegen.overwriteConfirmMessage')}
+          isShow={showConfirmOverwrite}
+          onConfirm={() => {
+            setShowConfirmOverwrite(false)
+            onFinished(res!)
+          }}
+          onCancel={() => setShowConfirmOverwrite(false)}
+        />
+      )}
+    </Modal>
+  )
+}
+
+export default React.memo(GetCodeGeneratorResModal)

web/app/components/workflow/nodes/_base/components/code-generator-button.tsx

@@ -0,0 +1,48 @@
+'use client'
+import type { FC } from 'react'
+import React, { useCallback } from 'react'
+import { useBoolean } from 'ahooks'
+import cn from 'classnames'
+import type { CodeLanguage } from '../../code/types'
+import { Generator } from '@/app/components/base/icons/src/vender/other'
+import { ActionButton } from '@/app/components/base/action-button'
+import { AppType } from '@/types/app'
+import type { CodeGenRes } from '@/service/debug'
+import { GetCodeGeneratorResModal } from '@/app/components/app/configuration/config/code-generator/get-code-generator-res'
+
+type Props = {
+  className?: string
+  onGenerated?: (prompt: string) => void
+  codeLanguages: CodeLanguage
+}
+
+const CodeGenerateBtn: FC<Props> = ({
+  className,
+  codeLanguages,
+  onGenerated,
+}) => {
+  const [showAutomatic, { setTrue: showAutomaticTrue, setFalse: showAutomaticFalse }] = useBoolean(false)
+  const handleAutomaticRes = useCallback((res: CodeGenRes) => {
+    onGenerated?.(res.code)
+    showAutomaticFalse()
+  }, [onGenerated, showAutomaticFalse])
+  return (
+    <div className={cn(className)}>
+      <ActionButton
+        className='hover:bg-[#155EFF]/8'
+        onClick={showAutomaticTrue}>
+        <Generator className='w-4 h-4 text-primary-600' />
+      </ActionButton>
+      {showAutomatic && (
+        <GetCodeGeneratorResModal
+          mode={AppType.chat}
+          isShow={showAutomatic}
+          codeLanguages={codeLanguages}
+          onClose={showAutomaticFalse}
+          onFinished={handleAutomaticRes}
+        />
+      )}
+    </div>
+  )
+}
+export default React.memo(CodeGenerateBtn)

web/app/components/workflow/nodes/_base/components/editor/base.tsx

@@ -2,6 +2,9 @@
 import type { FC } from 'react'
 import React, { useCallback, useRef, useState } from 'react'
 import copy from 'copy-to-clipboard'
+import ToggleExpandBtn from '../toggle-expand-btn'
+import CodeGeneratorButton from '../code-generator-button'
+import type { CodeLanguage } from '../../../code/types'
 import Wrap from './wrap'
 import cn from '@/utils/classnames'
 import PromptEditorHeightResizeWrap from '@/app/components/app/configuration/config-prompt/prompt-editor-height-resize-wrap'
@@ -9,7 +12,6 @@ import {
   Clipboard,
   ClipboardCheck,
 } from '@/app/components/base/icons/src/vender/line/files'
-import ToggleExpandBtn from '@/app/components/workflow/nodes/_base/components/toggle-expand-btn'
 import useToggleExpend from '@/app/components/workflow/nodes/_base/hooks/use-toggle-expend'
 import type { FileEntity } from '@/app/components/base/file-uploader/types'
 import FileListInLog from '@/app/components/base/file-uploader/file-list-in-log'
@@ -23,6 +25,8 @@ type Props = {
   value: string
   isFocus: boolean
   isInNode?: boolean
+  onGenerated?: (prompt: string) => void
+  codeLanguages: CodeLanguage
   fileList?: FileEntity[]
   showFileList?: boolean
 }
@@ -36,6 +40,8 @@ const Base: FC<Props> = ({
   value,
   isFocus,
   isInNode,
+  onGenerated,
+  codeLanguages,
   fileList = [],
   showFileList,
 }) => {
@@ -70,6 +76,9 @@ const Base: FC<Props> = ({
             e.stopPropagation()
           }}>
             {headerRight}
+            <div className='ml-1'>
+              <CodeGeneratorButton onGenerated={onGenerated} codeLanguages={codeLanguages}/>
+            </div>
             {!isCopied
               ? (
                 <Clipboard className='mx-1 w-3.5 h-3.5 text-gray-500 cursor-pointer' onClick={handleCopy} />
@@ -78,6 +87,7 @@ const Base: FC<Props> = ({
                 <ClipboardCheck className='mx-1 w-3.5 h-3.5 text-gray-500' />
               )
             }
+
             <div className='ml-1'>
               <ToggleExpandBtn isExpand={isExpand} onExpandChange={setIsExpand} />
             </div>

web/app/components/workflow/nodes/_base/components/editor/code-editor/index.tsx

@@ -33,7 +33,7 @@ export type Props = {
   showFileList?: boolean
 }
 
-const languageMap = {
+export const languageMap = {
   [CodeLanguage.javascript]: 'javascript',
   [CodeLanguage.python3]: 'python',
   [CodeLanguage.json]: 'json',
@@ -149,6 +149,9 @@ const CodeEditor: FC<Props> = ({
 
     return isFocus ? 'focus-theme' : 'blur-theme'
   })()
+  const handleGenerated = (code: string) => {
+    handleEditorChange(code)
+  }
 
   const main = (
     <>
@@ -200,6 +203,8 @@ const CodeEditor: FC<Props> = ({
             isFocus={isFocus && !readOnly}
             minHeight={minHeight}
             isInNode={isInNode}
+            onGenerated={handleGenerated}
+            codeLanguages={language}
             fileList={fileList}
             showFileList={showFileList}
           >

web/i18n/en-US/app-debug.ts

@@ -219,6 +219,20 @@ const translation = {
       manage: 'Manage',
     },
   },
+  codegen: {
+    title: 'Code Generator',
+    description: 'The Code Generator uses configured models to generate high-quality code based on your instructions. Please provide clear and detailed instructions.',
+    instruction: 'Instructions',
+    instructionPlaceholder: 'Enter detailed description of the code you want to generate.',
+    generate: 'Generate',
+    generatedCodeTitle: 'Generated Code',
+    loading: 'Generating code...',
+    apply: 'Apply',
+    applyChanges: 'Apply Changes',
+    resTitle: 'Generated Code',
+    overwriteConfirmTitle: 'Overwrite existing code?',
+    overwriteConfirmMessage: 'This action will overwrite the existing code. Do you want to continue?',
+  },
   generate: {
     title: 'Prompt Generator',
     description: 'The Prompt Generator uses the configured model to optimize prompts for higher quality and better structure. Please write clear and detailed instructions.',

web/i18n/ja-JP/app-debug.ts

@@ -199,6 +199,20 @@ const translation = {
       },
     },
   },
+  codegen: {
+    title: 'コードジェネレーター',
+    description: 'コードジェネレーターは、設定されたモデルを使用して指示に基づいて高品質なコードを生成します。明確で詳細な指示を提供してください。',
+    instruction: '指示',
+    instructionPlaceholder: '生成したいコードの詳細な説明を入力してください。',
+    generate: '生成',
+    generatedCodeTitle: '生成されたコード',
+    loading: 'コードを生成中...',
+    apply: '適用',
+    applyChanges: '変更を適用',
+    resTitle: '生成されたコード',
+    overwriteConfirmTitle: '既存のコードを上書きしますか？',
+    overwriteConfirmMessage: 'この操作は既存のコードを上書きします。続行しますか？',
+  },
   generate: {
     title: 'プロンプト生成器',
     description: 'プロンプト生成器は、設定済みのモデルを使って、高品質で構造的に優れたプロンプトを作成するための最適化を行います。具体的で詳細な指示をお書きください。',

web/i18n/zh-Hans/app-debug.ts

@@ -219,6 +219,20 @@ const translation = {
       manage: '管理',
     },
   },
+  codegen: {
+    title: '代码生成器',
+    description: '代码生成器使用配置的模型根据您的指令生成高质量的代码。请提供清晰详细的说明。',
+    instruction: '指令',
+    instructionPlaceholder: '请输入您想要生成的代码的详细描述。',
+    generate: '生成',
+    generatedCodeTitle: '生成的代码',
+    loading: '正在生成代码...',
+    apply: '应用',
+    applyChanges: '应用更改',
+    resTitle: '生成的代码',
+    overwriteConfirmTitle: '是否覆盖现有代码？',
+    overwriteConfirmMessage: '此操作将覆盖现有代码。您确定要继续吗？',
+  },
   generate: {
     title: '提示词生成器',
     description: '提示词生成器使用配置的模型来优化提示词，以获得更高的质量和更好的结构。请写出清晰详细的说明。',

web/service/debug.ts

@@ -9,6 +9,11 @@ export type AutomaticRes = {
   opening_statement: string
   error?: string
 }
+export type CodeGenRes = {
+  code: string
+  language: string[]
+  error?: string
+}
 
 export const sendChatMessage = async (appId: string, body: Record<string, any>, { onData, onCompleted, onThought, onFile, onError, getAbortController, onMessageEnd, onMessageReplace }: {
   onData: IOnData
@@ -71,6 +76,11 @@ export const generateRule = (body: Record<string, any>) => {
     body,
   })
 }
+export const generateRuleCode = (body: Record<string, any>) => {
+  return post<CodeGenRes>('/rule-code-generate', {
+    body,
+  })
+}
 
 export const fetchModelParams = (providerName: string, modelId: string) => {
   return get(`workspaces/current/model-providers/${providerName}/models/parameter-rules`, {