feat(rag): enable multi-document and page-level reasoning

Signed-off-by: Cesar Berrospi Ramis <ceb@zurich.ibm.com>

Author: ceberam

docling_agent/agent/rag.py

@@ -59,6 +59,11 @@ class DoclingRAGAgent(BaseDoclingAgent):
 
     max_iterations: int = 5
     verbose: bool = False
+    enable_document_selection: bool = False  # NEW: Enable document filtering before RAG
+    use_page_level: bool = False  # NEW: Use pages instead of sections
+    use_batch_selection: bool = False  # NEW: Batch-based selection instead of iterative
+    batch_size: int = 30  # NEW: Pages/sections per batch
+    top_k: int = 10  # NEW: Maximum pages/sections to select
 
     def __init__(
         self,
@@ -67,6 +72,11 @@ def __init__(
         backend=None,
         max_iterations: int = 5,
         verbose: bool = False,
+        enable_document_selection: bool = False,
+        use_page_level: bool = False,
+        use_batch_selection: bool = False,
+        batch_size: int = 30,
+        top_k: int = 10,
     ):
         super().__init__(
             agent_type=DoclingAgentType.DOCLING_DOCUMENT_RAG,
@@ -75,6 +85,11 @@ def __init__(
         )
         self.max_iterations = max_iterations
         self.verbose = verbose
+        self.enable_document_selection = enable_document_selection
+        self.use_page_level = use_page_level
+        self.use_batch_selection = use_batch_selection
+        self.batch_size = batch_size
+        self.top_k = top_k
         self._console = Console(highlight=False) if verbose else None
 
     def _rprint(self, renderable: Any) -> None:
@@ -95,6 +110,31 @@ def run(
         if not docs:
             raise ValueError("DoclingRAGAgent requires at least one DoclingDocument.")
 
+        # Phase 1: Optional document selection
+        if self.enable_document_selection and len(docs) > 1:
+            self._rprint(Rule(f"[bold cyan]Selecting relevant documents from {len(docs)} candidates[/bold cyan]"))
+
+            # Build document summaries
+            documents_dict = {doc.name: doc for doc in docs}
+            doc_summaries = {doc.name: self._extract_document_summary(doc) for doc in docs}
+
+            # Select relevant documents
+            selected_doc_names = self._select_relevant_documents(
+                query=task,
+                documents=documents_dict,
+                doc_summaries=doc_summaries,
+            )
+
+            # Filter docs to only selected ones
+            docs = [documents_dict[name] for name in selected_doc_names if name in documents_dict]
+            self._rprint(
+                Panel(
+                    f"Selected {len(docs)} document(s): {[d.name for d in docs]}",
+                    title="[cyan]Document Selection[/cyan]",
+                    border_style="cyan",
+                )
+            )
+
         per_doc_answers: list[str] = []
         all_iterations: list[RAGIteration] = []
 
@@ -261,18 +301,111 @@ def _rag_loop(self, *, query: str, doc: DoclingDocument) -> RAGResult:
     # ------------------------------------------------------------------
 
     def _build_outline(self, doc: DoclingDocument) -> str:
+        """Build outline or page list based on mode."""
+        if self.use_page_level:
+            # Page-level mode: build a list of pages with summaries
+            return self._build_page_list(doc)
+
+        # Section-level mode: use standard outline
         serializer = OutlineDocSerializer(
             doc=doc,
             params=OutlineParams(mode=OutlineMode.OUTLINE),
         )
         return serializer.serialize().text
 
+    def _build_page_list(self, doc: DoclingDocument) -> str:
+        """Build a formatted list of pages with their summaries."""
+        lines = []
+        num_pages = len(doc.pages) if hasattr(doc, "pages") and doc.pages else 0
+
+        for i in range(num_pages):
+            page_ref = f"#/pages/{i}"
+            page_num = i + 1  # 1-indexed for display
+
+            # Get page summary if available
+            summary = "No summary available"
+            page_node = get_item_by_ref(doc, page_ref)
+            if page_node and hasattr(page_node, "children") and page_node.children:
+                # Look for summary in first child's meta
+                first_child = page_node.children[0]
+                if hasattr(first_child, "meta") and first_child.meta:
+                    meta_dict = first_child.meta.model_dump() if hasattr(first_child.meta, "model_dump") else {}
+                    if "summary" in meta_dict and isinstance(meta_dict["summary"], dict):
+                        summary = meta_dict["summary"].get("text", summary)
+
+            # Use page_ref directly in the display (e.g., "#/pages/0")
+            lines.append(f"Page {page_num} ({page_ref}): {summary}")
+
+        return "\n".join(lines)
+
     def _extract_section_refs(self, doc: DoclingDocument) -> set[str]:
-        refs: set[str] = set()
-        for item, _ in doc.iterate_items():
-            if isinstance(item, TitleItem | SectionHeaderItem):
-                refs.add(item.self_ref)
-        return refs
+        """Extract section or page references based on mode.
+
+        Returns:
+            Set of section refs (e.g., "#/body/0") or page refs (e.g., "#/pages/0")
+        """
+        if self.use_page_level:
+            # Page-level mode: return page references using JSON pointer format
+            refs: set[str] = set()
+            num_pages = len(doc.pages) if hasattr(doc, "pages") and doc.pages else 0
+            for i in range(num_pages):
+                refs.add(f"#/pages/{i}")  # Use JSON pointer format
+            return refs
+        else:
+            # Section-level mode: return section header references
+            refs = set()
+            for item, _ in doc.iterate_items():
+                if isinstance(item, TitleItem | SectionHeaderItem):
+                    refs.add(item.self_ref)
+            return refs
+
+    def _extract_page_summaries(self, doc: DoclingDocument) -> dict[int, str]:
+        """Extract page-level summaries from enriched document.
+
+        For page-level enrichment, the summary is stored in the meta field of the
+        first document item on each page.
+
+        Args:
+            doc: The enriched DoclingDocument
+
+        Returns:
+            Dictionary mapping page number (1-indexed) to summary text
+        """
+        page_summaries: dict[int, str] = {}
+
+        try:
+            for item, _ in doc.iterate_items():
+                # Check if item has prov (provenance) with page number
+                if hasattr(item, "prov") and item.prov:
+                    for prov in item.prov:
+                        if hasattr(prov, "page_no") and prov.page_no is not None:
+                            page_num = prov.page_no + 1  # Convert to 1-indexed
+
+                            # Only process if we haven't seen this page yet
+                            if page_num not in page_summaries:
+                                # Check if item has meta with summary
+                                if hasattr(item, "meta") and item.meta:
+                                    meta = item.meta
+                                    # Handle both dict and Pydantic object
+                                    if isinstance(meta, dict):
+                                        summary_data = meta.get("summary", {})
+                                    else:
+                                        summary_data = getattr(meta, "summary", {})
+
+                                    if isinstance(summary_data, dict):
+                                        summary_text = summary_data.get("text", "")
+                                    elif hasattr(summary_data, "text"):
+                                        summary_text = summary_data.text
+                                    else:
+                                        summary_text = ""
+
+                                    if summary_text:
+                                        page_summaries[page_num] = summary_text
+                            break  # Only need first prov entry
+        except Exception as e:
+            log_warning(f"Error extracting page summaries: {e}")
+
+        return page_summaries
 
     # ------------------------------------------------------------------
     # Section selection
@@ -336,7 +469,11 @@ def _validate(content: str) -> bool:
     # ------------------------------------------------------------------
 
     def _get_section_content(self, doc: DoclingDocument, section_ref: str) -> str:
-        """Return all text belonging to the given section node."""
+        """Return all text belonging to the given section node or page."""
+        # Handle page-level mode
+        if self.use_page_level and section_ref.startswith("#/pages/"):
+            return self._get_page_content(doc, section_ref)
+
         node = get_item_by_ref(doc, section_ref)
         if node is None:
             log_warning(f"Could not resolve section ref {section_ref!r}")
@@ -353,6 +490,16 @@ def _get_section_content(self, doc: DoclingDocument, section_ref: str) -> str:
 
         return subtree
 
+    def _get_page_content(self, doc: DoclingDocument, page_ref: str) -> str:
+        """Extract all text content from a specific page."""
+        page_node = get_item_by_ref(doc, page_ref)
+        if page_node is None:
+            log_warning(f"Could not resolve page ref {page_ref!r}")
+            return ""
+
+        # Collect all text from the page using subtree traversal
+        return collect_subtree_text(page_node, doc)
+
     def _collect_flat_section_text(self, doc: DoclingDocument, section_ref: str) -> str:
         """Scan iterate_items for the section and collect siblings until next section."""
         texts: list[str] = []
@@ -440,3 +587,110 @@ def _merge_answers(self, *, query: str, answers: list[str]) -> str:
             retry_budget=3,
         )
         return answer.strip()
+
+    # ------------------------------------------------------------------
+    # Document selection (Phase 1: Multi-document support)
+    # ------------------------------------------------------------------
+
+    def _extract_document_summary(self, doc: DoclingDocument) -> str:
+        """Extract document-level summary from enriched document.
+
+        The summary is stored in the meta field of the root body item.
+        """
+        try:
+            # Find the root body item (self_ref == "#/body")
+            for item, _ in doc.iterate_items():
+                if hasattr(item, "self_ref") and item.self_ref == "#/body":
+                    # Check if meta exists and has summary
+                    if hasattr(item, "meta") and item.meta:
+                        meta = item.meta
+                        # Handle both dict and Pydantic object
+                        if isinstance(meta, dict):
+                            summary_data = meta.get("summary", {})
+                        else:
+                            summary_data = getattr(meta, "summary", {})
+
+                        if isinstance(summary_data, dict):
+                            return summary_data.get("text", "")
+                        elif hasattr(summary_data, "text"):
+                            return summary_data.text
+
+            log_warning(f"No document summary found for {doc.name}")
+            return f"Document: {doc.name}"
+        except Exception as e:
+            log_warning(f"Error extracting document summary: {e}")
+            return f"Document: {doc.name}"
+
+    def _select_relevant_documents(
+        self,
+        *,
+        query: str,
+        documents: dict[str, DoclingDocument],
+        doc_summaries: dict[str, str],
+    ) -> list[str]:
+        """Select which documents are relevant for answering the query.
+
+        This removes evaluation bias by not assuming we know which document
+        contains the answer. The model must decide based on document summaries.
+
+        Args:
+            query: The query to answer
+            documents: Dictionary mapping doc_id to DoclingDocument
+            doc_summaries: Dictionary mapping doc_id to document summary
+
+        Returns:
+            List of relevant doc_ids
+        """
+        log_info(f"Selecting relevant documents from {len(documents)} candidates")
+
+        # Build prompt with all document summaries
+        doc_list = "\n".join([f"Document '{doc_id}':\n{summary}" for doc_id, summary in doc_summaries.items()])
+
+        prompt = f"""You are analyzing a collection of documents to find which ones are relevant for answering a query.
+
+QUERY:
+{query}
+
+AVAILABLE DOCUMENTS:
+{doc_list}
+
+TASK:
+Identify the MOST relevant document(s) for answering the query. Be selective and precise.
+
+IMPORTANT GUIDELINES:
+- Prefer selecting 1-2 documents that are highly relevant
+- Only select additional documents if they provide essential complementary information
+- Do NOT select documents just because they might be tangentially related
+- Quality over quantity: it's better to select fewer, highly relevant documents than many loosely related ones
+- If the query is specific to one company/topic, typically only 1 document is needed
+- If the query compares multiple entities, select only the documents for those specific entities
+
+Format your response as:
+Document 'doc_id': [reason]
+
+Only include documents that are actually relevant to the query.
+"""
+
+        try:
+            # Get reasoning model
+            m = self._create_reasoning_session(system_prompt=self._RAG_SYSTEM_PROMPT)
+            response = m.instruct(prompt, retry_budget=3)
+
+            # Parse response to extract doc_ids
+            selected_docs = []
+            for doc_id in documents.keys():
+                # Look for the doc_id in the response
+                if f"'{doc_id}'" in response or f'"{doc_id}"' in response or doc_id in response:
+                    selected_docs.append(doc_id)
+
+            if not selected_docs:
+                log_warning("No documents selected by model, using all documents as fallback")
+                selected_docs = list(documents.keys())
+
+            log_info(f"Selected {len(selected_docs)} relevant document(s): {selected_docs}")
+            return selected_docs
+
+        except Exception as e:
+            log_warning(f"Error selecting documents: {e}")
+            # Fallback to all documents
+            return list(documents.keys())