diff --git a/api/controllers/console/app/workflow_draft_variable.py b/api/controllers/console/app/workflow_draft_variable.py index 8b4841ed00da0a..e536a27d8ce9bc 100644 --- a/api/controllers/console/app/workflow_draft_variable.py +++ b/api/controllers/console/app/workflow_draft_variable.py @@ -7,12 +7,13 @@ from flask import Response from flask_restx import Resource from pydantic import BaseModel, ConfigDict, Field, model_validator -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.common.fields import SimpleResultResponse from controllers.common.rbac import PlainApp, RBACCheck from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models +from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.app.error import ( DraftWorkflowNotExist, @@ -41,6 +42,8 @@ WorkflowDraftVariableListWithoutValueResponse, WorkflowDraftVariableResponse, WorkflowDraftVariableWithoutValueResponse, + draft_variable_list_response_source, + draft_variable_response_source, ) from graphon.variables.types import SegmentType from libs.helper import dump_response @@ -246,12 +249,12 @@ def get(self, req_data: WorkflowDraftVariableListQuery, current_user: Account, a @console_ns.doc(description="Delete all draft workflow variables") @console_ns.response(204, "Workflow variables deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, app_model: App): + @with_session + def delete(self, session: Session, current_user: Account, app_model: App): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) draft_var_srv.delete_user_workflow_variables(app_model.id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -284,25 +287,27 @@ class NodeVariableCollectionApi(Resource): ) @_api_prerequisite @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - def get(self, current_user: Account, app_model: App, node_id: str): + @with_session(write=False) + def get(self, session: Session, current_user: Account, app_model: App, node_id: str): validate_node_id(node_id) - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - draft_var_srv = WorkflowDraftVariableService( - session=session, - ) - node_vars = draft_var_srv.list_node_variables(app_model.id, node_id, user_id=current_user.id) - - return dump_response(WorkflowDraftVariableListResponse, node_vars) + draft_var_srv = WorkflowDraftVariableService( + session=session, + ) + node_vars = draft_var_srv.list_node_variables(app_model.id, node_id, user_id=current_user.id) + return dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(node_vars, session=session), + ) @console_ns.doc("delete_node_variables") @console_ns.doc(description="Delete all variables for a specific node") @console_ns.response(204, "Node variables deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, app_model: App, node_id: str): + @with_session + def delete(self, session: Session, current_user: Account, app_model: App, node_id: str): validate_node_id(node_id) - srv = WorkflowDraftVariableService(db.session()) + srv = WorkflowDraftVariableService(session) srv.delete_node_variables(app_model.id, node_id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -322,9 +327,10 @@ class VariableApi(Resource): @console_ns.response(404, "Variable not found") @_api_prerequisite @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - def get(self, current_user: Account, app_model: App, variable_id: UUID): + @with_session(write=False) + def get(self, session: Session, current_user: Account, app_model: App, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) variable = ensure_variable_access( @@ -333,7 +339,7 @@ def get(self, current_user: Account, app_model: App, variable_id: UUID): variable_id=variable_id_str, current_user_id=current_user.id, ) - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.doc("update_variable") @console_ns.doc(description="Update a workflow variable") @@ -346,8 +352,10 @@ def get(self, current_user: Account, app_model: App, variable_id: UUID): @console_ns.response(404, "Variable not found") @_api_prerequisite @model_validate(WorkflowDraftVariableUpdatePayload) + @with_session def patch( self, + session: Session, req_data: WorkflowDraftVariableUpdatePayload, current_user: Account, app_model: App, @@ -375,7 +383,7 @@ def patch( # } draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) @@ -389,7 +397,9 @@ def patch( new_name = req_data.name raw_value = req_data.value if new_name is None and raw_value is None: - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session) + ) new_value = None if raw_value is not None: @@ -416,17 +426,17 @@ def patch( pass new_value = build_segment_with_type(variable.value_type, raw_value) draft_var_srv.update_variable(variable, name=new_name, value=new_value) - db.session.commit() - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.doc("delete_variable") @console_ns.doc(description="Delete a workflow variable") @console_ns.response(204, "Variable deleted successfully") @console_ns.response(404, "Variable not found") @_api_prerequisite - def delete(self, current_user: Account, app_model: App, variable_id: UUID): + @with_session + def delete(self, session: Session, current_user: Account, app_model: App, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) variable = ensure_variable_access( @@ -436,7 +446,6 @@ def delete(self, current_user: Account, app_model: App, variable_id: UUID): current_user_id=current_user.id, ) draft_var_srv.delete_variable(variable) - db.session.commit() return Response("", 204) @@ -453,13 +462,14 @@ class VariableResetApi(Resource): @console_ns.response(204, "Variable reset (no content)") @console_ns.response(404, "Variable not found") @_api_prerequisite - def put(self, current_user: Account, app_model: App, variable_id: UUID): + @with_session + def put(self, session: Session, current_user: Account, app_model: App, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) workflow_srv = WorkflowService() - draft_workflow = workflow_srv.get_draft_workflow(app_model, session=db.session()) + draft_workflow = workflow_srv.get_draft_workflow(app_model, session=session) if draft_workflow is None: raise NotFoundError( f"Draft workflow not found, app_id={app_model.id}", @@ -473,27 +483,27 @@ def put(self, current_user: Account, app_model: App, variable_id: UUID): ) resetted = draft_var_srv.reset_variable(draft_workflow, variable) - db.session.commit() if resetted is None: return Response("", 204) - return dump_response(WorkflowDraftVariableResponse, resetted) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session)) -def _get_variable_list(app_model: App, node_id: str, current_user_id: str) -> WorkflowDraftVariableList: - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - draft_var_srv = WorkflowDraftVariableService( - session=session, +def _get_variable_list( + app_model: App, node_id: str, current_user_id: str, *, session: Session +) -> WorkflowDraftVariableList: + draft_var_srv = WorkflowDraftVariableService( + session=session, + ) + if node_id == CONVERSATION_VARIABLE_NODE_ID: + draft_vars = draft_var_srv.list_conversation_variables(app_model.id, user_id=current_user_id) + elif node_id == SYSTEM_VARIABLE_NODE_ID: + draft_vars = draft_var_srv.list_system_variables(app_model.id, user_id=current_user_id) + else: + draft_vars = draft_var_srv.list_node_variables( + app_id=app_model.id, + node_id=node_id, + user_id=current_user_id, ) - if node_id == CONVERSATION_VARIABLE_NODE_ID: - draft_vars = draft_var_srv.list_conversation_variables(app_model.id, user_id=current_user_id) - elif node_id == SYSTEM_VARIABLE_NODE_ID: - draft_vars = draft_var_srv.list_system_variables(app_model.id, user_id=current_user_id) - else: - draft_vars = draft_var_srv.list_node_variables( - app_id=app_model.id, - node_id=node_id, - user_id=current_user_id, - ) return draft_vars @@ -510,19 +520,22 @@ class ConversationVariableCollectionApi(Resource): @console_ns.response(404, "Draft workflow not found") @_api_prerequisite @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - def get(self, current_user: Account, app_model: App): + @with_session + def get(self, session: Session, current_user: Account, app_model: App): # NOTE(QuantumGhost): Prefill conversation variables into the draft variables table # so their IDs can be returned to the caller. workflow_srv = WorkflowService() - draft_workflow = workflow_srv.get_draft_workflow(app_model, session=db.session()) + draft_workflow = workflow_srv.get_draft_workflow(app_model, session=session) if draft_workflow is None: raise NotFoundError(description=f"draft workflow not found, id={app_model.id}") - draft_var_srv = WorkflowDraftVariableService(db.session()) + draft_var_srv = WorkflowDraftVariableService(session) draft_var_srv.prefill_conversation_variable_default_values(draft_workflow, user_id=current_user.id) - db.session.commit() return dump_response( WorkflowDraftVariableListResponse, - _get_variable_list(app_model, CONVERSATION_VARIABLE_NODE_ID, current_user.id), + draft_variable_list_response_source( + _get_variable_list(app_model, CONVERSATION_VARIABLE_NODE_ID, current_user.id, session=session), + session=session, + ), ) @console_ns.expect(console_ns.models[ConversationVariableUpdatePayload.__name__]) @@ -575,10 +588,14 @@ class SystemVariableCollectionApi(Resource): ) @_api_prerequisite @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - def get(self, current_user: Account, app_model: App): + @with_session(write=False) + def get(self, session: Session, current_user: Account, app_model: App): return dump_response( WorkflowDraftVariableListResponse, - _get_variable_list(app_model, SYSTEM_VARIABLE_NODE_ID, current_user.id), + draft_variable_list_response_source( + _get_variable_list(app_model, SYSTEM_VARIABLE_NODE_ID, current_user.id, session=session), + session=session, + ), ) diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py index 336eab18c5de5d..bccc2b3b59ebb0 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py @@ -7,11 +7,12 @@ from flask import Response from flask_restx import Resource from pydantic import BaseModel, Field -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.common.rbac import DatasetByPipeline, RBACCheck from controllers.common.schema import query_params_from_model, register_schema_models +from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.app.error import ( DraftWorkflowNotExist, @@ -40,6 +41,8 @@ WorkflowDraftVariableListResponse, WorkflowDraftVariableListWithoutValueResponse, WorkflowDraftVariableResponse, + draft_variable_list_response_source, + draft_variable_response_source, ) from graphon.variables.types import SegmentType from libs.helper import dump_response @@ -130,12 +133,12 @@ def get(self, req_data: PaginationQuery, current_user: Account, pipeline: Pipeli @console_ns.response(204, "Workflow variables deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, pipeline: Pipeline): + @with_session + def delete(self, session: Session, current_user: Account, pipeline: Pipeline): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) draft_var_srv.delete_user_workflow_variables(pipeline.id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -165,23 +168,25 @@ class RagPipelineNodeVariableCollectionApi(Resource): console_ns.models[WorkflowDraftVariableListResponse.__name__], ) @_api_prerequisite - def get(self, current_user: Account, pipeline: Pipeline, node_id: str): + @with_session(write=False) + def get(self, session: Session, current_user: Account, pipeline: Pipeline, node_id: str): validate_node_id(node_id) - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - draft_var_srv = WorkflowDraftVariableService( - session=session, - ) - node_vars = draft_var_srv.list_node_variables(pipeline.id, node_id, user_id=current_user.id) - - return dump_response(WorkflowDraftVariableListResponse, node_vars) + draft_var_srv = WorkflowDraftVariableService( + session=session, + ) + node_vars = draft_var_srv.list_node_variables(pipeline.id, node_id, user_id=current_user.id) + return dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(node_vars, session=session), + ) @console_ns.response(204, "Node variables deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, pipeline: Pipeline, node_id: str): + @with_session + def delete(self, session: Session, current_user: Account, pipeline: Pipeline, node_id: str): validate_node_id(node_id) - srv = WorkflowDraftVariableService(db.session()) + srv = WorkflowDraftVariableService(session) srv.delete_node_variables(pipeline.id, node_id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -196,9 +201,10 @@ class RagPipelineVariableApi(Resource): console_ns.models[WorkflowDraftVariableResponse.__name__], ) @_api_prerequisite - def get(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): + @with_session(write=False) + def get(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) variable = ensure_variable_access( @@ -207,7 +213,7 @@ def get(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): variable_id=variable_id_str, current_user_id=current_user.id, ) - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.response( 200, @@ -217,8 +223,10 @@ def get(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): @_api_prerequisite @console_ns.expect(console_ns.models[WorkflowDraftVariablePatchPayload.__name__]) @model_validate(WorkflowDraftVariablePatchPayload) + @with_session def patch( self, + session: Session, req_data: WorkflowDraftVariablePatchPayload, current_user: Account, pipeline: Pipeline, @@ -246,7 +254,7 @@ def patch( # } draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) args = req_data.model_dump(exclude_none=True) @@ -261,7 +269,9 @@ def patch( new_name = args.get(self._PATCH_NAME_FIELD, None) raw_value = args.get(self._PATCH_VALUE_FIELD, None) if new_name is None and raw_value is None: - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session) + ) new_value = None if raw_value is not None: @@ -288,14 +298,14 @@ def patch( pass new_value = build_segment_with_type(variable.value_type, raw_value) draft_var_srv.update_variable(variable, name=new_name, value=new_value) - db.session.commit() - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.response(204, "Variable deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): + @with_session + def delete(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) variable = ensure_variable_access( @@ -305,7 +315,6 @@ def delete(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): current_user_id=current_user.id, ) draft_var_srv.delete_variable(variable) - db.session.commit() return Response("", 204) @@ -318,12 +327,13 @@ class RagPipelineVariableResetApi(Resource): ) @console_ns.response(204, "Variable reset (no content)") @_api_prerequisite - def put(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): + @with_session + def put(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) - rag_pipeline_service = RagPipelineService(db.session()) + rag_pipeline_service = RagPipelineService(session) draft_workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline) if draft_workflow is None: raise NotFoundError( @@ -338,23 +348,23 @@ def put(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): ) resetted = draft_var_srv.reset_variable(draft_workflow, variable) - db.session.commit() if resetted is None: return Response("", 204) - return dump_response(WorkflowDraftVariableResponse, resetted) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session)) -def _get_variable_list(pipeline: Pipeline, node_id: str, current_user_id: str) -> WorkflowDraftVariableList: - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - draft_var_srv = WorkflowDraftVariableService( - session=session, - ) - if node_id == CONVERSATION_VARIABLE_NODE_ID: - draft_vars = draft_var_srv.list_conversation_variables(pipeline.id, user_id=current_user_id) - elif node_id == SYSTEM_VARIABLE_NODE_ID: - draft_vars = draft_var_srv.list_system_variables(pipeline.id, user_id=current_user_id) - else: - draft_vars = draft_var_srv.list_node_variables(app_id=pipeline.id, node_id=node_id, user_id=current_user_id) +def _get_variable_list( + pipeline: Pipeline, node_id: str, current_user_id: str, *, session: Session +) -> WorkflowDraftVariableList: + draft_var_srv = WorkflowDraftVariableService( + session=session, + ) + if node_id == CONVERSATION_VARIABLE_NODE_ID: + draft_vars = draft_var_srv.list_conversation_variables(pipeline.id, user_id=current_user_id) + elif node_id == SYSTEM_VARIABLE_NODE_ID: + draft_vars = draft_var_srv.list_system_variables(pipeline.id, user_id=current_user_id) + else: + draft_vars = draft_var_srv.list_node_variables(app_id=pipeline.id, node_id=node_id, user_id=current_user_id) return draft_vars @@ -366,10 +376,14 @@ class RagPipelineSystemVariableCollectionApi(Resource): console_ns.models[WorkflowDraftVariableListResponse.__name__], ) @_api_prerequisite - def get(self, current_user: Account, pipeline: Pipeline): + @with_session(write=False) + def get(self, session: Session, current_user: Account, pipeline: Pipeline): return dump_response( WorkflowDraftVariableListResponse, - _get_variable_list(pipeline, SYSTEM_VARIABLE_NODE_ID, current_user.id), + draft_variable_list_response_source( + _get_variable_list(pipeline, SYSTEM_VARIABLE_NODE_ID, current_user.id, session=session), + session=session, + ), ) diff --git a/api/controllers/console/snippets/snippet_workflow_draft_variable.py b/api/controllers/console/snippets/snippet_workflow_draft_variable.py index f3cf9c0a284828..e3aa23042b7223 100644 --- a/api/controllers/console/snippets/snippet_workflow_draft_variable.py +++ b/api/controllers/console/snippets/snippet_workflow_draft_variable.py @@ -48,6 +48,8 @@ WorkflowDraftVariableListResponse, WorkflowDraftVariableListWithoutValueResponse, WorkflowDraftVariableResponse, + draft_variable_list_response_source, + draft_variable_response_source, ) from graphon.variables.types import SegmentType from libs.helper import dump_response @@ -158,7 +160,9 @@ def get(self, session: Session, current_user: Account, snippet: CustomizedSnippe draft_var_srv = WorkflowDraftVariableService(session=session) node_vars = draft_var_srv.list_node_variables(snippet.id, node_id, user_id=current_user.id) - return dump_response(WorkflowDraftVariableListResponse, node_vars) + return dump_response( + WorkflowDraftVariableListResponse, draft_variable_list_response_source(node_vars, session=session) + ) @console_ns.doc("delete_snippet_node_variables") @console_ns.doc(description="Delete all variables for a specific node (snippet draft workflow)") @@ -195,7 +199,7 @@ def get( current_user_id=current_user.id, ) _ensure_snippet_draft_variable_row_allowed(variable=variable, variable_id=variable_id) - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.doc("update_snippet_workflow_variable") @console_ns.doc(description="Update a draft workflow variable (snippet scope)") @@ -230,7 +234,9 @@ def patch( new_name = req_data.name raw_value = req_data.value if new_name is None and raw_value is None: - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session) + ) new_value = None if raw_value is not None: @@ -254,7 +260,7 @@ def patch( ) new_value = build_segment_with_type(variable.value_type, raw_value) draft_var_srv.update_variable(variable, name=new_name, value=new_value) - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.doc("delete_snippet_workflow_variable") @console_ns.doc(description="Delete a draft workflow variable (snippet scope)") @@ -309,7 +315,7 @@ def put( resetted = draft_var_srv.reset_variable(draft_workflow, variable) if resetted is None: return Response("", 204) - return dump_response(WorkflowDraftVariableResponse, resetted) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session)) @console_ns.route("/snippets//workflows/draft/conversation-variables") diff --git a/api/fields/workflow_draft_variable_fields.py b/api/fields/workflow_draft_variable_fields.py index 354082edc1935a..9415f057915654 100644 --- a/api/fields/workflow_draft_variable_fields.py +++ b/api/fields/workflow_draft_variable_fields.py @@ -3,8 +3,9 @@ from typing import Any, override from pydantic import model_validator +from sqlalchemy.orm import Session -from fields.base import ResponseModel +from fields.base import ResponseModel, SessionResponseSource from graphon.file import helpers as file_helpers from graphon.variables.segment_group import SegmentGroup from graphon.variables.segments import ArrayFileSegment, FileSegment, Segment @@ -25,7 +26,27 @@ def _convert_values_to_json_serializable_object(value: Segment) -> JSONValue: return value.value -def _serialize_var_value(variable: WorkflowDraftVariable) -> JSONValue: +class WorkflowDraftVariableResponseSource(SessionResponseSource[WorkflowDraftVariable]): + """Expose the session-backed value decoding during response validation.""" + + def get_value(self) -> Segment: + return self._source.get_value(session=self._session) + + +def draft_variable_response_source( + variable: WorkflowDraftVariable, *, session: Session +) -> WorkflowDraftVariableResponseSource: + return WorkflowDraftVariableResponseSource(variable, session=session) + + +def draft_variable_list_response_source(variable_list: Any, *, session: Session) -> dict[str, Any]: + """Wrap each variable in a list payload so value decoding resolves via the session.""" + return { + "items": [draft_variable_response_source(variable, session=session) for variable in variable_list.variables], + } + + +def _serialize_var_value(variable: WorkflowDraftVariableResponseSource) -> JSONValue: value = variable.get_value() # Create a copy to avoid mutating the model's cached deserialized value. value = value.model_copy(deep=True) @@ -105,12 +126,18 @@ class WorkflowDraftVariableResponse(WorkflowDraftVariableWithoutValueResponse): @classmethod @override def _from_workflow_draft_variable(cls, value: Any) -> Any: - if isinstance(value, WorkflowDraftVariable): + if isinstance(value, WorkflowDraftVariableResponseSource): return { - **_serialize_without_value(value), + **_serialize_without_value(value._source), "value": _serialize_var_value(value), - "full_content": _serialize_full_content(value), + "full_content": _serialize_full_content(value._source), } + if isinstance(value, WorkflowDraftVariable): + # Decoding the value needs a database session; a bare model would fall through to + # from_attributes and expose the raw serialized JSON string instead of the decoded value. + raise TypeError( + "WorkflowDraftVariableResponse requires draft_variable_response_source(variable, session=...)" + ) return value diff --git a/api/models/workflow.py b/api/models/workflow.py index bc595df4cc27b2..78d1cdd5b83114 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -20,7 +20,7 @@ orm, select, ) -from sqlalchemy.orm import Mapped, mapped_column +from sqlalchemy.orm import Mapped, Session, mapped_column from typing_extensions import deprecated from core.trigger.constants import TRIGGER_PLUGIN_NODE_TYPE @@ -89,10 +89,10 @@ SerializedWorkflowVariables = dict[str, SerializedWorkflowValue] -def _resolve_workflow_app_tenant_id(app_id: str) -> str: +def _resolve_workflow_app_tenant_id(session: Session, app_id: str) -> str: from .model import App - tenant_id = db.session.scalar(select(App.tenant_id).where(App.id == app_id)) + tenant_id = session.scalar(select(App.tenant_id).where(App.id == app_id)) if not tenant_id: raise ValueError(f"Unable to resolve tenant_id for app {app_id}") return tenant_id @@ -1712,11 +1712,11 @@ def get_selector(self) -> list[str]: def _set_selector(self, value: list[str]): self.selector = json.dumps(value) - def _loads_value(self) -> Segment: + def _loads_value(self, *, session: Session) -> Segment: value = json.loads(self.value) - return self.build_segment_from_serialized_value(self.value_type, value) + return self.build_segment_from_serialized_value(self.value_type, value, session=session) - def _rebuild_file_types(self, value: Any): + def _rebuild_file_types(self, value: Any, *, session: Session): # NOTE(QuantumGhost): Temporary workaround for structured data handling. # By this point, `output` has been converted to dict by # `WorkflowEntry.handle_special_values`, so we need to @@ -1731,7 +1731,7 @@ def _rebuild_file_types(self, value: Any): case dict(): if not maybe_file_object(value): return cast(Any, value) - tenant_id = _resolve_workflow_app_tenant_id(self.app_id) + tenant_id = _resolve_workflow_app_tenant_id(session, self.app_id) return build_file_from_stored_mapping( file_mapping=cast(dict[str, Any], value), tenant_id=tenant_id, @@ -1741,7 +1741,7 @@ def _rebuild_file_types(self, value: Any): first: Any = value_list[0] if not maybe_file_object(first): return cast(Any, value) - tenant_id = _resolve_workflow_app_tenant_id(self.app_id) + tenant_id = _resolve_workflow_app_tenant_id(session, self.app_id) file_list: list[File] = [] for item in value_list: file_list.append( @@ -1754,7 +1754,9 @@ def _rebuild_file_types(self, value: Any): case _: return cast(Any, value) - def build_segment_from_serialized_value(self, segment_type: SegmentType, value: Any) -> Segment: + def build_segment_from_serialized_value( + self, segment_type: SegmentType, value: Any, *, session: Session + ) -> Segment: # Persisted draft variable rows may contain historical file payloads. # Rebuild them through the file factory so tenant ownership, signed URLs, # and storage-backed metadata come from canonical records instead of the @@ -1765,14 +1767,14 @@ def build_segment_from_serialized_value(self, segment_type: SegmentType, value: case File(): return build_segment_with_type(segment_type, value) case dict(): - file = self._rebuild_file_types(value) + file = self._rebuild_file_types(value, session=session) return build_segment_with_type(segment_type, file) case _: raise TypeMismatchError(f"expected dict or File for FileSegment, got {type(value)}") case SegmentType.ARRAY_FILE: if not isinstance(value, list): raise TypeMismatchError(f"expected list for ArrayFileSegment, got {type(value)}") - file_list = self._rebuild_file_types(value) + file_list = self._rebuild_file_types(value, session=session) return build_segment_with_type(segment_type=segment_type, value=file_list) case _: return build_segment_with_type(segment_type=segment_type, value=value) @@ -1827,7 +1829,7 @@ def build_segment_with_type(cls, segment_type: SegmentType, value: Any) -> Segme case _: return build_segment_with_type(segment_type=segment_type, value=value) - def get_value(self) -> Segment: + def get_value(self, *, session: Session) -> Segment: """Decode the serialized value into its corresponding `Segment` object. This method caches the result, so repeated calls will return the same @@ -1845,7 +1847,7 @@ def get_value(self) -> Segment: if self.__value is not None: return self.__value - value = self._loads_value() + value = self._loads_value(session=session) self.__value = value return value diff --git a/api/services/workflow_draft_variable_service.py b/api/services/workflow_draft_variable_service.py index e0a0c18d5c6fbe..2613c28b56ef9f 100644 --- a/api/services/workflow_draft_variable_service.py +++ b/api/services/workflow_draft_variable_service.py @@ -115,21 +115,37 @@ def load_variables(self, selectors: list[list[str]]) -> list[VariableBase]: # Map each selector (as a tuple via `_selector_to_tuple`) to its corresponding variable instance. variable_by_selector: dict[tuple[str, str], VariableBase] = {} + files: list[File] = [] + offloaded_draft_vars: list[WorkflowDraftVariable] = [] with Session(bind=self._engine, expire_on_commit=False) as session: srv = WorkflowDraftVariableService(session) draft_vars = srv.get_draft_variables_by_selectors(self._app_id, selectors, user_id=self._user_id) - # Important: - files: list[File] = [] - # FileSegment and ArrayFileSegment are not subject to offloading, so their values - # can be safely accessed before any offloading logic is applied. - for draft_var in draft_vars: - value = draft_var.get_value() - match value: - case FileSegment(): - files.append(value.value) - case ArrayFileSegment(): - files.extend(value.value) + for draft_var in draft_vars: + if draft_var.is_truncated(): + offloaded_draft_vars.append(draft_var) + continue + + # FileSegment and ArrayFileSegment are not subject to offloading, so their values + # can be safely accessed before any offloading logic is applied. + segment = draft_var.get_value(session=session) + match segment: + case FileSegment(): + files.append(segment.value) + case ArrayFileSegment(): + files.extend(segment.value) + variable = segment_to_variable( + segment=segment, + selector=draft_var.get_selector(), + variable_id=draft_var.id, + name=draft_var.name, + description=draft_var.description, + ) + selector_tuple = self._selector_to_tuple(variable.selector) + variable_by_selector[selector_tuple] = variable + + # The files above are the same instances held by the variables, so loading their + # storage keys here is visible through `variable_by_selector` as well. with Session(bind=self._engine) as session: storage_key_loader = StorageKeyLoader( session, @@ -138,23 +154,6 @@ def load_variables(self, selectors: list[list[str]]) -> list[VariableBase]: ) storage_key_loader.load_storage_keys(files) - offloaded_draft_vars = [] - for draft_var in draft_vars: - if draft_var.is_truncated(): - offloaded_draft_vars.append(draft_var) - continue - - segment = draft_var.get_value() - variable = segment_to_variable( - segment=segment, - selector=draft_var.get_selector(), - variable_id=draft_var.id, - name=draft_var.name, - description=draft_var.description, - ) - selector_tuple = self._selector_to_tuple(variable.selector) - variable_by_selector[selector_tuple] = variable - # Load offloaded variables using multithreading. # This approach reduces loading time by querying external systems concurrently. with ThreadPoolExecutor(max_workers=10) as executor: @@ -189,7 +188,11 @@ def _load_offloaded_variable(self, draft_var: WorkflowDraftVariable) -> tuple[tu return (draft_var.node_id, draft_var.name), variable deserialized = json.loads(content) - segment = draft_var.build_segment_from_serialized_value(variable_file.value_type, deserialized) + # Runs on a worker thread, so it opens its own session instead of sharing one. + with Session(bind=self._engine, expire_on_commit=False) as session: + segment = draft_var.build_segment_from_serialized_value( + variable_file.value_type, deserialized, session=session + ) variable = segment_to_variable( segment=segment, selector=draft_var.get_selector(), @@ -574,7 +577,7 @@ def _get_conversation_id_from_draft_variable(self, app_id: str, user_id: str) -> ) if draft_var is None: return None - segment = draft_var.get_value() + segment = draft_var.get_value(session=self._session) if not isinstance(segment, StringSegment): logger.warning( "sys.conversation_id variable is not a string: app_id=%s, user_id=%s, id=%s", diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_workflow_draft_variable.py b/api/tests/test_containers_integration_tests/controllers/console/app/test_workflow_draft_variable.py index a071d22ee9112e..7902bc36ba0985 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_workflow_draft_variable.py +++ b/api/tests/test_containers_integration_tests/controllers/console/app/test_workflow_draft_variable.py @@ -288,6 +288,9 @@ def test_variable_api_patch_success( assert payload["id"] == variable.id assert payload["name"] == "renamed_var" + # The handler commits through its own decorator-injected session, so expire this + # session's cached instance to read the persisted row instead of the identity map. + db_session_with_containers.expire_all() refreshed = db_session_with_containers.scalar( select(WorkflowDraftVariable).where(WorkflowDraftVariable.id == variable.id) ) diff --git a/api/tests/test_containers_integration_tests/services/test_workflow_draft_variable_service.py b/api/tests/test_containers_integration_tests/services/test_workflow_draft_variable_service.py index 82fe391b08af52..613c87061401f9 100644 --- a/api/tests/test_containers_integration_tests/services/test_workflow_draft_variable_service.py +++ b/api/tests/test_containers_integration_tests/services/test_workflow_draft_variable_service.py @@ -213,7 +213,7 @@ def test_get_variable_success(self, db_session_with_containers: Session, mock_ex assert retrieved_variable.id == variable.id assert retrieved_variable.name == "test_var" assert retrieved_variable.app_id == app.id - assert retrieved_variable.get_value().value == test_value.value + assert retrieved_variable.get_value(session=db_session_with_containers).value == test_value.value def test_get_variable_not_found(self, db_session_with_containers: Session, mock_external_service_dependencies): """ @@ -274,11 +274,11 @@ def test_get_draft_variables_by_selectors_success( assert "var3" in var_names for var in retrieved_variables: if var.name == "var1": - assert var.get_value().value == var1_value.value + assert var.get_value(session=db_session_with_containers).value == var1_value.value elif var.name == "var2": - assert var.get_value().value == var2_value.value + assert var.get_value(session=db_session_with_containers).value == var2_value.value elif var.name == "var3": - assert var.get_value().value == var3_value.value + assert var.get_value(session=db_session_with_containers).value == var3_value.value def test_list_variables_without_values_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -431,12 +431,12 @@ def test_update_variable_success(self, db_session_with_containers: Session, mock service = WorkflowDraftVariableService(db_session_with_containers) updated_variable = service.update_variable(variable, name="new_name", value=new_value) assert updated_variable.name == "new_name" - assert updated_variable.get_value().value == new_value.value + assert updated_variable.get_value(session=db_session_with_containers).value == new_value.value assert updated_variable.last_edited_at is not None db_session_with_containers.refresh(variable) assert variable.name == "new_name" - assert variable.get_value().value == new_value.value + assert variable.get_value(session=db_session_with_containers).value == new_value.value assert variable.last_edited_at is not None def test_update_variable_not_editable( @@ -509,10 +509,10 @@ def test_reset_conversation_variable_success( service = WorkflowDraftVariableService(db_session_with_containers) reset_variable = service.reset_variable(workflow, variable) assert reset_variable is not None - assert reset_variable.get_value().value == "default_value" + assert reset_variable.get_value(session=db_session_with_containers).value == "default_value" assert reset_variable.last_edited_at is None db_session_with_containers.refresh(variable) - assert variable.get_value().value == "default_value" + assert variable.get_value(session=db_session_with_containers).value == "default_value" assert variable.last_edited_at is None def test_delete_variable_success(self, db_session_with_containers: Session, mock_external_service_dependencies): diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_draft_variable_handlers.py b/api/tests/unit_tests/controllers/console/app/test_workflow_draft_variable_handlers.py new file mode 100644 index 00000000000000..241452c305e63e --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_draft_variable_handlers.py @@ -0,0 +1,222 @@ +"""Direct-call tests for the app draft-variable handlers converted to `@with_session`.""" + +import json +from inspect import unwrap +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask +from sqlalchemy.orm import Session + +from controllers.console.app import workflow_draft_variable as module +from graphon.variables.segments import StringSegment +from models.account import Account, AccountStatus +from models.model import App, AppMode +from models.workflow import Workflow, WorkflowDraftVariable + +_APP_ID = "app-1" +_USER_ID = "user-1" + + +def _make_account() -> Account: + account = Account( + name="tester", + email="tester@example.com", + status=AccountStatus.ACTIVE, + ) + account.id = _USER_ID # type: ignore[assignment] + return account + + +def _make_app_model() -> App: + return App( + id=_APP_ID, + tenant_id="tenant-1", + name="App", + mode=AppMode.ADVANCED_CHAT, + enable_site=True, + enable_api=True, + ) + + +def _make_draft_workflow() -> Workflow: + workflow = Workflow.new( + tenant_id="tenant-1", + app_id=_APP_ID, + type="workflow", + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by=_USER_ID, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow-1" + return workflow + + +def _make_node_variable(variable_id: str, *, node_id: str = "llm-1") -> WorkflowDraftVariable: + variable = WorkflowDraftVariable.new_node_variable( + app_id=_APP_ID, + user_id=_USER_ID, + node_id=node_id, + name=variable_id, + value=StringSegment(value=f"value-{variable_id}"), + node_execution_id="execution-1", + ) + variable.id = variable_id + return variable + + +def _make_system_variable(variable_id: str, *, name: str) -> WorkflowDraftVariable: + variable = WorkflowDraftVariable.new_sys_variable( + app_id=_APP_ID, + user_id=_USER_ID, + name=name, + value=StringSegment(value=f"value-{name}"), + node_execution_id="execution-1", + ) + variable.id = variable_id + return variable + + +@pytest.fixture +def app() -> Flask: + app = Flask("test_workflow_draft_variable_handlers") + app.config["TESTING"] = True + return app + + +def test_node_variable_collection_get_returns_decoded_values(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.NodeVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), node_id="llm-1") + + assert [item["id"] for item in result["items"]] == ["var-1"] + assert result["items"][0]["value"] == "value-var-1" + + +def test_variable_get_returns_decoded_value(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.VariableApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), variable_id="var-1") + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_patch_without_changes_returns_variable(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.VariableApi() + handler = unwrap(api.patch) + + with app.test_request_context("/", method="PATCH", json={}): + result = handler( + api, + sqlite_session, + module.WorkflowDraftVariableUpdatePayload(), + _make_account(), + app_model=_make_app_model(), + variable_id="var-1", + ) + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_patch_renames_variable(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.VariableApi() + handler = unwrap(api.patch) + + with app.test_request_context("/", method="PATCH", json={"name": "renamed"}): + result = handler( + api, + sqlite_session, + module.WorkflowDraftVariableUpdatePayload(name="renamed"), + _make_account(), + app_model=_make_app_model(), + variable_id="var-1", + ) + + assert result["name"] == "renamed" + stored = sqlite_session.get(WorkflowDraftVariable, "var-1") + assert stored is not None + assert stored.name == "renamed" + + +def test_variable_reset_returns_default_value(app: Flask, sqlite_session: Session) -> None: + workflow = _make_draft_workflow() + variable = _make_node_variable("var-1") + variable.node_execution_id = None + sqlite_session.add_all([workflow, variable]) + sqlite_session.commit() + api = module.VariableResetApi() + handler = unwrap(api.put) + + workflow_service = MagicMock() + workflow_service.get_draft_workflow.return_value = workflow + with ( + app.test_request_context("/", method="PUT"), + patch.object(module, "WorkflowService", return_value=workflow_service), + ): + response = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), variable_id="var-1") + + # A node variable without an associated execution has no default to restore, so the + # handler answers 204 after removing the row. + assert response.status_code == 204 + + +def test_conversation_variable_collection_get_lists_conversation_variables(app: Flask, sqlite_session: Session) -> None: + workflow = _make_draft_workflow() + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id=_APP_ID, + name="conv_var", + value=StringSegment(value="conv-value"), + ) + conv_var.id = "conv-1" + conv_var.user_id = _USER_ID + sqlite_session.add_all([workflow, conv_var]) + sqlite_session.commit() + api = module.ConversationVariableCollectionApi() + handler = unwrap(api.get) + + workflow_service = MagicMock() + workflow_service.get_draft_workflow.return_value = workflow + with ( + app.test_request_context("/"), + patch.object(module, "WorkflowService", return_value=workflow_service), + ): + result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model()) + + assert [item["id"] for item in result["items"]] == ["conv-1"] + assert result["items"][0]["value"] == "conv-value" + + +def test_system_variable_collection_get_lists_system_variables(app: Flask, sqlite_session: Session) -> None: + sys_var = _make_system_variable("sys-1", name="query") + sqlite_session.add(sys_var) + sqlite_session.commit() + api = module.SystemVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model()) + + assert [item["id"] for item in result["items"]] == ["sys-1"] + assert result["items"][0]["value"] == "value-query" diff --git a/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py b/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py index 709ae1e72a4d62..44fd4a10d1ae64 100644 --- a/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py +++ b/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py @@ -4,6 +4,7 @@ from unittest.mock import patch import pytest +from sqlalchemy.orm import Session from core.workflow.variable_prefixes import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID from factories.variable_factory import build_segment @@ -13,6 +14,8 @@ WorkflowDraftVariableResponse, WorkflowDraftVariableWithoutValueResponse, _serialize_full_content, + draft_variable_list_response_source, + draft_variable_response_source, ) from graphon.variables.types import SegmentType from libs.datetime_utils import naive_utc_now @@ -76,7 +79,7 @@ def test_serialize_full_content_should_raises_when_file_id_exists_but_file_is_no with pytest.raises(AssertionError): result = _serialize_full_content(draft_var) - def test_conversation_variable(self): + def test_conversation_variable(self, sqlite_session: Session): conv_var = WorkflowDraftVariable.new_conversation_variable( app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1) ) @@ -102,9 +105,14 @@ def test_conversation_variable(self): expected_with_value = expected_without_value.copy() expected_with_value["value"] = 1 expected_with_value["full_content"] = None - assert dump_response(WorkflowDraftVariableResponse, conv_var) == expected_with_value + assert ( + dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(conv_var, session=sqlite_session) + ) + == expected_with_value + ) - def test_create_sys_variable(self): + def test_create_sys_variable(self, sqlite_session: Session): sys_var = WorkflowDraftVariable.new_sys_variable( app_id=_TEST_APP_ID, name="sys_var", @@ -134,9 +142,14 @@ def test_create_sys_variable(self): expected_with_value = expected_without_value.copy() expected_with_value["value"] = "a" expected_with_value["full_content"] = None - assert dump_response(WorkflowDraftVariableResponse, sys_var) == expected_with_value + assert ( + dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(sys_var, session=sqlite_session) + ) + == expected_with_value + ) - def test_node_variable(self): + def test_node_variable(self, sqlite_session: Session): node_var = WorkflowDraftVariable.new_node_variable( app_id=_TEST_APP_ID, node_id="test_node", @@ -167,9 +180,14 @@ def test_node_variable(self): expected_with_value = expected_without_value.copy() expected_with_value["value"] = [1, "a"] expected_with_value["full_content"] = None - assert dump_response(WorkflowDraftVariableResponse, node_var) == expected_with_value + assert ( + dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(node_var, session=sqlite_session) + ) + == expected_with_value + ) - def test_node_variable_with_file(self): + def test_node_variable_with_file(self, sqlite_session: Session): node_var = WorkflowDraftVariable.new_node_variable( app_id=_TEST_APP_ID, node_id="test_node", @@ -219,7 +237,12 @@ def test_node_variable_with_file(self): "length": 10, "download_url": "http://example.com/signed-url", } - assert dump_response(WorkflowDraftVariableResponse, node_var) == expected_with_value + assert ( + dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(node_var, session=sqlite_session) + ) + == expected_with_value + ) class TestWorkflowDraftVariableList: @@ -301,11 +324,42 @@ class TestCase(NamedTuple): ) -def test_workflow_node_variables_fields(): +def test_value_response_rejects_bare_model(): + """A bare model would expose the raw serialized JSON string as the value, so it must fail loudly.""" + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1) + ) + with pytest.raises(TypeError, match="draft_variable_response_source"): + dump_response(WorkflowDraftVariableResponse, conv_var) + + +def test_value_response_accepts_plain_mapping(): + resp = WorkflowDraftVariableResponse.model_validate( + { + "id": "variable-id", + "type": "conversation", + "name": "conv_var", + "description": "", + "selector": ["conversation", "conv_var"], + "value_type": "number", + "edited": False, + "visible": True, + "is_truncated": False, + "value": 1, + "full_content": None, + } + ) + assert resp.value == 1 + + +def test_workflow_node_variables_fields(sqlite_session: Session): conv_var = WorkflowDraftVariable.new_conversation_variable( app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1) ) - resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[conv_var])) + resp = dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(WorkflowDraftVariableList(variables=[conv_var]), session=sqlite_session), + ) assert isinstance(resp, dict) assert len(resp["items"]) == 1 item_dict = resp["items"][0] @@ -313,7 +367,7 @@ def test_workflow_node_variables_fields(): assert item_dict["value"] == 1 -def test_workflow_file_variable_with_signed_url(): +def test_workflow_file_variable_with_signed_url(sqlite_session: Session): """Test that File type variables include signed URLs in API responses.""" from graphon.file import File, FileTransferMethod, FileType @@ -339,7 +393,10 @@ def test_workflow_file_variable_with_signed_url(): ) # Marshal the variable using the API fields - resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[file_var])) + resp = dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(WorkflowDraftVariableList(variables=[file_var]), session=sqlite_session), + ) # Verify the response structure assert isinstance(resp, dict) @@ -369,7 +426,7 @@ def test_workflow_file_variable_with_signed_url(): assert "sign=" in remote_url -def test_workflow_file_variable_remote_url(): +def test_workflow_file_variable_remote_url(sqlite_session: Session): """Test that File type variables with REMOTE_URL transfer method return the remote URL.""" from graphon.file import File, FileTransferMethod, FileType @@ -395,7 +452,10 @@ def test_workflow_file_variable_remote_url(): ) # Marshal the variable using the API fields - resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[file_var])) + resp = dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(WorkflowDraftVariableList(variables=[file_var]), session=sqlite_session), + ) # Verify the response structure assert isinstance(resp, dict) diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py index c5f43180ceee89..74236448460070 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py @@ -178,7 +178,7 @@ def test_delete_variables_success(self, app: Flask, sqlite_db, editor_user): patch("controllers.console.datasets.rag_pipeline.rag_pipeline_draft_variable.db", sqlite_db), patch("controllers.console.datasets.rag_pipeline.rag_pipeline_draft_variable.WorkflowDraftVariableService"), ): - result = method(api, editor_user, pipeline) + result = method(api, sqlite_db.session(), editor_user, pipeline) assert isinstance(result, Response) assert result.status_code == 204 @@ -206,12 +206,12 @@ def test_get_node_variables_success(self, app: Flask, sqlite_db, editor_user, re return_value=srv, ), ): - result = method(api, editor_user, pipeline, "node1") + result = method(api, sqlite_db.session(), editor_user, pipeline, "node1") assert result == {"items": []} srv.list_node_variables.assert_called_once_with("p1", "node1", user_id="account-1") - def test_get_node_variables_invalid_node(self, app: Flask, editor_user): + def test_get_node_variables_invalid_node(self, app: Flask, sqlite_db, editor_user): api = RagPipelineNodeVariableCollectionApi() method = unwrap(api.get) @@ -219,7 +219,7 @@ def test_get_node_variables_invalid_node(self, app: Flask, editor_user): app.test_request_context("/"), ): with pytest.raises(InvalidArgumentError): - method(api, editor_user, _pipeline(), SYSTEM_VARIABLE_NODE_ID) + method(api, sqlite_db.session(), editor_user, _pipeline(), SYSTEM_VARIABLE_NODE_ID) class TestRagPipelineVariableApi: @@ -239,7 +239,7 @@ def test_get_variable_not_found(self, app: Flask, sqlite_db, editor_user): ), ): with pytest.raises(NotFoundError): - method(api, editor_user, _pipeline(), "v1") + method(api, sqlite_db.session(), editor_user, _pipeline(), "v1") def test_patch_variable_invalid_file_payload(self, app: Flask, sqlite_db, editor_user): api = RagPipelineVariableApi() @@ -263,7 +263,14 @@ def test_patch_variable_invalid_file_payload(self, app: Flask, sqlite_db, editor ), ): with pytest.raises(InvalidArgumentError): - method(api, WorkflowDraftVariablePatchPayload.model_validate(payload), editor_user, pipeline, "v1") + method( + api, + sqlite_db.session(), + WorkflowDraftVariablePatchPayload.model_validate(payload), + editor_user, + pipeline, + "v1", + ) def test_delete_variable_success(self, app: Flask, sqlite_db, editor_user): api = RagPipelineVariableApi() @@ -283,7 +290,7 @@ def test_delete_variable_success(self, app: Flask, sqlite_db, editor_user): return_value=srv, ), ): - result = method(api, editor_user, pipeline, "v1") + result = method(api, sqlite_db.session(), editor_user, pipeline, "v1") assert result.status_code == 204 @@ -318,9 +325,9 @@ def test_direct_variable_access_rejects_different_user( rag_service = MagicMock() rag_service.get_draft_workflow.return_value = _workflow(pipeline) if payload is not None: - call_args = (api, payload, editor_user, pipeline, "v1") + call_args = (api, sqlite_db.session(), payload, editor_user, pipeline, "v1") else: - call_args = (api, editor_user, pipeline, "v1") + call_args = (api, sqlite_db.session(), editor_user, pipeline, "v1") with ( app.test_request_context("/"), @@ -376,7 +383,7 @@ def test_reset_variable_success(self, app: Flask, sqlite_db, editor_user): return_value={"id": "v1"}, ), ): - result = method(api, editor_user, pipeline, "v1") + result = method(api, sqlite_db.session(), editor_user, pipeline, "v1") assert result == {"id": "v1"} @@ -403,7 +410,7 @@ def test_system_variables_success(self, app: Flask, sqlite_db, editor_user, rest return_value=srv, ), ): - result = method(api, editor_user, pipeline) + result = method(api, sqlite_db.session(), editor_user, pipeline) assert result == {"items": []} srv.list_system_variables.assert_called_once_with("p1", user_id="account-1") diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable_handlers.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable_handlers.py new file mode 100644 index 00000000000000..ac0b29c05c42e7 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable_handlers.py @@ -0,0 +1,164 @@ +"""Direct-call tests for the rag-pipeline draft-variable handlers converted to `@with_session`.""" + +import json +from inspect import unwrap +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask +from sqlalchemy.orm import Session + +from controllers.console.datasets.rag_pipeline import rag_pipeline_draft_variable as module +from graphon.variables.segments import StringSegment +from models.account import Account, AccountStatus +from models.dataset import Pipeline +from models.workflow import Workflow, WorkflowDraftVariable + +_PIPELINE_ID = "pipeline-1" +_USER_ID = "user-1" + + +def _make_account() -> Account: + account = Account( + name="tester", + email="tester@example.com", + status=AccountStatus.ACTIVE, + ) + account.id = _USER_ID # type: ignore[assignment] + return account + + +def _make_pipeline() -> Pipeline: + pipeline = Pipeline( + tenant_id="tenant-1", + name="Pipeline", + ) + pipeline.id = _PIPELINE_ID + return pipeline + + +def _make_draft_workflow() -> Workflow: + workflow = Workflow.new( + tenant_id="tenant-1", + app_id=_PIPELINE_ID, + type="rag-pipeline", + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by=_USER_ID, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow-1" + return workflow + + +def _make_node_variable(variable_id: str, *, node_id: str = "extract-1") -> WorkflowDraftVariable: + variable = WorkflowDraftVariable.new_node_variable( + app_id=_PIPELINE_ID, + user_id=_USER_ID, + node_id=node_id, + name=variable_id, + value=StringSegment(value=f"value-{variable_id}"), + node_execution_id="execution-1", + ) + variable.id = variable_id + return variable + + +@pytest.fixture +def app() -> Flask: + app = Flask("test_rag_pipeline_draft_variable_handlers") + app.config["TESTING"] = True + return app + + +def test_node_variable_collection_get_returns_decoded_values(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.RagPipelineNodeVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), node_id="extract-1") + + assert [item["id"] for item in result["items"]] == ["var-1"] + assert result["items"][0]["value"] == "value-var-1" + + +def test_variable_get_returns_decoded_value(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.RagPipelineVariableApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), variable_id="var-1") + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_patch_without_changes_returns_variable(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.RagPipelineVariableApi() + handler = unwrap(api.patch) + + with app.test_request_context("/", method="PATCH", json={}): + result = handler( + api, + sqlite_session, + module.WorkflowDraftVariablePatchPayload(), + _make_account(), + pipeline=_make_pipeline(), + variable_id="var-1", + ) + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_reset_removes_variable_without_default(app: Flask, sqlite_session: Session) -> None: + workflow = _make_draft_workflow() + variable = _make_node_variable("var-1") + variable.node_execution_id = None + sqlite_session.add_all([workflow, variable]) + sqlite_session.commit() + api = module.RagPipelineVariableResetApi() + handler = unwrap(api.put) + + rag_pipeline_service = MagicMock() + rag_pipeline_service.get_draft_workflow.return_value = workflow + with ( + app.test_request_context("/", method="PUT"), + patch.object(module, "RagPipelineService", return_value=rag_pipeline_service), + ): + response = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), variable_id="var-1") + + assert response.status_code == 204 + + +def test_system_variable_collection_get_lists_system_variables(app: Flask, sqlite_session: Session) -> None: + sys_var = WorkflowDraftVariable.new_sys_variable( + app_id=_PIPELINE_ID, + user_id=_USER_ID, + name="query", + value=StringSegment(value="sys-value"), + node_execution_id="execution-1", + ) + sys_var.id = "sys-1" + sqlite_session.add(sys_var) + sqlite_session.commit() + api = module.RagPipelineSystemVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline()) + + assert [item["id"] for item in result["items"]] == ["sys-1"] + assert result["items"][0]["value"] == "sys-value" diff --git a/api/tests/unit_tests/models/test_workflow.py b/api/tests/unit_tests/models/test_workflow.py index 4e77e3e80cb16c..c58dc0e8fd8084 100644 --- a/api/tests/unit_tests/models/test_workflow.py +++ b/api/tests/unit_tests/models/test_workflow.py @@ -15,11 +15,13 @@ from graphon.variables import FloatVariable, IntegerVariable, SecretVariable, StringVariable from graphon.variables.segments import IntegerSegment, Segment from models.account import Account +from models.model import App, AppMode from models.tools import WorkflowToolProvider from models.workflow import ( Workflow, WorkflowDraftVariable, WorkflowNodeExecutionModel, + _resolve_workflow_app_tenant_id, is_system_variable_editable, ) @@ -302,7 +304,7 @@ def test_is_system_variable(self): class TestWorkflowDraftVariableGetValue: - def test_get_value_by_case(self): + def test_get_value_by_case(self, sqlite_session: Session): @dataclasses.dataclass class TestCase: name: str @@ -374,9 +376,9 @@ class TestCase: fail_msg = f"test case {c.name} failed, index={idx}" draft_var = WorkflowDraftVariable() draft_var.set_value(c.value) - assert c.value == draft_var.get_value(), fail_msg + assert c.value == draft_var.get_value(session=sqlite_session), fail_msg - def test_file_variable_preserves_all_fields(self): + def test_file_variable_preserves_all_fields(self, sqlite_session: Session): """Test that File type variables preserve all fields during encoding/decoding.""" tenant_id = "test_tenant_id" @@ -399,7 +401,7 @@ def test_file_variable_preserves_all_fields(self): draft_var.set_value(file_segment) # Retrieve the value and verify all fields are preserved - retrieved_segment = draft_var.get_value() + retrieved_segment = draft_var.get_value(session=sqlite_session) retrieved_file = retrieved_segment.value # Verify all important fields are preserved @@ -416,7 +418,7 @@ def test_file_variable_preserves_all_fields(self): # Verify the segments have the same type and the important fields match assert file_segment.value_type == retrieved_segment.value_type - def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self): + def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self, sqlite_session: Session): persisted_file = File( file_id="test_file_id", file_type=FileType.DOCUMENT, @@ -448,15 +450,71 @@ def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self): mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"), mock.patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file, ): - retrieved_segment = draft_var.get_value() + retrieved_segment = draft_var.get_value(session=sqlite_session) assert retrieved_segment.value == rebuilt_file rebuild_file.assert_called_once() assert rebuild_file.call_args.kwargs["tenant_id"] == "tenant-1" - def test_get_and_set_value(self): + def test_get_and_set_value(self, sqlite_session: Session): draft_var = WorkflowDraftVariable() int_var = IntegerSegment(value=1) draft_var.set_value(int_var) - value = draft_var.get_value() + value = draft_var.get_value(session=sqlite_session) assert value == int_var + + def test_array_file_variable_rebuilds_each_stored_payload_with_app_tenant(self, sqlite_session: Session): + files = [ + File( + file_id=f"test_file_id_{index}", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=build_file_reference(record_id=f"upload-{index}", storage_key="legacy-storage-key"), + filename=f"test_{index}.txt", + extension=".txt", + mime_type="text/plain", + size=12, + ) + for index in range(2) + ] + rebuilt_files = [file.model_copy(update={"storage_key": "canonical-storage-key"}, deep=True) for file in files] + encoded = WorkflowDraftVariable( + app_id="app-1", + ) + encoded.set_value(build_segment(files)) + # A fresh instance holding only the persisted columns, so get_value decodes from scratch. + draft_var = WorkflowDraftVariable( + app_id="app-1", + ) + draft_var.value = encoded.value + draft_var.value_type = encoded.value_type + + with ( + mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1") as resolve, + mock.patch("models.workflow.build_file_from_stored_mapping", side_effect=rebuilt_files) as rebuild_file, + ): + retrieved_segment = draft_var.get_value(session=sqlite_session) + + assert retrieved_segment.value == rebuilt_files + resolve.assert_called_once_with(sqlite_session, "app-1") + assert rebuild_file.call_count == 2 + assert all(call.kwargs["tenant_id"] == "tenant-1" for call in rebuild_file.call_args_list) + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +def test_resolve_workflow_app_tenant_id_reads_through_caller_session(sqlite_session: Session): + app = App( + id="app-1", + tenant_id="tenant-1", + name="Workflow", + mode=AppMode.WORKFLOW, + enable_site=True, + enable_api=True, + ) + sqlite_session.add(app) + sqlite_session.flush() + + assert _resolve_workflow_app_tenant_id(sqlite_session, "app-1") == "tenant-1" + + with pytest.raises(ValueError, match="Unable to resolve tenant_id"): + _resolve_workflow_app_tenant_id(sqlite_session, "missing-app") diff --git a/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py b/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py index 2fe76f7becccd9..45361d9b66d7b8 100644 --- a/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py +++ b/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py @@ -11,7 +11,7 @@ from core.workflow.file_reference import build_file_reference from extensions.storage.storage_type import StorageType from graphon.file import File, FileTransferMethod, FileType -from graphon.variables.segments import StringSegment +from graphon.variables.segments import ArrayFileSegment, FileSegment, StringSegment from graphon.variables.types import SegmentType from models.enums import CreatorUserRole from models.model import UploadFile @@ -159,6 +159,70 @@ def test_load_offloaded_variable_missing_upload_file_unit(self, draft_var_loader with pytest.raises(AssertionError): draft_var_loader._load_offloaded_variable(draft_var) + @pytest.mark.parametrize( + "sqlite_session", + [(WorkflowDraftVariable, WorkflowDraftVariableFile, UploadFile)], + indirect=True, + ) + def test_load_variables_collects_file_values_for_storage_key_loading( + self, + draft_var_loader: DraftVarLoader, + sqlite_session: Session, + ): + """File and array-file variables surface their File objects to the storage-key loader.""" + + def _make_file(index: int) -> File: + return File( + file_id=f"file-{index}", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=build_file_reference(record_id=f"upload-{index}", storage_key="legacy-storage-key"), + filename=f"test_{index}.txt", + extension=".txt", + mime_type="text/plain", + size=12, + ) + + single_file = _make_file(0) + array_files = [_make_file(1), _make_file(2)] + file_var = WorkflowDraftVariable.new_node_variable( + app_id="test-app-id", + user_id="test-user-id", + node_id="node1", + name="file_var", + value=FileSegment(value=single_file), + node_execution_id="execution-node1", + ) + array_var = WorkflowDraftVariable.new_node_variable( + app_id="test-app-id", + user_id="test-user-id", + node_id="node1", + name="array_var", + value=ArrayFileSegment(value=array_files), + node_execution_id="execution-node1", + ) + sqlite_session.add_all([file_var, array_var]) + sqlite_session.commit() + + rebuilt = { + file.id: file.model_copy(update={"storage_key": "canonical"}, deep=True) + for file in [single_file, *array_files] + } + + with ( + patch("models.workflow._resolve_workflow_app_tenant_id", return_value="test-tenant-id"), + patch( + "models.workflow.build_file_from_stored_mapping", + side_effect=lambda **kwargs: rebuilt[kwargs["file_mapping"]["id"]], + ), + patch("services.workflow_draft_variable_service.StorageKeyLoader") as storage_key_loader_cls, + ): + result = draft_var_loader.load_variables([["node1", "file_var"], ["node1", "array_var"]]) + + loaded_files = storage_key_loader_cls.return_value.load_storage_keys.call_args.args[0] + assert {file.id for file in loaded_files} == {"file-0", "file-1", "file-2"} + assert {variable.name for variable in result} == {"file_var", "array_var"} + def test_load_variables_empty_selectors_unit(self, draft_var_loader): """Test load_variables returns empty list for empty selectors.""" result = draft_var_loader.load_variables([]) diff --git a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py index ed4d78845d4c0f..6fb0dfd5ee37f4 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py @@ -167,7 +167,7 @@ def test_build_variables_from_start_mapping_rebuilds_system_files(self, sqlite_s draft_vars = saver._build_variables_from_start_mapping({"sys.files": [raw_file]}) sys_var = draft_vars[0] - assert sys_var.get_value().value[0] == rebuilt_file + assert sys_var.get_value(session=sqlite_session).value[0] == rebuilt_file rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1") @pytest.fixture diff --git a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service_sqlite.py b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service_sqlite.py index 2f66d12e6e2666..e5c6078a51ecd1 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service_sqlite.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service_sqlite.py @@ -160,7 +160,7 @@ def test_get_node_variable(self) -> None: assert node_var is not None assert node_var.id == self._node1_str_var_id assert node_var.name == "str_var" - assert node_var.get_value() == build_segment("str_value") + assert node_var.get_value(session=self._session) == build_segment("str_value") def test_get_system_variable(self) -> None: srv = self._get_test_srv() @@ -168,7 +168,7 @@ def test_get_system_variable(self) -> None: assert sys_var is not None assert sys_var.id == self._sys_var_id assert sys_var.name == "sys_var" - assert sys_var.get_value() == build_segment("sys_value") + assert sys_var.get_value(session=self._session) == build_segment("sys_value") def test_get_conversation_variable(self) -> None: srv = self._get_test_srv() @@ -176,7 +176,7 @@ def test_get_conversation_variable(self) -> None: assert conv_var is not None assert conv_var.id == self._conv_var_id assert conv_var.name == "conv_var" - assert conv_var.get_value() == build_segment("conv_value") + assert conv_var.get_value(session=self._session) == build_segment("conv_value") def test_delete_node_variables(self) -> None: srv = self._get_test_srv()