diff --git a/api/models/dataset.py b/api/models/dataset.py index 01c17fa243..33d396aeb9 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -21,6 +21,7 @@ from configs import dify_config from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource from core.rag.retrieval.retrieval_methods import RetrievalMethod from extensions.ext_storage import storage +from models.base import TypeBase from services.entities.knowledge_entities.knowledge_entities import ParentMode, Rule from .account import Account @@ -906,17 +907,21 @@ class ChildChunk(Base): return db.session.query(DocumentSegment).where(DocumentSegment.id == self.segment_id).first() -class AppDatasetJoin(Base): +class AppDatasetJoin(TypeBase): __tablename__ = "app_dataset_joins" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="app_dataset_join_pkey"), sa.Index("app_dataset_join_app_dataset_idx", "dataset_id", "app_id"), ) - id = mapped_column(StringUUID, primary_key=True, nullable=False, server_default=sa.text("uuid_generate_v4()")) - app_id = mapped_column(StringUUID, nullable=False) - dataset_id = mapped_column(StringUUID, nullable=False) - created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=sa.func.current_timestamp()) + id: Mapped[str] = mapped_column( + StringUUID, primary_key=True, nullable=False, server_default=sa.text("uuid_generate_v4()"), init=False + ) + app_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) + created_at: Mapped[datetime] = mapped_column( + DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False + ) @property def app(self): diff --git a/api/services/workflow/workflow_converter.py b/api/services/workflow/workflow_converter.py index e70b2b5c95..067feb994f 100644 --- a/api/services/workflow/workflow_converter.py +++ b/api/services/workflow/workflow_converter.py @@ -1,5 +1,5 @@ import json -from typing import Any +from typing import Any, TypedDict from core.app.app_config.entities import ( DatasetEntity, @@ -28,6 +28,12 @@ from models.model import App, AppMode, AppModelConfig from models.workflow import Workflow, WorkflowType +class _NodeType(TypedDict): + id: str + position: None + data: dict[str, Any] + + class WorkflowConverter: """ App Convert to Workflow Mode @@ -217,7 +223,7 @@ class WorkflowConverter: return app_config - def _convert_to_start_node(self, variables: list[VariableEntity]): + def _convert_to_start_node(self, variables: list[VariableEntity]) -> _NodeType: """ Convert to Start Node :param variables: list of variables @@ -235,7 +241,7 @@ class WorkflowConverter: def _convert_to_http_request_node( self, app_model: App, variables: list[VariableEntity], external_data_variables: list[ExternalDataVariableEntity] - ) -> tuple[list[dict], dict[str, str]]: + ) -> tuple[list[_NodeType], dict[str, str]]: """ Convert API Based Extension to HTTP Request Node :param app_model: App instance @@ -285,7 +291,7 @@ class WorkflowConverter: request_body_json = json.dumps(request_body) request_body_json = request_body_json.replace(r"\{\{", "{{").replace(r"\}\}", "}}") - http_request_node = { + http_request_node: _NodeType = { "id": f"http_request_{index}", "position": None, "data": { @@ -303,7 +309,7 @@ class WorkflowConverter: nodes.append(http_request_node) # append code node for response body parsing - code_node: dict[str, Any] = { + code_node: _NodeType = { "id": f"code_{index}", "position": None, "data": { @@ -326,7 +332,7 @@ class WorkflowConverter: def _convert_to_knowledge_retrieval_node( self, new_app_mode: AppMode, dataset_config: DatasetEntity, model_config: ModelConfigEntity - ) -> dict | None: + ) -> _NodeType | None: """ Convert datasets to Knowledge Retrieval Node :param new_app_mode: new app mode @@ -384,7 +390,7 @@ class WorkflowConverter: prompt_template: PromptTemplateEntity, file_upload: FileUploadConfig | None = None, external_data_variable_node_mapping: dict[str, str] | None = None, - ): + ) -> _NodeType: """ Convert to LLM Node :param original_app_mode: original app mode @@ -561,7 +567,7 @@ class WorkflowConverter: return template - def _convert_to_end_node(self): + def _convert_to_end_node(self) -> _NodeType: """ Convert to End Node :return: @@ -577,7 +583,7 @@ class WorkflowConverter: }, } - def _convert_to_answer_node(self): + def _convert_to_answer_node(self) -> _NodeType: """ Convert to Answer Node :return: @@ -598,7 +604,7 @@ class WorkflowConverter: """ return {"id": f"{source}-{target}", "source": source, "target": target} - def _append_node(self, graph: dict, node: dict): + def _append_node(self, graph: dict[str, Any], node: _NodeType): """ Append Node to Graph diff --git a/api/tests/unit_tests/services/workflow/test_workflow_converter.py b/api/tests/unit_tests/services/workflow/test_workflow_converter.py index 63ce4c0c3c..8ea5754363 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_converter.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_converter.py @@ -199,6 +199,7 @@ def test__convert_to_knowledge_retrieval_node_for_chatbot(): node = WorkflowConverter()._convert_to_knowledge_retrieval_node( new_app_mode=new_app_mode, dataset_config=dataset_config, model_config=model_config ) + assert node is not None assert node["data"]["type"] == "knowledge-retrieval" assert node["data"]["query_variable_selector"] == ["sys", "query"] @@ -231,6 +232,7 @@ def test__convert_to_knowledge_retrieval_node_for_workflow_app(): node = WorkflowConverter()._convert_to_knowledge_retrieval_node( new_app_mode=new_app_mode, dataset_config=dataset_config, model_config=model_config ) + assert node is not None assert node["data"]["type"] == "knowledge-retrieval" assert node["data"]["query_variable_selector"] == ["start", dataset_config.retrieve_config.query_variable]