Add pre-loaded documents support and improve RLM test coverage

haiku_rag_slim/haiku/rag/agents/rlm/prompts.py

@@ -32,6 +32,16 @@ Call an LLM directly with the given prompt. Returns the response as a string.
 Use this for classification, summarization, extraction, or any task where you
 already have the content and just need LLM reasoning.
 
+## Pre-loaded Documents Variable
+
+If documents were pre-loaded for this session, a `documents` variable is available:
+```python
+# documents is a list of dicts with keys: id, title, uri, content
+for doc in documents:
+    print(doc['title'], len(doc['content']))
+```
+Check if it exists with: `if 'documents' in dir(): ...`
+
 ## Standard Library Modules
 You can import: json, re, collections, math, statistics, itertools, functools, datetime, typing
 

tests/agents/rlm/test_agent.py

@@ -343,3 +343,38 @@ class TestClientRLMIntegration:
                 assert (
                     label in answer_lower or label.replace("-", " ") in answer_lower
                 ), f"Missing label: {label}"
+
+    @pytest.mark.asyncio
+    @pytest.mark.vcr()
+    async def test_rlm_with_preloaded_documents(
+        self, allow_model_requests, temp_db_path
+    ):
+        """Test RLM agent can use pre-loaded documents variable.
+
+        Agent program:
+            if 'documents' in dir():
+                for doc in documents:
+                    print(doc['title'], len(doc['content']))
+            else:
+                print('No preloaded documents')
+        """
+        from haiku.rag.client import HaikuRAG
+
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            await client.create_document(
+                "The company was founded in 1985 by Jane Smith.",
+                title="Company History",
+            )
+            await client.create_document(
+                "Our mission is to make technology accessible to everyone.",
+                title="Mission Statement",
+            )
+
+            answer = await client.rlm(
+                "Using the pre-loaded documents variable, "
+                "tell me when was the company founded and what is their mission?",
+                documents=["Company History", "Mission Statement"],
+            )
+
+            assert "1985" in answer
+            assert "accessible" in answer.lower() or "technology" in answer.lower()

tests/agents/rlm/test_sandbox.py

@@ -482,6 +482,96 @@ class TestContextFilter:
             )
 
 
+class TestPreloadedDocuments:
+    """Test pre-loaded documents context variable."""
+
+    @pytest.mark.asyncio
+    async def test_documents_variable_available_when_preloaded(self, temp_db_path):
+        """documents variable is available when context.documents is set."""
+        from haiku.rag.agents.rlm.dependencies import RLMContext
+        from haiku.rag.agents.rlm.sandbox import REPLEnvironment
+        from haiku.rag.client import HaikuRAG
+        from haiku.rag.config.models import RLMConfig
+        from haiku.rag.store.models import Document
+
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            preloaded = [
+                Document(
+                    id="doc-1",
+                    title="First Doc",
+                    uri="test://first",
+                    content="Content of first document about cats.",
+                ),
+                Document(
+                    id="doc-2",
+                    title="Second Doc",
+                    uri="test://second",
+                    content="Content of second document about dogs.",
+                ),
+            ]
+            context = RLMContext(documents=preloaded)
+            repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
+
+            result = await repl.execute_async(
+                "print(len(documents))\n"
+                "print([d['title'] for d in documents])\n"
+                "print('cats' in documents[0]['content'])"
+            )
+            assert result.success
+            assert "2" in result.stdout
+            assert "First Doc" in result.stdout
+            assert "Second Doc" in result.stdout
+            assert "True" in result.stdout
+
+    @pytest.mark.asyncio
+    async def test_documents_variable_not_available_without_preload(self, temp_db_path):
+        """documents variable is not available when context.documents is None."""
+        from haiku.rag.agents.rlm.dependencies import RLMContext
+        from haiku.rag.agents.rlm.sandbox import REPLEnvironment
+        from haiku.rag.client import HaikuRAG
+        from haiku.rag.config.models import RLMConfig
+
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            context = RLMContext()
+            repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
+
+            result = await repl.execute_async("print(documents)")
+            assert not result.success
+            assert "NameError" in result.stderr
+
+    @pytest.mark.asyncio
+    async def test_documents_has_expected_fields(self, temp_db_path):
+        """documents variable contains expected dict fields."""
+        from haiku.rag.agents.rlm.dependencies import RLMContext
+        from haiku.rag.agents.rlm.sandbox import REPLEnvironment
+        from haiku.rag.client import HaikuRAG
+        from haiku.rag.config.models import RLMConfig
+        from haiku.rag.store.models import Document
+
+        async with HaikuRAG(temp_db_path, create=True) as client:
+            preloaded = [
+                Document(
+                    id="doc-1",
+                    title="Test Doc",
+                    uri="test://doc",
+                    content="Test content",
+                ),
+            ]
+            context = RLMContext(documents=preloaded)
+            repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
+
+            result = await repl.execute_async(
+                "d = documents[0]\n"
+                "print(sorted(d.keys()))\n"
+                "print(d['id'], d['title'], d['uri'])"
+            )
+            assert result.success
+            assert "['content', 'id', 'title', 'uri']" in result.stdout
+            assert "doc-1" in result.stdout
+            assert "Test Doc" in result.stdout
+            assert "test://doc" in result.stdout
+
+
 class TestSecurityEscapes:
     """Test that common security escape attempts are blocked."""