|
import json |
|
import os |
|
import time |
|
import uuid |
|
from collections.abc import Generator |
|
from unittest.mock import MagicMock |
|
|
|
import pytest |
|
|
|
from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity |
|
from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle |
|
from core.entities.provider_entities import CustomConfiguration, CustomProviderConfiguration, SystemConfiguration |
|
from core.model_manager import ModelInstance |
|
from core.model_runtime.entities.model_entities import ModelType |
|
from core.model_runtime.model_providers import ModelProviderFactory |
|
from core.workflow.entities.variable_pool import VariablePool |
|
from core.workflow.enums import SystemVariableKey |
|
from core.workflow.graph_engine.entities.graph import Graph |
|
from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams |
|
from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState |
|
from core.workflow.nodes.event import RunCompletedEvent |
|
from core.workflow.nodes.llm.node import LLMNode |
|
from extensions.ext_database import db |
|
from models.enums import UserFrom |
|
from models.provider import ProviderType |
|
from models.workflow import WorkflowNodeExecutionStatus, WorkflowType |
|
|
|
"""FOR MOCK FIXTURES, DO NOT REMOVE""" |
|
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock |
|
from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock |
|
|
|
|
|
def init_llm_node(config: dict) -> LLMNode: |
|
graph_config = { |
|
"edges": [ |
|
{ |
|
"id": "start-source-next-target", |
|
"source": "start", |
|
"target": "llm", |
|
}, |
|
], |
|
"nodes": [{"data": {"type": "start"}, "id": "start"}, config], |
|
} |
|
|
|
graph = Graph.init(graph_config=graph_config) |
|
|
|
init_params = GraphInitParams( |
|
tenant_id="1", |
|
app_id="1", |
|
workflow_type=WorkflowType.WORKFLOW, |
|
workflow_id="1", |
|
graph_config=graph_config, |
|
user_id="1", |
|
user_from=UserFrom.ACCOUNT, |
|
invoke_from=InvokeFrom.DEBUGGER, |
|
call_depth=0, |
|
) |
|
|
|
|
|
variable_pool = VariablePool( |
|
system_variables={ |
|
SystemVariableKey.QUERY: "what's the weather today?", |
|
SystemVariableKey.FILES: [], |
|
SystemVariableKey.CONVERSATION_ID: "abababa", |
|
SystemVariableKey.USER_ID: "aaa", |
|
}, |
|
user_inputs={}, |
|
environment_variables=[], |
|
conversation_variables=[], |
|
) |
|
variable_pool.add(["abc", "output"], "sunny") |
|
|
|
node = LLMNode( |
|
id=str(uuid.uuid4()), |
|
graph_init_params=init_params, |
|
graph=graph, |
|
graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), |
|
config=config, |
|
) |
|
|
|
return node |
|
|
|
|
|
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True) |
|
def test_execute_llm(setup_openai_mock): |
|
node = init_llm_node( |
|
config={ |
|
"id": "llm", |
|
"data": { |
|
"title": "123", |
|
"type": "llm", |
|
"model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}}, |
|
"prompt_template": [ |
|
{"role": "system", "text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}."}, |
|
{"role": "user", "text": "{{#sys.query#}}"}, |
|
], |
|
"memory": None, |
|
"context": {"enabled": False}, |
|
"vision": {"enabled": False}, |
|
}, |
|
}, |
|
) |
|
|
|
credentials = {"openai_api_key": os.environ.get("OPENAI_API_KEY")} |
|
|
|
provider_instance = ModelProviderFactory().get_provider_instance("openai") |
|
model_type_instance = provider_instance.get_model_instance(ModelType.LLM) |
|
provider_model_bundle = ProviderModelBundle( |
|
configuration=ProviderConfiguration( |
|
tenant_id="1", |
|
provider=provider_instance.get_provider_schema(), |
|
preferred_provider_type=ProviderType.CUSTOM, |
|
using_provider_type=ProviderType.CUSTOM, |
|
system_configuration=SystemConfiguration(enabled=False), |
|
custom_configuration=CustomConfiguration(provider=CustomProviderConfiguration(credentials=credentials)), |
|
model_settings=[], |
|
), |
|
provider_instance=provider_instance, |
|
model_type_instance=model_type_instance, |
|
) |
|
model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo") |
|
model_schema = model_type_instance.get_model_schema("gpt-3.5-turbo") |
|
assert model_schema is not None |
|
model_config = ModelConfigWithCredentialsEntity( |
|
model="gpt-3.5-turbo", |
|
provider="openai", |
|
mode="chat", |
|
credentials=credentials, |
|
parameters={}, |
|
model_schema=model_schema, |
|
provider_model_bundle=provider_model_bundle, |
|
) |
|
|
|
|
|
db.session.close = MagicMock() |
|
|
|
node._fetch_model_config = MagicMock(return_value=(model_instance, model_config)) |
|
|
|
|
|
result = node._run() |
|
assert isinstance(result, Generator) |
|
|
|
for item in result: |
|
if isinstance(item, RunCompletedEvent): |
|
assert item.run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED |
|
assert item.run_result.process_data is not None |
|
assert item.run_result.outputs is not None |
|
assert item.run_result.outputs.get("text") is not None |
|
assert item.run_result.outputs.get("usage", {})["total_tokens"] > 0 |
|
|
|
|
|
@pytest.mark.parametrize("setup_code_executor_mock", [["none"]], indirect=True) |
|
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True) |
|
def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_openai_mock): |
|
""" |
|
Test execute LLM node with jinja2 |
|
""" |
|
node = init_llm_node( |
|
config={ |
|
"id": "llm", |
|
"data": { |
|
"title": "123", |
|
"type": "llm", |
|
"model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}}, |
|
"prompt_config": { |
|
"jinja2_variables": [ |
|
{"variable": "sys_query", "value_selector": ["sys", "query"]}, |
|
{"variable": "output", "value_selector": ["abc", "output"]}, |
|
] |
|
}, |
|
"prompt_template": [ |
|
{ |
|
"role": "system", |
|
"text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}", |
|
"jinja2_text": "you are a helpful assistant.\ntoday's weather is {{output}}.", |
|
"edition_type": "jinja2", |
|
}, |
|
{ |
|
"role": "user", |
|
"text": "{{#sys.query#}}", |
|
"jinja2_text": "{{sys_query}}", |
|
"edition_type": "basic", |
|
}, |
|
], |
|
"memory": None, |
|
"context": {"enabled": False}, |
|
"vision": {"enabled": False}, |
|
}, |
|
}, |
|
) |
|
|
|
credentials = {"openai_api_key": os.environ.get("OPENAI_API_KEY")} |
|
|
|
provider_instance = ModelProviderFactory().get_provider_instance("openai") |
|
model_type_instance = provider_instance.get_model_instance(ModelType.LLM) |
|
provider_model_bundle = ProviderModelBundle( |
|
configuration=ProviderConfiguration( |
|
tenant_id="1", |
|
provider=provider_instance.get_provider_schema(), |
|
preferred_provider_type=ProviderType.CUSTOM, |
|
using_provider_type=ProviderType.CUSTOM, |
|
system_configuration=SystemConfiguration(enabled=False), |
|
custom_configuration=CustomConfiguration(provider=CustomProviderConfiguration(credentials=credentials)), |
|
model_settings=[], |
|
), |
|
provider_instance=provider_instance, |
|
model_type_instance=model_type_instance, |
|
) |
|
|
|
model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo") |
|
model_schema = model_type_instance.get_model_schema("gpt-3.5-turbo") |
|
assert model_schema is not None |
|
model_config = ModelConfigWithCredentialsEntity( |
|
model="gpt-3.5-turbo", |
|
provider="openai", |
|
mode="chat", |
|
credentials=credentials, |
|
parameters={}, |
|
model_schema=model_schema, |
|
provider_model_bundle=provider_model_bundle, |
|
) |
|
|
|
|
|
db.session.close = MagicMock() |
|
|
|
node._fetch_model_config = MagicMock(return_value=(model_instance, model_config)) |
|
|
|
|
|
result = node._run() |
|
|
|
for item in result: |
|
if isinstance(item, RunCompletedEvent): |
|
assert item.run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED |
|
assert item.run_result.process_data is not None |
|
assert "sunny" in json.dumps(item.run_result.process_data) |
|
assert "what's the weather today?" in json.dumps(item.run_result.process_data) |
|
|
