✨add_image_retrieval

Author: yzAiden

test/backend/database/test_attachment_db.py

@@ -17,6 +17,8 @@
 # Mock consts module
 consts_mock = MagicMock()
 consts_mock.const = MagicMock()
+# Ensure constants are real strings to avoid startswith TypeError
+consts_mock.const.S3_URL_PREFIX = "s3://"
 # Environment variables are now configured in conftest.py
 
 sys.modules['consts'] = consts_mock

test/backend/services/test_data_process_service.py

@@ -1667,14 +1667,18 @@ async def async_test_create_batch_tasks_impl_success(self, mock_process, mock_fo
                 'source_type': 'url',
                 'chunking_strategy': 'semantic',
                 'index_name': 'test_index_1',
-                'original_filename': 'doc1.pdf'
+                'original_filename': 'doc1.pdf',
+                'embedding_model_id': None,
+                'tenant_id': None
             },
             {
                 'source': 'http://example.com/doc2.pdf',
                 'source_type': 'url',
                 'chunking_strategy': 'fixed',
                 'index_name': 'test_index_2',
-                'original_filename': 'doc2.pdf'
+                'original_filename': 'doc2.pdf',
+                'embedding_model_id': None,
+                'tenant_id': None
             }
         ]
         actual_process_calls = [kwargs for args,

test/backend/services/test_vectordatabase_service.py

@@ -512,7 +512,9 @@ def test_create_knowledge_base_with_embedding_model_name(self, mock_get_embeddin
         self.assertEqual(result["knowledge_id"], 10)
 
         # Verify get_embedding_model was called with the model name
-        mock_get_embedding.assert_called_once_with("tenant-1", "text-embedding-3-small")
+        mock_get_embedding.assert_called_once_with(
+            "tenant-1", is_multimodal=False, model_name="text-embedding-3-small"
+        )
 
         # Verify knowledge record was created with the embedding model name
         mock_create_knowledge.assert_called_once()
@@ -559,7 +561,7 @@ def test_create_knowledge_base_without_embedding_model_name_uses_default(self, m
         self.assertEqual(result["status"], "success")
 
         # Verify get_embedding_model was called with None (no specific model)
-        mock_get_embedding.assert_called_once_with("tenant-1", None)
+        mock_get_embedding.assert_called_once_with("tenant-1", is_multimodal=False)
 
         # Verify knowledge record was created with the model's display name
         mock_create_knowledge.assert_called_once()

test/sdk/data_process/test_core.py

@@ -207,7 +207,8 @@ def test_validate_parameters_invalid_processor(self, core):
     )
     def test_select_processor_by_filename(self, core, filename, expected_processor, expected_extractor):
         """Test processor selection based on filename"""
-        processor_name, extractor = core._select_processor_by_filename(filename)
+        params = {"model_type": "multi_embedding"} if expected_extractor else {}
+        processor_name, extractor = core._select_processor_by_filename(filename, params)
         assert processor_name == expected_processor
         assert extractor == expected_extractor
 
@@ -349,7 +350,7 @@ def test_file_process_returns_images_when_extractor_available(self, core, mocker
         core.processors["UniversalImageExtractor"] = mock_extractor
 
         result = core.file_process(
-            b"data", "sample.pdf", chunking_strategy="basic"
+            b"data", "sample.pdf", chunking_strategy="basic", model_type="multi_embedding"
         )
 
         chunks = _unpack_chunks(result)
@@ -366,7 +367,11 @@ def test_file_process_with_explicit_processor_still_extracts_images(self, core):
         )
 
         result = core.file_process(
-            b"data", "report.pdf", chunking_strategy="basic", processor="Unstructured"
+            b"data",
+            "report.pdf",
+            chunking_strategy="basic",
+            processor="Unstructured",
+            model_type="multi_embedding",
         )
 
         chunks = _unpack_chunks(result)