diff --git a/docling_agent/agent/editor.py b/docling_agent/agent/editor.py index 2835bde..38a4b38 100644 --- a/docling_agent/agent/editor.py +++ b/docling_agent/agent/editor.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Annotated, ClassVar, Literal -from docling_core.experimental.serializer.outline import OutlineFormat +from docling_core.experimental.serializer.outline import OutlineFormat, OutlineMode from docling_core.types.base import _JSON_POINTER_REGEX from docling_core.types.doc.document import ( ContentLayer, @@ -173,7 +173,9 @@ def _identify_document_items( log_info(f"task: {task}") # TODO: check best format for describing the outline (MARKDOWN vs JSON) for short and long documents - outline = create_document_outline(doc=document, format=OutlineFormat.MARKDOWN) + # The full outline (not only the table of contents) lists the refs of text items and tables, + # which the update_content and rewrite_content operations need to target. + outline = create_document_outline(doc=document, mode=OutlineMode.OUTLINE, format=OutlineFormat.MARKDOWN) log_debug(f"outline: {outline}") context = rf"""Given the current outline of the document: diff --git a/tests/test_editor_validation.py b/tests/test_editor_validation.py index 8c56bba..38c78df 100644 --- a/tests/test_editor_validation.py +++ b/tests/test_editor_validation.py @@ -6,6 +6,8 @@ DocItemLabel, DoclingDocument, SectionHeaderItem, + TableCell, + TableData, TextItem, TitleItem, ) @@ -65,6 +67,40 @@ def test_invalid_ref_in_list(self): assert "String should match pattern" in str(exc_info.value) +class TestIdentifyDocumentItems: + """Test the operation planning step of the DoclingEditingAgent.""" + + def test_outline_exposes_table_refs(self, monkeypatch): + """The planning prompt lists tables, so the model can target them with update_content.""" + document = DoclingDocument(name="form") + document.add_title(text="Assessment Form") + cells = [ + TableCell( + text=text, + start_row_offset_idx=0, + end_row_offset_idx=1, + start_col_offset_idx=col, + end_col_offset_idx=col + 1, + ) + for col, text in enumerate(["Field", "Value"]) + ] + table = document.add_table(data=TableData(num_rows=1, num_cols=2, table_cells=cells)) + prompts = [] + + def answer_with_table_ref(self, prompt, *, requirements=None, retry_budget=1): + prompts.append(prompt) + return f'```json\n{{"operation": "update_content", "ref": "{table.self_ref}"}}\n```' + + monkeypatch.setattr(MockSession, "instruct", answer_with_table_ref) + + op = DoclingEditingAgent(backend=MockBackend(), tools=[])._identify_document_items( + task="Fill in the Value column of the table", document=document + ) + + assert table.self_ref in prompts[0] + assert op == UpdateContentOperation(operation="update_content", ref=table.self_ref) + + class TestRewriteContent: @staticmethod def _rewrite(monkeypatch, document, refs):