mirror of
https://github.com/langgenius/dify.git
synced 2026-08-01 09:50:50 +08:00
test: use SQLite sessions in core llm_generator (#39070)
This commit is contained in:
Asuka Minato
GitHub
parent
bab3db7f77
commit
8c4c1dc656
@ -1,14 +1,170 @@
|
||||
"""Tests for LLM generation and database-backed instruction modification."""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
from collections.abc import Iterator
|
||||
from datetime import datetime
|
||||
from decimal import Decimal
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session, scoped_session, sessionmaker
|
||||
|
||||
from core.app.app_config.entities import ModelConfig
|
||||
from core.llm_generator import llm_generator as llm_generator_module
|
||||
from core.llm_generator.entities import RuleCodeGeneratePayload, RuleGeneratePayload, RuleStructuredOutputPayload
|
||||
from core.llm_generator.llm_generator import LLMGenerator
|
||||
from graphon.model_runtime.entities.llm_entities import LLMMode, LLMResult
|
||||
from core.model_manager import ModelInstance, ModelManager
|
||||
from graphon.enums import WorkflowNodeExecutionStatus
|
||||
from graphon.model_runtime.entities.llm_entities import LLMMode, LLMResult, LLMUsage
|
||||
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage
|
||||
from graphon.model_runtime.entities.model_entities import ModelType
|
||||
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError
|
||||
from models.enums import ConversationFromSource, CreatorUserRole
|
||||
from models.model import App, AppMode, Message
|
||||
from models.workflow import (
|
||||
Workflow,
|
||||
WorkflowNodeExecutionModel,
|
||||
WorkflowNodeExecutionTriggeredFrom,
|
||||
)
|
||||
from services.workflow_service import WorkflowService
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def database(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Iterator[Session]:
|
||||
"""Bind the shared SQLite session to the generator's scoped-session interface."""
|
||||
|
||||
registry = scoped_session(lambda: sqlite_session)
|
||||
monkeypatch.setattr(llm_generator_module.db, "session", registry)
|
||||
try:
|
||||
yield sqlite_session
|
||||
finally:
|
||||
registry.remove()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def recording_model_instance(monkeypatch: pytest.MonkeyPatch) -> Mock:
|
||||
model_instance = Mock(spec=ModelInstance)
|
||||
model_instance.invoke_llm.return_value = _llm_result('{"modified": "workflow"}')
|
||||
model_manager = Mock(spec=ModelManager)
|
||||
model_manager.get_model_instance.return_value = model_instance
|
||||
monkeypatch.setattr(
|
||||
llm_generator_module.ModelManager,
|
||||
"for_tenant",
|
||||
Mock(return_value=model_manager),
|
||||
)
|
||||
return model_instance
|
||||
|
||||
|
||||
def _llm_result(content: str) -> LLMResult:
|
||||
return LLMResult(
|
||||
model="test-model",
|
||||
message=AssistantPromptMessage(content=content),
|
||||
usage=LLMUsage.empty_usage(),
|
||||
)
|
||||
|
||||
|
||||
def _persist_app(database: Session, *, tenant_id: str | None = None) -> App:
|
||||
app = App(
|
||||
id=str(uuid4()),
|
||||
tenant_id=tenant_id or str(uuid4()),
|
||||
name="Generator app",
|
||||
description="",
|
||||
mode=AppMode.WORKFLOW,
|
||||
icon_type=None,
|
||||
icon="",
|
||||
icon_background=None,
|
||||
enable_site=True,
|
||||
enable_api=True,
|
||||
)
|
||||
database.add(app)
|
||||
database.commit()
|
||||
return app
|
||||
|
||||
|
||||
def _persist_message(
|
||||
database: Session,
|
||||
app: App,
|
||||
*,
|
||||
query: str = "q",
|
||||
answer: str = "a",
|
||||
created_at: datetime | None = None,
|
||||
) -> Message:
|
||||
message = Message(
|
||||
id=str(uuid4()),
|
||||
app_id=app.id,
|
||||
conversation_id=str(uuid4()),
|
||||
_inputs={},
|
||||
query=query,
|
||||
message={},
|
||||
message_unit_price=Decimal(0),
|
||||
answer=answer,
|
||||
answer_unit_price=Decimal(0),
|
||||
currency="USD",
|
||||
from_source=ConversationFromSource.API,
|
||||
error="e",
|
||||
created_at=created_at,
|
||||
)
|
||||
database.add(message)
|
||||
database.commit()
|
||||
return message
|
||||
|
||||
|
||||
def _persist_workflow(database: Session, app: App, *, node_type: str | None) -> Workflow:
|
||||
nodes = [] if node_type is None else [{"id": "node", "data": {"type": node_type}}]
|
||||
workflow = Workflow.new(
|
||||
tenant_id=app.tenant_id,
|
||||
app_id=app.id,
|
||||
type="workflow",
|
||||
version=Workflow.VERSION_DRAFT,
|
||||
graph=json.dumps({"graph": {"nodes": nodes}}),
|
||||
features="{}",
|
||||
created_by=str(uuid4()),
|
||||
environment_variables=[],
|
||||
conversation_variables=[],
|
||||
rag_pipeline_variables=[],
|
||||
)
|
||||
database.add(workflow)
|
||||
database.commit()
|
||||
return workflow
|
||||
|
||||
|
||||
def _persist_node_execution(
|
||||
database: Session,
|
||||
app: App,
|
||||
workflow: Workflow,
|
||||
*,
|
||||
agent_log: list[dict[str, object]],
|
||||
inputs: dict[str, object] | None = None,
|
||||
) -> WorkflowNodeExecutionModel:
|
||||
execution = WorkflowNodeExecutionModel(
|
||||
id=str(uuid4()),
|
||||
tenant_id=app.tenant_id,
|
||||
app_id=app.id,
|
||||
workflow_id=workflow.id,
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.SINGLE_STEP,
|
||||
workflow_run_id=None,
|
||||
index=1,
|
||||
predecessor_node_id=None,
|
||||
node_execution_id=str(uuid4()),
|
||||
node_id="node",
|
||||
node_type="llm",
|
||||
title="LLM",
|
||||
inputs=json.dumps(inputs or {}),
|
||||
process_data=None,
|
||||
outputs=None,
|
||||
status=WorkflowNodeExecutionStatus.SUCCEEDED,
|
||||
error="",
|
||||
elapsed_time=0.1,
|
||||
execution_metadata=json.dumps({"agent_log": agent_log}),
|
||||
created_at=datetime(2026, 1, 1),
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=str(uuid4()),
|
||||
finished_at=datetime(2026, 1, 1),
|
||||
)
|
||||
database.add(execution)
|
||||
database.commit()
|
||||
return execution
|
||||
|
||||
|
||||
class TestLLMGenerator:
|
||||
@ -409,298 +565,296 @@ class TestLLMGenerator:
|
||||
result = LLMGenerator.generate_structured_output("tenant_id", payload)
|
||||
assert "An unexpected error occurred" in result["error"]
|
||||
|
||||
def test_instruction_modify_legacy_no_last_run(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session.scalar") as mock_scalar:
|
||||
mock_scalar.return_value = None
|
||||
def test_instruction_modify_legacy_without_last_run_uses_real_empty_query(
|
||||
self,
|
||||
database: Session,
|
||||
recording_model_instance: Mock,
|
||||
model_config_entity: ModelConfig,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
recording_model_instance.invoke_llm.return_value = _llm_result('{"modified": "prompt"}')
|
||||
|
||||
# Mock __instruction_modify_common call via invoke_llm
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = '{"modified": "prompt"}'
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
app.tenant_id,
|
||||
app.id,
|
||||
"current_val",
|
||||
"Test {{#last_run#}} and {{#current#}} and {{#error_message#}}",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
)
|
||||
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
"tenant_id", "flow_id", "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
assert result == {"modified": "prompt"}
|
||||
stmt = mock_scalar.call_args.args[0]
|
||||
compiled = stmt.compile()
|
||||
statement = str(compiled)
|
||||
assert "messages.app_id" in statement
|
||||
assert "apps.tenant_id" in statement
|
||||
assert "flow_id" in compiled.params.values()
|
||||
assert "tenant_id" in compiled.params.values()
|
||||
assert result == {"modified": "prompt"}
|
||||
user_payload = json.loads(recording_model_instance.invoke_llm.call_args.kwargs["prompt_messages"][1].content)
|
||||
assert "null" in user_payload["instruction"]
|
||||
assert "current_val" in user_payload["instruction"]
|
||||
|
||||
def test_instruction_modify_legacy_with_last_run(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session.scalar") as mock_scalar:
|
||||
last_run = MagicMock()
|
||||
last_run.query = "q"
|
||||
last_run.answer = "a"
|
||||
last_run.error = "e"
|
||||
mock_scalar.return_value = last_run
|
||||
def test_instruction_modify_legacy_reads_latest_tenant_scoped_message(
|
||||
self,
|
||||
database: Session,
|
||||
recording_model_instance: Mock,
|
||||
model_config_entity: ModelConfig,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
_persist_message(
|
||||
database,
|
||||
app,
|
||||
query="older question",
|
||||
answer="older answer",
|
||||
created_at=datetime(2026, 1, 1),
|
||||
)
|
||||
_persist_message(
|
||||
database,
|
||||
app,
|
||||
query="latest question",
|
||||
answer="latest answer",
|
||||
created_at=datetime(2026, 1, 2),
|
||||
)
|
||||
other_app = _persist_app(database)
|
||||
_persist_message(
|
||||
database,
|
||||
other_app,
|
||||
query="other tenant question",
|
||||
created_at=datetime(2026, 1, 3),
|
||||
)
|
||||
recording_model_instance.invoke_llm.return_value = _llm_result('{"modified": "prompt"}')
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = '{"modified": "prompt"}'
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
"tenant_id", "flow_id", "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
assert result == {"modified": "prompt"}
|
||||
stmt = mock_scalar.call_args.args[0]
|
||||
compiled = stmt.compile()
|
||||
statement = str(compiled)
|
||||
assert "messages.app_id" in statement
|
||||
assert "apps.tenant_id" in statement
|
||||
assert "flow_id" in compiled.params.values()
|
||||
assert "tenant_id" in compiled.params.values()
|
||||
assert result == {"modified": "prompt"}
|
||||
user_payload = json.loads(recording_model_instance.invoke_llm.call_args.kwargs["prompt_messages"][1].content)
|
||||
assert user_payload["last_run"]["query"] == "latest question"
|
||||
assert user_payload["last_run"]["answer"] == "latest answer"
|
||||
assert "older question" not in json.dumps(user_payload)
|
||||
assert "other tenant question" not in json.dumps(user_payload)
|
||||
|
||||
def test_instruction_modify_workflow_app_not_found(self):
|
||||
with patch("extensions.ext_database.db.session") as mock_session:
|
||||
mock_session.return_value.scalar.return_value = None
|
||||
with pytest.raises(ValueError, match="App not found."):
|
||||
LLMGenerator.instruction_modify_workflow("t", "f", "n", "c", "i", MagicMock(), "o", MagicMock())
|
||||
stmt = mock_session.return_value.scalar.call_args.args[0]
|
||||
compiled = stmt.compile()
|
||||
statement = str(compiled)
|
||||
assert "apps.id" in statement
|
||||
assert "apps.tenant_id" in statement
|
||||
assert "f" in compiled.params.values()
|
||||
assert "t" in compiled.params.values()
|
||||
def test_instruction_modify_workflow_app_not_found(
|
||||
self,
|
||||
database: Session,
|
||||
sqlite_session_factory: sessionmaker[Session],
|
||||
model_config_entity: ModelConfig,
|
||||
):
|
||||
workflow_service = WorkflowService(sqlite_session_factory)
|
||||
|
||||
def test_instruction_modify_workflow_no_workflow(self):
|
||||
with patch("extensions.ext_database.db.session") as mock_session:
|
||||
mock_session.return_value.scalar.return_value = MagicMock()
|
||||
workflow_service = MagicMock()
|
||||
workflow_service.get_draft_workflow.return_value = None
|
||||
with pytest.raises(ValueError, match="Workflow not found for the given app model."):
|
||||
LLMGenerator.instruction_modify_workflow("t", "f", "n", "c", "i", MagicMock(), "o", workflow_service)
|
||||
|
||||
def test_instruction_modify_workflow_success(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session") as mock_session:
|
||||
mock_session.return_value.scalar.return_value = MagicMock()
|
||||
workflow = MagicMock()
|
||||
workflow.graph_dict = {"graph": {"nodes": [{"id": "node_id", "data": {"type": "llm"}}]}}
|
||||
|
||||
workflow_service = MagicMock()
|
||||
workflow_service.get_draft_workflow.return_value = workflow
|
||||
|
||||
last_run = MagicMock()
|
||||
last_run.node_type = "llm"
|
||||
last_run.status = "s"
|
||||
last_run.error = "e"
|
||||
# Return regular values, not Mocks
|
||||
last_run.execution_metadata_dict = {"agent_log": [{"status": "s", "error": "e", "data": {}}]}
|
||||
last_run.load_full_inputs.return_value = {"in": "val"}
|
||||
|
||||
workflow_service.get_node_last_run.return_value = last_run
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = '{"modified": "workflow"}'
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
|
||||
result = LLMGenerator.instruction_modify_workflow(
|
||||
"tenant_id",
|
||||
"flow_id",
|
||||
"node_id",
|
||||
with pytest.raises(ValueError, match="App not found"):
|
||||
LLMGenerator.instruction_modify_workflow(
|
||||
str(uuid4()),
|
||||
str(uuid4()),
|
||||
"node",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
assert result == {"modified": "workflow"}
|
||||
|
||||
def test_instruction_modify_workflow_no_last_run_fallback(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session") as mock_session:
|
||||
mock_session.return_value.scalar.return_value = MagicMock()
|
||||
workflow = MagicMock()
|
||||
workflow.graph_dict = {"graph": {"nodes": [{"id": "node_id", "data": {"type": "code"}}]}}
|
||||
def test_instruction_modify_workflow_rejects_app_from_another_tenant(
|
||||
self,
|
||||
database: Session,
|
||||
sqlite_session_factory: sessionmaker[Session],
|
||||
model_config_entity: ModelConfig,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
workflow_service = WorkflowService(sqlite_session_factory)
|
||||
|
||||
workflow_service = MagicMock()
|
||||
workflow_service.get_draft_workflow.return_value = workflow
|
||||
workflow_service.get_node_last_run.return_value = None
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = '{"modified": "fallback"}'
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
|
||||
result = LLMGenerator.instruction_modify_workflow(
|
||||
"tenant_id",
|
||||
"flow_id",
|
||||
"node_id",
|
||||
with pytest.raises(ValueError, match="App not found"):
|
||||
LLMGenerator.instruction_modify_workflow(
|
||||
str(uuid4()),
|
||||
app.id,
|
||||
"node",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
assert result == {"modified": "fallback"}
|
||||
|
||||
def test_instruction_modify_workflow_node_type_fallback(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session") as mock_session:
|
||||
mock_session.return_value.scalar.return_value = MagicMock()
|
||||
workflow = MagicMock()
|
||||
# Cause exception in node_type logic
|
||||
workflow.graph_dict = {"graph": {"nodes": []}}
|
||||
def test_instruction_modify_workflow_requires_draft_workflow(
|
||||
self,
|
||||
database: Session,
|
||||
sqlite_session_factory: sessionmaker[Session],
|
||||
model_config_entity: ModelConfig,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
workflow_service = WorkflowService(sqlite_session_factory)
|
||||
|
||||
workflow_service = MagicMock()
|
||||
workflow_service.get_draft_workflow.return_value = workflow
|
||||
workflow_service.get_node_last_run.return_value = None
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = '{"modified": "fallback"}'
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
|
||||
result = LLMGenerator.instruction_modify_workflow(
|
||||
"tenant_id",
|
||||
"flow_id",
|
||||
"node_id",
|
||||
with pytest.raises(ValueError, match="Workflow not found"):
|
||||
LLMGenerator.instruction_modify_workflow(
|
||||
app.tenant_id,
|
||||
app.id,
|
||||
"node",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
assert result == {"modified": "fallback"}
|
||||
|
||||
def test_instruction_modify_workflow_empty_agent_log(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session") as mock_session:
|
||||
mock_session.return_value.scalar.return_value = MagicMock()
|
||||
workflow = MagicMock()
|
||||
workflow.graph_dict = {"graph": {"nodes": [{"id": "node_id", "data": {"type": "llm"}}]}}
|
||||
def test_instruction_modify_workflow_uses_last_run(
|
||||
self,
|
||||
database: Session,
|
||||
sqlite_session_factory: sessionmaker[Session],
|
||||
recording_model_instance: Mock,
|
||||
model_config_entity: ModelConfig,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
workflow = _persist_workflow(database, app, node_type="llm")
|
||||
_persist_node_execution(
|
||||
database,
|
||||
app,
|
||||
workflow,
|
||||
inputs={"input": "value"},
|
||||
agent_log=[{"status": "s", "error": "", "data": {"step": 1}}],
|
||||
)
|
||||
workflow_service = WorkflowService(sqlite_session_factory)
|
||||
|
||||
workflow_service = MagicMock()
|
||||
workflow_service.get_draft_workflow.return_value = workflow
|
||||
result = LLMGenerator.instruction_modify_workflow(
|
||||
app.tenant_id,
|
||||
app.id,
|
||||
"node",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
|
||||
last_run = MagicMock()
|
||||
last_run.node_type = "llm"
|
||||
last_run.status = "s"
|
||||
last_run.error = "e"
|
||||
# Return regular empty list, not a Mock
|
||||
last_run.execution_metadata_dict = {"agent_log": []}
|
||||
last_run.load_full_inputs.return_value = {}
|
||||
assert result == {"modified": "workflow"}
|
||||
user_payload = json.loads(recording_model_instance.invoke_llm.call_args.kwargs["prompt_messages"][1].content)
|
||||
assert user_payload["last_run"]["inputs"] == {"input": "value"}
|
||||
assert user_payload["last_run"]["agent_log"][0]["data"] == {"step": 1}
|
||||
|
||||
workflow_service.get_node_last_run.return_value = last_run
|
||||
def test_instruction_modify_workflow_accepts_empty_agent_log(
|
||||
self,
|
||||
database: Session,
|
||||
sqlite_session_factory: sessionmaker[Session],
|
||||
recording_model_instance: Mock,
|
||||
model_config_entity: ModelConfig,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
workflow = _persist_workflow(database, app, node_type="llm")
|
||||
_persist_node_execution(database, app, workflow, agent_log=[])
|
||||
workflow_service = WorkflowService(sqlite_session_factory)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = '{"modified": "workflow"}'
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
result = LLMGenerator.instruction_modify_workflow(
|
||||
app.tenant_id,
|
||||
app.id,
|
||||
"node",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
|
||||
result = LLMGenerator.instruction_modify_workflow(
|
||||
"tenant_id",
|
||||
"flow_id",
|
||||
"node_id",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
assert result == {"modified": "workflow"}
|
||||
assert result == {"modified": "workflow"}
|
||||
user_payload = json.loads(recording_model_instance.invoke_llm.call_args.kwargs["prompt_messages"][1].content)
|
||||
assert user_payload["last_run"]["agent_log"] == []
|
||||
|
||||
def test_instruction_modify_common_placeholders(self, mock_model_instance, model_config_entity):
|
||||
# Testing placeholders replacement via instruction_modify_legacy for convenience
|
||||
with patch("extensions.ext_database.db.session.scalar") as mock_scalar:
|
||||
mock_scalar.return_value = None
|
||||
@pytest.mark.parametrize(
|
||||
"node_type",
|
||||
[
|
||||
"code",
|
||||
None,
|
||||
],
|
||||
)
|
||||
def test_instruction_modify_workflow_falls_back_without_last_run(
|
||||
self,
|
||||
database: Session,
|
||||
sqlite_session_factory: sessionmaker[Session],
|
||||
recording_model_instance: Mock,
|
||||
model_config_entity: ModelConfig,
|
||||
node_type: str | None,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
_persist_workflow(database, app, node_type=node_type)
|
||||
workflow_service = WorkflowService(sqlite_session_factory)
|
||||
recording_model_instance.invoke_llm.return_value = _llm_result('{"modified": "fallback"}')
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = '{"ok": true}'
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
result = LLMGenerator.instruction_modify_workflow(
|
||||
app.tenant_id,
|
||||
app.id,
|
||||
"node",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
|
||||
instruction = "Test {{#last_run#}} and {{#current#}} and {{#error_message#}}"
|
||||
LLMGenerator.instruction_modify_legacy(
|
||||
"tenant_id", "flow_id", "current_val", instruction, model_config_entity, "ideal"
|
||||
)
|
||||
assert result == {"modified": "fallback"}
|
||||
|
||||
# Verify the call to invoke_llm contains replaced instruction
|
||||
args, kwargs = mock_model_instance.invoke_llm.call_args
|
||||
prompt_messages = kwargs["prompt_messages"]
|
||||
user_msg = prompt_messages[1].content
|
||||
user_msg_dict = json.loads(user_msg)
|
||||
assert "null" in user_msg_dict["instruction"] # because last_run is None and current is current_val etc.
|
||||
assert "current_val" in user_msg_dict["instruction"]
|
||||
def test_instruction_modify_workflow_falls_back_for_unknown_node_type(
|
||||
self,
|
||||
database: Session,
|
||||
sqlite_session_factory: sessionmaker[Session],
|
||||
recording_model_instance: Mock,
|
||||
model_config_entity: ModelConfig,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
_persist_workflow(database, app, node_type="unknown")
|
||||
workflow_service = WorkflowService(sqlite_session_factory)
|
||||
|
||||
def test_instruction_modify_common_no_braces(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session.scalar") as mock_scalar:
|
||||
mock_scalar.return_value = None
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = "No braces here"
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
"tenant_id", "flow_id", "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
assert "An unexpected error occurred" in result["error"]
|
||||
assert "Could not find a valid JSON object" in result["error"]
|
||||
result = LLMGenerator.instruction_modify_workflow(
|
||||
app.tenant_id,
|
||||
app.id,
|
||||
"node",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
|
||||
def test_instruction_modify_common_not_dict(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session.scalar") as mock_scalar:
|
||||
mock_scalar.return_value = None
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = "[1, 2, 3]"
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
"tenant_id", "flow_id", "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
# The exception message is "Expected a JSON object, but got list"
|
||||
assert "An unexpected error occurred" in result["error"]
|
||||
assert result == {"modified": "workflow"}
|
||||
system_prompt = recording_model_instance.invoke_llm.call_args.kwargs["prompt_messages"][0].content
|
||||
assert system_prompt == llm_generator_module.LLM_MODIFY_PROMPT_SYSTEM
|
||||
|
||||
def test_instruction_modify_common_other_node_type(self, mock_model_instance, model_config_entity):
|
||||
with patch("core.llm_generator.llm_generator.ModelManager.for_tenant") as mock_manager:
|
||||
instance = MagicMock()
|
||||
mock_manager.return_value.get_model_instance.return_value = instance
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = '{"ok": true}'
|
||||
instance.invoke_llm.return_value = mock_response
|
||||
@pytest.mark.parametrize(
|
||||
("raw_output", "error_fragment"),
|
||||
[
|
||||
("No braces here", "Could not find a valid JSON object"),
|
||||
("[1, 2, 3]", "Could not find a valid JSON object"),
|
||||
],
|
||||
)
|
||||
def test_instruction_modify_rejects_invalid_model_output(
|
||||
self,
|
||||
database: Session,
|
||||
mock_model_instance: MagicMock,
|
||||
model_config_entity: ModelConfig,
|
||||
raw_output: str,
|
||||
error_fragment: str,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
response = MagicMock()
|
||||
response.message.get_text_content.return_value = raw_output
|
||||
mock_model_instance.invoke_llm.return_value = response
|
||||
|
||||
with patch("extensions.ext_database.db.session") as mock_session:
|
||||
mock_session.return_value.scalar.return_value = MagicMock()
|
||||
workflow = MagicMock()
|
||||
workflow.graph_dict = {"graph": {"nodes": [{"id": "node_id", "data": {"type": "other"}}]}}
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
|
||||
workflow_service = MagicMock()
|
||||
workflow_service.get_draft_workflow.return_value = workflow
|
||||
workflow_service.get_node_last_run.return_value = None
|
||||
assert "An unexpected error occurred" in result["error"]
|
||||
assert error_fragment in result["error"]
|
||||
|
||||
LLMGenerator.instruction_modify_workflow(
|
||||
"tenant_id",
|
||||
"flow_id",
|
||||
"node_id",
|
||||
"current",
|
||||
"instruction",
|
||||
model_config_entity,
|
||||
"ideal",
|
||||
workflow_service,
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("model_error", "error_fragment"),
|
||||
[(InvokeError("invoke failed"), "Failed to generate code"), (RuntimeError("boom"), "unexpected error")],
|
||||
)
|
||||
def test_instruction_modify_handles_model_errors(
|
||||
self,
|
||||
database: Session,
|
||||
mock_model_instance: MagicMock,
|
||||
model_config_entity: ModelConfig,
|
||||
model_error: Exception,
|
||||
error_fragment: str,
|
||||
):
|
||||
app = _persist_app(database)
|
||||
mock_model_instance.invoke_llm.side_effect = model_error
|
||||
|
||||
def test_instruction_modify_common_invoke_error(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session.scalar") as mock_scalar:
|
||||
mock_scalar.return_value = None
|
||||
mock_model_instance.invoke_llm.side_effect = InvokeError("Invoke Failed")
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
app.tenant_id, app.id, "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
"tenant_id", "flow_id", "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
assert "Failed to generate code" in result["error"]
|
||||
|
||||
def test_instruction_modify_common_exception(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session.scalar") as mock_scalar:
|
||||
mock_scalar.return_value = None
|
||||
mock_model_instance.invoke_llm.side_effect = Exception("Random error")
|
||||
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
"tenant_id", "flow_id", "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
assert "An unexpected error occurred" in result["error"]
|
||||
|
||||
def test_instruction_modify_common_json_error(self, mock_model_instance, model_config_entity):
|
||||
with patch("extensions.ext_database.db.session.scalar") as mock_scalar:
|
||||
mock_scalar.return_value = None
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.message.get_text_content.return_value = "No JSON here"
|
||||
mock_model_instance.invoke_llm.return_value = mock_response
|
||||
|
||||
result = LLMGenerator.instruction_modify_legacy(
|
||||
"tenant_id", "flow_id", "current", "instruction", model_config_entity, "ideal"
|
||||
)
|
||||
assert "An unexpected error occurred" in result["error"]
|
||||
assert error_fragment.lower() in result["error"].lower()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user