from unittest.mock import MagicMock, patch import pytest from langchain_core.language_models import BaseLanguageModel from langflow.components.helpers.structured_output import StructuredOutputComponent from langflow.schema.data import Data from pydantic import BaseModel from typing_extensions import override class TestStructuredOutputComponent: # Ensure that the structured output is successfully generated with the correct BaseModel instance returned by # the mock function def test_successful_structured_output_generation_with_patch_with_config(self): from unittest.mock import patch class MockLanguageModel(BaseLanguageModel): @override def with_structured_output(self, *args, **kwargs): return self @override def with_config(self, *args, **kwargs): return self @override def invoke(self, *args, **kwargs): return self @override def generate_prompt(self, *args, **kwargs): raise NotImplementedError @override async def agenerate_prompt(self, *args, **kwargs): raise NotImplementedError @override def predict(self, *args, **kwargs): raise NotImplementedError @override def predict_messages(self, *args, **kwargs): raise NotImplementedError @override async def apredict(self, *args, **kwargs): raise NotImplementedError @override async def apredict_messages(self, *args, **kwargs): raise NotImplementedError def mock_get_chat_result(runnable, input_value, config): # noqa: ARG001 class MockBaseModel(BaseModel): @override def model_dump(self, **kwargs): return {"field": "value"} return MockBaseModel() component = StructuredOutputComponent( llm=MockLanguageModel(), input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, ) with patch("langflow.components.helpers.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output() assert isinstance(result, Data) assert result.data == {"field": "value"} # Raises ValueError when the language model does not support structured output def test_raises_value_error_for_unsupported_language_model(self): # Mocking an incompatible language model class MockLanguageModel: pass # Creating an instance of StructuredOutputComponent component = StructuredOutputComponent( llm=MockLanguageModel(), input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, ) with pytest.raises(TypeError, match="Language model does not support structured output."): component.build_structured_output() # Correctly builds the output model from the provided schema def test_correctly_builds_output_model(self): # Import internal organization modules, packages, and libraries from langflow.helpers.base_model import build_model_from_schema from langflow.inputs.inputs import TableInput # Setup component = StructuredOutputComponent() schema = [ { "name": "name", "display_name": "Name", "type": "str", "description": "Specify the name of the output field.", }, { "name": "description", "display_name": "Description", "type": "str", "description": "Describe the purpose of the output field.", }, { "name": "type", "display_name": "Type", "type": "str", "description": ( "Indicate the data type of the output field " "(e.g., str, int, float, bool, list, dict)." ), }, { "name": "multiple", "display_name": "Multiple", "type": "boolean", "description": "Set to True if this output field should be a list of the specified type.", }, ] component.output_schema = TableInput(name="output_schema", display_name="Output Schema", table_schema=schema) # Assertion output_model = build_model_from_schema(schema) assert isinstance(output_model, type) # Properly handles multiple outputs when 'multiple' is set to True def test_handles_multiple_outputs(self): # Import internal organization modules, packages, and libraries from langflow.helpers.base_model import build_model_from_schema from langflow.inputs.inputs import TableInput # Setup component = StructuredOutputComponent() schema = [ { "name": "name", "display_name": "Name", "type": "str", "description": "Specify the name of the output field.", }, { "name": "description", "display_name": "Description", "type": "str", "description": "Describe the purpose of the output field.", }, { "name": "type", "display_name": "Type", "type": "str", "description": ( "Indicate the data type of the output field " "(e.g., str, int, float, bool, list, dict)." ), }, { "name": "multiple", "display_name": "Multiple", "type": "boolean", "description": "Set to True if this output field should be a list of the specified type.", }, ] component.output_schema = TableInput(name="output_schema", display_name="Output Schema", table_schema=schema) component.multiple = True # Assertion output_model = build_model_from_schema(schema) assert isinstance(output_model, type) def test_empty_output_schema(self): component = StructuredOutputComponent( llm=MagicMock(), input_value="Test input", schema_name="EmptySchema", output_schema=[], multiple=False, ) with pytest.raises(ValueError, match="Output schema cannot be empty"): component.build_structured_output() def test_invalid_output_schema_type(self): component = StructuredOutputComponent( llm=MagicMock(), input_value="Test input", schema_name="InvalidSchema", output_schema=[{"name": "field", "type": "invalid_type", "description": "Invalid field"}], multiple=False, ) with pytest.raises(ValueError, match="Invalid type: invalid_type"): component.build_structured_output() @patch("langflow.components.helpers.structured_output.get_chat_result") def test_nested_output_schema(self, mock_get_chat_result): class ChildModel(BaseModel): child: str = "value" class ParentModel(BaseModel): parent: ChildModel = ChildModel() mock_llm = MagicMock() mock_llm.with_structured_output.return_value = mock_llm mock_get_chat_result.return_value = ParentModel(parent=ChildModel(child="value")) component = StructuredOutputComponent( llm=mock_llm, input_value="Test input", schema_name="NestedSchema", output_schema=[ { "name": "parent", "type": "dict", "description": "Parent field", "fields": [{"name": "child", "type": "str", "description": "Child field"}], } ], multiple=False, ) result = component.build_structured_output() assert isinstance(result, Data) assert result.data == {"parent": {"child": "value"}} @patch("langflow.components.helpers.structured_output.get_chat_result") def test_large_input_value(self, mock_get_chat_result): large_input = "Test input " * 1000 class MockBaseModel(BaseModel): field: str = "value" mock_get_chat_result.return_value = MockBaseModel(field="value") component = StructuredOutputComponent( llm=MagicMock(), input_value=large_input, schema_name="LargeInputSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, ) result = component.build_structured_output() assert isinstance(result, Data) assert result.data == {"field": "value"} mock_get_chat_result.assert_called_once() def test_invalid_llm_config(self): component = StructuredOutputComponent( llm="invalid_llm", # Not a proper LLM instance input_value="Test input", schema_name="InvalidLLMSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, ) with pytest.raises(TypeError, match="Language model does not support structured output."): component.build_structured_output()