Underground-Digital/Workflow-Engine

1import json2import os3import time4import uuid5from collections.abc import Generator6from unittest.mock import MagicMock7 8import pytest9 10from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity11from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle12from core.entities.provider_entities import CustomConfiguration, CustomProviderConfiguration, SystemConfiguration13from core.model_manager import ModelInstance14from core.model_runtime.entities.model_entities import ModelType15from core.model_runtime.model_providers import ModelProviderFactory16from core.workflow.entities.variable_pool import VariablePool17from core.workflow.enums import SystemVariableKey18from core.workflow.graph_engine.entities.graph import Graph19from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams20from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState21from core.workflow.nodes.event import RunCompletedEvent22from core.workflow.nodes.llm.node import LLMNode23from extensions.ext_database import db24from models.enums import UserFrom25from models.provider import ProviderType26from models.workflow import WorkflowNodeExecutionStatus, WorkflowType27 28"""FOR MOCK FIXTURES, DO NOT REMOVE"""29from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock30from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock31 32 33def init_llm_node(config: dict) -> LLMNode:34    graph_config = {35        "edges": [36            {37                "id": "start-source-next-target",38                "source": "start",39                "target": "llm",40            },41        ],42        "nodes": [{"data": {"type": "start"}, "id": "start"}, config],43    }44 45    graph = Graph.init(graph_config=graph_config)46 47    init_params = GraphInitParams(48        tenant_id="1",49        app_id="1",50        workflow_type=WorkflowType.WORKFLOW,51        workflow_id="1",52        graph_config=graph_config,53        user_id="1",54        user_from=UserFrom.ACCOUNT,55        invoke_from=InvokeFrom.DEBUGGER,56        call_depth=0,57    )58 59    # construct variable pool60    variable_pool = VariablePool(61        system_variables={62            SystemVariableKey.QUERY: "what's the weather today?",63            SystemVariableKey.FILES: [],64            SystemVariableKey.CONVERSATION_ID: "abababa",65            SystemVariableKey.USER_ID: "aaa",66        },67        user_inputs={},68        environment_variables=[],69        conversation_variables=[],70    )71    variable_pool.add(["abc", "output"], "sunny")72 73    node = LLMNode(74        id=str(uuid.uuid4()),75        graph_init_params=init_params,76        graph=graph,77        graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()),78        config=config,79    )80 81    return node82 83 84@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)85def test_execute_llm(setup_openai_mock):86    node = init_llm_node(87        config={88            "id": "llm",89            "data": {90                "title": "123",91                "type": "llm",92                "model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}},93                "prompt_template": [94                    {"role": "system", "text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}."},95                    {"role": "user", "text": "{{#sys.query#}}"},96                ],97                "memory": None,98                "context": {"enabled": False},99                "vision": {"enabled": False},100            },101        },102    )103 104    credentials = {"openai_api_key": os.environ.get("OPENAI_API_KEY")}105 106    provider_instance = ModelProviderFactory().get_provider_instance("openai")107    model_type_instance = provider_instance.get_model_instance(ModelType.LLM)108    provider_model_bundle = ProviderModelBundle(109        configuration=ProviderConfiguration(110            tenant_id="1",111            provider=provider_instance.get_provider_schema(),112            preferred_provider_type=ProviderType.CUSTOM,113            using_provider_type=ProviderType.CUSTOM,114            system_configuration=SystemConfiguration(enabled=False),115            custom_configuration=CustomConfiguration(provider=CustomProviderConfiguration(credentials=credentials)),116            model_settings=[],117        ),118        provider_instance=provider_instance,119        model_type_instance=model_type_instance,120    )121    model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo")122    model_schema = model_type_instance.get_model_schema("gpt-3.5-turbo")123    assert model_schema is not None124    model_config = ModelConfigWithCredentialsEntity(125        model="gpt-3.5-turbo",126        provider="openai",127        mode="chat",128        credentials=credentials,129        parameters={},130        model_schema=model_schema,131        provider_model_bundle=provider_model_bundle,132    )133 134    # Mock db.session.close()135    db.session.close = MagicMock()136 137    node._fetch_model_config = MagicMock(return_value=(model_instance, model_config))138 139    # execute node140    result = node._run()141    assert isinstance(result, Generator)142 143    for item in result:144        if isinstance(item, RunCompletedEvent):145            assert item.run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED146            assert item.run_result.process_data is not None147            assert item.run_result.outputs is not None148            assert item.run_result.outputs.get("text") is not None149            assert item.run_result.outputs.get("usage", {})["total_tokens"] > 0150 151 152@pytest.mark.parametrize("setup_code_executor_mock", [["none"]], indirect=True)153@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)154def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_openai_mock):155    """156    Test execute LLM node with jinja2157    """158    node = init_llm_node(159        config={160            "id": "llm",161            "data": {162                "title": "123",163                "type": "llm",164                "model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}},165                "prompt_config": {166                    "jinja2_variables": [167                        {"variable": "sys_query", "value_selector": ["sys", "query"]},168                        {"variable": "output", "value_selector": ["abc", "output"]},169                    ]170                },171                "prompt_template": [172                    {173                        "role": "system",174                        "text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}",175                        "jinja2_text": "you are a helpful assistant.\ntoday's weather is {{output}}.",176                        "edition_type": "jinja2",177                    },178                    {179                        "role": "user",180                        "text": "{{#sys.query#}}",181                        "jinja2_text": "{{sys_query}}",182                        "edition_type": "basic",183                    },184                ],185                "memory": None,186                "context": {"enabled": False},187                "vision": {"enabled": False},188            },189        },190    )191 192    credentials = {"openai_api_key": os.environ.get("OPENAI_API_KEY")}193 194    provider_instance = ModelProviderFactory().get_provider_instance("openai")195    model_type_instance = provider_instance.get_model_instance(ModelType.LLM)196    provider_model_bundle = ProviderModelBundle(197        configuration=ProviderConfiguration(198            tenant_id="1",199            provider=provider_instance.get_provider_schema(),200            preferred_provider_type=ProviderType.CUSTOM,201            using_provider_type=ProviderType.CUSTOM,202            system_configuration=SystemConfiguration(enabled=False),203            custom_configuration=CustomConfiguration(provider=CustomProviderConfiguration(credentials=credentials)),204            model_settings=[],205        ),206        provider_instance=provider_instance,207        model_type_instance=model_type_instance,208    )209 210    model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo")211    model_schema = model_type_instance.get_model_schema("gpt-3.5-turbo")212    assert model_schema is not None213    model_config = ModelConfigWithCredentialsEntity(214        model="gpt-3.5-turbo",215        provider="openai",216        mode="chat",217        credentials=credentials,218        parameters={},219        model_schema=model_schema,220        provider_model_bundle=provider_model_bundle,221    )222 223    # Mock db.session.close()224    db.session.close = MagicMock()225 226    node._fetch_model_config = MagicMock(return_value=(model_instance, model_config))227 228    # execute node229    result = node._run()230 231    for item in result:232        if isinstance(item, RunCompletedEvent):233            assert item.run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED234            assert item.run_result.process_data is not None235            assert "sunny" in json.dumps(item.run_result.process_data)236            assert "what's the weather today?" in json.dumps(item.run_result.process_data)237