dify/api/controllers/console/app/generator.py

from collections.abc import Sequence

from flask_login import current_user
from flask_restx import Resource, fields, reqparse

from controllers.console import api, console_ns
from controllers.console.app.error import (
    CompletionRequestError,
    ProviderModelCurrentlyNotSupportError,
    ProviderNotInitializeError,
    ProviderQuotaExceededError,
)
from controllers.console.wraps import account_initialization_required, setup_required
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
from core.helper.code_executor.javascript.javascript_code_provider import JavascriptCodeProvider
from core.helper.code_executor.python3.python3_code_provider import Python3CodeProvider
from core.llm_generator.llm_generator import LLMGenerator
from core.model_runtime.errors.invoke import InvokeError
from libs.login import login_required


@console_ns.route("/rule-generate")
class RuleGenerateApi(Resource):
    @api.doc("generate_rule_config")
    @api.doc(description="Generate rule configuration using LLM")
    @api.expect(
        api.model(
            "RuleGenerateRequest",
            {
                "instruction": fields.String(required=True, description="Rule generation instruction"),
                "model_config": fields.Raw(required=True, description="Model configuration"),
                "no_variable": fields.Boolean(required=True, default=False, description="Whether to exclude variables"),
            },
        )
    )
    @api.response(200, "Rule configuration generated successfully")
    @api.response(400, "Invalid request parameters")
    @api.response(402, "Provider quota exceeded")
    @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")
        args = parser.parse_args()

        account = current_user
        try:
            rules = LLMGenerator.generate_rule_config(
                tenant_id=account.current_tenant_id,
                instruction=args["instruction"],
                model_config=args["model_config"],
                no_variable=args["no_variable"],
            )
        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 rules


@console_ns.route("/rule-code-generate")
class RuleCodeGenerateApi(Resource):
    @api.doc("generate_rule_code")
    @api.doc(description="Generate code rules using LLM")
    @api.expect(
        api.model(
            "RuleCodeGenerateRequest",
            {
                "instruction": fields.String(required=True, description="Code generation instruction"),
                "model_config": fields.Raw(required=True, description="Model configuration"),
                "no_variable": fields.Boolean(required=True, default=False, description="Whether to exclude variables"),
                "code_language": fields.String(
                    default="javascript", description="Programming language for code generation"
                ),
            },
        )
    )
    @api.response(200, "Code rules generated successfully")
    @api.response(400, "Invalid request parameters")
    @api.response(402, "Provider quota exceeded")
    @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
        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"],
            )
        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


@console_ns.route("/rule-structured-output-generate")
class RuleStructuredOutputGenerateApi(Resource):
    @api.doc("generate_structured_output")
    @api.doc(description="Generate structured output rules using LLM")
    @api.expect(
        api.model(
            "StructuredOutputGenerateRequest",
            {
                "instruction": fields.String(required=True, description="Structured output generation instruction"),
                "model_config": fields.Raw(required=True, description="Model configuration"),
            },
        )
    )
    @api.response(200, "Structured output generated successfully")
    @api.response(400, "Invalid request parameters")
    @api.response(402, "Provider quota exceeded")
    @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")
        args = parser.parse_args()

        account = current_user
        try:
            structured_output = LLMGenerator.generate_structured_output(
                tenant_id=account.current_tenant_id,
                instruction=args["instruction"],
                model_config=args["model_config"],
            )
        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 structured_output


@console_ns.route("/instruction-generate")
class InstructionGenerateApi(Resource):
    @api.doc("generate_instruction")
    @api.doc(description="Generate instruction for workflow nodes or general use")
    @api.expect(
        api.model(
            "InstructionGenerateRequest",
            {
                "flow_id": fields.String(required=True, description="Workflow/Flow ID"),
                "node_id": fields.String(description="Node ID for workflow context"),
                "current": fields.String(description="Current instruction text"),
                "language": fields.String(default="javascript", description="Programming language (javascript/python)"),
                "instruction": fields.String(required=True, description="Instruction for generation"),
                "model_config": fields.Raw(required=True, description="Model configuration"),
                "ideal_output": fields.String(description="Expected ideal output"),
            },
        )
    )
    @api.response(200, "Instruction generated successfully")
    @api.response(400, "Invalid request parameters or flow/workflow not found")
    @api.response(402, "Provider quota exceeded")
    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument("flow_id", type=str, required=True, default="", location="json")
        parser.add_argument("node_id", type=str, required=False, default="", location="json")
        parser.add_argument("current", type=str, required=False, default="", location="json")
        parser.add_argument("language", type=str, required=False, default="javascript", location="json")
        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("ideal_output", type=str, required=False, default="", location="json")
        args = parser.parse_args()
        code_template = (
            Python3CodeProvider.get_default_code()
            if args["language"] == "python"
            else (JavascriptCodeProvider.get_default_code())
            if args["language"] == "javascript"
            else ""
        )
        try:
            # Generate from nothing for a workflow node
            if (args["current"] == code_template or args["current"] == "") and args["node_id"] != "":
                from models import App, db
                from services.workflow_service import WorkflowService

                app = db.session.query(App).where(App.id == args["flow_id"]).first()
                if not app:
                    return {"error": f"app {args['flow_id']} not found"}, 400
                workflow = WorkflowService().get_draft_workflow(app_model=app)
                if not workflow:
                    return {"error": f"workflow {args['flow_id']} not found"}, 400
                nodes: Sequence = workflow.graph_dict["nodes"]
                node = [node for node in nodes if node["id"] == args["node_id"]]
                if len(node) == 0:
                    return {"error": f"node {args['node_id']} not found"}, 400
                node_type = node[0]["data"]["type"]
                match node_type:
                    case "llm":
                        return LLMGenerator.generate_rule_config(
                            current_user.current_tenant_id,
                            instruction=args["instruction"],
                            model_config=args["model_config"],
                            no_variable=True,
                        )
                    case "agent":
                        return LLMGenerator.generate_rule_config(
                            current_user.current_tenant_id,
                            instruction=args["instruction"],
                            model_config=args["model_config"],
                            no_variable=True,
                        )
                    case "code":
                        return LLMGenerator.generate_code(
                            tenant_id=current_user.current_tenant_id,
                            instruction=args["instruction"],
                            model_config=args["model_config"],
                            code_language=args["language"],
                        )
                    case _:
                        return {"error": f"invalid node type: {node_type}"}
            if args["node_id"] == "" and args["current"] != "":  # For legacy app without a workflow
                return LLMGenerator.instruction_modify_legacy(
                    tenant_id=current_user.current_tenant_id,
                    flow_id=args["flow_id"],
                    current=args["current"],
                    instruction=args["instruction"],
                    model_config=args["model_config"],
                    ideal_output=args["ideal_output"],
                )
            if args["node_id"] != "" and args["current"] != "":  # For workflow node
                return LLMGenerator.instruction_modify_workflow(
                    tenant_id=current_user.current_tenant_id,
                    flow_id=args["flow_id"],
                    node_id=args["node_id"],
                    current=args["current"],
                    instruction=args["instruction"],
                    model_config=args["model_config"],
                    ideal_output=args["ideal_output"],
                )
            return {"error": "incompatible parameters"}, 400
        except ProviderTokenNotInitError as ex:
            raise ProviderNotInitializeError(ex.description)
        except QuotaExceededError:
            raise ProviderQuotaExceededError()
        except ModelCurrentlyNotSupportError:
            raise ProviderModelCurrentlyNotSupportError()
        except InvokeError as e:
            raise CompletionRequestError(e.description)


@console_ns.route("/instruction-generate/template")
class InstructionGenerationTemplateApi(Resource):
    @api.doc("get_instruction_template")
    @api.doc(description="Get instruction generation template")
    @api.expect(
        api.model(
            "InstructionTemplateRequest",
            {
                "instruction": fields.String(required=True, description="Template instruction"),
                "ideal_output": fields.String(description="Expected ideal output"),
            },
        )
    )
    @api.response(200, "Template retrieved successfully")
    @api.response(400, "Invalid request parameters")
    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument("type", type=str, required=True, default=False, location="json")
        args = parser.parse_args()
        match args["type"]:
            case "prompt":
                from core.llm_generator.prompts import INSTRUCTION_GENERATE_TEMPLATE_PROMPT

                return {"data": INSTRUCTION_GENERATE_TEMPLATE_PROMPT}
            case "code":
                from core.llm_generator.prompts import INSTRUCTION_GENERATE_TEMPLATE_CODE

                return {"data": INSTRUCTION_GENERATE_TEMPLATE_CODE}
            case _:
                raise ValueError(f"Invalid type: {args['type']}")