refactor: resolve completion model via manager (#40756)

api/core/app/apps/completion/app_runner.py

@@ -12,11 +12,12 @@ from core.app.entities.app_invoke_entities import (
 )
 from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
 from core.db.session_factory import create_session
-from core.model_manager import ModelInstance
+from core.model_manager import ModelManager
 from core.moderation.base import ModerationError
 from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
 from graphon.file import File
 from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent
+from graphon.model_runtime.entities.model_entities import ModelType
 from models.model import App, Message
 
 logger = logging.getLogger(__name__)
@@ -184,8 +185,10 @@ class CompletionAppRunner(AppRunner):
         self.recalc_llm_max_tokens(model_config=application_generate_entity.model_conf, prompt_messages=prompt_messages)
 
         # Invoke model
-        model_instance = ModelInstance(
-            provider_model_bundle=application_generate_entity.model_conf.provider_model_bundle,
+        model_instance = ModelManager.for_tenant(tenant_id=app_config.tenant_id).get_model_instance(
+            tenant_id=app_config.tenant_id,
+            provider=application_generate_entity.model_conf.provider,
+            model_type=ModelType.LLM,
             model=application_generate_entity.model_conf.model,
         )
 

api/tests/unit_tests/core/app/apps/completion/test_app_runner.py

@@ -9,6 +9,7 @@ import core.app.apps.completion.app_runner as module
 from core.app.apps.completion.app_runner import CompletionAppRunner
 from core.moderation.base import ModerationError
 from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent
+from graphon.model_runtime.entities.model_entities import ModelType
 from models.model import App, AppMode, IconType
 
 APP_ID = "00000000-0000-0000-0000-000000000001"
@@ -34,6 +35,7 @@ def _build_app_config(dataset=None, external_tools=None, additional_features=Non
 
 def _build_generate_entity(app_config, file_upload_config=None):
     model_conf = MagicMock(
+        provider="provider",
         provider_model_bundle="bundle",
         model="model",
         parameters={"max_tokens": 10},
@@ -138,7 +140,9 @@ class TestCompletionAppRunner:
 
         model_instance = MagicMock()
         model_instance.invoke_llm.return_value = "invoke_result"
-        mocker.patch.object(module, "ModelInstance", return_value=model_instance)
+        model_manager = MagicMock()
+        model_manager.get_model_instance.return_value = model_instance
+        mocker.patch.object(module.ModelManager, "for_tenant", return_value=model_manager)
 
         runner.run(app_generate_entity, MagicMock(), MagicMock(id="msg", tenant_id=TENANT_ID), sqlite_session)
 
@@ -186,7 +190,9 @@ class TestCompletionAppRunner:
             return invoke_stream()
 
         model_instance.invoke_llm.side_effect = invoke_llm
-        mocker.patch.object(module, "ModelInstance", return_value=model_instance)
+        model_manager = MagicMock()
+        model_manager.get_model_instance.return_value = model_instance
+        mocker.patch.object(module.ModelManager, "for_tenant", return_value=model_manager)
 
         runner.run(app_generate_entity, queue_manager, MagicMock(id="msg"), session)
 
@@ -200,6 +206,48 @@ class TestCompletionAppRunner:
             tenant_id=TENANT_ID,
         )
 
+    @pytest.mark.parametrize("stream", [False, True])
+    def test_run_invokes_model_resolved_by_model_manager(
+        self,
+        runner,
+        mocker: MockerFixture,
+        sqlite_session: Session,
+        stream: bool,
+    ):
+        _persist_app(sqlite_session)
+        app_config = _build_app_config()
+        app_generate_entity = _build_generate_entity(app_config)
+        app_generate_entity.stream = stream
+
+        runner.organize_prompt_messages = MagicMock(return_value=(["prompt"], ["stop"]))
+        runner.moderation_for_inputs = MagicMock(return_value=(None, app_generate_entity.inputs, "query"))
+        runner.check_hosting_moderation = MagicMock(return_value=False)
+        runner.recalc_llm_max_tokens = MagicMock()
+        runner._handle_invoke_result = MagicMock()
+
+        model_instance = MagicMock()
+        model_instance.invoke_llm.return_value = "invoke_result"
+        model_manager = MagicMock()
+        model_manager.get_model_instance.return_value = model_instance
+        model_manager_factory = mocker.patch.object(module.ModelManager, "for_tenant", return_value=model_manager)
+
+        runner.run(app_generate_entity, MagicMock(), MagicMock(id="msg"), sqlite_session)
+
+        model_manager_factory.assert_called_once_with(tenant_id=TENANT_ID)
+        model_manager.get_model_instance.assert_called_once_with(
+            tenant_id=TENANT_ID,
+            provider="provider",
+            model_type=ModelType.LLM,
+            model="model",
+        )
+        model_instance.invoke_llm.assert_called_once_with(
+            prompt_messages=["prompt"],
+            model_parameters={"max_tokens": 10},
+            stop=["stop"],
+            stream=stream,
+            request_metadata={"app_id": APP_ID},
+        )
+
     def test_run_uses_low_image_detail_default(self, runner, mocker: MockerFixture, sqlite_session: Session):
         _persist_app(sqlite_session)