From c647cfd6a38f4091dfd4012bda0cbcc98f6947e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E7=8E=AE=20=28Jade=20Lin=29?= Date: Thu, 13 Aug 2026 06:22:22 +0000 Subject: [PATCH] refactor: resolve summary models through manager (#40696) --- .../processor/paragraph_index_processor.py | 23 ++--- .../test_paragraph_index_processor.py | 87 ++++++++++--------- 2 files changed, 52 insertions(+), 58 deletions(-) diff --git a/api/core/rag/index_processor/processor/paragraph_index_processor.py b/api/core/rag/index_processor/processor/paragraph_index_processor.py index 6e4a1fef757..518d70eb284 100644 --- a/api/core/rag/index_processor/processor/paragraph_index_processor.py +++ b/api/core/rag/index_processor/processor/paragraph_index_processor.py @@ -9,12 +9,10 @@ from sqlalchemy import select from sqlalchemy.orm import Session from core.app.file_access import DatabaseFileAccessController -from core.app.llm import deduct_llm_quota from core.db.session_factory import session_factory from core.entities.knowledge_entities import PreviewDetail from core.llm_generator.prompts import DEFAULT_GENERATOR_SUMMARY_PROMPT -from core.model_manager import ModelInstance -from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager +from core.model_manager import ModelManager from core.rag.cleaner.clean_processor import CleanProcessor from core.rag.datasource.keyword.keyword_factory import Keyword from core.rag.datasource.vdb.vector_factory import Vector @@ -388,7 +386,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor): session: Session, ) -> tuple[str, LLMUsage]: """ - Generate summary for the given text using ModelInstance.invoke_llm and the default or custom summary prompt, + Generate summary for the given text using the configured LLM and the default or custom summary prompt, and supports vision models by including images from the segment attachments or text content. Args: @@ -431,11 +429,13 @@ class ParagraphIndexProcessor(BaseIndexProcessor): # If default prompt doesn't have {language} placeholder, use it as-is pass - provider_manager = create_plugin_provider_manager(tenant_id=tenant_id) - provider_model_bundle = provider_manager.get_provider_model_bundle( - tenant_id, model_provider_name, ModelType.LLM + model_manager = ModelManager.for_tenant(tenant_id=tenant_id) + model_instance = model_manager.get_model_instance( + tenant_id=tenant_id, + provider=model_provider_name, + model_type=ModelType.LLM, + model=model_name, ) - model_instance = ModelInstance(provider_model_bundle, model_name) # Get model schema to check if vision is supported model_schema = model_instance.model_type_instance.get_model_schema(model_name, model_instance.credentials) @@ -496,13 +496,6 @@ class ParagraphIndexProcessor(BaseIndexProcessor): summary_content = result.message.get_text_content() usage = result.usage - # Deduct quota for summary generation (same as workflow nodes) - try: - deduct_llm_quota(tenant_id=tenant_id, model_instance=model_instance, usage=usage) - except Exception as e: - # Log but don't fail summary generation if quota deduction fails - logger.warning("Failed to deduct quota for summary generation: %s", str(e)) - return summary_content, usage @staticmethod diff --git a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py index 228aaf1a531..40c14147f53 100644 --- a/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py +++ b/api/tests/unit_tests/core/rag/indexing/processor/test_paragraph_index_processor.py @@ -16,7 +16,7 @@ from core.rag.models.document import AttachmentDocument, Document from extensions.storage.storage_type import StorageType from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage, ImagePromptMessageContent -from graphon.model_runtime.entities.model_entities import ModelFeature +from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType from models.dataset import DocumentSegment, SegmentAttachmentBinding from models.enums import CreatorUserRole from models.model import UploadFile @@ -519,41 +519,51 @@ class TestParagraphIndexProcessor: with pytest.raises(ValueError, match="model_name and model_provider_name"): ParagraphIndexProcessor.generate_summary("tenant-1", "text", {"enable": True}, session=self.session) - def test_generate_summary_text_only_flow(self, caplog: pytest.LogCaptureFixture) -> None: + def test_generate_summary_text_only_flow(self) -> None: model_instance = Mock() model_instance.credentials = {"k": "v"} model_instance.model_type_instance.get_model_schema.return_value = SimpleNamespace(features=[]) model_instance.invoke_llm.return_value = self._llm_result("text summary") - with ( - patch( - "core.rag.index_processor.processor.paragraph_index_processor.create_plugin_provider_manager" - ) as mock_provider_manager, - patch( - "core.rag.index_processor.processor.paragraph_index_processor.ModelInstance", - return_value=model_instance, - ), - patch( - "core.rag.index_processor.processor.paragraph_index_processor.deduct_llm_quota", - side_effect=RuntimeError("quota"), - ), - ): - mock_provider_manager.return_value.get_provider_model_bundle.return_value = Mock() - with caplog.at_level( - logging.WARNING, logger="core.rag.index_processor.processor.paragraph_index_processor" - ): - summary, usage = ParagraphIndexProcessor.generate_summary( - "tenant-1", - "text content", - {"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"}, - document_language="English", - session=self.session, - ) + with patch( + "core.rag.index_processor.processor.paragraph_index_processor.ModelManager.for_tenant" + ) as mock_model_manager: + mock_model_manager.return_value.get_model_instance.return_value = model_instance + summary, usage = ParagraphIndexProcessor.generate_summary( + "tenant-1", + "text content", + {"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"}, + document_language="English", + session=self.session, + ) assert summary == "text summary" assert isinstance(usage, LLMUsage) - assert sum(1 for r in caplog.records if r.levelno == logging.WARNING) == 1 - assert any("Failed to deduct quota for summary generation" in record.message for record in caplog.records) + mock_model_manager.assert_called_once_with(tenant_id="tenant-1") + mock_model_manager.return_value.get_model_instance.assert_called_once_with( + tenant_id="tenant-1", + provider="provider-a", + model_type=ModelType.LLM, + model="model-a", + ) + + def test_generate_summary_propagates_model_invocation_errors(self) -> None: + model_instance = Mock() + model_instance.credentials = {"k": "v"} + model_instance.model_type_instance.get_model_schema.return_value = SimpleNamespace(features=[]) + model_instance.invoke_llm.side_effect = RuntimeError("invocation failed") + + with patch( + "core.rag.index_processor.processor.paragraph_index_processor.ModelManager.for_tenant" + ) as mock_model_manager: + mock_model_manager.return_value.get_model_instance.return_value = model_instance + with pytest.raises(RuntimeError, match="invocation failed"): + ParagraphIndexProcessor.generate_summary( + "tenant-1", + "text content", + {"enable": True, "model_name": "model-a", "model_provider_name": "provider-a"}, + session=self.session, + ) def test_generate_summary_handles_vision_and_image_conversion(self) -> None: model_instance = Mock() @@ -567,12 +577,8 @@ class TestParagraphIndexProcessor: with ( patch( - "core.rag.index_processor.processor.paragraph_index_processor.create_plugin_provider_manager" - ) as mock_provider_manager, - patch( - "core.rag.index_processor.processor.paragraph_index_processor.ModelInstance", - return_value=model_instance, - ), + "core.rag.index_processor.processor.paragraph_index_processor.ModelManager.for_tenant" + ) as mock_model_manager, patch.object( ParagraphIndexProcessor, "_extract_images_from_segment_attachments", return_value=[image_file] ), @@ -581,9 +587,8 @@ class TestParagraphIndexProcessor: "core.rag.index_processor.processor.paragraph_index_processor.file_manager.to_prompt_message_content", return_value=image_content, ), - patch("core.rag.index_processor.processor.paragraph_index_processor.deduct_llm_quota"), ): - mock_provider_manager.return_value.get_provider_model_bundle.return_value = Mock() + mock_model_manager.return_value.get_model_instance.return_value = model_instance summary, _ = ParagraphIndexProcessor.generate_summary( "tenant-1", "text content", @@ -606,12 +611,8 @@ class TestParagraphIndexProcessor: with ( patch( - "core.rag.index_processor.processor.paragraph_index_processor.create_plugin_provider_manager" - ) as mock_provider_manager, - patch( - "core.rag.index_processor.processor.paragraph_index_processor.ModelInstance", - return_value=model_instance, - ), + "core.rag.index_processor.processor.paragraph_index_processor.ModelManager.for_tenant" + ) as mock_model_manager, patch( "core.rag.index_processor.processor.paragraph_index_processor.DEFAULT_GENERATOR_SUMMARY_PROMPT", "Prompt {missing}", @@ -623,7 +624,7 @@ class TestParagraphIndexProcessor: side_effect=RuntimeError("bad image"), ), ): - mock_provider_manager.return_value.get_provider_model_bundle.return_value = Mock() + mock_model_manager.return_value.get_model_instance.return_value = model_instance with pytest.raises(ValueError, match="Expected LLMResult"): with caplog.at_level( logging.WARNING, logger="core.rag.index_processor.processor.paragraph_index_processor"