From b9b5d43dc69f731457acf1be6f31183f4c859642 Mon Sep 17 00:00:00 2001 From: Heyang Wang Date: Mon, 26 May 2025 16:41:57 +0800 Subject: [PATCH] fix: add 'floatfmt' when extract number from excel ( #20153 ) (#20193) Co-authored-by: wangheyang Co-authored-by: crazywoola <427733928@qq.com> --- .../workflow/nodes/document_extractor/node.py | 2 +- .../nodes/test_document_extractor_node.py | 179 ++++++++++++++++++ 2 files changed, 180 insertions(+), 1 deletion(-) diff --git a/api/core/workflow/nodes/document_extractor/node.py b/api/core/workflow/nodes/document_extractor/node.py index 8fb1baec89..84abac7b15 100644 --- a/api/core/workflow/nodes/document_extractor/node.py +++ b/api/core/workflow/nodes/document_extractor/node.py @@ -366,7 +366,7 @@ def _extract_text_from_excel(file_content: bytes) -> str: df = excel_file.parse(sheet_name=sheet_name) df.dropna(how="all", inplace=True) # Create Markdown table two times to separate tables with a newline - markdown_table += df.to_markdown(index=False) + "\n\n" + markdown_table += df.to_markdown(index=False, floatfmt="") + "\n\n" except Exception as e: continue return markdown_table diff --git a/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py b/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py index de739ce9f5..1e8aec7f88 100644 --- a/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/test_document_extractor_node.py @@ -10,6 +10,7 @@ from core.workflow.entities.node_entities import NodeRunResult from core.workflow.nodes.document_extractor import DocumentExtractorNode, DocumentExtractorNodeData from core.workflow.nodes.document_extractor.node import ( _extract_text_from_docx, + _extract_text_from_excel, _extract_text_from_pdf, _extract_text_from_plain_text, ) @@ -182,3 +183,181 @@ def test_extract_text_from_docx(mock_document): def test_node_type(document_extractor_node): assert document_extractor_node._node_type == NodeType.DOCUMENT_EXTRACTOR + + +@patch("pandas.ExcelFile") +def test_extract_text_from_excel_single_sheet(mock_excel_file): + """Test extracting text from Excel file with single sheet.""" + # Mock DataFrame + mock_df = Mock() + mock_df.dropna = Mock() + mock_df.to_markdown.return_value = "| Name | Age |\n|------|-----|\n| John | 25 |" + + # Mock ExcelFile + mock_excel_instance = Mock() + mock_excel_instance.sheet_names = ["Sheet1"] + mock_excel_instance.parse.return_value = mock_df + mock_excel_file.return_value = mock_excel_instance + + file_content = b"fake_excel_content" + result = _extract_text_from_excel(file_content) + + expected = "| Name | Age |\n|------|-----|\n| John | 25 |\n\n" + assert result == expected + mock_excel_file.assert_called_once() + mock_df.dropna.assert_called_once_with(how="all", inplace=True) + mock_df.to_markdown.assert_called_once_with(index=False, floatfmt="") + + +@patch("pandas.ExcelFile") +def test_extract_text_from_excel_multiple_sheets(mock_excel_file): + """Test extracting text from Excel file with multiple sheets.""" + # Mock DataFrames for different sheets + mock_df1 = Mock() + mock_df1.dropna = Mock() + mock_df1.to_markdown.return_value = "| Product | Price |\n|---------|-------|\n| Apple | 1.50 |" + + mock_df2 = Mock() + mock_df2.dropna = Mock() + mock_df2.to_markdown.return_value = "| City | Population |\n|------|------------|\n| NYC | 8000000 |" + + # Mock ExcelFile + mock_excel_instance = Mock() + mock_excel_instance.sheet_names = ["Products", "Cities"] + mock_excel_instance.parse.side_effect = [mock_df1, mock_df2] + mock_excel_file.return_value = mock_excel_instance + + file_content = b"fake_excel_content_multiple_sheets" + result = _extract_text_from_excel(file_content) + + expected = ( + "| Product | Price |\n|---------|-------|\n| Apple | 1.50 |\n\n" + "| City | Population |\n|------|------------|\n| NYC | 8000000 |\n\n" + ) + assert result == expected + assert mock_excel_instance.parse.call_count == 2 + + +@patch("pandas.ExcelFile") +def test_extract_text_from_excel_empty_sheets(mock_excel_file): + """Test extracting text from Excel file with empty sheets.""" + # Mock empty DataFrame + mock_df = Mock() + mock_df.dropna = Mock() + mock_df.to_markdown.return_value = "" + + # Mock ExcelFile + mock_excel_instance = Mock() + mock_excel_instance.sheet_names = ["EmptySheet"] + mock_excel_instance.parse.return_value = mock_df + mock_excel_file.return_value = mock_excel_instance + + file_content = b"fake_excel_empty_content" + result = _extract_text_from_excel(file_content) + + expected = "\n\n" + assert result == expected + + +@patch("pandas.ExcelFile") +def test_extract_text_from_excel_sheet_parse_error(mock_excel_file): + """Test handling of sheet parsing errors - should continue with other sheets.""" + # Mock DataFrames - one successful, one that raises exception + mock_df_success = Mock() + mock_df_success.dropna = Mock() + mock_df_success.to_markdown.return_value = "| Data | Value |\n|------|-------|\n| Test | 123 |" + + # Mock ExcelFile + mock_excel_instance = Mock() + mock_excel_instance.sheet_names = ["GoodSheet", "BadSheet"] + mock_excel_instance.parse.side_effect = [mock_df_success, Exception("Parse error")] + mock_excel_file.return_value = mock_excel_instance + + file_content = b"fake_excel_mixed_content" + result = _extract_text_from_excel(file_content) + + expected = "| Data | Value |\n|------|-------|\n| Test | 123 |\n\n" + assert result == expected + + +@patch("pandas.ExcelFile") +def test_extract_text_from_excel_file_error(mock_excel_file): + """Test handling of Excel file reading errors.""" + mock_excel_file.side_effect = Exception("Invalid Excel file") + + file_content = b"invalid_excel_content" + + with pytest.raises(Exception) as exc_info: + _extract_text_from_excel(file_content) + + # Note: The function should raise TextExtractionError, but since it's not imported in the test, + # we check for the general Exception pattern + assert "Failed to extract text from Excel file" in str(exc_info.value) + + +@patch("pandas.ExcelFile") +def test_extract_text_from_excel_io_bytesio_usage(mock_excel_file): + """Test that BytesIO is properly used with the file content.""" + import io + + # Mock DataFrame + mock_df = Mock() + mock_df.dropna = Mock() + mock_df.to_markdown.return_value = "| Test | Data |\n|------|------|\n| 1 | A |" + + # Mock ExcelFile + mock_excel_instance = Mock() + mock_excel_instance.sheet_names = ["TestSheet"] + mock_excel_instance.parse.return_value = mock_df + mock_excel_file.return_value = mock_excel_instance + + file_content = b"test_excel_bytes" + result = _extract_text_from_excel(file_content) + + # Verify that ExcelFile was called with a BytesIO object + mock_excel_file.assert_called_once() + call_args = mock_excel_file.call_args[0][0] + assert isinstance(call_args, io.BytesIO) + + expected = "| Test | Data |\n|------|------|\n| 1 | A |\n\n" + assert result == expected + + +@patch("pandas.ExcelFile") +def test_extract_text_from_excel_all_sheets_fail(mock_excel_file): + """Test when all sheets fail to parse - should return empty string.""" + # Mock ExcelFile + mock_excel_instance = Mock() + mock_excel_instance.sheet_names = ["BadSheet1", "BadSheet2"] + mock_excel_instance.parse.side_effect = [Exception("Error 1"), Exception("Error 2")] + mock_excel_file.return_value = mock_excel_instance + + file_content = b"fake_excel_all_bad_sheets" + result = _extract_text_from_excel(file_content) + + # Should return empty string when all sheets fail + assert result == "" + + +@patch("pandas.ExcelFile") +def test_extract_text_from_excel_markdown_formatting(mock_excel_file): + """Test that markdown formatting parameters are correctly applied.""" + # Mock DataFrame + mock_df = Mock() + mock_df.dropna = Mock() + mock_df.to_markdown.return_value = "| Float | Int |\n|-------|-----|\n| 123456.78 | 42 |" + + # Mock ExcelFile + mock_excel_instance = Mock() + mock_excel_instance.sheet_names = ["NumberSheet"] + mock_excel_instance.parse.return_value = mock_df + mock_excel_file.return_value = mock_excel_instance + + file_content = b"fake_excel_numbers" + result = _extract_text_from_excel(file_content) + + # Verify to_markdown was called with correct parameters + mock_df.to_markdown.assert_called_once_with(index=False, floatfmt="") + + expected = "| Float | Int |\n|-------|-----|\n| 123456.78 | 42 |\n\n" + assert result == expected