import os import re from unittest.mock import MagicMock, patch import openai import pytest from langflow.helpers.base_model import build_model_from_schema from langflow.inputs.inputs import TableInput from lfx.components.llm_operations.structured_output import StructuredOutputComponent from pydantic import BaseModel from tests.base import ComponentTestBaseWithoutClient from tests.unit.mock_language_model import MockLanguageModel class TestStructuredOutputComponent(ComponentTestBaseWithoutClient): @pytest.fixture def component_class(self): """Return the component class to test.""" return StructuredOutputComponent @pytest.fixture def mock_llm(self): """Create a mock LLM for testing.""" return MockLanguageModel() @pytest.fixture def default_kwargs(self): """Return the default kwargs for the component with proper model metadata.""" return { "model": [ { "name": "gpt-3.5-turbo", "provider": "OpenAI", "metadata": { "model_class": "MockLanguageModel", "model_name_param": "model", "api_key_param": "api_key", }, } ], "api_key": "test-api-key", "input_value": "Test input", "schema_name": "TestSchema", "output_schema": [{"name": "field", "type": "str", "description": "A test field"}], "multiple": False, "system_prompt": "Test system prompt", } @pytest.fixture def file_names_mapping(self): """Return the file names mapping for version-specific files.""" @pytest.fixture def mock_model_classes(self): """Helper fixture to create a mock model class factory for MockLanguageModel.""" def _create_mock_model_class(mock_llm_instance): """Create a mock model class that returns the provided mock LLM.""" return MagicMock(return_value=mock_llm_instance) return _create_mock_model_class @pytest.fixture def model_metadata(self): """Helper fixture that returns standard model metadata structure.""" return [ { "name": "gpt-3.5-turbo", "provider": "OpenAI", "metadata": { "model_class": "MockLanguageModel", "model_name_param": "model", "api_key_param": "api_key", }, } ] @patch("lfx.base.models.unified_models.get_model_class") def test_successful_structured_output_generation_with_patch_with_config(self, mock_get_model_class, mock_llm): def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return {"objects": [{"field": "value"}]} # Return trustcall-style response structure return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } # Mock get_model_class to return MockLanguageModel factory mock_model_class = MagicMock(return_value=mock_llm) mock_get_model_class.return_value = mock_model_class component = StructuredOutputComponent( model=[ { "name": "gpt-3.5-turbo", "provider": "OpenAI", "metadata": { "model_class": "MockLanguageModel", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, system_prompt="Test system prompt", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output_base() assert isinstance(result, list) assert result == [{"field": "value"}] @patch("lfx.base.models.unified_models.get_model_class") def test_raises_value_error_for_unsupported_language_model(self, mock_get_model_class): # Mocking an incompatible language model that doesn't support with_structured_output class IncompatibleModel: pass # Mock get_model_class to return IncompatibleModel factory mock_model_class = MagicMock(return_value=IncompatibleModel()) mock_get_model_class.return_value = mock_model_class # Creating an instance of StructuredOutputComponent component = StructuredOutputComponent( model=[ { "name": "test-model", "provider": "Test", "metadata": { "model_class": "MockLanguageModel", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, ) with pytest.raises(TypeError, match=re.escape("Language model does not support structured output.")): component.build_structured_output() def test_correctly_builds_output_model(self): # 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) def test_handles_multiple_outputs(self): # 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) @patch("lfx.base.models.unified_models.get_model_class") def test_empty_output_schema(self, mock_get_model_class, mock_llm, mock_model_classes): mock_get_model_class.return_value = mock_model_classes(mock_llm) component = StructuredOutputComponent( model=[ { "name": "gpt-3.5-turbo", "provider": "OpenAI", "metadata": { "model_class": "MockLanguageModel", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key="test-api-key", 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() @patch("lfx.base.models.unified_models.get_model_class") def test_invalid_output_schema_type(self, mock_get_model_class, mock_llm, mock_model_classes): mock_get_model_class.return_value = mock_model_classes(mock_llm) component = StructuredOutputComponent( model=[ { "name": "gpt-3.5-turbo", "provider": "OpenAI", "metadata": { "model_class": "MockLanguageModel", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key="test-api-key", 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("lfx.components.llm_operations.structured_output.get_chat_result") @patch("lfx.base.models.unified_models.get_model_class") def test_nested_output_schema( self, mock_get_model_class, mock_get_chat_result, mock_llm, mock_model_classes, model_metadata ): mock_get_model_class.return_value = mock_model_classes(mock_llm) class ChildModel(BaseModel): child: str = "value" class ParentModel(BaseModel): objects: list[dict] = [{"parent": {"child": "value"}}] def model_dump(self, **__): return {"objects": self.objects} # Update to return trustcall-style response mock_get_chat_result.return_value = { "messages": ["mock_message"], "responses": [ParentModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", 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, system_prompt="Test system prompt", ) result = component.build_structured_output_base() assert isinstance(result, list) assert result == [{"parent": {"child": "value"}}] @patch("lfx.components.llm_operations.structured_output.get_chat_result") @patch("lfx.base.models.unified_models.get_model_class") def test_large_input_value( self, mock_get_model_class, mock_get_chat_result, mock_llm, mock_model_classes, model_metadata ): mock_get_model_class.return_value = mock_model_classes(mock_llm) large_input = "Test input " * 1000 class MockBaseModel(BaseModel): objects: list[dict] = [{"field": "value"}] def model_dump(self, **__): return {"objects": self.objects} # Update to return trustcall-style response mock_get_chat_result.return_value = { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value=large_input, schema_name="LargeInputSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, system_prompt="Test system prompt", ) result = component.build_structured_output_base() assert isinstance(result, list) assert result == [{"field": "value"}] mock_get_chat_result.assert_called_once() @pytest.mark.skipif( not (os.getenv("OPENAI_API_KEY") or "").strip(), reason="OPENAI_API_KEY is not set or is empty", ) def test_with_real_openai_model_simple_schema(self): # Create a component with a simple schema using real OpenAI model component = StructuredOutputComponent( model=[ { "name": "gpt-4o-mini", "provider": "OpenAI", "metadata": { "model_class": "ChatOpenAI", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key=os.getenv("OPENAI_API_KEY"), input_value="Extract the name and age from this text: John Doe is 30 years old.", schema_name="PersonInfo", output_schema=[ {"name": "name", "type": "str", "description": "The person's name"}, {"name": "age", "type": "int", "description": "The person's age"}, ], multiple=False, system_prompt="Extract structured information from the input text.", ) # Get the structured output result = component.build_structured_output_base() # Verify the result assert isinstance(result, list) assert len(result) > 0 assert "name" in result[0] assert "age" in result[0] assert result[0]["name"] == "John Doe" assert result[0]["age"] == 30 @pytest.mark.skipif( not (os.getenv("OPENAI_API_KEY") or "").strip(), reason="OPENAI_API_KEY environment variable not set", ) def test_with_real_openai_model_multiple_patterns(self): # Create a component with multiple people in the input using real OpenAI model component = StructuredOutputComponent( model=[ { "name": "gpt-4o-mini", "provider": "OpenAI", "metadata": { "model_class": "ChatOpenAI", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key=os.getenv("OPENAI_API_KEY"), input_value=( "Extract all people from this text: John Doe is 30 years old, Jane Smith is 25, and Bob Johnson is 35." ), schema_name="PersonInfo", output_schema=[ {"name": "name", "type": "str", "description": "The person's name"}, {"name": "age", "type": "int", "description": "The person's age"}, ], multiple=False, system_prompt=( "You are an AI that extracts structured JSON objects from unstructured text. " "Use a predefined schema with expected types (str, int, float, bool, dict). " "Extract ALL relevant instances that match the schema - if multiple patterns exist, capture them all. " "Fill missing or ambiguous values with defaults: null for missing values. " "Remove exact duplicates but keep variations that have different field values. " "Always return valid JSON in the expected format, never throw errors. " "If multiple objects can be extracted, return them all in the structured format." ), ) # Get the structured output result = component.build_structured_output_base() # Verify the result contains multiple people assert isinstance(result, list) assert len(result) >= 3 # Should extract all three people # Check that we have names and ages for multiple people names = [item["name"] for item in result if "name" in item] ages = [item["age"] for item in result if "age" in item] assert len(names) >= 3 assert len(ages) >= 3 # Check that we extracted the expected people (order may vary) expected_names = ["John Doe", "Jane Smith", "Bob Johnson"] expected_ages = [30, 25, 35] for expected_name in expected_names: assert any(expected_name in name for name in names) for expected_age in expected_ages: assert expected_age in ages @patch("lfx.base.models.unified_models.get_model_class") def test_multiple_patterns_with_duplicates_and_variations( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that multiple patterns are extracted while removing exact duplicates but keeping variations.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return { "objects": [ {"product": "iPhone", "price": 999.99}, {"product": "iPhone", "price": 1099.99}, # Variation - different price {"product": "Samsung", "price": 899.99}, {"product": "iPhone", "price": 999.99}, # Exact duplicate - should be removed ] } return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Products: iPhone $999.99, iPhone $1099.99, Samsung $899.99, iPhone $999.99", schema_name="ProductSchema", output_schema=[ {"name": "product", "type": "str", "description": "Product name"}, {"name": "price", "type": "float", "description": "Product price"}, ], multiple=False, system_prompt="Remove exact duplicates but keep variations that have different field values.", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output() # Check that result is a Data object from lfx.schema.data import Data assert isinstance(result, Data) # Should have multiple results due to multiple patterns assert isinstance(result.data, dict) assert "results" in result.data assert ( len(result.data["results"]) == 4 ) # All items returned (duplicate handling is expected to be done by LLM) # Verify the expected products are present products = [item["product"] for item in result.data["results"]] prices = [item["price"] for item in result.data["results"]] assert "iPhone" in products assert "Samsung" in products assert 999.99 in prices assert 1099.99 in prices assert 899.99 in prices @pytest.mark.skipif( not (os.getenv("OPENAI_API_KEY") or "").strip(), reason="OPENAI_API_KEY environment variable not set", ) def test_with_real_openai_model_simple_schema_fail(self): # Create a component with a simple schema using real OpenAI model with very low max_tokens # Note: max_tokens parameter should be passed through model metadata if supported component = StructuredOutputComponent( model=[ { "name": "gpt-4o-mini", "provider": "OpenAI", "metadata": { "model_class": "ChatOpenAI", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key=os.getenv("OPENAI_API_KEY"), input_value="Extract the name and age from this text: John Doe is 30 years old.", schema_name="PersonInfo", output_schema=[ {"name": "name", "type": "str", "description": "The person's name"}, {"name": "age", "type": "int", "description": "The person's age"}, ], multiple=False, system_prompt="Extract structured information from the input text.", ) # Attempt to build structured output with max_tokens=1 # Note: Some API versions may handle this gracefully and succeed, # while others may raise BadRequestError. This test verifies error handling when it occurs. try: result = component.build_structured_output_base() # If the API succeeds with max_tokens=1, that's also valid behavior # Just verify we got some result back assert result is not None, "Expected a result even with max_tokens=1" except (ValueError, openai.BadRequestError) as exc_info: # If an error is raised, verify it's related to max_tokens error_message = str(exc_info.value) assert any( phrase in error_message for phrase in [ "max_tokens was reached", "max_tokens or model output limit was reached", "Could not finish the message because max_tokens", "BadRequestError", "with_structured_output also failed", # Langchain fallback error message ] ), f"Expected max_tokens error but got: {error_message}" @pytest.mark.skipif( not (os.getenv("OPENAI_API_KEY") or "").strip(), reason="OPENAI_API_KEY environment variable not set", ) def test_with_real_openai_model_complex_schema(self): # Create a component with a more complex schema using real OpenAI model component = StructuredOutputComponent( model=[ { "name": "gpt-4o-mini", "provider": "OpenAI", "metadata": { "model_class": "ChatOpenAI", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key=os.getenv("OPENAI_API_KEY"), input_value=""" Product Review: I purchased the XYZ Wireless Headphones last month. The sound quality is excellent, and the battery lasts about 8 hours. However, they're a bit uncomfortable after wearing them for a long time. The price was $129.99, which I think is reasonable for the quality. Overall rating: 4/5. """, schema_name="ProductReview", output_schema=[ {"name": "product_name", "type": "str", "description": "The name of the product"}, {"name": "sound_quality", "type": "str", "description": "Description of sound quality"}, {"name": "comfort", "type": "str", "description": "Description of comfort"}, {"name": "battery_life", "type": "str", "description": "Description of battery life"}, {"name": "price", "type": "float", "description": "The price of the product"}, {"name": "rating", "type": "float", "description": "The overall rating out of 5"}, ], multiple=False, system_prompt="Extract detailed product review information from the input text.", ) # Get the structured output result = component.build_structured_output_base() # Verify the result assert isinstance(result, list) assert len(result) > 0 assert "product_name" in result[0] assert "sound_quality" in result[0] assert "comfort" in result[0] assert "battery_life" in result[0] assert "price" in result[0] assert "rating" in result[0] assert result[0]["product_name"] == "XYZ Wireless Headphones" assert result[0]["price"] == 129.99 assert result[0]["rating"] == 4.0 @pytest.mark.skipif( not (os.getenv("OPENAI_API_KEY") or "").strip(), reason="OPENAI_API_KEY environment variable not set", ) def test_with_real_openai_model_nested_schema(self): # Create a component with a flattened schema (no nested structures) using real OpenAI model component = StructuredOutputComponent( model=[ { "name": "gpt-4o-mini", "provider": "OpenAI", "metadata": { "model_class": "ChatOpenAI", "model_name_param": "model", "api_key_param": "api_key", }, } ], api_key=os.getenv("OPENAI_API_KEY"), input_value=""" Restaurant: Bella Italia Address: 123 Main St, Anytown, CA 12345 Visited: June 15, 2023 Ordered: - Margherita Pizza ($14.99) - Delicious with fresh basil - Tiramisu ($8.50) - Perfect sweetness Service was excellent, atmosphere was cozy. Total bill: $35.49 including tip. Would definitely visit again! """, schema_name="RestaurantReview", output_schema=[ {"name": "restaurant_name", "type": "str", "description": "The name of the restaurant"}, {"name": "street", "type": "str", "description": "Street address"}, {"name": "city", "type": "str", "description": "City"}, {"name": "state", "type": "str", "description": "State"}, {"name": "zip", "type": "str", "description": "ZIP code"}, {"name": "first_item_name", "type": "str", "description": "Name of first item ordered"}, {"name": "first_item_price", "type": "float", "description": "Price of first item"}, {"name": "second_item_name", "type": "str", "description": "Name of second item ordered"}, {"name": "second_item_price", "type": "float", "description": "Price of second item"}, {"name": "total_bill", "type": "float", "description": "Total bill amount"}, {"name": "would_return", "type": "bool", "description": "Whether the reviewer would return"}, ], multiple=False, system_prompt="Extract detailed restaurant review information from the input text.", ) # Get the structured output result = component.build_structured_output_base() # Verify the result assert isinstance(result, list) assert len(result) > 0 assert "restaurant_name" in result[0] assert "street" in result[0] assert "city" in result[0] assert "state" in result[0] assert "zip" in result[0] assert "first_item_name" in result[0] assert "first_item_price" in result[0] assert "total_bill" in result[0] assert "would_return" in result[0] assert result[0]["restaurant_name"] == "Bella Italia" assert result[0]["street"] == "123 Main St" assert result[0]["total_bill"] == 35.49 assert result[0]["would_return"] is True @pytest.mark.skipif( not (os.getenv("NVIDIA_API_KEY") or "").strip(), reason="NVIDIA_API_KEY environment variable not set", ) def test_with_real_nvidia_model_simple_schema(self): # Create a component with a simple schema using real NVIDIA model # Note: ChatNVIDIA is not imported as it's resolved dynamically via model_class string component = StructuredOutputComponent( model=[ { "name": "meta/llama-3.2-3b-instruct", "provider": "NVIDIA", "metadata": { "model_class": "ChatNVIDIA", "model_name_param": "model", "api_key_param": "nvidia_api_key", }, } ], api_key=os.getenv("NVIDIA_API_KEY"), input_value="Extract the name and age from this text: John Doe is 30 years old.", schema_name="PersonInfo", output_schema=[ {"name": "name", "type": "str", "description": "The person's name"}, {"name": "age", "type": "int", "description": "The person's age"}, ], multiple=False, system_prompt="Extract structured information from the input text.", ) # Test that it now works with NVIDIA models (previously expected to fail but now supports structured output) try: result = component.build_structured_output_base() # If it succeeds, verify it returns a valid runnable assert result is not None except (TypeError, ValueError, RuntimeError) as e: # If it still fails, verify it's with a known error message error_msg = str(e) assert any( msg in error_msg for msg in [ "Language model does not support structured output", "400 Bad Request", "not supported", ] ), f"Unexpected error: {error_msg}" @patch("lfx.base.models.unified_models.get_model_class") def test_structured_output_returns_dict_when_no_objects_key( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that when trustcall returns a dict without 'objects' key, we return the dict directly.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 # Return trustcall-style response but without BaseModel that creates "objects" key return { "messages": ["mock_message"], "responses": [{"field": "value", "another_field": "another_value"}], # Direct dict, not BaseModel "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, system_prompt="Test system prompt", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output_base() # Should return the dict directly since there's no "objects" key assert isinstance(result, dict) assert result == {"field": "value", "another_field": "another_value"} @patch("lfx.base.models.unified_models.get_model_class") def test_structured_output_returns_direct_response_when_not_dict( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that when trustcall returns a non-dict response, we return it directly.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 # Return a string response (edge case) return "Simple string response" component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, system_prompt="Test system prompt", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output_base() # Should return the string directly assert isinstance(result, str) assert result == "Simple string response" @patch("lfx.base.models.unified_models.get_model_class") def test_structured_output_handles_empty_responses_array( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that when trustcall returns empty responses array, we return the result dict.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 # Return trustcall-style response with empty responses return { "messages": ["mock_message"], "responses": [], # Empty responses array "response_metadata": [], "attempts": 1, "fallback_data": {"field": "fallback_value"}, # Some other data in the result } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, system_prompt="Test system prompt", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output_base() # Should return the entire result dict when responses is empty assert isinstance(result, dict) assert "messages" in result assert "responses" in result assert "fallback_data" in result @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_output_fails_when_base_returns_non_list( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that build_structured_output() fails when base method returns non-list.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 # Return a dict instead of list with objects return { "messages": ["mock_message"], "responses": [{"single_item": "value"}], # Dict without "objects" key "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, system_prompt="Test system prompt", ) with ( patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result), pytest.raises(ValueError, match="No structured output returned"), ): component.build_structured_output() @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_output_returns_data_with_dict( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that build_structured_output() returns Data object with dict data.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return {"objects": [{"field": "value2", "number": 24}]} # Return only one object # Return trustcall-style response structure return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[ {"name": "field", "type": "str", "description": "A test field"}, {"name": "number", "type": "int", "description": "A test number"}, ], multiple=False, system_prompt="Test system prompt", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output() # Check that result is a Data object from lfx.schema.data import Data assert isinstance(result, Data) # Check that result.data is a dict assert isinstance(result.data, dict) # Check the content of the dict assert result.data == {"field": "value2", "number": 24} # Verify the data has the expected keys assert "field" in result.data assert "number" in result.data assert result.data["field"] == "value2" assert result.data["number"] == 24 @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_output_returns_multiple_objects( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that build_structured_output() returns Data object with multiple objects wrapped in results.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return { "objects": [ {"name": "John", "age": 30}, {"name": "Jane", "age": 25}, {"name": "Bob", "age": 35}, ] } return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Extract multiple people: John is 30, Jane is 25, Bob is 35", schema_name="PersonSchema", output_schema=[ {"name": "name", "type": "str", "description": "Person's name"}, {"name": "age", "type": "int", "description": "Person's age"}, ], multiple=False, system_prompt="Extract ALL relevant instances that match the schema", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output() # Check that result is a Data object from lfx.schema.data import Data assert isinstance(result, Data) # Check that result.data is a dict with results key assert isinstance(result.data, dict) assert "results" in result.data assert len(result.data["results"]) == 3 # Check the content of each result assert result.data["results"][0] == {"name": "John", "age": 30} assert result.data["results"][1] == {"name": "Jane", "age": 25} assert result.data["results"][2] == {"name": "Bob", "age": 35} @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_output_returns_data_with_single_item( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that build_structured_output() returns Data object when only one item in objects.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return {"objects": [{"name": "John Doe", "age": 30}]} return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Extract name and age from: John Doe is 30 years old", schema_name="PersonInfo", output_schema=[ {"name": "name", "type": "str", "description": "Person's name"}, {"name": "age", "type": "int", "description": "Person's age"}, ], multiple=False, system_prompt="Extract person info", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output() # Check that result is a Data object from lfx.schema.data import Data assert isinstance(result, Data) # Check that result.data is a dict assert isinstance(result.data, dict) # Check the content matches exactly assert result.data == {"name": "John Doe", "age": 30} @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_output_data_object_properties( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that the returned Data object has proper properties.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return {"objects": [{"product": "iPhone", "price": 999.99, "available": True}]} return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Product info: iPhone costs $999.99 and is available", schema_name="ProductInfo", output_schema=[ {"name": "product", "type": "str", "description": "Product name"}, {"name": "price", "type": "float", "description": "Product price"}, {"name": "available", "type": "bool", "description": "Product availability"}, ], multiple=False, system_prompt="Extract product info", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_output() # Check that result is a Data object from lfx.schema.data import Data assert isinstance(result, Data) # Check that result.data is a dict with correct types assert isinstance(result.data, dict) assert isinstance(result.data["product"], str) assert isinstance(result.data["price"], float) assert isinstance(result.data["available"], bool) # Check values assert result.data["product"] == "iPhone" assert result.data["price"] == 999.99 assert result.data["available"] is True # Test Data object methods if they exist if hasattr(result, "get_text"): # Data object should be able to represent itself as text text_repr = result.get_text() assert isinstance(text_repr, str) @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_dataframe_returns_dataframe_with_single_data( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that build_structured_dataframe() returns DataFrame object with single Data item.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return {"objects": [{"field": "value2", "number": 24}]} return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[ {"name": "field", "type": "str", "description": "A test field"}, {"name": "number", "type": "int", "description": "A test number"}, ], multiple=False, system_prompt="Test system prompt", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_dataframe() # Check that result is a DataFrame object from lfx.schema.dataframe import DataFrame assert isinstance(result, DataFrame) assert len(result) == 1 assert result.iloc[0]["field"] == "value2" assert result.iloc[0]["number"] == 24 # Test conversion back to Data list data_list = result.to_data_list() assert len(data_list) == 1 assert data_list[0].data == {"field": "value2", "number": 24} @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_dataframe_returns_dataframe_with_multiple_data( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that build_structured_dataframe() returns DataFrame object with multiple Data items.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return { "objects": [ {"name": "John", "age": 30}, {"name": "Jane", "age": 25}, {"name": "Bob", "age": 35}, ] } return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input with multiple people", schema_name="PersonSchema", output_schema=[ {"name": "name", "type": "str", "description": "Person's name"}, {"name": "age", "type": "int", "description": "Person's age"}, ], multiple=False, system_prompt="Test system prompt", ) with patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result): result = component.build_structured_dataframe() # Check that result is a DataFrame object from lfx.schema.dataframe import DataFrame assert isinstance(result, DataFrame) assert len(result) == 3 assert result.iloc[0]["name"] == "John" assert result.iloc[0]["age"] == 30 assert result.iloc[1]["name"] == "Jane" assert result.iloc[1]["age"] == 25 assert result.iloc[2]["name"] == "Bob" assert result.iloc[2]["age"] == 35 # Test conversion back to Data list data_list = result.to_data_list() assert len(data_list) == 3 assert data_list[0].data == {"name": "John", "age": 30} assert data_list[1].data == {"name": "Jane", "age": 25} assert data_list[2].data == {"name": "Bob", "age": 35} @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_dataframe_fails_when_base_returns_non_list( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that build_structured_dataframe() fails when base method returns non-list.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 return { "messages": ["mock_message"], "responses": [{"single_item": "value"}], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, system_prompt="Test system prompt", ) with ( patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result), pytest.raises(ValueError, match="No structured output returned"), ): component.build_structured_dataframe() @patch("lfx.base.models.unified_models.get_model_class") def test_build_structured_dataframe_fails_when_empty_output( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that build_structured_dataframe() fails when base method returns empty list.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return {"objects": []} return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], multiple=False, system_prompt="Test system prompt", ) with ( patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result), pytest.raises(ValueError, match="No structured output returned"), ): component.build_structured_dataframe() @patch("lfx.base.models.unified_models.get_model_class") def test_fallback_to_langchain_on_trustcall_generic_exception( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that when trustcall fails with a generic exception, it falls back to langchain.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], system_prompt="Test system prompt", ) with ( patch.object(component, "_extract_output_with_trustcall", return_value=None) as mock_trustcall, patch.object( component, "_extract_output_with_langchain", return_value=[{"field": "langchain_value"}] ) as mock_langchain, ): result = component.build_structured_output_base() # Verify fallback was successful assert isinstance(result, list) assert result == [{"field": "langchain_value"}] # Verify both methods were called mock_trustcall.assert_called_once() mock_langchain.assert_called_once() @patch("lfx.base.models.unified_models.get_model_class") def test_fallback_both_methods_fail_raises_value_error( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that when both trustcall and langchain fail, a ValueError is raised.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], system_prompt="Test system prompt", ) # Mock trustcall to return None (indicating failure) # Mock langchain to raise an exception with ( patch.object(component, "_extract_output_with_trustcall", return_value=None), patch.object( component, "_extract_output_with_langchain", side_effect=ValueError( "Model does not support tool calling (trustcall failed) and fallback " "with_structured_output also failed: Langchain parsing error" ), ), ): with pytest.raises(ValueError, match="trustcall failed") as exc_info: component.build_structured_output_base() error_msg = str(exc_info.value) # Verify error message mentions both failures assert "Model does not support tool calling" in error_msg assert "trustcall failed" in error_msg assert "fallback with_structured_output also failed" in error_msg assert "Langchain parsing error" in error_msg @patch("lfx.base.models.unified_models.get_model_class") def test_langchain_fallback_processes_basemodel_response( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that langchain fallback correctly processes BaseModel responses.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], system_prompt="Test system prompt", ) # Mock trustcall to return None (fail) # Mock langchain to return list with dict with ( patch.object(component, "_extract_output_with_trustcall", return_value=None), patch.object(component, "_extract_output_with_langchain", return_value=[{"field": "test_value"}]), ): result = component.build_structured_output_base() # Verify it extracted the objects assert isinstance(result, list) assert result == [{"field": "test_value"}] @patch("lfx.base.models.unified_models.get_model_class") def test_langchain_fallback_processes_dict_response( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that langchain fallback correctly processes dict responses without BaseModel conversion.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], system_prompt="Test system prompt", ) # Mock trustcall to return None (fail) # Mock langchain to return dict directly with ( patch.object(component, "_extract_output_with_trustcall", return_value=None), patch.object(component, "_extract_output_with_langchain", return_value={"field": "dict_value"}), ): result = component.build_structured_output_base() # When langchain returns dict, it's returned as-is assert result == {"field": "dict_value"} @patch("lfx.base.models.unified_models.get_model_class") def test_fallback_error_message_includes_both_errors( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that the error message when both methods fail includes context about both failures.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], system_prompt="Test system prompt", ) # Mock trustcall to return None (fail) # Mock langchain to raise an exception with full error message with ( patch.object(component, "_extract_output_with_trustcall", return_value=None), patch.object( component, "_extract_output_with_langchain", side_effect=ValueError( "Model does not support tool calling (trustcall failed) and fallback " "with_structured_output also failed: Langchain parsing error" ), ), ): with pytest.raises(ValueError, match="trustcall failed") as exc_info: component.build_structured_output_base() error_msg = str(exc_info.value) # Verify error message mentions both failures assert "Model does not support tool calling" in error_msg assert "trustcall failed" in error_msg assert "fallback with_structured_output also failed" in error_msg assert "Langchain parsing error" in error_msg @patch("lfx.base.models.unified_models.get_model_class") def test_trustcall_success_no_fallback_attempted( self, mock_get_model_class, mock_llm, mock_model_classes, model_metadata ): """Test that when trustcall succeeds, langchain fallback is not attempted.""" mock_get_model_class.return_value = mock_model_classes(mock_llm) def mock_get_chat_result(runnable, system_message, input_value, config, **kwargs): # noqa: ARG001 class MockBaseModel(BaseModel): def model_dump(self, **__): return {"objects": [{"field": "trustcall_value"}]} return { "messages": ["mock_message"], "responses": [MockBaseModel()], "response_metadata": [{"id": "mock_id"}], "attempts": 1, } component = StructuredOutputComponent( model=model_metadata, api_key="test-api-key", input_value="Test input", schema_name="TestSchema", output_schema=[{"name": "field", "type": "str", "description": "A test field"}], system_prompt="Test system prompt", ) with ( patch("lfx.components.llm_operations.structured_output.get_chat_result", mock_get_chat_result), patch.object(component, "_extract_output_with_langchain") as mock_lc_fallback, ): result = component.build_structured_output_base() # Verify trustcall succeeded assert isinstance(result, list) assert result == [{"field": "trustcall_value"}] # Verify langchain was NOT called mock_lc_fallback.assert_not_called()