diff --git a/api/controllers/inner_api/workspace/workspace.py b/api/controllers/inner_api/workspace/workspace.py index 77568b75f1..276990f09d 100644 --- a/api/controllers/inner_api/workspace/workspace.py +++ b/api/controllers/inner_api/workspace/workspace.py @@ -5,10 +5,13 @@ from flask_restful import Resource, reqparse from controllers.console.wraps import setup_required from controllers.inner_api import api from controllers.inner_api.wraps import enterprise_inner_api_only +from core.tools.entities.tool_entities import CredentialType from events.tenant_event import tenant_was_created from extensions.ext_database import db +from libs.helper import StrLen from models.account import Account from services.account_service import TenantService +from services.tools.builtin_tools_manage_service import BuiltinToolManageService class EnterpriseWorkspace(Resource): @@ -73,5 +76,58 @@ class EnterpriseWorkspaceNoOwnerEmail(Resource): } +class EnterpriseWorkspaceCreateToolCredential(Resource): + @setup_required + @enterprise_inner_api_only + def post(self): + parser = reqparse.RequestParser() + parser.add_argument("tenant_id", type=str, required=True, location="json") + parser.add_argument("user_id", type=str, required=True, location="json") + parser.add_argument("provider", type=str, required=True, location="json") + parser.add_argument("credentials", type=dict, required=True, nullable=False, location="json") + parser.add_argument("name", type=StrLen(30), required=False, nullable=False, location="json") + parser.add_argument("type", type=str, required=True, nullable=False, location="json") + args = parser.parse_args() + + if args["type"] not in CredentialType.values(): + raise ValueError(f"Invalid credential type: {args['type']}") + + return BuiltinToolManageService.add_builtin_tool_provider( + user_id=args["user_id"], + tenant_id=args["tenant_id"], + provider=args["provider"], + credentials=args["credentials"], + name=args["name"], + api_type=CredentialType.of(args["type"]), + ) + + +class EnterpriseWorkspaceUpdateToolCredential(Resource): + @setup_required + @enterprise_inner_api_only + def post(self): + parser = reqparse.RequestParser() + parser.add_argument("credential_id", type=str, required=True, nullable=False, location="json") + parser.add_argument("credentials", type=dict, required=False, nullable=True, location="json") + parser.add_argument("name", type=StrLen(30), required=False, nullable=True, location="json") + parser.add_argument("tenant_id", type=str, required=True, nullable=False, location="json") + parser.add_argument("user_id", type=str, required=True, nullable=False, location="json") + parser.add_argument("provider", type=str, required=True, nullable=False, location="json") + args = parser.parse_args() + + args = parser.parse_args() + result = BuiltinToolManageService.update_builtin_tool_provider( + user_id=args["user_id"], + tenant_id=args["tenant_id"], + provider=args["provider"], + credential_id=args["credential_id"], + credentials=args.get("credentials", None), + name=args.get("name", ""), + ) + return result + + api.add_resource(EnterpriseWorkspace, "/enterprise/workspace") api.add_resource(EnterpriseWorkspaceNoOwnerEmail, "/enterprise/workspace/ownerless") +api.add_resource(EnterpriseWorkspaceCreateToolCredential, "/enterprise/workspace/tool/credential") +api.add_resource(EnterpriseWorkspaceUpdateToolCredential, "/enterprise/workspace/tool/credential/update") diff --git a/api/services/tools/builtin_tools_manage_service.py b/api/services/tools/builtin_tools_manage_service.py index 841eeb4333..b8c56f9355 100644 --- a/api/services/tools/builtin_tools_manage_service.py +++ b/api/services/tools/builtin_tools_manage_service.py @@ -278,10 +278,11 @@ class BuiltinToolManageService: session.add(db_provider) session.commit() + provider_id = db_provider.id except Exception as e: session.rollback() raise ValueError(str(e)) - return {"result": "success"} + return {"result": "success", "id": provider_id} @staticmethod def create_tool_encrypter(