feat(agent): support new Agent DSL import and export (#38849)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/controllers/console/agent/composer.py

@@ -2,6 +2,7 @@ from uuid import UUID
 
 from flask import request
 from flask_restx import Resource
+from werkzeug.exceptions import NotFound
 
 from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
 from controllers.console import console_ns
@@ -34,6 +35,7 @@ from services.entities.agent_entities import (
     WorkflowAgentComposerQuery,
     WorkflowComposerCopyFromRosterPayload,
 )
+from services.snippet_service import SnippetService
 
 register_schema_models(
     console_ns, ComposerSavePayload, WorkflowAgentComposerQuery, WorkflowComposerCopyFromRosterPayload
@@ -237,6 +239,205 @@ class WorkflowAgentComposerSaveToRosterApi(Resource):
         )
 
 
+def _require_snippet_app_id(*, tenant_id: str, snippet_id: UUID) -> str:
+    snippet = SnippetService(session=db.session()).get_snippet_by_id(
+        snippet_id=str(snippet_id),
+        tenant_id=tenant_id,
+    )
+    if snippet is None:
+        raise NotFound("Snippet not found")
+    return snippet.id
+
+
+@console_ns.route("/snippets/<uuid:snippet_id>/workflows/draft/nodes/<string:node_id>/agent-composer")
+class SnippetAgentComposerApi(Resource):
+    @console_ns.response(200, "Snippet agent composer state", console_ns.models[WorkflowAgentComposerResponse.__name__])
+    @console_ns.doc(params=query_params_from_model(WorkflowAgentComposerQuery))
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @with_current_user_id
+    @with_current_tenant_id
+    def get(self, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
+        query = WorkflowAgentComposerQuery.model_validate(request.args.to_dict(flat=True))
+        return dump_response(
+            WorkflowAgentComposerResponse,
+            AgentComposerService.load_workflow_composer(
+                tenant_id=tenant_id,
+                app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
+                node_id=node_id,
+                account_id=account_id,
+                snapshot_id=query.snapshot_id,
+                session=db.session(),
+            ),
+        )
+
+    @console_ns.expect(console_ns.models[ComposerSavePayload.__name__])
+    @console_ns.response(200, "Snippet agent composer saved", console_ns.models[WorkflowAgentComposerResponse.__name__])
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @edit_permission_required
+    @rbac_permission_required(
+        RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False
+    )
+    @with_current_user_id
+    @with_current_tenant_id
+    def put(self, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
+        payload = ComposerSavePayload.model_validate(console_ns.payload or {})
+        return dump_response(
+            WorkflowAgentComposerResponse,
+            AgentComposerService.save_workflow_composer(
+                tenant_id=tenant_id,
+                app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
+                node_id=node_id,
+                account_id=account_id,
+                payload=payload,
+                session=db.session(),
+            ),
+        )
+
+
+@console_ns.route("/snippets/<uuid:snippet_id>/workflows/draft/nodes/<string:node_id>/agent-composer/copy-from-roster")
+class SnippetAgentComposerCopyFromRosterApi(Resource):
+    @console_ns.expect(console_ns.models[WorkflowComposerCopyFromRosterPayload.__name__])
+    @console_ns.response(
+        200, "Roster agent copied into snippet", console_ns.models[WorkflowAgentComposerResponse.__name__]
+    )
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @edit_permission_required
+    @rbac_permission_required(
+        RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False
+    )
+    @with_current_user_id
+    @with_current_tenant_id
+    def post(self, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
+        payload = WorkflowComposerCopyFromRosterPayload.model_validate(console_ns.payload or {})
+        return dump_response(
+            WorkflowAgentComposerResponse,
+            AgentComposerService.copy_workflow_composer_from_roster(
+                tenant_id=tenant_id,
+                app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
+                node_id=node_id,
+                account_id=account_id,
+                source_agent_id=payload.source_agent_id,
+                source_snapshot_id=payload.source_snapshot_id,
+                idempotency_key=payload.idempotency_key,
+                session=db.session(),
+            ),
+        )
+
+
+@console_ns.route("/snippets/<uuid:snippet_id>/workflows/draft/nodes/<string:node_id>/agent-composer/validate")
+class SnippetAgentComposerValidateApi(Resource):
+    @console_ns.expect(console_ns.models[ComposerSavePayload.__name__])
+    @console_ns.response(
+        200, "Snippet agent composer validation", console_ns.models[AgentComposerValidateResponse.__name__]
+    )
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @with_current_tenant_id
+    def post(self, tenant_id: str, snippet_id: UUID, node_id: str):
+        app_id = _require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id)
+        payload = ComposerSavePayload.model_validate(console_ns.payload or {})
+        ComposerConfigValidator.validate_publish_payload(payload)
+        AgentComposerService.validate_knowledge_datasets(tenant_id=tenant_id, agent_soul=payload.agent_soul)
+        findings = AgentComposerService.collect_validation_findings(
+            tenant_id=tenant_id,
+            payload=payload,
+            agent_id=AgentComposerService.resolve_workflow_node_agent_id(
+                tenant_id=tenant_id,
+                app_id=app_id,
+                node_id=node_id,
+                session=db.session(),
+            ),
+            session=db.session(),
+        )
+        return dump_response(AgentComposerValidateResponse, {"result": "success", "errors": [], **findings})
+
+
+@console_ns.route("/snippets/<uuid:snippet_id>/workflows/draft/nodes/<string:node_id>/agent-composer/candidates")
+class SnippetAgentComposerCandidatesApi(Resource):
+    @console_ns.response(
+        200, "Snippet agent composer candidates", console_ns.models[AgentComposerCandidatesResponse.__name__]
+    )
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @with_current_user_id
+    @with_current_tenant_id
+    def get(self, tenant_id: str, current_user_id: str, snippet_id: UUID, node_id: str):
+        return dump_response(
+            AgentComposerCandidatesResponse,
+            AgentComposerService.get_workflow_candidates(
+                tenant_id=tenant_id,
+                app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
+                node_id=node_id,
+                user_id=current_user_id,
+                session=db.session(),
+            ),
+        )
+
+
+@console_ns.route("/snippets/<uuid:snippet_id>/workflows/draft/nodes/<string:node_id>/agent-composer/impact")
+class SnippetAgentComposerImpactApi(Resource):
+    @console_ns.expect(console_ns.models[ComposerSavePayload.__name__])
+    @console_ns.response(200, "Snippet agent composer impact", console_ns.models[AgentComposerImpactResponse.__name__])
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @with_current_tenant_id
+    def post(self, tenant_id: str, snippet_id: UUID, node_id: str):
+        _require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id)
+        payload = ComposerSavePayload.model_validate(console_ns.payload or {})
+        current_snapshot_id = payload.binding.current_snapshot_id if payload.binding else None
+        if not current_snapshot_id:
+            return dump_response(
+                AgentComposerImpactResponse, {"current_snapshot_id": None, "workflow_node_count": 0, "bindings": []}
+            )
+        return dump_response(
+            AgentComposerImpactResponse,
+            AgentComposerService.calculate_impact(
+                tenant_id=tenant_id,
+                current_snapshot_id=current_snapshot_id,
+                session=db.session(),
+            ),
+        )
+
+
+@console_ns.route("/snippets/<uuid:snippet_id>/workflows/draft/nodes/<string:node_id>/agent-composer/save-to-roster")
+class SnippetAgentComposerSaveToRosterApi(Resource):
+    @console_ns.expect(console_ns.models[ComposerSavePayload.__name__])
+    @console_ns.response(
+        200, "Snippet agent saved to roster", console_ns.models[WorkflowAgentComposerResponse.__name__]
+    )
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @edit_permission_required
+    @rbac_permission_required(
+        RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False
+    )
+    @with_current_user_id
+    @with_current_tenant_id
+    def post(self, tenant_id: str, account_id: str, snippet_id: UUID, node_id: str):
+        payload = ComposerSavePayload.model_validate(console_ns.payload or {})
+        return dump_response(
+            WorkflowAgentComposerResponse,
+            AgentComposerService.save_workflow_composer(
+                tenant_id=tenant_id,
+                app_id=_require_snippet_app_id(tenant_id=tenant_id, snippet_id=snippet_id),
+                node_id=node_id,
+                account_id=account_id,
+                payload=payload,
+                session=db.session(),
+            ),
+        )
+
+
 @console_ns.route("/agent/<uuid:agent_id>/composer")
 class AgentComposerApi(Resource):
     @console_ns.response(200, "Agent app composer state", console_ns.models[AgentAppComposerResponse.__name__])

api/controllers/console/app/app.py

@@ -55,7 +55,7 @@ from services.app_dsl_service import AppDslService
 from services.app_service import AppListParams, AppListSortBy, AppService, CreateAppParams, StarredAppListParams
 from services.enterprise import rbac_service as enterprise_rbac_service
 from services.enterprise.enterprise_service import EnterpriseService
-from services.entities.dsl_entities import ImportMode, ImportStatus
+from services.entities.dsl_entities import DslImportWarning, ImportMode, ImportStatus
 from services.entities.knowledge_entities.knowledge_entities import (
     DataSource,
     InfoList,
@@ -461,6 +461,7 @@ class AppImportResponse(ResponseModel):
     current_dsl_version: str
     imported_dsl_version: str = ""
     error: str = ""
+    warnings: list[DslImportWarning] = Field(default_factory=list)
 
 
 def _enrich_app_list_items(session: Session, *, apps: Sequence[App], tenant_id: str) -> None:

api/controllers/console/snippets/snippet_workflow.py

@@ -56,6 +56,7 @@ from libs.helper import TimestampField
 from libs.login import current_account_with_tenant, login_required
 from models import Account
 from models.snippet import CustomizedSnippet
+from services.agent.workflow_publish_service import WorkflowAgentPublishService
 from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError
 from services.snippet_generate_service import SnippetGenerateService
 from services.snippet_service import SnippetService
@@ -177,6 +178,10 @@ class SnippetDraftWorkflowApi(Resource):
 
         workflow.conversation_variables = []
         response = SnippetWorkflowResponse.model_validate(workflow, from_attributes=True).model_dump(mode="json")
+        response["graph"] = WorkflowAgentPublishService.project_draft_bindings_to_graph(
+            session=db.session(),
+            draft_workflow=workflow,
+        )
         response["input_fields"] = snippet.input_fields_list
         return response
 

api/controllers/console/workspace/snippets.py

@@ -14,6 +14,7 @@ from controllers.common.schema import (
     register_schema_models,
 )
 from controllers.console import console_ns
+from controllers.console.app.wraps import with_session
 from controllers.console.snippets.payloads import (
     CreateSnippetPayload,
     IncludeSecretQuery,
@@ -39,6 +40,7 @@ from libs.helper import dump_response
 from libs.login import login_required
 from models import Account
 from models.snippet import SnippetType
+from services.entities.dsl_entities import DslImportWarning
 from services.snippet_dsl_service import ImportStatus, SnippetDslService
 from services.snippet_service import SnippetService
 
@@ -52,6 +54,7 @@ class SnippetImportResponse(ResponseModel):
     current_dsl_version: str
     imported_dsl_version: str
     error: str
+    warnings: list[DslImportWarning]
 
 
 class SnippetDependencyCheckResponse(ResponseModel):
@@ -256,8 +259,9 @@ class CustomizedSnippetDetailApi(Resource):
     @account_initialization_required
     @edit_permission_required
     @rbac_permission_required(RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_MANAGE, resource_required=False)
+    @with_current_user
     @with_current_tenant_id
-    def delete(self, current_tenant_id: str, snippet_id: str):
+    def delete(self, current_tenant_id: str, current_user: Account, snippet_id: str):
         """Delete customized snippet."""
         snippet_service = _snippet_service()
         snippet = snippet_service.get_snippet_by_id(
@@ -273,6 +277,7 @@ class CustomizedSnippetDetailApi(Resource):
             SnippetService.delete_snippet(
                 session=session,
                 snippet=snippet,
+                account_id=current_user.id,
             )
             session.commit()
 
@@ -343,22 +348,21 @@ class CustomizedSnippetImportApi(Resource):
         RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False
     )
     @with_current_user
-    def post(self, current_user: Account):
+    @with_session
+    def post(self, session: Session, current_user: Account):
         """Import snippet from DSL."""
         payload = SnippetImportPayload.model_validate(console_ns.payload or {})
 
-        with Session(db.engine) as session:
-            import_service = SnippetDslService(session)
-            result = import_service.import_snippet(
-                account=current_user,
-                import_mode=payload.mode,
-                yaml_content=payload.yaml_content,
-                yaml_url=payload.yaml_url,
-                snippet_id=payload.snippet_id,
-                name=payload.name,
-                description=payload.description,
-            )
-            session.commit()
+        import_service = SnippetDslService(session)
+        result = import_service.import_snippet(
+            account=current_user,
+            import_mode=payload.mode,
+            yaml_content=payload.yaml_content,
+            yaml_url=payload.yaml_url,
+            snippet_id=payload.snippet_id,
+            name=payload.name,
+            description=payload.description,
+        )
 
         # Return appropriate status code based on result
         status = result.status
@@ -384,12 +388,11 @@ class CustomizedSnippetImportConfirmApi(Resource):
         RBACResourceScope.WORKSPACE, RBACPermission.SNIPPETS_CREATE_AND_MODIFY, resource_required=False
     )
     @with_current_user
-    def post(self, current_user: Account, import_id: str):
+    @with_session
+    def post(self, session: Session, current_user: Account, import_id: str):
         """Confirm a pending snippet import."""
-        with Session(db.engine) as session:
-            import_service = SnippetDslService(session)
-            result = import_service.confirm_import(import_id=import_id, account=current_user)
-            session.commit()
+        import_service = SnippetDslService(session)
+        result = import_service.confirm_import(import_id=import_id, account=current_user)
 
         if result.status == ImportStatus.FAILED:
             return result.model_dump(mode="json"), 400