feat: add code generator (#9051)

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

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"