update tests for ollama

Author: HzaRashid

src/backend/tests/integration/components/languagemodels/test_chatollama_integration.py

@@ -28,7 +28,11 @@ async def test_end_to_end_structured_output_to_data(self, mock_chat_ollama):
 
         # Create component with schema format
         component = ChatOllamaComponent(
-            base_url="http://localhost:11434", model_name="llama3.1", format=json_schema, temperature=0.1
+            base_url="http://localhost:11434",
+            model_name="llama3.1",
+            format=json_schema,
+            temperature=0.1,
+            enable_structured_output=True
         )
 
         # Set up input message
@@ -67,7 +71,11 @@ async def test_end_to_end_structured_output_to_dataframe(self, mock_chat_ollama)
 
         # Create component with JSON format
         component = ChatOllamaComponent(
-            base_url="http://localhost:11434", model_name="llama3.1", format="json", temperature=0.1
+            base_url="http://localhost:11434",
+            model_name="llama3.1",
+            format="json",
+            temperature=0.1,
+            enable_structured_output=True
         )
 
         # Set up input message
@@ -117,7 +125,11 @@ class PersonInfo(BaseModel):
 
         # Create component with Pydantic schema
         component = ChatOllamaComponent(
-            base_url="http://localhost:11434", model_name="llama3.1", format=pydantic_schema, temperature=0.1
+            base_url="http://localhost:11434",
+            model_name="llama3.1",
+            format=pydantic_schema,
+            temperature=0.1,
+            enable_structured_output=True
         )
 
         component.input_value = "Extract person info"

src/backend/tests/unit/components/languagemodels/test_chatollama_component.py

@@ -41,6 +41,7 @@ def default_kwargs(self):
             "system": "",
             "tool_model_enabled": True,
             "template": "",
+            "enable_structured_output": True
         }
 
     @pytest.fixture
@@ -1037,3 +1038,119 @@ def test_build_model_cloud_with_v1_suffix_stripped(self, mock_chat_ollama):
         assert call_args["base_url"] == DEFAULT_OLLAMA_API_URL
         assert "/v1" not in call_args["base_url"]
         assert model == mock_model
+
+    @patch("lfx.components.ollama.ollama.ChatOllama")
+    def test_build_model_with_structured_output_disabled(self, mock_chat_ollama, component_class, default_kwargs):
+        """Test that format field is NOT passed when enable_structured_output is False (default)."""
+        mock_instance = MagicMock()
+        mock_chat_ollama.return_value = mock_instance
+
+        # Remove enable_structured_output to use default (False)
+        kwargs = default_kwargs.copy()
+        kwargs.pop("enable_structured_output", None)
+        kwargs["format"] = "json"  # Set format but it should be ignored
+
+        component = component_class(**kwargs)
+        model = component.build_model()
+
+        # Verify ChatOllama was called WITHOUT format parameter
+        call_args = mock_chat_ollama.call_args[1]
+        assert "format" not in call_args, "format should not be passed when enable_structured_output is False"
+        assert model == mock_instance
+
+
+    @patch("lfx.components.ollama.ollama.ChatOllama")
+    def test_build_model_with_structured_output_enabled_string_format(self, mock_chat_ollama, component_class, default_kwargs):
+        """Test that format field IS passed when enable_structured_output is True with string format."""
+        mock_instance = MagicMock()
+        mock_chat_ollama.return_value = mock_instance
+
+        kwargs = default_kwargs.copy()
+        kwargs["enable_structured_output"] = True
+        kwargs["format"] = "json"
+
+        component = component_class(**kwargs)
+        model = component.build_model()
+
+        # Verify ChatOllama was called WITH format parameter
+        call_args = mock_chat_ollama.call_args[1]
+        assert "format" in call_args, "format should be passed when enable_structured_output is True"
+        assert call_args["format"] == "json"
+        assert model == mock_instance
+
+
+    @patch("lfx.components.ollama.ollama.ChatOllama")
+    def test_build_model_with_structured_output_enabled_dict_format(self, mock_chat_ollama, component_class, default_kwargs):
+        """Test that JSON schema format IS passed when enable_structured_output is True."""
+        mock_instance = MagicMock()
+        mock_chat_ollama.return_value = mock_instance
+
+        json_schema = {
+            "type": "object",
+            "properties": {"name": {"type": "string"}},
+            "required": ["name"],
+        }
+
+        kwargs = default_kwargs.copy()
+        kwargs["enable_structured_output"] = True
+        kwargs["format"] = json_schema
+
+        component = component_class(**kwargs)
+        model = component.build_model()
+
+        # Verify ChatOllama was called WITH format parameter as dict
+        call_args = mock_chat_ollama.call_args[1]
+        assert "format" in call_args
+        assert call_args["format"] == json_schema
+        assert call_args["format"]["type"] == "object"
+        assert model == mock_instance
+
+
+    @patch("lfx.components.ollama.ollama.ChatOllama")
+    def test_build_model_with_structured_output_enabled_no_format(self, mock_chat_ollama, component_class, default_kwargs):
+        """Test that format is not passed when enable_structured_output is True but format is None/empty."""
+        mock_instance = MagicMock()
+        mock_chat_ollama.return_value = mock_instance
+
+        kwargs = default_kwargs.copy()
+        kwargs["enable_structured_output"] = True
+        kwargs["format"] = None  # No format specified
+
+        component = component_class(**kwargs)
+        model = component.build_model()
+
+        # Verify ChatOllama was called WITHOUT format parameter
+        call_args = mock_chat_ollama.call_args[1]
+        assert "format" not in call_args, "format should not be passed when it's None"
+        assert model == mock_instance
+
+
+    @patch("lfx.components.ollama.ollama.ChatOllama")
+    def test_build_model_structured_output_toggle_behavior(self, mock_chat_ollama, component_class, default_kwargs):
+        """Test toggling enable_structured_output affects format parameter passing."""
+        mock_instance = MagicMock()
+        mock_chat_ollama.return_value = mock_instance
+
+        # First: Test with structured output disabled
+        kwargs = default_kwargs.copy()
+        kwargs["enable_structured_output"] = False
+        kwargs["format"] = "json"
+
+        component = component_class(**kwargs)
+        model = component.build_model()
+
+        call_args = mock_chat_ollama.call_args[1]
+        assert "format" not in call_args, "format should not be in call when disabled"
+
+        # Reset mock
+        mock_chat_ollama.reset_mock()
+
+        # Second: Test with structured output enabled
+        kwargs["enable_structured_output"] = True
+        component = component_class(**kwargs)
+        model = component.build_model()
+
+        call_args = mock_chat_ollama.call_args[1]
+        assert "format" in call_args, "format should be in call when enabled"
+        assert call_args["format"] == "json"
+        assert model == mock_instance

src/lfx/src/lfx/_assets/component_index.json

@@ -101,6 +101,7 @@ class ChatOllamaComponent(LCModelComponent):
             info="Refer to https://ollama.com/library for more models.",
             refresh_button=True,
             real_time_refresh=True,
+            required=True,
         ),
         SecretStrInput(
             name="api_key",
@@ -122,9 +123,9 @@ class ChatOllamaComponent(LCModelComponent):
             name="format",
             display_name="Format",
             info="Specify the format of the output.",
-            advanced=False,
             table_schema=TABLE_SCHEMA,
             value=default_table_row,
+            show=False,
         ),
         DictInput(name="metadata", display_name="Metadata", info="Metadata to add to the run trace.", advanced=True),
         DropdownInput(
@@ -215,6 +216,14 @@ class ChatOllamaComponent(LCModelComponent):
         MessageTextInput(
             name="template", display_name="Template", info="Template to use for generating text.", advanced=True
         ),
+        BoolInput(
+            name="enable_structured_output",
+            display_name="Enable Structured Output",
+            info="Whether to enable structured output in the model.",
+            value=False,
+            advanced=False,
+            real_time_refresh=True,
+        ),
         *LCModelComponent.get_base_inputs(),
     ]
 
@@ -254,7 +263,7 @@ def build_model(self) -> LanguageModel:  # type: ignore[type-var]
             )
 
         try:
-            output_format = self._parse_format_field(self.format)
+            output_format = self._parse_format_field(self.format) if self.enable_structured_output else None
         except Exception as e:
             msg = f"Failed to parse the format field: {e}"
             raise ValueError(msg) from e
@@ -264,7 +273,7 @@ def build_model(self) -> LanguageModel:  # type: ignore[type-var]
             "base_url": transformed_base_url,
             "model": self.model_name,
             "mirostat": mirostat_value,
-            "format": output_format,
+            "format": output_format or None,
             "metadata": self.metadata,
             "tags": self.tags.split(",") if self.tags else None,
             "mirostat_eta": mirostat_eta,
@@ -319,6 +328,9 @@ async def is_valid_ollama_url(self, url: str) -> bool:
             return False
 
     async def update_build_config(self, build_config: dict, field_value: Any, field_name: str | None = None):
+        if field_name == "enable_structured_output":  # bind enable_structured_output boolean to format show value
+            build_config["format"]["show"] = field_value
+
         if field_name == "mirostat":
             if field_value == "Disabled":
                 build_config["mirostat_eta"]["advanced"] = True