agent-flow / src /backend /tests /unit /components /helpers /test_structured_output_component.py

Upload folder using huggingface_hub

4b0794d verified almost 2 years ago

9.88 kB

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()

	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()