From 5e7832a6ea234a3292b0c607f887197bca0623ca Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 29 Jan 2026 18:41:07 +0200
Subject: [PATCH] Fix REPL for docling document, add test

---
 .../haiku/rag/agents/rlm/sandbox.py           |  6 ++---
 tests/agents/rlm/test_agent.py                | 26 +++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py b/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py
index 5daf7e2a..e10a5410 100644
--- a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py
+++ b/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py
@@ -229,17 +229,17 @@ class REPLEnvironment:
             async def _get():
                 doc = await self.client.get_document_by_id(id_or_title)
                 if doc:
-                    return doc.docling_document
+                    return doc.get_docling_document()
                 docs = await self.client.list_documents(
                     filter=f"title = '{id_or_title}'"
                 )
                 if docs and docs[0].id:
                     full_doc = await self.client.get_document_by_id(docs[0].id)
-                    return full_doc.docling_document if full_doc else None
+                    return full_doc.get_docling_document() if full_doc else None
                 docs = await self.client.list_documents(filter=f"uri = '{id_or_title}'")
                 if docs and docs[0].id:
                     full_doc = await self.client.get_document_by_id(docs[0].id)
-                    return full_doc.docling_document if full_doc else None
+                    return full_doc.get_docling_document() if full_doc else None
                 return None
 
             return self._run_async_from_thread(_get())
diff --git a/tests/agents/rlm/test_agent.py b/tests/agents/rlm/test_agent.py
index 2d5cd060..f8bb2356 100644
--- a/tests/agents/rlm/test_agent.py
+++ b/tests/agents/rlm/test_agent.py
@@ -171,3 +171,29 @@ class TestClientRLMIntegration:
             )
 
             assert "1" in answer
+
+    @pytest.mark.asyncio
+    @pytest.mark.vcr()
+    async def test_rlm_docling_document_structure(
+        self, allow_model_requests, temp_db_path
+    ):
+        """Test RLM agent can analyze document structure using DoclingDocument."""
+        from pathlib import Path
+
+        from haiku.rag.client import HaikuRAG
+        from haiku.rag.config import AppConfig
+
+        pdf_path = Path("tests/data/doclaynet.pdf")
+        config = AppConfig()
+        config.processing.conversion_options.do_ocr = False
+
+        async with HaikuRAG(temp_db_path, config=config, create=True) as client:
+            await client.create_document_from_source(pdf_path)
+
+            answer = await client.rlm(
+                "How many tables are in the document? "
+                "Also tell me how many pictures/figures it contains."
+            )
+
+            # The doclaynet.pdf has 1 table and 1 picture
+            assert "1" in answer