import os import time import uuid from unittest.mock import MagicMock from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.model_manager import ModelInstance from core.workflow.node_runtime import DifyPromptMessageSerializer from core.workflow.system_variables import build_system_variables from extensions.ext_database import db from graphon.enums import WorkflowNodeExecutionStatus from graphon.model_runtime.entities import AssistantPromptMessage, UserPromptMessage from graphon.nodes.llm.protocols import CredentialsProvider, ModelFactory from graphon.nodes.parameter_extractor.entities import ParameterExtractorNodeData from graphon.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode from graphon.runtime import GraphRuntimeState, VariablePool from tests.integration_tests.workflow.nodes.__mock.model import get_mocked_fetch_model_instance from tests.workflow_test_utils import build_test_graph_init_params pytest_plugins = ("tests.integration_tests.model_runtime.__mock.plugin_daemon",) def get_mocked_fetch_memory(memory_text: str): class MemoryMock: def get_history_prompt_messages( self, max_token_limit: int = 2000, message_limit: int | None = None, ): return [UserPromptMessage(content=memory_text), AssistantPromptMessage(content="mocked answer")] return MagicMock(return_value=MemoryMock()) def init_parameter_extractor_node(config: dict, memory=None): graph_config = { "edges": [ { "id": "start-source-next-target", "source": "start", "target": "llm", }, ], "nodes": [{"data": {"type": "start", "title": "Start"}, "id": "start"}, config], } init_params = build_test_graph_init_params( workflow_id="1", graph_config=graph_config, tenant_id="1", app_id="1", user_id="1", 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", files=[], query="what's the weather in SF", conversation_id="abababa" ), user_inputs={}, environment_variables=[], conversation_variables=[], ) variable_pool.add(["a", "args1"], 1) variable_pool.add(["a", "args2"], 2) graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()) node = ParameterExtractorNode( node_id=str(uuid.uuid4()), data=ParameterExtractorNodeData.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), memory=memory, prompt_message_serializer=DifyPromptMessageSerializer(), ) return node def _mock_db_session_close(monkeypatch) -> None: monkeypatch.setattr(db.session, "close", MagicMock()) def test_function_calling_parameter_extractor(setup_model_mock, monkeypatch): """ Test function calling for parameter extractor. """ node = init_parameter_extractor_node( config={ "id": "llm", "data": { "title": "123", "type": "parameter-extractor", "model": { "provider": "langgenius/openai/openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}, }, "query": ["sys", "query"], "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], "instruction": "", "reasoning_mode": "function_call", "memory": None, }, } ) node._model_instance = get_mocked_fetch_model_instance( provider="langgenius/openai/openai", model="gpt-3.5-turbo", mode="chat", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, )() _mock_db_session_close(monkeypatch) result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED assert result.outputs is not None assert result.outputs.get("location") == "kawaii" assert result.outputs.get("__reason") == None def test_instructions(setup_model_mock, monkeypatch): """ Test chat parameter extractor. """ node = init_parameter_extractor_node( config={ "id": "llm", "data": { "title": "123", "type": "parameter-extractor", "model": { "provider": "langgenius/openai/openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}, }, "query": ["sys", "query"], "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], "reasoning_mode": "function_call", "instruction": "{{#sys.query#}}", "memory": None, }, }, ) node._model_instance = get_mocked_fetch_model_instance( provider="langgenius/openai/openai", model="gpt-3.5-turbo", mode="chat", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, )() _mock_db_session_close(monkeypatch) result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED assert result.outputs is not None assert result.outputs.get("location") == "kawaii" assert result.outputs.get("__reason") == None process_data = result.process_data assert process_data is not None process_data.get("prompts") for prompt in process_data.get("prompts", []): if prompt.get("role") == "system": assert "what's the weather in SF" in prompt.get("text") def test_chat_parameter_extractor(setup_model_mock, monkeypatch): """ Test chat parameter extractor. """ node = init_parameter_extractor_node( config={ "id": "llm", "data": { "title": "123", "type": "parameter-extractor", "model": { "provider": "langgenius/openai/openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}, }, "query": ["sys", "query"], "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], "reasoning_mode": "prompt", "instruction": "", "memory": None, }, }, ) node._model_instance = get_mocked_fetch_model_instance( provider="langgenius/openai/openai", model="gpt-3.5-turbo", mode="chat", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, )() _mock_db_session_close(monkeypatch) result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED assert result.outputs is not None assert result.outputs.get("location") == "" assert ( result.outputs.get("__reason") == "Failed to extract result from function call or text response, using empty result." ) assert result.process_data is not None prompts = result.process_data.get("prompts", []) for prompt in prompts: if prompt.get("role") == "user": if "" in prompt.get("text"): assert '\n{"type": "object"' in prompt.get("text") def test_completion_parameter_extractor(setup_model_mock, monkeypatch): """ Test completion parameter extractor. """ node = init_parameter_extractor_node( config={ "id": "llm", "data": { "title": "123", "type": "parameter-extractor", "model": { "provider": "langgenius/openai/openai", "name": "gpt-3.5-turbo-instruct", "mode": "completion", "completion_params": {}, }, "query": ["sys", "query"], "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], "reasoning_mode": "prompt", "instruction": "{{#sys.query#}}", "memory": None, }, }, ) node._model_instance = get_mocked_fetch_model_instance( provider="langgenius/openai/openai", model="gpt-3.5-turbo-instruct", mode="completion", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, )() _mock_db_session_close(monkeypatch) result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED assert result.outputs is not None assert result.outputs.get("location") == "" assert ( result.outputs.get("__reason") == "Failed to extract result from function call or text response, using empty result." ) assert result.process_data is not None assert len(result.process_data.get("prompts", [])) == 1 assert "SF" in result.process_data.get("prompts", [])[0].get("text") def test_extract_json_response(): """ Test extract json response. """ node = init_parameter_extractor_node( config={ "id": "llm", "data": { "title": "123", "type": "parameter-extractor", "model": { "provider": "langgenius/openai/openai", "name": "gpt-3.5-turbo-instruct", "mode": "completion", "completion_params": {}, }, "query": ["sys", "query"], "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], "reasoning_mode": "prompt", "instruction": "{{#sys.query#}}", "memory": None, }, }, ) result = node._extract_complete_json_response(""" uwu{ovo} { "location": "kawaii" } hello world. """) assert result is not None assert result["location"] == "kawaii" def test_extract_json_from_tool_call(): """ Test extract json response. """ node = init_parameter_extractor_node( config={ "id": "llm", "data": { "title": "123", "type": "parameter-extractor", "model": { "provider": "langgenius/openai/openai", "name": "gpt-3.5-turbo-instruct", "mode": "completion", "completion_params": {}, }, "query": ["sys", "query"], "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], "reasoning_mode": "prompt", "instruction": "{{#sys.query#}}", "memory": None, }, }, ) result = node._extract_json_from_tool_call( AssistantPromptMessage.ToolCall( id="llm", type="parameter-extractor", function=AssistantPromptMessage.ToolCall.ToolCallFunction( name="foo", arguments="""{"location":"kawaii"}{"location": 1}""" ), ) ) assert result is not None assert result["location"] == "kawaii" def test_chat_parameter_extractor_with_memory(setup_model_mock, monkeypatch): """ Test chat parameter extractor with memory. """ node = init_parameter_extractor_node( config={ "id": "llm", "data": { "title": "123", "type": "parameter-extractor", "model": { "provider": "langgenius/openai/openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}, }, "query": ["sys", "query"], "parameters": [{"name": "location", "type": "string", "description": "location", "required": True}], "reasoning_mode": "prompt", "instruction": "", "memory": {"window": {"enabled": True, "size": 50}}, }, }, memory=get_mocked_fetch_memory("customized memory")(), ) node._model_instance = get_mocked_fetch_model_instance( provider="langgenius/openai/openai", model="gpt-3.5-turbo", mode="chat", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}, )() _mock_db_session_close(monkeypatch) result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED assert result.outputs is not None assert result.outputs.get("location") == "" assert ( result.outputs.get("__reason") == "Failed to extract result from function call or text response, using empty result." ) assert result.process_data is not None prompts = result.process_data.get("prompts", []) latest_role = None for prompt in prompts: if prompt.get("role") == "user": if "" in prompt.get("text"): assert '\n{"type": "object"' in prompt.get("text") elif prompt.get("role") == "system": assert "customized memory" in prompt.get("text") if latest_role is not None: assert latest_role != prompt.get("role") if prompt.get("role") in {"user", "assistant"}: latest_role = prompt.get("role")