import json import time import uuid from collections.abc import Generator from unittest.mock import MagicMock, patch from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.llm_generator.output_parser.structured_output import _parse_structured_output from core.model_manager import ModelInstance from core.workflow.system_variables import build_system_variables from extensions.ext_database import db from graphon.enums import WorkflowNodeExecutionStatus from graphon.node_events import StreamCompletedEvent from graphon.nodes.llm.entities import LLMNodeData from graphon.nodes.llm.file_saver import LLMFileSaver from graphon.nodes.llm.node import LLMNode from graphon.nodes.llm.protocols import CredentialsProvider, ModelFactory from graphon.nodes.llm.runtime_protocols import PromptMessageSerializerProtocol from graphon.nodes.protocols import HttpClientProtocol from graphon.runtime import GraphRuntimeState, VariablePool from tests.workflow_test_utils import build_test_graph_init_params """FOR MOCK FIXTURES, DO NOT REMOVE""" def init_llm_node(config: dict) -> LLMNode: graph_config = { "edges": [ { "id": "start-source-next-target", "source": "start", "target": "llm", }, ], "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config], } # Use proper UUIDs for database compatibility tenant_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056b" app_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056c" workflow_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056d" user_id = "9d2074fc-6f86-45a9-b09d-6ecc63b9056e" init_params = build_test_graph_init_params( workflow_id=workflow_id, graph_config=graph_config, tenant_id=tenant_id, app_id=app_id, user_id=user_id, user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, call_depth=0, ) # construct variable pool variable_pool = VariablePool.from_bootstrap( system_variables=build_system_variables( user_id="aaa", app_id=app_id, workflow_id=workflow_id, files=[], query="what's the weather today?", conversation_id="abababa", ), user_inputs={}, environment_variables=[], conversation_variables=[], ) variable_pool.add(["abc", "output"], "sunny") graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()) prompt_message_serializer = MagicMock(spec=PromptMessageSerializerProtocol) prompt_message_serializer.serialize.side_effect = lambda *, model_mode, prompt_messages: [ message.model_dump(mode="json") for message in prompt_messages ] llm_file_saver = MagicMock(spec=LLMFileSaver) node = LLMNode( node_id=str(uuid.uuid4()), data=LLMNodeData.model_validate(config["data"]), graph_init_params=init_params, graph_runtime_state=graph_runtime_state, credentials_provider=MagicMock(spec=CredentialsProvider), model_factory=MagicMock(spec=ModelFactory), model_instance=MagicMock(spec=ModelInstance), llm_file_saver=llm_file_saver, prompt_message_serializer=prompt_message_serializer, http_client=MagicMock(spec=HttpClientProtocol), ) return node def _mock_db_session_close(monkeypatch) -> None: monkeypatch.setattr(db.session, "close", MagicMock()) def test_execute_llm(monkeypatch): 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}, }, }, ) _mock_db_session_close(monkeypatch) def build_mock_model_instance() -> MagicMock: from decimal import Decimal from unittest.mock import MagicMock from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage # Create mock model instance mock_model_instance = MagicMock(spec=ModelInstance) mock_model_instance.provider = "openai" mock_model_instance.model_name = "gpt-3.5-turbo" mock_model_instance.credentials = {} mock_model_instance.parameters = {} mock_model_instance.stop = [] mock_model_instance.model_type_instance = MagicMock() mock_model_instance.model_type_instance.get_model_schema.return_value = MagicMock( model_properties={}, parameter_rules=[], features=[], ) mock_model_instance.provider_model_bundle = MagicMock() mock_model_instance.provider_model_bundle.configuration.using_provider_type = "custom" mock_usage = LLMUsage( prompt_tokens=30, prompt_unit_price=Decimal("0.001"), prompt_price_unit=Decimal(1000), prompt_price=Decimal("0.00003"), completion_tokens=20, completion_unit_price=Decimal("0.002"), completion_price_unit=Decimal(1000), completion_price=Decimal("0.00004"), total_tokens=50, total_price=Decimal("0.00007"), currency="USD", latency=0.5, ) mock_message = AssistantPromptMessage(content="Test response from mock") mock_llm_result = LLMResult( model="gpt-3.5-turbo", prompt_messages=[], message=mock_message, usage=mock_usage, ) mock_model_instance.invoke_llm.return_value = mock_llm_result return mock_model_instance # Mock fetch_prompt_messages to avoid database calls def mock_fetch_prompt_messages_1(**_kwargs): from graphon.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage return [ SystemPromptMessage(content="you are a helpful assistant. today's weather is sunny."), UserPromptMessage(content="what's the weather today?"), ], [] node._model_instance = build_mock_model_instance() with patch.object(LLMNode, "fetch_prompt_messages", mock_fetch_prompt_messages_1): # execute node result = node._run() assert isinstance(result, Generator) for item in result: if isinstance(item, StreamCompletedEvent): if item.node_run_result.status != WorkflowNodeExecutionStatus.SUCCEEDED: print(f"Error: {item.node_run_result.error}") print(f"Error type: {item.node_run_result.error_type}") assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED assert item.node_run_result.process_data is not None assert item.node_run_result.outputs is not None assert item.node_run_result.outputs.get("text") is not None assert item.node_run_result.outputs.get("usage", {})["total_tokens"] > 0 def test_execute_llm_with_jinja2(monkeypatch): """ 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}, }, }, ) _mock_db_session_close(monkeypatch) def build_mock_model_instance() -> MagicMock: from decimal import Decimal from unittest.mock import MagicMock from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage # Create mock model instance mock_model_instance = MagicMock(spec=ModelInstance) mock_model_instance.provider = "openai" mock_model_instance.model_name = "gpt-3.5-turbo" mock_model_instance.credentials = {} mock_model_instance.parameters = {} mock_model_instance.stop = [] mock_model_instance.model_type_instance = MagicMock() mock_model_instance.model_type_instance.get_model_schema.return_value = MagicMock( model_properties={}, parameter_rules=[], features=[], ) mock_model_instance.provider_model_bundle = MagicMock() mock_model_instance.provider_model_bundle.configuration.using_provider_type = "custom" mock_usage = LLMUsage( prompt_tokens=30, prompt_unit_price=Decimal("0.001"), prompt_price_unit=Decimal(1000), prompt_price=Decimal("0.00003"), completion_tokens=20, completion_unit_price=Decimal("0.002"), completion_price_unit=Decimal(1000), completion_price=Decimal("0.00004"), total_tokens=50, total_price=Decimal("0.00007"), currency="USD", latency=0.5, ) mock_message = AssistantPromptMessage(content="Test response: sunny weather and what's the weather today?") mock_llm_result = LLMResult( model="gpt-3.5-turbo", prompt_messages=[], message=mock_message, usage=mock_usage, ) mock_model_instance.invoke_llm.return_value = mock_llm_result return mock_model_instance # Mock fetch_prompt_messages to avoid database calls def mock_fetch_prompt_messages_2(**_kwargs): from graphon.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage return [ SystemPromptMessage(content="you are a helpful assistant. today's weather is sunny."), UserPromptMessage(content="what's the weather today?"), ], [] node._model_instance = build_mock_model_instance() with patch.object(LLMNode, "fetch_prompt_messages", mock_fetch_prompt_messages_2): # execute node result = node._run() for item in result: if isinstance(item, StreamCompletedEvent): assert item.node_run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED assert item.node_run_result.process_data is not None assert "sunny" in json.dumps(item.node_run_result.process_data) assert "what's the weather today?" in json.dumps(item.node_run_result.process_data) def test_extract_json(): llm_texts = [ '\n\n{"name": "test", "age": 123', # resoning model (deepseek-r1) '{"name":"test","age":123}', # json schema model (gpt-4o) '{\n "name": "test",\n "age": 123\n}', # small model (llama-3.2-1b) '```json\n{"name": "test", "age": 123}\n```', # json markdown (deepseek-chat) '{"name":"test",age:123}', # without quotes (qwen-2.5-0.5b) ] result = {"name": "test", "age": 123} assert all(_parse_structured_output(item) == result for item in llm_texts)