From 08c7a2c5f3d5ada4c58b27ab6bf4334e7950cc8f Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 4 Sep 2026 12:16:31 +0000 Subject: [PATCH] test: migrate application generation sessions and ORM models to SQLite (#40509) --- .../app/apps/agent_app/test_app_generator.py | 283 ++++++++++++------ .../agent_chat/test_agent_chat_app_runner.py | 88 +++--- .../chat/test_app_generator_and_runner.py | 256 +++++++++------- .../services/test_app_generate_service.py | 96 ++++-- 4 files changed, 448 insertions(+), 275 deletions(-) diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py index f4a0c4e90f3..410efadba33 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_generator.py @@ -12,9 +12,13 @@ from __future__ import annotations import contextlib import inspect import json +from decimal import Decimal +from types import SimpleNamespace import pytest from pytest_mock import MockerFixture +from sqlalchemy import event +from sqlalchemy.orm import Session import core.app.apps.agent_app.app_generator as module from core.app.apps.agent_app.app_generator import ( @@ -27,14 +31,107 @@ from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.app.entities.queue_entities import QueueAnnotationReplyEvent from core.workflow.file_reference import build_file_reference from models import Account, AppModelConfig +from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource, AgentStatus +from models.agent_config_entities import AgentSoulConfig +from models.enums import AppStatus, ConversationFromSource +from models.model import App, AppMode, Conversation, Message, MessageAnnotation MODULE = "core.app.apps.agent_app.app_generator" -class DummyAccount: - def __init__(self, user_id: str) -> None: - self.id = user_id - self.session_id = f"session-{user_id}" +def _account(user_id: str = "user") -> Account: + account = Account(name="User", email=f"{user_id}@example.com") + account.id = user_id + return account + + +def _app(*, app_model_config_id: str | None = None) -> App: + return App( + id="app1", + tenant_id="tenant", + name="Agent App", + description="", + mode=AppMode.AGENT, + app_model_config_id=app_model_config_id, + status=AppStatus.NORMAL, + enable_site=False, + enable_api=False, + api_rpm=0, + api_rph=0, + ) + + +def _agent(*, agent_id: str = "agent1") -> Agent: + return Agent( + id=agent_id, + tenant_id="tenant", + name="Agent", + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=AgentStatus.ACTIVE, + app_id="app1", + ) + + +def _snapshot(*, snapshot_id: str = "snap1", agent_id: str = "agent1") -> AgentConfigSnapshot: + return AgentConfigSnapshot( + id=snapshot_id, + tenant_id="tenant", + agent_id=agent_id, + version=1, + config_snapshot=AgentSoulConfig(), + home_snapshot_id="home-1", + ) + + +def _conversation(*, invoke_from: InvokeFrom = InvokeFrom.WEB_APP) -> Conversation: + conversation = Conversation( + id="conv", + app_id="app1", + mode=AppMode.AGENT, + name="Conversation", + invoke_from=invoke_from, + from_source=ConversationFromSource.CONSOLE, + from_account_id="user", + is_deleted=False, + ) + conversation._inputs = {} + return conversation + + +def _message(*, query: str = "query") -> Message: + message = Message( + id="msg", + app_id="app1", + conversation_id="conv", + query=query, + message={"role": "user", "content": query}, + answer="", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + from_account_id="user", + ) + message._inputs = {} + return message + + +_CURRENT_SESSION: Session | None = None + + +@pytest.fixture(autouse=True) +def _bind_real_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch): + global _CURRENT_SESSION + _CURRENT_SESSION = sqlite_session + monkeypatch.setattr(module.db, "session", sqlite_session) + yield + _CURRENT_SESSION = None + + +def _session() -> Session: + assert _CURRENT_SESSION is not None + return _CURRENT_SESSION @pytest.fixture @@ -46,35 +143,35 @@ def generator(mocker: MockerFixture) -> AgentAppGenerator: class TestGenerateGuards: - def test_rejects_blocking_mode(self, generator: AgentAppGenerator, mocker: MockerFixture): + def test_rejects_blocking_mode(self, generator: AgentAppGenerator): with pytest.raises(AgentAppGeneratorError, match="only supports streaming"): generator.generate( - app_model=mocker.MagicMock(), - user=DummyAccount("u"), + app_model=_app(), + user=_account("u"), args={}, invoke_from=InvokeFrom.WEB_APP, - session=mocker.MagicMock(), + session=_session(), streaming=False, ) - def test_requires_query(self, generator: AgentAppGenerator, mocker: MockerFixture): + def test_requires_query(self, generator: AgentAppGenerator): with pytest.raises(AgentAppGeneratorError, match="query is required"): generator.generate( - app_model=mocker.MagicMock(), - user=DummyAccount("u"), + app_model=_app(), + user=_account("u"), args={"inputs": {}}, invoke_from=InvokeFrom.WEB_APP, - session=mocker.MagicMock(), + session=_session(), ) - def test_rejects_blank_query(self, generator: AgentAppGenerator, mocker: MockerFixture): + def test_rejects_blank_query(self, generator: AgentAppGenerator): with pytest.raises(AgentAppGeneratorError, match="query is required"): generator.generate( - app_model=mocker.MagicMock(), - user=DummyAccount("u"), + app_model=_app(), + user=_account("u"), args={"query": " ", "inputs": {}}, invoke_from=InvokeFrom.WEB_APP, - session=mocker.MagicMock(), + session=_session(), ) @@ -94,18 +191,17 @@ class TestGenerateSuccess: ) def test_generate_orchestrates_and_starts_worker(self, generator, mocker: MockerFixture): - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") - app_model.app_model_config_id = "config-1" - user = DummyAccount("user") - session = mocker.MagicMock() + config = AppModelConfig(app_id="app1") + config.id = "config-1" + session = _session() + session.add(config) + session.commit() + app_model = _app(app_model_config_id=config.id) + user = _account() - generator._resolve_agent = mocker.MagicMock( - return_value=(mocker.MagicMock(id="agent1"), "snap1", "snapshot", mocker.MagicMock()) - ) + generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "snapshot", AgentSoulConfig())) generator._prepare_user_inputs = mocker.MagicMock(return_value={"x": 1}) - generator._init_generate_records = mocker.MagicMock( - return_value=(mocker.MagicMock(id="conv", mode="agent"), mocker.MagicMock(id="msg")) - ) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) generator._handle_response = mocker.MagicMock(return_value="raw-response") mocker.patch( @@ -151,23 +247,19 @@ class TestGenerateSuccess: session=session, conversation=None, ) - session.get.assert_called_once_with(AppModelConfig, "config-1") + assert session.get(AppModelConfig, "config-1") is config assert generate_entity.call_args.kwargs["prompt_file_mappings"] == file_mappings assert "agent_runtime_exit_intent" not in generate_entity.call_args.kwargs def test_generate_loads_existing_conversation(self, generator: AgentAppGenerator, mocker: MockerFixture): - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") + app_model = _app() generator._resolve_agent = mocker.MagicMock( - return_value=(mocker.MagicMock(id="a"), "snap1", "snapshot", mocker.MagicMock()) + return_value=(_agent(agent_id="a"), "snap1", "snapshot", AgentSoulConfig()) ) generator._prepare_user_inputs = mocker.MagicMock(return_value={}) - generator._init_generate_records = mocker.MagicMock( - return_value=(mocker.MagicMock(id="conv", mode="agent"), mocker.MagicMock(id="msg")) - ) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) generator._handle_response = mocker.MagicMock(return_value="raw") - get_conv = mocker.patch( - f"{MODULE}.ConversationService.get_conversation", return_value=mocker.MagicMock(id="conv") - ) + get_conv = mocker.patch(f"{MODULE}.ConversationService.get_conversation", return_value=_conversation()) mocker.patch(f"{MODULE}.AgentAppConfigManager.get_app_config", return_value=mocker.MagicMock(variables=[])) mocker.patch(f"{MODULE}.load_annotation_reply_config", return_value={"enabled": False}) mocker.patch(f"{MODULE}.ModelConfigConverter.convert", return_value=mocker.MagicMock()) @@ -176,8 +268,8 @@ class TestGenerateSuccess: mocker.patch(f"{MODULE}.MessageBasedAppQueueManager", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.threading.Thread", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.AgentAppGenerateResponseConverter.convert", return_value={"result": "ok"}) - session = mocker.MagicMock() - user = DummyAccount("user") + session = _session() + user = _account() generator.generate( app_model=app_model, @@ -200,16 +292,12 @@ class TestGenerateSuccess: def test_generate_does_not_include_trace_session_id_in_extras( self, generator: AgentAppGenerator, mocker: MockerFixture ): - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") - user = DummyAccount("user") + app_model = _app() + user = _account() - generator._resolve_agent = mocker.MagicMock( - return_value=(mocker.MagicMock(id="agent1"), "snap1", "snapshot", mocker.MagicMock()) - ) + generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "snapshot", AgentSoulConfig())) generator._prepare_user_inputs = mocker.MagicMock(return_value={}) - generator._init_generate_records = mocker.MagicMock( - return_value=(mocker.MagicMock(id="conv", mode="agent"), mocker.MagicMock(id="msg")) - ) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) generator._handle_response = mocker.MagicMock(return_value="raw-response") mocker.patch( @@ -230,7 +318,7 @@ class TestGenerateSuccess: user=user, args={"query": "hello", "inputs": {}, "trace_session_id": "session-1"}, invoke_from=InvokeFrom.WEB_APP, - session=mocker.MagicMock(), + session=_session(), streaming=True, ) @@ -255,25 +343,22 @@ class TestGenerateWorker: handled=False, guard_query="query", ): - generator._get_conversation = mocker.MagicMock(return_value=mocker.MagicMock(id="conv")) - generator._get_message = mocker.MagicMock(return_value=mocker.MagicMock(id="msg")) + generator._get_conversation = mocker.MagicMock(return_value=_conversation()) + generator._get_message = mocker.MagicMock(return_value=_message()) generator._run_input_guards = mocker.MagicMock(return_value=(handled, guard_query, None)) - resolved_agent = mocker.MagicMock(id="a") - resolved_config = mocker.MagicMock(id="s", home_snapshot_id="home-1") - generator._resolve_agent_by_id = mocker.MagicMock( - return_value=(resolved_agent, resolved_config, mocker.MagicMock()) - ) - session = mocker.MagicMock() - session.get.return_value = mocker.MagicMock(id="app1") - session_context = mocker.MagicMock() - session_context.__enter__.return_value = session - session_maker = mocker.patch(f"{MODULE}.session_factory.get_session_maker").return_value - session_maker.begin.return_value = session_context - resolver_session = mocker.MagicMock() - resolver_context = mocker.MagicMock() - resolver_context.__enter__.return_value = resolver_session - mocker.patch(f"{MODULE}.session_factory.create_session", return_value=resolver_context) - mocker.patch(f"{MODULE}.db.session.close") + resolved_agent = _agent(agent_id="a") + resolved_config = _snapshot(snapshot_id="s", agent_id="a") + resolver_sessions: list[Session] = [] + + def resolve_agent_by_id(**kwargs): + resolver_sessions.append(kwargs["session"]) + return resolved_agent, resolved_config, AgentSoulConfig() + + generator._resolve_agent_by_id = mocker.MagicMock(side_effect=resolve_agent_by_id) + session = _session() + if session.get(App, "app1") is None: + session.add(_app()) + session.commit() mocker.patch(f"{MODULE}.DifyRunContext", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.AgentAppRuntimeRequestBuilder", return_value=mocker.MagicMock()) mocker.patch(f"{MODULE}.create_agent_backend_run_client", return_value=mocker.MagicMock()) @@ -283,7 +368,7 @@ class TestGenerateWorker: if run_side_effect is not None: runner.run.side_effect = run_side_effect mocker.patch(f"{MODULE}.AgentAppRunner", return_value=runner) - return runner, resolver_session + return runner, resolver_sessions def _call( self, @@ -300,6 +385,7 @@ class TestGenerateWorker: flask_app=mocker.MagicMock(), context=mocker.MagicMock(), application_generate_entity=mocker.MagicMock( + app_config=SimpleNamespace(app_id="app1", tenant_id="tenant"), agent_id="a", agent_config_snapshot_id="s", agent_session_scope_config_version_id=session_scope_config_version_id, @@ -315,11 +401,12 @@ class TestGenerateWorker: ) def test_happy_path_runs_backend(self, generator: AgentAppGenerator, mocker: MockerFixture): - runner, resolver_session = self._wire(generator, mocker) + runner, resolver_sessions = self._wire(generator, mocker) queue_manager = mocker.MagicMock() self._call(generator, mocker, queue_manager) runner.run.assert_called_once() - assert generator._resolve_agent_by_id.call_args.kwargs["session"] is resolver_session + assert resolver_sessions == [generator._resolve_agent_by_id.call_args.kwargs["session"]] + assert resolver_sessions[0].get_bind() is not None assert runner.run.call_args.kwargs["home_snapshot_id"] == "home-1" assert "home_snapshot_ref" not in runner.run.call_args.kwargs queue_manager.publish_error.assert_not_called() @@ -384,20 +471,30 @@ class TestGenerateWorker: def test_annotation_reply_publishes_after_guard_transaction_commits(self, generator, mocker: MockerFixture): runner, _ = self._wire(generator, mocker, handled=True) - annotation_reply = mocker.MagicMock(id="annotation-1", content="annotated answer") + annotation_reply = MessageAnnotation( + app_id="app1", + question="query", + content="annotated answer", + account_id="user", + ) generator._run_input_guards.return_value = (True, "query", annotation_reply) events: list[str] = [] - guard_context = module.session_factory.get_session_maker.return_value.begin.return_value - guard_context.__exit__.side_effect = lambda *args: events.append("commit") or False queue_manager = mocker.MagicMock() + def record_commit(_session: Session) -> None: + events.append("commit") + def publish(event, *_args): if isinstance(event, QueueAnnotationReplyEvent): events.append("publish") queue_manager.publish.side_effect = publish - self._call(generator, mocker, queue_manager) + event.listen(type(_session()), "after_commit", record_commit) + try: + self._call(generator, mocker, queue_manager) + finally: + event.remove(type(_session()), "after_commit", record_commit) assert events == ["commit", "publish"] runner.run.assert_not_called() @@ -451,16 +548,12 @@ class TestResumeAfterFormSubmission: composition's user-prompt layer matches the suspended snapshot (never blank).""" def _wire(self, generator, mocker: MockerFixture): - generator._resolve_agent = mocker.MagicMock( - return_value=(mocker.MagicMock(id="agent1"), "snap1", "draft", mocker.MagicMock()) - ) - generator._init_generate_records = mocker.MagicMock( - return_value=(mocker.MagicMock(id="conv", mode="agent"), mocker.MagicMock(id="msg")) - ) + generator._resolve_agent = mocker.MagicMock(return_value=(_agent(), "snap1", "draft", AgentSoulConfig())) + generator._init_generate_records = mocker.MagicMock(return_value=(_conversation(), _message())) generator._handle_response = mocker.MagicMock(return_value=None) get_conversation = mocker.patch( f"{MODULE}.ConversationService.get_conversation", - return_value=mocker.MagicMock(id="conv", invoke_from=InvokeFrom.WEB_APP), + return_value=_conversation(), ) mocker.patch(f"{MODULE}.AgentAppConfigManager.get_app_config", return_value=mocker.MagicMock(variables=[])) mocker.patch(f"{MODULE}.load_annotation_reply_config", return_value={"enabled": False}) @@ -478,12 +571,13 @@ class TestResumeAfterFormSubmission: def test_resume_resends_paused_turn_query(self, generator, mocker: MockerFixture): entity, get_conversation = self._wire(generator, mocker) - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") - app_model.app_model_config_id = "config-1" - user = DummyAccount("user") - session = mocker.MagicMock() - session.get.return_value = mocker.MagicMock() - session.scalar.return_value = mocker.MagicMock(query="original question") + session = _session() + config = AppModelConfig(app_id="app1") + config.id = "config-1" + session.add_all([config, _conversation(), _message(query="original question")]) + session.commit() + app_model = _app(app_model_config_id=config.id) + user = _account() generator.resume_after_form_submission( app_model=app_model, @@ -504,17 +598,16 @@ class TestResumeAfterFormSubmission: session=session, ) assert generator._init_generate_records.call_args.kwargs["session"] is session - session.get.assert_called_once_with(AppModelConfig, "config-1") + assert session.get(AppModelConfig, "config-1") is config assert generator._resolve_agent.call_args.kwargs["session"] is session def test_resume_falls_back_to_placeholder_when_no_paused_message(self, generator, mocker: MockerFixture): entity, _ = self._wire(generator, mocker) - session = mocker.MagicMock() - session.scalar.return_value = None + session = _session() generator.resume_after_form_submission( - app_model=mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent"), - user=DummyAccount("user"), + app_model=_app(), + user=_account(), conversation_id="conv", form_id="form-1", invoke_from=InvokeFrom.WEB_APP, @@ -526,15 +619,17 @@ class TestResumeAfterFormSubmission: def test_resume_uses_build_draft_for_debugger_conversation(self, generator, mocker: MockerFixture): self._wire(generator, mocker) - conversation = mocker.MagicMock(id="conv", invoke_from=InvokeFrom.DEBUGGER) + conversation = _conversation(invoke_from=InvokeFrom.DEBUGGER) mocker.patch(f"{MODULE}.ConversationService.get_conversation", return_value=conversation) generator._resolve_resume_draft.return_value = ("debug_build", "draft-build-1") account_user = Account(name="Test Account", email="test@example.com") account_user.id = "user" - app_model = mocker.MagicMock(id="app1", tenant_id="tenant", mode="agent") - app_model.app_model_config_id = "config-1" - session = mocker.MagicMock() - session.scalar.return_value = mocker.MagicMock(query="original question") + session = _session() + config = AppModelConfig(app_id="app1") + config.id = "config-1" + session.add_all([config, conversation, _message(query="original question")]) + session.commit() + app_model = _app(app_model_config_id=config.id) generator.resume_after_form_submission( app_model=app_model, diff --git a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_runner.py index 1bbfb9e5f1f..1c1ee3663dd 100644 --- a/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/agent_chat/test_agent_chat_app_runner.py @@ -3,7 +3,7 @@ from datetime import datetime import pytest from pytest_mock import MockerFixture from sqlalchemy import event -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from core.agent.entities import AgentEntity from core.app.apps.agent_chat.app_runner import AgentChatAppRunner @@ -12,7 +12,7 @@ from core.moderation.base import ModerationError from graphon.model_runtime.entities.llm_entities import LLMMode from graphon.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey from models.enums import ConversationFromSource -from models.model import App, AppMode, Conversation, Message +from models.model import App, AppMode, Conversation, Message, MessageAnnotation @pytest.fixture @@ -74,9 +74,12 @@ def runner(sqlite_session: Session): return AgentChatAppRunner() -@pytest.fixture(autouse=True) -def _patch_create_session(mocker: MockerFixture, sqlite_session_factory: sessionmaker[Session]) -> None: - mocker.patch("core.app.apps.agent_chat.app_runner.create_session", side_effect=sqlite_session_factory) +def _records(session: Session) -> tuple[Conversation, Message]: + conversation = session.get(Conversation, "conv") + message = session.get(Message, "msg") + assert conversation is not None + assert message is not None + return conversation, message class TestAgentChatAppRunnerRun: @@ -88,14 +91,14 @@ class TestAgentChatAppRunnerRun: assert app is not None sqlite_session.delete(app) sqlite_session.commit() + conversation, message = _records(sqlite_session) with pytest.raises(ValueError): - runner.run(generate_entity, mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock(), sqlite_session) + runner.run(generate_entity, mocker.MagicMock(), conversation, message, sqlite_session) def test_run_moderation_error_direct_output( self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session ): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = mocker.MagicMock() generate_entity = mocker.MagicMock( @@ -111,15 +114,15 @@ class TestAgentChatAppRunnerRun: mocker.patch.object(runner, "organize_prompt_messages", return_value=([], None)) mocker.patch.object(runner, "moderation_for_inputs", side_effect=ModerationError("bad")) mocker.patch.object(runner, "direct_output") + conversation, message = _records(sqlite_session) - runner.run(generate_entity, mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock(), sqlite_session) + runner.run(generate_entity, mocker.MagicMock(), conversation, message, sqlite_session) runner.direct_output.assert_called_once() def test_run_annotation_reply_short_circuits( self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session ): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = mocker.MagicMock() generate_entity = mocker.MagicMock( @@ -136,13 +139,19 @@ class TestAgentChatAppRunnerRun: mocker.patch.object(runner, "organize_prompt_messages", return_value=([], None)) mocker.patch.object(runner, "moderation_for_inputs", return_value=(None, {}, "q")) - annotation = mocker.MagicMock(id="anno", content="answer") + annotation = MessageAnnotation( + app_id="app1", + question="q", + content="answer", + account_id="user", + ) annotation_query = mocker.patch.object(runner, "query_app_annotations_to_reply", return_value=annotation) mocker.patch.object(runner, "direct_output") queue_manager = mocker.MagicMock() write_session = sqlite_session - runner.run(generate_entity, queue_manager, mocker.MagicMock(), mocker.MagicMock(), write_session) + conversation, message = _records(sqlite_session) + runner.run(generate_entity, queue_manager, conversation, message, write_session) queue_manager.publish.assert_called_once() assert annotation_query.call_args.kwargs["session"] is write_session @@ -151,7 +160,6 @@ class TestAgentChatAppRunnerRun: def test_run_hosting_moderation_short_circuits( self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session ): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = mocker.MagicMock() generate_entity = mocker.MagicMock( @@ -170,11 +178,11 @@ class TestAgentChatAppRunnerRun: mocker.patch.object(runner, "moderation_for_inputs", return_value=(None, {}, "q")) mocker.patch.object(runner, "query_app_annotations_to_reply", return_value=None) mocker.patch.object(runner, "check_hosting_moderation", return_value=True) + conversation, message = _records(sqlite_session) - runner.run(generate_entity, mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock(), sqlite_session) + runner.run(generate_entity, mocker.MagicMock(), conversation, message, sqlite_session) def test_run_model_schema_missing(self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = AgentEntity(provider="p", model="m", strategy=AgentEntity.Strategy.CHAIN_OF_THOUGHT) @@ -203,9 +211,10 @@ class TestAgentChatAppRunnerRun: llm_instance = mocker.MagicMock() llm_instance.model_type_instance.get_model_schema.return_value = None mocker.patch("core.app.apps.agent_chat.app_runner.ModelInstance", return_value=llm_instance) + conversation, message = _records(sqlite_session) with pytest.raises(ValueError): - runner.run(generate_entity, mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock(), sqlite_session) + runner.run(generate_entity, mocker.MagicMock(), conversation, message, sqlite_session) @pytest.mark.parametrize( ("mode", "expected_runner"), @@ -222,7 +231,6 @@ class TestAgentChatAppRunnerRun: expected_runner, sqlite_session: Session, ): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = AgentEntity(provider="p", model="m", strategy=AgentEntity.Strategy.CHAIN_OF_THOUGHT) @@ -256,8 +264,7 @@ class TestAgentChatAppRunnerRun: llm_instance.model_type_instance.get_model_schema.return_value = model_schema mocker.patch("core.app.apps.agent_chat.app_runner.ModelInstance", return_value=llm_instance) - conversation = mocker.MagicMock(id="conv") - message = mocker.MagicMock(id="msg") + conversation, message = _records(sqlite_session) runner_cls = mocker.MagicMock() mocker.patch(f"core.app.apps.agent_chat.app_runner.{expected_runner}", runner_cls) @@ -268,25 +275,23 @@ class TestAgentChatAppRunnerRun: runner_instance.run.side_effect = lambda **_kwargs: events.append("agent-run") or [] mocker.patch.object(runner, "_handle_invoke_result") session = sqlite_session - event.listen(session, "after_commit", lambda _session: events.append("commit")) - original_close = session.close - def close_session() -> None: - events.append("close") - original_close() + def record_commit(_session: Session) -> None: + events.append("commit") - session.close = close_session + event.listen(session, "after_commit", record_commit) + try: + runner.run(generate_entity, mocker.MagicMock(), conversation, message, session) + finally: + event.remove(session, "after_commit", record_commit) - runner.run(generate_entity, mocker.MagicMock(), conversation, message, session) - - assert events == ["commit", "close", "commit", "close", "agent-run"] + assert events == ["commit", "commit", "agent-run"] runner_instance.run.assert_called_once() runner._handle_invoke_result.assert_called_once() def test_run_invalid_llm_mode_raises( self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session ): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = AgentEntity(provider="p", model="m", strategy=AgentEntity.Strategy.CHAIN_OF_THOUGHT) @@ -320,8 +325,7 @@ class TestAgentChatAppRunnerRun: llm_instance.model_type_instance.get_model_schema.return_value = model_schema mocker.patch("core.app.apps.agent_chat.app_runner.ModelInstance", return_value=llm_instance) - conversation = mocker.MagicMock(id="conv") - message = mocker.MagicMock(id="msg") + conversation, message = _records(sqlite_session) with pytest.raises(ValueError): runner.run(generate_entity, mocker.MagicMock(), conversation, message, sqlite_session) @@ -329,7 +333,6 @@ class TestAgentChatAppRunnerRun: def test_run_function_calling_strategy_selected_by_features( self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session ): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = AgentEntity(provider="p", model="m", strategy=AgentEntity.Strategy.CHAIN_OF_THOUGHT) @@ -363,8 +366,7 @@ class TestAgentChatAppRunnerRun: llm_instance.model_type_instance.get_model_schema.return_value = model_schema mocker.patch("core.app.apps.agent_chat.app_runner.ModelInstance", return_value=llm_instance) - conversation = mocker.MagicMock(id="conv") - message = mocker.MagicMock(id="msg") + conversation, message = _records(sqlite_session) runner_cls = mocker.MagicMock() mocker.patch("core.app.apps.agent_chat.app_runner.FunctionCallAgentRunner", runner_cls) @@ -382,7 +384,6 @@ class TestAgentChatAppRunnerRun: def test_run_conversation_not_found( self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session ): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = AgentEntity(provider="p", model="m", strategy=AgentEntity.Strategy.FUNCTION_CALLING) @@ -403,8 +404,7 @@ class TestAgentChatAppRunnerRun: user_id="user", ) - conversation_record = sqlite_session.get(Conversation, "conv") - assert conversation_record is not None + conversation_record, message_record = _records(sqlite_session) sqlite_session.delete(conversation_record) sqlite_session.commit() mocker.patch.object(runner, "organize_prompt_messages", return_value=([], None)) @@ -416,13 +416,12 @@ class TestAgentChatAppRunnerRun: runner.run( generate_entity, mocker.MagicMock(), - mocker.MagicMock(id="conv"), - mocker.MagicMock(id="msg"), + conversation_record, + message_record, sqlite_session, ) def test_run_message_not_found(self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = AgentEntity(provider="p", model="m", strategy=AgentEntity.Strategy.FUNCTION_CALLING) @@ -443,8 +442,7 @@ class TestAgentChatAppRunnerRun: user_id="user", ) - message_record = sqlite_session.get(Message, "msg") - assert message_record is not None + conversation_record, message_record = _records(sqlite_session) sqlite_session.delete(message_record) sqlite_session.commit() mocker.patch.object(runner, "organize_prompt_messages", return_value=([], None)) @@ -456,15 +454,14 @@ class TestAgentChatAppRunnerRun: runner.run( generate_entity, mocker.MagicMock(), - mocker.MagicMock(id="conv"), - mocker.MagicMock(id="msg"), + conversation_record, + message_record, sqlite_session, ) def test_run_invalid_agent_strategy_raises( self, runner: AgentChatAppRunner, mocker: MockerFixture, sqlite_session: Session ): - app_record = mocker.MagicMock(id="app1", tenant_id="tenant") app_config = mocker.MagicMock(app_id="app1", tenant_id="tenant", prompt_template=mocker.MagicMock()) app_config.agent = mocker.MagicMock(strategy="invalid", provider="p", model="m") @@ -498,8 +495,7 @@ class TestAgentChatAppRunnerRun: llm_instance.model_type_instance.get_model_schema.return_value = model_schema mocker.patch("core.app.apps.agent_chat.app_runner.ModelInstance", return_value=llm_instance) - conversation = mocker.MagicMock(id="conv") - message = mocker.MagicMock(id="msg") + conversation, message = _records(sqlite_session) with pytest.raises(ValueError): runner.run(generate_entity, mocker.MagicMock(), conversation, message, sqlite_session) diff --git a/api/tests/unit_tests/core/app/apps/chat/test_app_generator_and_runner.py b/api/tests/unit_tests/core/app/apps/chat/test_app_generator_and_runner.py index 5f31cf6b49a..d98c39d9215 100644 --- a/api/tests/unit_tests/core/app/apps/chat/test_app_generator_and_runner.py +++ b/api/tests/unit_tests/core/app/apps/chat/test_app_generator_and_runner.py @@ -1,10 +1,12 @@ -from contextlib import contextmanager +from decimal import Decimal from types import SimpleNamespace from unittest.mock import ANY, MagicMock, Mock, patch import pytest -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy import event +from sqlalchemy.orm import Session +import core.app.apps.chat.app_generator as generator_module from core.app.apps.chat.app_generator import ChatAppGenerator from core.app.apps.chat.app_runner import ChatAppRunner from core.app.apps.exc import GenerateTaskStoppedError @@ -12,7 +14,9 @@ from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import QueueAnnotationReplyEvent from core.moderation.base import ModerationError from graphon.model_runtime.errors.invoke import InvokeAuthorizationError -from models.model import App, AppMode, IconType +from models import Account +from models.enums import ConversationFromSource +from models.model import App, AppMode, AppModelConfig, Conversation, IconType, Message, MessageAnnotation class DummyGenerateEntity: @@ -31,19 +35,8 @@ class DummyQueueManager: self.published.append((event, pub_from)) -@contextmanager -def patched_create_session(session_factory: sessionmaker[Session]): - @contextmanager - def create_session(): - with session_factory() as session: - yield session - - with patch("core.app.apps.chat.app_runner.create_session", create_session): - yield - - -def _persist_app(session: Session) -> App: - app = App( +def _app() -> App: + return App( id="app-1", tenant_id="tenant-1", name="Chat app", @@ -54,9 +47,75 @@ def _persist_app(session: Session) -> App: enable_site=False, enable_api=True, ) - session.add(app) + + +def _account() -> Account: + account = Account(name="User", email="user-1@example.com") + account.id = "user-1" + return account + + +def _conversation() -> Conversation: + conversation = Conversation( + id="c1", + app_id="app-1", + app_model_config_id=None, + model_provider=None, + override_model_configs=None, + model_id=None, + mode=AppMode.CHAT, + name="Conversation", + inputs={}, + introduction="", + system_instruction="", + system_instruction_tokens=0, + status="normal", + invoke_from=InvokeFrom.SERVICE_API, + from_source=ConversationFromSource.API, + from_end_user_id=None, + from_account_id="user-1", + ) + return conversation + + +def _message() -> Message: + return Message( + id="m1", + app_id="app-1", + conversation_id="c1", + inputs={}, + query="hi", + message={}, + message_tokens=0, + message_unit_price=Decimal(0), + message_price_unit=Decimal(0), + answer="", + answer_tokens=0, + answer_unit_price=Decimal(0), + answer_price_unit=Decimal(0), + provider_response_latency=0, + total_price=Decimal(0), + currency="USD", + invoke_from=InvokeFrom.SERVICE_API, + from_source=ConversationFromSource.API, + from_end_user_id=None, + from_account_id="user-1", + app_mode=AppMode.CHAT, + ) + + +def _persist_records(session: Session) -> tuple[App, Conversation, Message]: + app = _app() + conversation = _conversation() + message = _message() + session.add_all([app, conversation, message]) session.commit() - return app + return app, conversation, message + + +@pytest.fixture(autouse=True) +def _bind_db_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(generator_module.db, "session", sqlite_session) class TestChatAppGenerator: @@ -65,8 +124,8 @@ class TestChatAppGenerator: with pytest.raises(ValueError): generator.generate( session=unbound_session, - app_model=SimpleNamespace(), - user=SimpleNamespace(), + app_model=_app(), + user=_account(), args={"inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, @@ -77,8 +136,8 @@ class TestChatAppGenerator: with pytest.raises(ValueError): generator.generate( session=unbound_session, - app_model=SimpleNamespace(), - user=SimpleNamespace(), + app_model=_app(), + user=_account(), args={"query": 1, "inputs": {}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, @@ -86,8 +145,8 @@ class TestChatAppGenerator: def test_generate_debugger_overrides_model_config(self, unbound_session: Session): generator = ChatAppGenerator() - app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") - user = SimpleNamespace(id="user-1", session_id="session-1") + app_model = _app() + user = _account() args = { "query": "hi", "inputs": {}, @@ -119,12 +178,12 @@ class TestChatAppGenerator: patch( "core.app.apps.chat.app_generator.ChatAppGenerateResponseConverter.convert", return_value={"ok": True} ), - patch.object(ChatAppGenerator, "_get_app_model_config", return_value=SimpleNamespace(to_dict=lambda: {})), + patch.object(ChatAppGenerator, "_get_app_model_config", return_value=AppModelConfig(app_id="app-1")), patch.object(ChatAppGenerator, "_prepare_user_inputs", return_value={}), patch.object( ChatAppGenerator, "_init_generate_records", - return_value=(SimpleNamespace(id="c1", mode="chat"), SimpleNamespace(id="m1")), + return_value=(_conversation(), _message()), ), patch.object(ChatAppGenerator, "_handle_response", return_value={"response": True}), patch("core.app.apps.chat.app_generator.copy_current_request_context", side_effect=lambda f: f), @@ -141,13 +200,14 @@ class TestChatAppGenerator: def test_generate_uses_session_for_annotation_reply(self, unbound_session: Session): generator = ChatAppGenerator() - app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") - app_model_config = MagicMock(id="config-1", app_id="app-1") + app_model = _app() + app_model_config = AppModelConfig(app_id="app-1") annotation_reply = {"enabled": False} - user = SimpleNamespace(id="user-1", session_id="session-1") + user = _account() with ( patch.object(ChatAppGenerator, "_get_app_model_config", return_value=app_model_config), + patch.object(AppModelConfig, "to_dict", return_value={}) as to_dict, patch( "core.app.apps.chat.app_generator.load_annotation_reply_config", return_value=annotation_reply, @@ -168,40 +228,34 @@ class TestChatAppGenerator: ) load_annotation_reply_config.assert_called_once_with(unbound_session, "app-1") - app_model_config.to_dict.assert_called_once_with(annotation_reply=annotation_reply) + to_dict.assert_called_once_with(annotation_reply=annotation_reply) assert get_app_config.call_args.kwargs["annotation_reply"] is annotation_reply def test_generate_rejects_model_config_override_for_non_debugger(self, unbound_session: Session): generator = ChatAppGenerator() with pytest.raises(ValueError): with ( - patch.object( - ChatAppGenerator, "_get_app_model_config", return_value=SimpleNamespace(to_dict=lambda: {}) - ), + patch.object(ChatAppGenerator, "_get_app_model_config", return_value=AppModelConfig(app_id="app-1")), + patch.object(AppModelConfig, "to_dict", return_value={}), ): generator.generate( session=unbound_session, - app_model=SimpleNamespace(tenant_id="t1", id="a1", mode=AppMode.CHAT.value), - user=SimpleNamespace(id="u1", session_id="s1"), + app_model=_app(), + user=_account(), args={"query": "hi", "inputs": {}, "model_config": {"foo": "bar"}}, invoke_from=InvokeFrom.SERVICE_API, streaming=False, ) - def test_generate_worker_handles_exceptions(self, unbound_session_factory: sessionmaker[Session]): + def test_generate_worker_handles_exceptions(self): generator = ChatAppGenerator() queue_manager = DummyQueueManager() entity = DummyGenerateEntity(task_id="t1", user_id="u1") with ( - patch.object(ChatAppGenerator, "_get_conversation", return_value=SimpleNamespace()), - patch.object(ChatAppGenerator, "_get_message", return_value=SimpleNamespace()), + patch.object(ChatAppGenerator, "_get_conversation", return_value=_conversation()), + patch.object(ChatAppGenerator, "_get_message", return_value=_message()), patch("core.app.apps.chat.app_generator.ChatAppRunner.run", side_effect=InvokeAuthorizationError()), - patch( - "core.app.apps.chat.app_generator.session_factory", - SimpleNamespace(create_session=unbound_session_factory), - ), - patch("core.app.apps.chat.app_generator.db.session.close"), ): generator._generate_worker( flask_app=Mock(app_context=Mock(return_value=Mock(__enter__=Mock(), __exit__=Mock()))), @@ -214,14 +268,9 @@ class TestChatAppGenerator: assert queue_manager.published with ( - patch.object(ChatAppGenerator, "_get_conversation", return_value=SimpleNamespace()), - patch.object(ChatAppGenerator, "_get_message", return_value=SimpleNamespace()), + patch.object(ChatAppGenerator, "_get_conversation", return_value=_conversation()), + patch.object(ChatAppGenerator, "_get_message", return_value=_message()), patch("core.app.apps.chat.app_generator.ChatAppRunner.run", side_effect=GenerateTaskStoppedError()), - patch( - "core.app.apps.chat.app_generator.session_factory", - SimpleNamespace(create_session=unbound_session_factory), - ), - patch("core.app.apps.chat.app_generator.db.session.close"), ): generator._generate_worker( flask_app=Mock(app_context=Mock(return_value=Mock(__enter__=Mock(), __exit__=Mock()))), @@ -233,7 +282,7 @@ class TestChatAppGenerator: class TestChatAppRunner: - def test_run_raises_when_app_missing(self, sqlite_session_factory: sessionmaker[Session], unbound_session: Session): + def test_run_raises_when_app_missing(self, sqlite_session: Session): runner = ChatAppRunner() app_config = SimpleNamespace( app_id="app-1", tenant_id="tenant-1", prompt_template=None, external_data_variables=[] @@ -251,20 +300,17 @@ class TestChatAppRunner: invoke_from=InvokeFrom.SERVICE_API, ) - with patched_create_session(sqlite_session_factory): - with pytest.raises(ValueError): - runner.run( - app_generate_entity, - DummyQueueManager(), - SimpleNamespace(), - SimpleNamespace(id="m1"), - unbound_session, - ) + with pytest.raises(ValueError): + runner.run( + app_generate_entity, + DummyQueueManager(), + _conversation(), + _message(), + sqlite_session, + ) - def test_run_moderation_error_direct_output( - self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] - ): - _persist_app(sqlite_session) + def test_run_moderation_error_direct_output(self, sqlite_session: Session): + _, conversation, message = _persist_records(sqlite_session) runner = ChatAppRunner() app_config = SimpleNamespace( app_id="app-1", @@ -289,7 +335,6 @@ class TestChatAppRunner: ) with ( - patched_create_session(sqlite_session_factory), patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])), patch.object(ChatAppRunner, "moderation_for_inputs", side_effect=ModerationError("blocked")), patch.object(ChatAppRunner, "direct_output") as mock_direct, @@ -297,17 +342,15 @@ class TestChatAppRunner: runner.run( app_generate_entity, DummyQueueManager(), - SimpleNamespace(), - SimpleNamespace(id="m1"), + conversation, + message, sqlite_session, ) mock_direct.assert_called_once() - def test_run_annotation_reply_short_circuits( - self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] - ): - _persist_app(sqlite_session) + def test_run_annotation_reply_short_circuits(self, sqlite_session: Session): + _, conversation, message = _persist_records(sqlite_session) runner = ChatAppRunner() app_config = SimpleNamespace( app_id="app-1", @@ -331,26 +374,28 @@ class TestChatAppRunner: invoke_from=InvokeFrom.SERVICE_API, ) - annotation = SimpleNamespace(id="ann-1", content="answer") + annotation = MessageAnnotation( + app_id="app-1", + question="hi", + content="answer", + account_id="user-1", + ) with ( - patched_create_session(sqlite_session_factory), patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])), patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")), patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=annotation) as annotation_query, patch.object(ChatAppRunner, "direct_output") as mock_direct, ): queue_manager = DummyQueueManager() - runner.run(app_generate_entity, queue_manager, SimpleNamespace(), SimpleNamespace(id="m1"), sqlite_session) + runner.run(app_generate_entity, queue_manager, conversation, message, sqlite_session) assert any(isinstance(item[0], QueueAnnotationReplyEvent) for item in queue_manager.published) assert annotation_query.call_args.kwargs["session"] is sqlite_session mock_direct.assert_called_once() - def test_run_returns_when_hosting_moderation_blocks( - self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] - ): - _persist_app(sqlite_session) + def test_run_returns_when_hosting_moderation_blocks(self, sqlite_session: Session): + _, conversation, message = _persist_records(sqlite_session) runner = ChatAppRunner() app_config = SimpleNamespace( app_id="app-1", @@ -374,7 +419,6 @@ class TestChatAppRunner: ) with ( - patched_create_session(sqlite_session_factory), patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])), patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")), patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=None), @@ -383,15 +427,13 @@ class TestChatAppRunner: runner.run( app_generate_entity, DummyQueueManager(), - SimpleNamespace(), - SimpleNamespace(id="m1"), + conversation, + message, sqlite_session, ) - def test_run_closes_explicit_session_before_stream_consumption( - self, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session] - ): - _persist_app(sqlite_session) + def test_run_closes_explicit_session_before_stream_consumption(self, sqlite_session: Session): + _, conversation, message = _persist_records(sqlite_session) runner = ChatAppRunner() app_config = SimpleNamespace( app_id="app-1", @@ -418,8 +460,9 @@ class TestChatAppRunner: events = [] queue_manager = DummyQueueManager() model_instance = MagicMock() - original_commit = sqlite_session.commit - original_close = sqlite_session.close + + def record_commit(_session: Session) -> None: + events.append("commit") def invoke_stream(): events.append("first-chunk") @@ -429,26 +472,27 @@ class TestChatAppRunner: events.append("invoke") return invoke_stream() - with ( - patched_create_session(sqlite_session_factory), - patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])), - patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")), - patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=None), - patch.object(ChatAppRunner, "check_hosting_moderation", return_value=False), - patch.object(ChatAppRunner, "recalc_llm_max_tokens"), - patch.object( - ChatAppRunner, - "_handle_invoke_result", - side_effect=lambda invoke_result, **kwargs: list(invoke_result), - ) as mock_handle, - patch("core.app.apps.chat.app_runner.ModelInstance", return_value=model_instance), - patch.object(sqlite_session, "commit", side_effect=lambda: (events.append("commit"), original_commit())[1]), - patch.object(sqlite_session, "close", side_effect=lambda: (events.append("close"), original_close())[1]), - ): - model_instance.invoke_llm.side_effect = invoke_llm - runner.run(app_generate_entity, queue_manager, SimpleNamespace(), SimpleNamespace(id="m1"), sqlite_session) + event.listen(sqlite_session, "after_commit", record_commit) + try: + with ( + patch.object(ChatAppRunner, "organize_prompt_messages", return_value=([], [])), + patch.object(ChatAppRunner, "moderation_for_inputs", return_value=(None, {}, "hi")), + patch.object(ChatAppRunner, "query_app_annotations_to_reply", return_value=None), + patch.object(ChatAppRunner, "check_hosting_moderation", return_value=False), + patch.object(ChatAppRunner, "recalc_llm_max_tokens"), + patch.object( + ChatAppRunner, + "_handle_invoke_result", + side_effect=lambda invoke_result, **kwargs: list(invoke_result), + ) as mock_handle, + patch("core.app.apps.chat.app_runner.ModelInstance", return_value=model_instance), + ): + model_instance.invoke_llm.side_effect = invoke_llm + runner.run(app_generate_entity, queue_manager, conversation, message, sqlite_session) + finally: + event.remove(sqlite_session, "after_commit", record_commit) - assert events == ["commit", "close", "commit", "close", "invoke", "first-chunk"] + assert events == ["commit", "commit", "invoke", "first-chunk"] mock_handle.assert_called_once_with( invoke_result=ANY, queue_manager=queue_manager, diff --git a/api/tests/unit_tests/services/test_app_generate_service.py b/api/tests/unit_tests/services/test_app_generate_service.py index 646a30d5cfb..5b20675c9c8 100644 --- a/api/tests/unit_tests/services/test_app_generate_service.py +++ b/api/tests/unit_tests/services/test_app_generate_service.py @@ -13,6 +13,7 @@ Covers: - get_response_generator (ended / non-ended workflow run) """ +import json import threading import uuid from collections.abc import Callable @@ -26,7 +27,11 @@ from sqlalchemy.orm import Session import services.app_generate_service as ags_module from core.app.entities.app_invoke_entities import InvokeFrom from enums import DeploymentEdition, QuotaType -from models.model import AppMode +from graphon.enums import WorkflowExecutionStatus +from models.account import Account +from models.enums import AppStatus, CreatorUserRole +from models.model import App, AppMode +from models.workflow import Workflow, WorkflowRun, WorkflowRunTriggeredFrom, WorkflowType from services.app_generate_service import AppGenerateService from services.errors.app import ( TriggerWorkflowServiceModeUnavailableError, @@ -67,19 +72,27 @@ class _DummyRateLimit: return generator -def _make_app(mode: AppMode | str, *, max_active_requests: int = 0, is_agent: bool = False) -> MagicMock: - app = MagicMock() - app.mode = mode - app.id = "app-id" - app.tenant_id = "tenant-id" - app.max_active_requests = max_active_requests - app.is_agent = is_agent - app.is_agent_with_session.return_value = is_agent +def _make_app(mode: AppMode | str, *, max_active_requests: int = 0) -> App: + app = App( + id="app-id", + tenant_id="tenant-id", + name="App", + description="", + mode=AppMode.CHAT if isinstance(mode, str) and mode == "invalid-mode" else mode, + status=AppStatus.NORMAL, + enable_site=False, + enable_api=False, + api_rpm=0, + api_rph=0, + max_active_requests=max_active_requests, + ) + if mode == "invalid-mode": + app.mode = mode # type: ignore[assignment] return app -def _make_user() -> MagicMock: - user = MagicMock() +def _make_user() -> Account: + user = Account(name="User", email="user@example.com") user.id = "user-id" return user @@ -95,14 +108,42 @@ def _make_workflow( workflow_id: str = "workflow-id", created_by: str = "owner-id", node_types: tuple[str, ...] = (), -) -> MagicMock: - workflow = MagicMock() - workflow.id = workflow_id - workflow.created_by = created_by - workflow.walk_nodes.return_value = [ - (f"node-{index}", {"type": node_type}) for index, node_type in enumerate(node_types) - ] - return workflow +) -> Workflow: + return Workflow( + id=workflow_id, + tenant_id="tenant-id", + app_id="app-id", + type=WorkflowType.WORKFLOW, + version=Workflow.VERSION_DRAFT, + graph=json.dumps( + { + "nodes": [ + {"id": f"node-{index}", "data": {"type": node_type}} for index, node_type in enumerate(node_types) + ], + "edges": [], + } + ), + features={}, + created_by=created_by, + environment_variables=[], + conversation_variables=[], + ) + + +def _make_workflow_run(*, run_id: str, ended: bool) -> WorkflowRun: + run = WorkflowRun( + tenant_id="tenant-id", + app_id="app-id", + workflow_id="workflow-id", + type=WorkflowType.WORKFLOW, + triggered_from=WorkflowRunTriggeredFrom.APP_RUN, + version="published", + status=WorkflowExecutionStatus.SUCCEEDED if ended else WorkflowExecutionStatus.RUNNING, + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-id", + ) + run.id = run_id + return run @contextmanager @@ -335,7 +376,8 @@ class TestGenerate(_RealSessionTest): "services.app_generate_service.AgentChatAppGenerator.convert_to_event_stream", side_effect=lambda x: x, ) - app = _make_app(AppMode.CHAT, is_agent=True) + app = _make_app(AppMode.CHAT) + is_agent = mocker.patch.object(App, "is_agent_with_session", return_value=True) session = self.session result = AppGenerateService.generate( app_model=app, @@ -347,7 +389,7 @@ class TestGenerate(_RealSessionTest): ) assert result == {"result": "agent-via-flag"} gen_spy.assert_called_once() - app.is_agent_with_session.assert_called_once_with(session=session) + is_agent.assert_called_once_with(session=session) # -- AGENT -------------------------------------------------------------- def test_agent_mode_passes_session(self, mocker: MockerFixture): @@ -383,7 +425,7 @@ class TestGenerate(_RealSessionTest): "services.app_generate_service.ChatAppGenerator.convert_to_event_stream", side_effect=lambda x: x, ) - app = _make_app(AppMode.CHAT, is_agent=False) + app = _make_app(AppMode.CHAT) result = AppGenerateService.generate( app_model=app, user=_make_user(), @@ -600,7 +642,7 @@ class TestGenerate(_RealSessionTest): # -- Invalid mode ------------------------------------------------------- def test_invalid_mode_raises(self, mocker: MockerFixture): - app = _make_app("invalid-mode", is_agent=False) + app = _make_app("invalid-mode") with pytest.raises(ValueError, match="Invalid app mode"): AppGenerateService.generate( app_model=app, @@ -1039,9 +1081,7 @@ class TestGenerateMoreLikeThis(_RealSessionTest): class TestGetResponseGenerator: def test_non_ended_workflow_run(self, mocker: MockerFixture): app = _make_app(AppMode.ADVANCED_CHAT) - workflow_run = MagicMock() - workflow_run.id = "run-1" - workflow_run.status.is_ended.return_value = False + workflow_run = _make_workflow_run(run_id="run-1", ended=False) gen_instance = MagicMock() gen_instance.retrieve_events.return_value = iter([{"event": "started"}]) @@ -1057,9 +1097,7 @@ class TestGetResponseGenerator: def test_ended_workflow_run_still_returns_generator(self, mocker: MockerFixture): """Even when the run is ended, the current code still returns a generator (TODO branch).""" app = _make_app(AppMode.WORKFLOW) - workflow_run = MagicMock() - workflow_run.id = "run-2" - workflow_run.status.is_ended.return_value = True + workflow_run = _make_workflow_run(run_id="run-2", ended=True) gen_instance = MagicMock() gen_instance.retrieve_events.return_value = iter([])