feat: lazy load imports for language model component

src/backend/tests/unit/components/llm_operations/test_lambda_filter.py

@@ -563,13 +563,13 @@ def test_should_format_json_when_result_is_dict(self, component_class):
 class TestProcessAsDataIntegration(TestLambdaFilterComponent):
     """Integration tests for process_as_data method."""
 
-    @patch("lfx.base.models.unified_models.get_model_classes")
+    @patch("lfx.base.models.unified_models.get_model_class")
     async def test_should_return_filtered_data_when_lambda_is_valid(
-        self, mock_get_model_classes, component_class, default_kwargs, mock_llm
+        self, mock_get_model_class, component_class, default_kwargs, mock_llm
     ):
         # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
-        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        mock_get_model_class.return_value = mock_model_class
         component = await self.component_setup(component_class, default_kwargs)
         mock_llm.ainvoke.return_value.content = "lambda x: [item for item in x['items'] if item['value'] > 15]"
 
@@ -584,13 +584,13 @@ async def test_should_return_filtered_data_when_lambda_is_valid(
         assert filtered_items[0]["name"] == "test2"
         assert filtered_items[0]["value"] == 20
 
-    @patch("lfx.base.models.unified_models.get_model_classes")
+    @patch("lfx.base.models.unified_models.get_model_class")
     async def test_should_raise_error_when_lambda_not_found_in_response(
-        self, mock_get_model_classes, component_class, default_kwargs, mock_llm
+        self, mock_get_model_class, component_class, default_kwargs, mock_llm
     ):
         # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
-        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        mock_get_model_class.return_value = mock_model_class
         component = await self.component_setup(component_class, default_kwargs)
         mock_llm.ainvoke.return_value.content = "invalid response without lambda"
 
@@ -614,13 +614,13 @@ def message_kwargs(self, model_metadata):
             "max_size": 30000,
         }
 
-    @patch("lfx.base.models.unified_models.get_model_classes")
+    @patch("lfx.base.models.unified_models.get_model_class")
     async def test_should_transform_message_when_input_is_message(
-        self, mock_get_model_classes, component_class, message_kwargs, mock_llm
+        self, mock_get_model_class, component_class, message_kwargs, mock_llm
     ):
         # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
-        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        mock_get_model_class.return_value = mock_model_class
         component = await self.component_setup(component_class, message_kwargs)
         mock_llm.ainvoke.return_value.content = "lambda text: text.upper()"
 
@@ -631,13 +631,13 @@ async def test_should_transform_message_when_input_is_message(
         assert isinstance(result, Message)
         assert result.text == "HELLO WORLD"
 
-    @patch("lfx.base.models.unified_models.get_model_classes")
+    @patch("lfx.base.models.unified_models.get_model_class")
     async def test_should_join_multiple_messages_when_input_is_list_of_messages(
-        self, mock_get_model_classes, component_class, model_metadata, mock_llm
+        self, mock_get_model_class, component_class, model_metadata, mock_llm
     ):
         # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
-        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        mock_get_model_class.return_value = mock_model_class
         kwargs = {
             "data": [Message(text="Hello"), Message(text="World")],
             "model": model_metadata,
@@ -660,13 +660,13 @@ async def test_should_join_multiple_messages_when_input_is_list_of_messages(
 class TestProcessAsDataframeIntegration(TestLambdaFilterComponent):
     """Integration tests for process_as_dataframe method."""
 
-    @patch("lfx.base.models.unified_models.get_model_classes")
+    @patch("lfx.base.models.unified_models.get_model_class")
     async def test_should_return_dataframe_when_lambda_returns_list_of_dicts(
-        self, mock_get_model_classes, component_class, default_kwargs, mock_llm
+        self, mock_get_model_class, component_class, default_kwargs, mock_llm
     ):
         # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
-        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        mock_get_model_class.return_value = mock_model_class
         component = await self.component_setup(component_class, default_kwargs)
         mock_llm.ainvoke.return_value.content = "lambda x: x['items']"
 
@@ -680,13 +680,13 @@ async def test_should_return_dataframe_when_lambda_returns_list_of_dicts(
 class TestLargeDataset(TestLambdaFilterComponent):
     """Tests for handling large datasets."""
 
-    @patch("lfx.base.models.unified_models.get_model_classes")
+    @patch("lfx.base.models.unified_models.get_model_class")
     async def test_should_filter_large_dataset_when_data_exceeds_max_size(
-        self, mock_get_model_classes, component_class, default_kwargs, mock_llm
+        self, mock_get_model_class, component_class, default_kwargs, mock_llm
     ):
         # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
-        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        mock_get_model_class.return_value = mock_model_class
         large_data = {"items": [{"name": f"test{i}", "value": i} for i in range(2000)]}
         default_kwargs["data"] = [Data(data=large_data)]
         default_kwargs["filter_instruction"] = "Filter items with value greater than 1500"
@@ -707,13 +707,13 @@ async def test_should_filter_large_dataset_when_data_exceeds_max_size(
 class TestComplexDataStructure(TestLambdaFilterComponent):
     """Tests for handling complex nested data structures."""
 
-    @patch("lfx.base.models.unified_models.get_model_classes")
+    @patch("lfx.base.models.unified_models.get_model_class")
     async def test_should_handle_nested_data_when_structure_is_complex(
-        self, mock_get_model_classes, component_class, default_kwargs, mock_llm
+        self, mock_get_model_class, component_class, default_kwargs, mock_llm
     ):
         # Arrange
         mock_model_class = MagicMock(return_value=mock_llm)
-        mock_get_model_classes.return_value = {"MockLanguageModel": mock_model_class}
+        mock_get_model_class.return_value = mock_model_class
         complex_data = {
             "categories": {
                 "A": [{"id": 1, "score": 90}, {"id": 2, "score": 85}],