diff --git a/docling_agent/agent/base_functions.py b/docling_agent/agent/base_functions.py index 6e71b41..f59baa5 100644 --- a/docling_agent/agent/base_functions.py +++ b/docling_agent/agent/base_functions.py @@ -199,7 +199,7 @@ def convert_html_to_docling_table(text: str) -> list[TableItem] | None: try: converter = DocumentConverter(allowed_formats=[InputFormat.HTML]) - buff = BytesIO(text.encode("utf-8")) + buff = BytesIO(text_.encode("utf-8")) doc_stream = DocumentStream(name="tmp.html", stream=buff) conv: ConversionResult = converter.convert(doc_stream) diff --git a/tests/test_base_functions.py b/tests/test_base_functions.py index c2d44d1..8bc6159 100644 --- a/tests/test_base_functions.py +++ b/tests/test_base_functions.py @@ -4,7 +4,12 @@ from docling_core.experimental.serializer.outline import OutlineMode from docling_core.types.doc.document import DoclingDocument, SectionHeaderItem -from docling_agent.agent.base_functions import create_document_outline, find_json_dicts, make_hierarchical_document +from docling_agent.agent.base_functions import ( + convert_html_to_docling_table, + create_document_outline, + find_json_dicts, + make_hierarchical_document, +) from .test_data_gen_flag import GEN_TEST_DATA @@ -120,6 +125,37 @@ def test_find_json_dicts_with_invalid_json(): assert result == [] +HTML_TABLE_BLOCK = """```html + + + +
YearRevenue
202412
+```""" + + +def test_convert_html_to_docling_table_ignores_prose_mentioning_tags(): + """Test that only the html code block is converted, not the prose around it.""" + text = f"I updated the below and removed the 2023 row.\n\n{HTML_TABLE_BLOCK}" + + tables = convert_html_to_docling_table(text) + + assert tables is not None + assert len(tables) == 1 + assert [cell.text for cell in tables[0].data.table_cells] == ["Year", "Revenue", "2024", "12"] + + +def test_convert_html_to_docling_table_ignores_tables_outside_code_block(): + """Test that a table quoted outside the html code block is not returned.""" + text = f"Original table:\n
old
\n\nUpdated table:\n{HTML_TABLE_BLOCK}" + + tables = convert_html_to_docling_table(text) + + assert tables is not None + assert len(tables) == 1 + assert tables[0].data.num_rows == 2 + assert tables[0].data.num_cols == 2 + + def test_make_hierarchical_document(): """Test that make_hierarchical_document creates proper hierarchical structure from flat document.""" # Load the JSON file