Simplify sandbox with run_monty_async, replace manual ThreadPoolExecutor start/resume loop

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

@@ -2,32 +2,33 @@ RLM_SYSTEM_PROMPT = """You are a Recursive Language Model (RLM) agent that solve
 
 IMPORTANT: You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside the execute_code tool - you cannot access them any other way. Always execute code to answer questions; do not just describe what code would do.
 
-CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just use them directly:
+CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
-- search("query")  ✓ CORRECT
+- results = await search("query")  ✓ CORRECT
-- from haiku.rag import search  ✗ WRONG - will fail
+- import search  ✗ WRONG - will fail
+- results = search("query")  ✗ WRONG - must use await
 
-You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly, no imports needed):
+You have access to a sandboxed Python interpreter with these functions (use them directly with `await`, no imports needed):
 
 ## Available Functions
 
-### search(query, limit=10) -> list[dict]
+### await search(query, limit=10) -> list[dict]
 Search the knowledge base using hybrid search (vector + full-text).
 Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings
 
-### list_documents(limit=10, offset=0) -> list[dict]
+### await list_documents(limit=10, offset=0) -> list[dict]
 List available documents in the knowledge base.
 Returns list of dicts with keys: id, title, uri, created_at
 
-### get_document(id_or_title) -> str | None
+### await get_document(id_or_title) -> str | None
 Get the full text content of a document by ID, title, or URI.
 Returns the document content as a string, or None if not found.
 
-### get_chunk(chunk_id) -> dict | None
+### await get_chunk(chunk_id) -> dict | None
 Get a specific chunk by its ID (from search results).
 Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
 Use this to retrieve full chunk details and metadata for citation.
 
-### llm(prompt) -> str
+### await llm(prompt) -> str
 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.
@@ -44,7 +45,7 @@ Check if it exists with: `if 'documents' in dir(): ...`
 
 ## Available Python Features
 
-The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, `map()`, `sorted()`/`.sort(key=...)`, try/except, and the `json` module.
+The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `sorted()`/`.sort(key=...)`, try/except, and the `json` module.
 
 Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
@@ -53,21 +54,21 @@ For pattern matching or text extraction, use string methods (`str.split`, `str.f
 ## Strategy Guide
 
 1. **Explore First**: Start by listing documents or searching to understand what's available. Document names may differ from filenames (e.g., "tbmed593.pdf" might be stored as "TB MED 593" or similar).
-2. **If get_document returns None**: Use `list_documents()` to see actual document titles, or `search()` to find relevant content.
+2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
 3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
 4. **Use print() Liberally**: The sandbox captures stdout - print intermediate results to see what you're working with.
 5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
-6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use llm().
+6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
 7. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
 
 ## Example Patterns
 
 ### Counting documents matching a condition
 ```python
-docs = list_documents(limit=100)
+docs = await list_documents(limit=100)
 count = 0
 for doc in docs:
-    content = get_document(doc['id'])
+    content = await get_document(doc['id'])
     if content and 'keyword' in content.lower():
         count += 1
         print(f"Found in: {doc['title']}")
@@ -77,9 +78,9 @@ print(f"Total: {count}")
 ### Extracting data with llm()
 ```python
 numbers = []
-results = search("financial data", limit=20)
+results = await search("financial data", limit=20)
 for r in results:
-    extracted = llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+    extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
     for part in extracted.split(','):
         part = part.strip().replace(',', '')
         if part.isdigit():
@@ -90,16 +91,16 @@ if numbers:
 
 ### Using search results with get_chunk for citations
 ```python
-results = search("safety requirements", limit=5)
+results = await search("safety requirements", limit=5)
 for r in results:
-    chunk = get_chunk(r['chunk_id'])
+    chunk = await get_chunk(r['chunk_id'])
     print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
 ```
 
 ### Using llm() for classification
 ```python
-content = get_document("Q1 Report")
+content = await get_document("Q1 Report")
-sentiment = llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
+sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
 print(sentiment)
 ```
 

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

@@ -1,7 +1,4 @@
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
-from functools import partial
 from typing import TYPE_CHECKING, Any, Literal
 
 import pydantic_monty
@@ -27,12 +24,10 @@ class Sandbox:
 
     Uses pydantic-monty, a minimal secure Python interpreter written in Rust.
     External functions (search, list_documents, etc.) are called by Monty code
-    and resolved asynchronously on the host.
+    using ``await`` and resolved asynchronously on the host.
 
-    Use as an async context manager:
+        sandbox = Sandbox(client, config, context)
-
+        result = await sandbox.execute("print('hello')")
-        async with Sandbox(client, config, context) as sandbox:
-            result = await sandbox.execute("print('hello')")
     """
 
     _client: "HaikuRAG"
@@ -49,14 +44,6 @@ class Sandbox:
         self._config = config
         self._context = context
 
-    async def __aenter__(self) -> "Sandbox":
-        return self
-
-    async def __aexit__(
-        self, exc_type: object, exc_val: object, exc_tb: object
-    ) -> None:
-        pass
-
     def _build_external_functions(self) -> dict[str, Any]:
         """Build async external functions for the Monty interpreter."""
         client = self._client
@@ -138,12 +125,7 @@ class Sandbox:
         }
 
     async def execute(self, code: str) -> SandboxResult:
-        """Execute Python code in the Monty interpreter.
+        """Execute Python code in the Monty interpreter."""
-
-        Uses a manual start/resume loop so that async external functions
-        are awaited on the host while Monty code calls them synchronously
-        (without ``await``).
-        """
         external_fns = self._build_external_functions()
 
         input_names: list[str] = []
@@ -181,45 +163,14 @@ class Sandbox:
             "max_duration_secs": self._config.rlm.code_timeout,
         }
 
-        loop = asyncio.get_running_loop()
-
         try:
-            with ThreadPoolExecutor() as pool:
+            output = await pydantic_monty.run_monty_async(
-
+                monty,
-                async def run_in_pool(func: Any) -> Any:
+                inputs=inputs,
-                    return await loop.run_in_executor(pool, func)
+                external_functions=external_fns,
-
+                limits=limits,
-                progress = await run_in_pool(
+                print_callback=print_callback,
-                    partial(
+            )
-                        monty.start,
-                        inputs=inputs,
-                        limits=limits,
-                        print_callback=print_callback,
-                    )
-                )
-
-                while not isinstance(progress, pydantic_monty.MontyComplete):
-                    assert isinstance(progress, pydantic_monty.MontySnapshot)
-                    fn = external_fns.get(progress.function_name)
-                    if fn is None:
-                        exc = KeyError(f"Function {progress.function_name} not found")
-                        progress = await run_in_pool(
-                            partial(progress.resume, exception=exc)
-                        )
-                        continue
-
-                    try:
-                        result = await fn(*progress.args, **progress.kwargs)
-                    except Exception as exc:
-                        progress = await run_in_pool(
-                            partial(progress.resume, exception=exc)
-                        )
-                    else:
-                        progress = await run_in_pool(
-                            partial(progress.resume, return_value=result)
-                        )
-
-                output = progress.output
         except pydantic_monty.MontyRuntimeError as e:
             stdout = "".join(stdout_lines)
             if len(stdout) > max_chars:

haiku_rag_slim/haiku/rag/client.py

@@ -1406,20 +1406,20 @@ class HaikuRAG:
                     loaded_docs.append(doc)
             context.documents = loaded_docs if loaded_docs else None
 
-        async with Sandbox(
+        sandbox = Sandbox(
             client=self,
             config=self._config,
             context=context,
-        ) as sandbox:
+        )
-            deps = RLMDeps(
+        deps = RLMDeps(
-                sandbox=sandbox,
+            sandbox=sandbox,
-                context=context,
+            context=context,
-            )
+        )
 
-            agent = create_rlm_agent(self._config)
+        agent = create_rlm_agent(self._config)
-            result = await agent.run(question, deps=deps)
+        result = await agent.run(question, deps=deps)
 
-            return result.output
+        return result.output
 
     async def visualize_chunk(self, chunk: Chunk) -> list:
         """Render page images with bounding box highlights for a chunk.

haiku_rag_slim/haiku/rag/skills/rag-rlm/SKILL.md

@@ -1,7 +1,7 @@
 ---
 name: rag-rlm
 description: >
-  Computational analysis of the knowledge base via code execution in a Docker sandbox.
+  Computational analysis of the knowledge base via code execution in a sandboxed Python interpreter.
   Use for questions requiring counting, aggregation, statistics, data traversal,
   comparison across documents, or any task best answered by writing Python code.
   Examples: "how many pages?", "compare table 3 across documents",
@@ -10,4 +10,4 @@ description: >
 
 # RLM Analysis
 
-Use the `analyze` tool for complex analytical questions. It writes and executes Python code against the knowledge base in an isolated Docker sandbox.
+Use the `analyze` tool for complex analytical questions. It writes and executes Python code against the knowledge base in a sandboxed Python interpreter.

haiku_rag_slim/haiku/rag/tools/analysis.py

@@ -62,25 +62,25 @@ def create_analysis_toolset(
 
         rlm_context = RLMContext(filter=effective_filter)
 
-        async with Sandbox(
+        sandbox = Sandbox(
             client=client,
             config=config,
             context=rlm_context,
-        ) as sandbox:
+        )
-            deps = RLMDeps(
+        deps = RLMDeps(
-                sandbox=sandbox,
+            sandbox=sandbox,
-                context=rlm_context,
+            context=rlm_context,
-            )
+        )
 
-            rlm_agent = create_rlm_agent(config)
+        rlm_agent = create_rlm_agent(config)
-            result = await rlm_agent.run(task, deps=deps)
+        result = await rlm_agent.run(task, deps=deps)
 
-            program = result.output.program
+        program = result.output.program
 
-            return AnalysisResult(
+        return AnalysisResult(
-                answer=result.output.answer,
+            answer=result.output.answer,
-                code_executed=bool(program),
+            code_executed=bool(program),
-            )
+        )
 
     toolset: FunctionToolset[RAGDeps] = FunctionToolset()
     toolset.add_function(analyze, name=tool_name)

tests/agents/rlm/conftest.py

@@ -18,5 +18,4 @@ async def sandbox(empty_client):
     """Create a Monty sandbox for testing."""
     config = AppConfig()
     context = RLMContext()
-    async with Sandbox(client=empty_client, config=config, context=context) as sandbox:
+    return Sandbox(client=empty_client, config=config, context=context)
-        yield sandbox

tests/agents/rlm/test_sandbox.py

@@ -74,7 +74,7 @@ class TestSandboxHaikuRAG:
     async def test_list_documents_empty(self, sandbox):
         """Test list_documents returns empty list for empty database."""
         result = await sandbox.execute(
-            "docs = list_documents()\nprint(type(docs).__name__, len(docs))"
+            "docs = await list_documents()\nprint(type(docs).__name__, len(docs))"
         )
         assert result.success
         assert "list 0" in result.stdout
@@ -92,13 +92,15 @@ class TestSandboxHaikuRAG:
             )
 
             context = RLMContext()
-            async with Sandbox(client=client, config=config, context=context) as sb:
+            sb = Sandbox(client=client, config=config, context=context)
-                result = await sb.execute(
+            result = await sb.execute(
-                    "docs = list_documents()\nprint(len(docs))\nprint(docs[0]['title'])"
+                "docs = await list_documents()\n"
-                )
+                "print(len(docs))\n"
-                assert result.success
+                "print(docs[0]['title'])"
-                assert "1" in result.stdout
+            )
-                assert "Test Document" in result.stdout
+            assert result.success
+            assert "1" in result.stdout
+            assert "Test Document" in result.stdout
 
     @pytest.mark.asyncio
     @pytest.mark.vcr()
@@ -113,15 +115,15 @@ class TestSandboxHaikuRAG:
             )
 
             context = RLMContext()
-            async with Sandbox(client=client, config=config, context=context) as sb:
+            sb = Sandbox(client=client, config=config, context=context)
-                result = await sb.execute(
+            result = await sb.execute(
-                    "results = search('fox', limit=5)\n"
+                "results = await search('fox', limit=5)\n"
-                    "print(len(results))\n"
+                "print(len(results))\n"
-                    "if results:\n"
+                "if results:\n"
-                    "    print('fox' in results[0]['content'].lower())"
+                "    print('fox' in results[0]['content'].lower())"
-                )
+            )
-                assert result.success
+            assert result.success
-                assert "True" in result.stdout or "1" in result.stdout
+            assert "True" in result.stdout or "1" in result.stdout
 
     @pytest.mark.asyncio
     @pytest.mark.vcr()
@@ -136,19 +138,19 @@ class TestSandboxHaikuRAG:
             )
 
             context = RLMContext()
-            async with Sandbox(client=client, config=config, context=context) as sb:
+            sb = Sandbox(client=client, config=config, context=context)
-                result = await sb.execute(
+            result = await sb.execute(
-                    f"content = get_document('{doc.id}')\n"
+                f"content = await get_document('{doc.id}')\n"
-                    "print('foxes' in content.lower() if content else 'None')"
+                "print('foxes' in content.lower() if content else 'None')"
-                )
+            )
-                assert result.success
+            assert result.success
-                assert "True" in result.stdout
+            assert "True" in result.stdout
 
     @pytest.mark.asyncio
     async def test_get_document_not_found(self, sandbox):
         """Test get_document returns None for missing document."""
         result = await sandbox.execute(
-            "content = get_document('nonexistent-id')\nprint(content is None)"
+            "content = await get_document('nonexistent-id')\nprint(content is None)"
         )
         assert result.success
         assert "True" in result.stdout
@@ -166,24 +168,24 @@ class TestSandboxHaikuRAG:
             )
 
             context = RLMContext()
-            async with Sandbox(client=client, config=config, context=context) as sb:
+            sb = Sandbox(client=client, config=config, context=context)
-                # First search to get a chunk_id
+            # First search to get a chunk_id
-                result = await sb.execute(
+            result = await sb.execute(
-                    "results = search('foxes', limit=1)\n"
+                "results = await search('foxes', limit=1)\n"
-                    "chunk_id = results[0]['chunk_id']\n"
+                "chunk_id = results[0]['chunk_id']\n"
-                    "chunk = get_chunk(chunk_id)\n"
+                "chunk = await get_chunk(chunk_id)\n"
-                    "print(chunk['document_title'])\n"
+                "print(chunk['document_title'])\n"
-                    "print('content' in chunk)"
+                "print('content' in chunk)"
-                )
+            )
-                assert result.success
+            assert result.success
-                assert "Fox Document" in result.stdout
+            assert "Fox Document" in result.stdout
-                assert "True" in result.stdout
+            assert "True" in result.stdout
 
     @pytest.mark.asyncio
     async def test_get_chunk_not_found(self, sandbox):
         """Test get_chunk returns None for missing chunk."""
         result = await sandbox.execute(
-            "chunk = get_chunk('nonexistent-id')\nprint(chunk is None)"
+            "chunk = await get_chunk('nonexistent-id')\nprint(chunk is None)"
         )
         assert result.success
         assert "True" in result.stdout
@@ -205,7 +207,11 @@ class TestSandboxExternalFunctionEdgeCases:
         sandbox._build_external_functions = patched_build
 
         result = await sandbox.execute(
-            "try:\n    search('hello')\nexcept:\n    print('caught')\nprint('done')"
+            "try:\n"
+            "    await search('hello')\n"
+            "except:\n"
+            "    print('caught')\n"
+            "print('done')"
         )
         assert result.success
         assert "caught" in result.stdout
@@ -213,7 +219,12 @@ class TestSandboxExternalFunctionEdgeCases:
 
     @pytest.mark.asyncio
     async def test_external_function_raises_exception(self, sandbox):
-        """Test that exceptions from external functions are propagated to Monty."""
+        """Test that exceptions from async external functions surface as errors.
+
+        With run_monty_async, exceptions from async external functions
+        propagate as MontyRuntimeError rather than being catchable inside
+        Monty's try/except.
+        """
         original_build = sandbox._build_external_functions
 
         def patched_build():
@@ -227,12 +238,9 @@ class TestSandboxExternalFunctionEdgeCases:
 
         sandbox._build_external_functions = patched_build
 
-        result = await sandbox.execute(
+        result = await sandbox.execute("await search('hello')")
-            "try:\n    search('hello')\nexcept:\n    print('caught')\nprint('done')"
+        assert not result.success
-        )
+        assert "external error" in result.stderr
-        assert result.success
-        assert "caught" in result.stdout
-        assert "done" in result.stdout
 
 
 class TestSandboxOutputTruncation:
@@ -244,12 +252,12 @@ class TestSandboxOutputTruncation:
         config = AppConfig()
         config.rlm.max_output_chars = 20
         context = RLMContext()
-        async with Sandbox(client=empty_client, config=config, context=context) as sb:
+        sb = Sandbox(client=empty_client, config=config, context=context)
-            result = await sb.execute("print('a' * 100)\nx = 1/0")
+        result = await sb.execute("print('a' * 100)\nx = 1/0")
-            assert not result.success
+        assert not result.success
-            assert "ZeroDivisionError" in result.stderr
+        assert "ZeroDivisionError" in result.stderr
-            assert result.stdout.endswith("... (output truncated)")
+        assert result.stdout.endswith("... (output truncated)")
-            assert len(result.stdout) < 100
+        assert len(result.stdout) < 100
 
     @pytest.mark.asyncio
     async def test_truncate_successful_output(self, empty_client):
@@ -257,11 +265,11 @@ class TestSandboxOutputTruncation:
         config = AppConfig()
         config.rlm.max_output_chars = 20
         context = RLMContext()
-        async with Sandbox(client=empty_client, config=config, context=context) as sb:
+        sb = Sandbox(client=empty_client, config=config, context=context)
-            result = await sb.execute("print('b' * 100)")
+        result = await sb.execute("print('b' * 100)")
-            assert result.success
+        assert result.success
-            assert result.stdout.endswith("... (output truncated)")
+        assert result.stdout.endswith("... (output truncated)")
-            assert len(result.stdout) < 100
+        assert len(result.stdout) < 100
 
 
 class TestSandboxContextFilter:
@@ -285,17 +293,17 @@ class TestSandboxContextFilter:
             )
 
             context = RLMContext(filter="uri LIKE 'public://%'")
-            async with Sandbox(client=client, config=config, context=context) as sb:
+            sb = Sandbox(client=client, config=config, context=context)
-                result = await sb.execute(
+            result = await sb.execute(
-                    "docs = list_documents()\n"
+                "docs = await list_documents()\n"
-                    "print(len(docs))\n"
+                "print(len(docs))\n"
-                    "if docs:\n"
+                "if docs:\n"
-                    "    print(docs[0]['title'])"
+                "    print(docs[0]['title'])"
-                )
+            )
-                assert result.success
+            assert result.success
-                assert "1" in result.stdout
+            assert "1" in result.stdout
-                assert "Public Doc" in result.stdout
+            assert "Public Doc" in result.stdout
-                assert "Private Doc" not in result.stdout
+            assert "Private Doc" not in result.stdout
 
 
 class TestSandboxPreloadedDocuments:
@@ -317,16 +325,16 @@ class TestSandboxPreloadedDocuments:
             Document(id="2", content="Content B", title="Doc B", uri="b://2"),
         ]
         context = RLMContext(documents=docs)
-        async with Sandbox(client=empty_client, config=config, context=context) as sb:
+        sb = Sandbox(client=empty_client, config=config, context=context)
-            result = await sb.execute(
+        result = await sb.execute(
-                "print(len(documents))\n"
+            "print(len(documents))\n"
-                "print(documents[0]['title'])\n"
+            "print(documents[0]['title'])\n"
-                "print(documents[1]['title'])"
+            "print(documents[1]['title'])"
-            )
+        )
-            assert result.success
+        assert result.success
-            assert "2" in result.stdout
+        assert "2" in result.stdout
-            assert "Doc A" in result.stdout
+        assert "Doc A" in result.stdout
-            assert "Doc B" in result.stdout
+        assert "Doc B" in result.stdout
 
 
 class TestSandboxLLM:
@@ -338,10 +346,10 @@ class TestSandboxLLM:
         """Test llm() calls the model and returns a string."""
         config = AppConfig()
         context = RLMContext()
-        async with Sandbox(client=empty_client, config=config, context=context) as sb:
+        sb = Sandbox(client=empty_client, config=config, context=context)
-            result = await sb.execute(
+        result = await sb.execute(
-                "answer = llm('What is 2 + 2? Reply with just the number.')\n"
+            "answer = await llm('What is 2 + 2? Reply with just the number.')\n"
-                "print(answer)"
+            "print(answer)"
-            )
+        )
-            assert result.success
+        assert result.success
-            assert "4" in result.stdout
+        assert "4" in result.stdout

tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_count_documents.yaml

@@ -128,7 +128,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '7083'
+      - '7325'
       content-type:
       - application/json
       host:
@@ -141,32 +141,33 @@ interactions:
 
           IMPORTANT: You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside the execute_code tool - you cannot access them any other way. Always execute code to answer questions; do not just describe what code would do.
 
-          CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just use them directly:
+          CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
-          - search("query")  ✓ CORRECT
+          - results = await search("query")  ✓ CORRECT
           - from haiku.rag import search  ✗ WRONG - will fail
+          - results = search("query")  ✗ WRONG - must use await
 
-          You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly, no imports needed):
+          You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly with `await`, no imports needed):
 
           ## Available Functions
 
-          ### search(query, limit=10) -> list[dict]
+          ### await search(query, limit=10) -> list[dict]
           Search the knowledge base using hybrid search (vector + full-text).
           Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings
 
-          ### list_documents(limit=10, offset=0) -> list[dict]
+          ### await list_documents(limit=10, offset=0) -> list[dict]
           List available documents in the knowledge base.
           Returns list of dicts with keys: id, title, uri, created_at
 
-          ### get_document(id_or_title) -> str | None
+          ### await get_document(id_or_title) -> str | None
           Get the full text content of a document by ID, title, or URI.
           Returns the document content as a string, or None if not found.
 
-          ### get_chunk(chunk_id) -> dict | None
+          ### await get_chunk(chunk_id) -> dict | None
           Get a specific chunk by its ID (from search results).
           Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
           Use this to retrieve full chunk details and metadata for citation.
 
-          ### llm(prompt) -> str
+          ### await llm(prompt) -> str
           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.
@@ -183,7 +184,7 @@ interactions:
 
           ## Available Python Features
 
-          The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, try/except, and the `json` module.
+          The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `sorted()`/`.sort(key=...)`, try/except, and the `json` module.
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
@@ -192,21 +193,21 @@ interactions:
           ## Strategy Guide
 
           1. **Explore First**: Start by listing documents or searching to understand what's available. Document names may differ from filenames (e.g., "tbmed593.pdf" might be stored as "TB MED 593" or similar).
-          2. **If get_document returns None**: Use `list_documents()` to see actual document titles, or `search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
           3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
           4. **Use print() Liberally**: The sandbox captures stdout - print intermediate results to see what you're working with.
           5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
-          6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use llm().
+          6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
           7. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
 
           ## Example Patterns
 
           ### Counting documents matching a condition
           ```python
-          docs = list_documents(limit=100)
+          docs = await list_documents(limit=100)
           count = 0
           for doc in docs:
-              content = get_document(doc['id'])
+              content = await get_document(doc['id'])
               if content and 'keyword' in content.lower():
                   count += 1
                   print(f"Found in: {doc['title']}")
@@ -216,9 +217,9 @@ interactions:
           ### Extracting data with llm()
           ```python
           numbers = []
-          results = search("financial data", limit=20)
+          results = await search("financial data", limit=20)
           for r in results:
-              extracted = llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
               for part in extracted.split(','):
                   part = part.strip().replace(',', '')
                   if part.isdigit():
@@ -229,16 +230,16 @@ interactions:
 
           ### Using search results with get_chunk for citations
           ```python
-          results = search("safety requirements", limit=5)
+          results = await search("safety requirements", limit=5)
           for r in results:
-              chunk = get_chunk(r['chunk_id'])
+              chunk = await get_chunk(r['chunk_id'])
               print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
           ```
 
           ### Using llm() for classification
           ```python
-          content = get_document("Q1 Report")
+          content = await get_document("Q1 Report")
-          sentiment = llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
+          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
           print(sentiment)
           ```
 
@@ -314,7 +315,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '552'
+      - '514'
       content-type:
       - application/json
     parsed_body:
@@ -323,24 +324,24 @@ interactions:
         index: 0
         message:
           content: ''
-          reasoning: Need to list documents.
+          reasoning: Need to list docs.
           role: assistant
           tool_calls:
           - function:
-              arguments: '{"code":"docs=list_documents(limit=1000)\nprint(len(docs))\nfor d in docs:\n    print(d[''title''])"}'
+              arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
               name: execute_code
-            id: call_tumky965
+            id: call_stp0fimx
             index: 0
             type: function
-      created: 1771336699
+      created: 1771924497
-      id: chatcmpl-910
+      id: chatcmpl-750
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 56
+        completion_tokens: 44
-        prompt_tokens: 1562
+        prompt_tokens: 1623
-        total_tokens: 1618
+        total_tokens: 1667
     status:
       code: 200
       message: OK
@@ -353,7 +354,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '7612'
+      - '7759'
       content-type:
       - application/json
       host:
@@ -366,32 +367,33 @@ interactions:
 
           IMPORTANT: You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside the execute_code tool - you cannot access them any other way. Always execute code to answer questions; do not just describe what code would do.
 
-          CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just use them directly:
+          CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
-          - search("query")  ✓ CORRECT
+          - results = await search("query")  ✓ CORRECT
           - from haiku.rag import search  ✗ WRONG - will fail
+          - results = search("query")  ✗ WRONG - must use await
 
-          You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly, no imports needed):
+          You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly with `await`, no imports needed):
 
           ## Available Functions
 
-          ### search(query, limit=10) -> list[dict]
+          ### await search(query, limit=10) -> list[dict]
           Search the knowledge base using hybrid search (vector + full-text).
           Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings
 
-          ### list_documents(limit=10, offset=0) -> list[dict]
+          ### await list_documents(limit=10, offset=0) -> list[dict]
           List available documents in the knowledge base.
           Returns list of dicts with keys: id, title, uri, created_at
 
-          ### get_document(id_or_title) -> str | None
+          ### await get_document(id_or_title) -> str | None
           Get the full text content of a document by ID, title, or URI.
           Returns the document content as a string, or None if not found.
 
-          ### get_chunk(chunk_id) -> dict | None
+          ### await get_chunk(chunk_id) -> dict | None
           Get a specific chunk by its ID (from search results).
           Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
           Use this to retrieve full chunk details and metadata for citation.
 
-          ### llm(prompt) -> str
+          ### await llm(prompt) -> str
           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.
@@ -408,7 +410,7 @@ interactions:
 
           ## Available Python Features
 
-          The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, try/except, and the `json` module.
+          The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `sorted()`/`.sort(key=...)`, try/except, and the `json` module.
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
@@ -417,21 +419,21 @@ interactions:
           ## Strategy Guide
 
           1. **Explore First**: Start by listing documents or searching to understand what's available. Document names may differ from filenames (e.g., "tbmed593.pdf" might be stored as "TB MED 593" or similar).
-          2. **If get_document returns None**: Use `list_documents()` to see actual document titles, or `search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
           3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
           4. **Use print() Liberally**: The sandbox captures stdout - print intermediate results to see what you're working with.
           5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
-          6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use llm().
+          6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
           7. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
 
           ## Example Patterns
 
           ### Counting documents matching a condition
           ```python
-          docs = list_documents(limit=100)
+          docs = await list_documents(limit=100)
           count = 0
           for doc in docs:
-              content = get_document(doc['id'])
+              content = await get_document(doc['id'])
               if content and 'keyword' in content.lower():
                   count += 1
                   print(f"Found in: {doc['title']}")
@@ -441,9 +443,9 @@ interactions:
           ### Extracting data with llm()
           ```python
           numbers = []
-          results = search("financial data", limit=20)
+          results = await search("financial data", limit=20)
           for r in results:
-              extracted = llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
               for part in extracted.split(','):
                   part = part.strip().replace(',', '')
                   if part.isdigit():
@@ -454,16 +456,16 @@ interactions:
 
           ### Using search results with get_chunk for citations
           ```python
-          results = search("safety requirements", limit=5)
+          results = await search("safety requirements", limit=5)
           for r in results:
-              chunk = get_chunk(r['chunk_id'])
+              chunk = await get_chunk(r['chunk_id'])
               print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
           ```
 
           ### Using llm() for classification
           ```python
-          content = get_document("Q1 Report")
+          content = await get_document("Q1 Report")
-          sentiment = llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
+          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
           print(sentiment)
           ```
 
@@ -490,18 +492,17 @@ interactions:
       - content: How many documents are in the database?
         role: user
       - content: null
-        reasoning: Need to list documents.
+        reasoning: Need to list docs.
         role: assistant
         tool_calls:
         - function:
-            arguments: '{"code":"docs=list_documents(limit=1000)\nprint(len(docs))\nfor d in docs:\n    print(d[''title''])"}'
+            arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
             name: execute_code
-          id: call_tumky965
+          id: call_stp0fimx
           type: function
-      - content: '{"code":"docs=list_documents(limit=1000)\nprint(len(docs))\nfor d in docs:\n    print(d[''title''])","stdout":"3\nDoc
+      - content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"3\n","stderr":"","success":true}'
-          1\nDoc 2\nDoc 3\n","stderr":"","success":true}'
         role: tool
-        tool_call_id: call_tumky965
+        tool_call_id: call_stp0fimx
       model: gpt-oss
       reasoning_effort: low
       stream: false
@@ -552,7 +553,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '517'
+      - '416'
       content-type:
       - application/json
     parsed_body:
@@ -560,18 +561,17 @@ interactions:
       - finish_reason: stop
         index: 0
         message:
-          content: '{"answer":"There are 3 documents in the database.","program":"docs = list_documents(limit=1000)\nprint(f\"Number
+          content: '{"answer":"There are 3 documents in the database.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
-            of documents: {len(docs)}\")\nfor doc in docs:\n    print(f\"- {doc[''title'']} (ID: {doc[''id'']})\")"}'
           role: assistant
-      created: 1771336701
+      created: 1771924498
-      id: chatcmpl-57
+      id: chatcmpl-945
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 82
+        completion_tokens: 38
-        prompt_tokens: 1684
+        prompt_tokens: 1709
-        total_tokens: 1766
+        total_tokens: 1747
     status:
       code: 200
       message: OK

tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_with_filter.yaml

@@ -128,7 +128,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '7077'
+      - '7319'
       content-type:
       - application/json
       host:
@@ -141,32 +141,33 @@ interactions:
 
           IMPORTANT: You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside the execute_code tool - you cannot access them any other way. Always execute code to answer questions; do not just describe what code would do.
 
-          CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just use them directly:
+          CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
-          - search("query")  ✓ CORRECT
+          - results = await search("query")  ✓ CORRECT
           - from haiku.rag import search  ✗ WRONG - will fail
+          - results = search("query")  ✗ WRONG - must use await
 
-          You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly, no imports needed):
+          You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly with `await`, no imports needed):
 
           ## Available Functions
 
-          ### search(query, limit=10) -> list[dict]
+          ### await search(query, limit=10) -> list[dict]
           Search the knowledge base using hybrid search (vector + full-text).
           Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings
 
-          ### list_documents(limit=10, offset=0) -> list[dict]
+          ### await list_documents(limit=10, offset=0) -> list[dict]
           List available documents in the knowledge base.
           Returns list of dicts with keys: id, title, uri, created_at
 
-          ### get_document(id_or_title) -> str | None
+          ### await get_document(id_or_title) -> str | None
           Get the full text content of a document by ID, title, or URI.
           Returns the document content as a string, or None if not found.
 
-          ### get_chunk(chunk_id) -> dict | None
+          ### await get_chunk(chunk_id) -> dict | None
           Get a specific chunk by its ID (from search results).
           Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
           Use this to retrieve full chunk details and metadata for citation.
 
-          ### llm(prompt) -> str
+          ### await llm(prompt) -> str
           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.
@@ -183,7 +184,7 @@ interactions:
 
           ## Available Python Features
 
-          The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, try/except, and the `json` module.
+          The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `sorted()`/`.sort(key=...)`, try/except, and the `json` module.
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
@@ -192,21 +193,21 @@ interactions:
           ## Strategy Guide
 
           1. **Explore First**: Start by listing documents or searching to understand what's available. Document names may differ from filenames (e.g., "tbmed593.pdf" might be stored as "TB MED 593" or similar).
-          2. **If get_document returns None**: Use `list_documents()` to see actual document titles, or `search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
           3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
           4. **Use print() Liberally**: The sandbox captures stdout - print intermediate results to see what you're working with.
           5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
-          6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use llm().
+          6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
           7. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
 
           ## Example Patterns
 
           ### Counting documents matching a condition
           ```python
-          docs = list_documents(limit=100)
+          docs = await list_documents(limit=100)
           count = 0
           for doc in docs:
-              content = get_document(doc['id'])
+              content = await get_document(doc['id'])
               if content and 'keyword' in content.lower():
                   count += 1
                   print(f"Found in: {doc['title']}")
@@ -216,9 +217,9 @@ interactions:
           ### Extracting data with llm()
           ```python
           numbers = []
-          results = search("financial data", limit=20)
+          results = await search("financial data", limit=20)
           for r in results:
-              extracted = llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
               for part in extracted.split(','):
                   part = part.strip().replace(',', '')
                   if part.isdigit():
@@ -229,16 +230,16 @@ interactions:
 
           ### Using search results with get_chunk for citations
           ```python
-          results = search("safety requirements", limit=5)
+          results = await search("safety requirements", limit=5)
           for r in results:
-              chunk = get_chunk(r['chunk_id'])
+              chunk = await get_chunk(r['chunk_id'])
               print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
           ```
 
           ### Using llm() for classification
           ```python
-          content = get_document("Q1 Report")
+          content = await get_document("Q1 Report")
-          sentiment = llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
+          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
           print(sentiment)
           ```
 
@@ -314,7 +315,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '526'
+      - '547'
       content-type:
       - application/json
     parsed_body:
@@ -323,24 +324,24 @@ interactions:
         index: 0
         message:
           content: ''
-          reasoning: Need to list_documents.
+          reasoning: Need to list documents.
           role: assistant
           tool_calls:
           - function:
-              arguments: '{"code":"docs=list_documents(limit=1000);print(len(docs)); print(docs[:3])"}'
+              arguments: '{"code":"docs = await list_documents(limit=1000, offset=0)\nprint(len(docs))\nprint(docs[:3])"}'
               name: execute_code
-            id: call_rpfcy176
+            id: call_pu34e9fx
             index: 0
             type: function
-      created: 1771336720
+      created: 1771924517
-      id: chatcmpl-390
+      id: chatcmpl-236
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 47
+        completion_tokens: 56
-        prompt_tokens: 1560
+        prompt_tokens: 1621
-        total_tokens: 1607
+        total_tokens: 1677
     status:
       code: 200
       message: OK
@@ -353,7 +354,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '7655'
+      - '7939'
       content-type:
       - application/json
       host:
@@ -366,32 +367,33 @@ interactions:
 
           IMPORTANT: You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside the execute_code tool - you cannot access them any other way. Always execute code to answer questions; do not just describe what code would do.
 
-          CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just use them directly:
+          CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
-          - search("query")  ✓ CORRECT
+          - results = await search("query")  ✓ CORRECT
           - from haiku.rag import search  ✗ WRONG - will fail
+          - results = search("query")  ✗ WRONG - must use await
 
-          You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly, no imports needed):
+          You have access to a sandboxed Python interpreter with these haiku.rag functions (use them directly with `await`, no imports needed):
 
           ## Available Functions
 
-          ### search(query, limit=10) -> list[dict]
+          ### await search(query, limit=10) -> list[dict]
           Search the knowledge base using hybrid search (vector + full-text).
           Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings
 
-          ### list_documents(limit=10, offset=0) -> list[dict]
+          ### await list_documents(limit=10, offset=0) -> list[dict]
           List available documents in the knowledge base.
           Returns list of dicts with keys: id, title, uri, created_at
 
-          ### get_document(id_or_title) -> str | None
+          ### await get_document(id_or_title) -> str | None
           Get the full text content of a document by ID, title, or URI.
           Returns the document content as a string, or None if not found.
 
-          ### get_chunk(chunk_id) -> dict | None
+          ### await get_chunk(chunk_id) -> dict | None
           Get a specific chunk by its ID (from search results).
           Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
           Use this to retrieve full chunk details and metadata for citation.
 
-          ### llm(prompt) -> str
+          ### await llm(prompt) -> str
           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.
@@ -408,7 +410,7 @@ interactions:
 
           ## Available Python Features
 
-          The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, try/except, and the `json` module.
+          The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `sorted()`/`.sort(key=...)`, try/except, and the `json` module.
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
@@ -417,21 +419,21 @@ interactions:
           ## Strategy Guide
 
           1. **Explore First**: Start by listing documents or searching to understand what's available. Document names may differ from filenames (e.g., "tbmed593.pdf" might be stored as "TB MED 593" or similar).
-          2. **If get_document returns None**: Use `list_documents()` to see actual document titles, or `search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
           3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
           4. **Use print() Liberally**: The sandbox captures stdout - print intermediate results to see what you're working with.
           5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
-          6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use llm().
+          6. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
           7. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
 
           ## Example Patterns
 
           ### Counting documents matching a condition
           ```python
-          docs = list_documents(limit=100)
+          docs = await list_documents(limit=100)
           count = 0
           for doc in docs:
-              content = get_document(doc['id'])
+              content = await get_document(doc['id'])
               if content and 'keyword' in content.lower():
                   count += 1
                   print(f"Found in: {doc['title']}")
@@ -441,9 +443,9 @@ interactions:
           ### Extracting data with llm()
           ```python
           numbers = []
-          results = search("financial data", limit=20)
+          results = await search("financial data", limit=20)
           for r in results:
-              extracted = llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
               for part in extracted.split(','):
                   part = part.strip().replace(',', '')
                   if part.isdigit():
@@ -454,16 +456,16 @@ interactions:
 
           ### Using search results with get_chunk for citations
           ```python
-          results = search("safety requirements", limit=5)
+          results = await search("safety requirements", limit=5)
           for r in results:
-              chunk = get_chunk(r['chunk_id'])
+              chunk = await get_chunk(r['chunk_id'])
               print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
           ```
 
           ### Using llm() for classification
           ```python
-          content = get_document("Q1 Report")
+          content = await get_document("Q1 Report")
-          sentiment = llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
+          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
           print(sentiment)
           ```
 
@@ -490,18 +492,18 @@ interactions:
       - content: How many documents are available?
         role: user
       - content: null
-        reasoning: Need to list_documents.
+        reasoning: Need to list documents.
         role: assistant
         tool_calls:
         - function:
-            arguments: '{"code":"docs=list_documents(limit=1000);print(len(docs)); print(docs[:3])"}'
+            arguments: '{"code":"docs = await list_documents(limit=1000, offset=0)\nprint(len(docs))\nprint(docs[:3])"}'
             name: execute_code
-          id: call_rpfcy176
+          id: call_pu34e9fx
           type: function
-      - content: '{"code":"docs=list_documents(limit=1000);print(len(docs)); print(docs[:3])","stdout":"1\n[{''id'': ''50c659ee-70c5-4854-a418-12ae2abdcf40'',
+      - content: '{"code":"docs = await list_documents(limit=1000, offset=0)\nprint(len(docs))\nprint(docs[:3])","stdout":"1\n[{''id'':
-          ''title'': ''Cats'', ''uri'': None, ''created_at'': ''2026-02-17 15:58:36.831575''}]\n","stderr":"","success":true}'
+          ''90fce4e5-0d68-4736-acaf-d075e3ed4822'', ''title'': ''Cats'', ''uri'': None, ''created_at'': ''2026-02-24 11:15:14.007278''}]\n","stderr":"","success":true}'
         role: tool
-        tool_call_id: call_rpfcy176
+        tool_call_id: call_pu34e9fx
       model: gpt-oss
       reasoning_effort: low
       stream: false
@@ -552,7 +554,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '453'
+      - '704'
       content-type:
       - application/json
     parsed_body:
@@ -560,18 +562,20 @@ interactions:
       - finish_reason: stop
         index: 0
         message:
-          content: '{"answer":"There is 1 document available in the knowledge base.","program":"docs=list_documents(limit=1000)\nprint(f\"Number
+          content: '{"answer":"There is 1 document available in the knowledge base.","program":"# List all documents\nimport
-            of documents: {len(docs)}\")"}'
+            json\nasync def main():\n    docs = await list_documents(limit=1000, offset=0)\n    print(\"Number of documents:\",
+            len(docs))\n    # Optional: print first few document titles for reference\n    for d in docs[:5]:\n        print(f\"ID:
+            {d[''id'']}, Title: {d[''title'']}\")\n\nawait main()"}'
           role: assistant
-      created: 1771336721
+      created: 1771924519
-      id: chatcmpl-97
+      id: chatcmpl-487
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 46
+        completion_tokens: 117
-        prompt_tokens: 1713
+        prompt_tokens: 1793
-        total_tokens: 1759
+        total_tokens: 1910
     status:
       code: 200
       message: OK

tests/cassettes/test_sandbox/TestSandboxLLM.test_llm_function.yaml

@@ -25,7 +25,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '307'
+      - '311'
       content-type:
       - application/json
     parsed_body:
@@ -34,17 +34,17 @@ interactions:
         index: 0
         message:
           content: '4'
-          reasoning: Answer 4.
+          reasoning: Just reply 4.
           role: assistant
-      created: 1771338974
+      created: 1771924616
-      id: chatcmpl-199
+      id: chatcmpl-525
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 15
+        completion_tokens: 16
         prompt_tokens: 81
-        total_tokens: 96
+        total_tokens: 97
     status:
       code: 200
       message: OK

uv.lock

@@ -180,7 +180,7 @@ wheels = [
 
 [[package]]
 name = "anthropic"
-version = "0.79.0"
+version = "0.83.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "anyio" },
@@ -192,9 +192,9 @@ dependencies = [
     { name = "sniffio" },
     { name = "typing-extensions" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/15/b1/91aea3f8fd180d01d133d931a167a78a3737b3fd39ccef2ae8d6619c24fd/anthropic-0.79.0.tar.gz", hash = "sha256:8707aafb3b1176ed6c13e2b1c9fb3efddce90d17aee5d8b83a86c70dcdcca871", size = 509825, upload-time = "2026-02-07T18:06:18.388Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/db/e5/02cd2919ec327b24234abb73082e6ab84c451182cc3cc60681af700f4c63/anthropic-0.83.0.tar.gz", hash = "sha256:a8732c68b41869266c3034541a31a29d8be0f8cd0a714f9edce3128b351eceb4", size = 534058, upload-time = "2026-02-19T19:26:38.904Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/95/b2/cc0b8e874a18d7da50b0fda8c99e4ac123f23bf47b471827c5f6f3e4a767/anthropic-0.79.0-py3-none-any.whl", hash = "sha256:04cbd473b6bbda4ca2e41dd670fe2f829a911530f01697d0a1e37321eb75f3cf", size = 405918, upload-time = "2026-02-07T18:06:20.246Z" },
+    { url = "https://files.pythonhosted.org/packages/5f/75/b9d58e4e2a4b1fc3e75ffbab978f999baf8b7c4ba9f96e60edb918ba386b/anthropic-0.83.0-py3-none-any.whl", hash = "sha256:f069ef508c73b8f9152e8850830d92bd5ef185645dbacf234bb213344a274810", size = 456991, upload-time = "2026-02-19T19:26:40.114Z" },
 ]
 
 [[package]]
@@ -1092,7 +1092,7 @@ name = "ffmpeg-python"
 version = "0.2.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "future" },
+    { name = "future", marker = "python_full_version < '3.14'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/dd/5e/d5f9105d59c1325759d838af4e973695081fbbc97182baf73afc78dec266/ffmpeg-python-0.2.0.tar.gz", hash = "sha256:65225db34627c578ef0e11c8b1eb528bb35e024752f6f10b78c011f6f64c4127", size = 21543, upload-time = "2019-07-06T00:19:08.989Z" }
 wheels = [
@@ -1306,30 +1306,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c4/ab/09169d5a4612a5f92490806649ac8d41e3ec9129c636754575b3553f4ea4/googleapis_common_protos-1.72.0-py3-none-any.whl", hash = "sha256:4299c5a82d5ae1a9702ada957347726b167f9f8d1fc352477702a1e851ff4038", size = 297515, upload-time = "2025-11-06T18:29:13.14Z" },
 ]
 
-[[package]]
-name = "griffe"
-version = "2.0.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "griffecli" },
-    { name = "griffelib" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/8b/94/ee21d41e7eb4f823b94603b9d40f86d3c7fde80eacc2c3c71845476dddaa/griffe-2.0.0-py3-none-any.whl", hash = "sha256:5418081135a391c3e6e757a7f3f156f1a1a746cc7b4023868ff7d5e2f9a980aa", size = 5214, upload-time = "2026-02-09T19:09:44.105Z" },
-]
-
-[[package]]
-name = "griffecli"
-version = "2.0.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "colorama" },
-    { name = "griffelib" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e6/ed/d93f7a447bbf7a935d8868e9617cbe1cadf9ee9ee6bd275d3040fbf93d60/griffecli-2.0.0-py3-none-any.whl", hash = "sha256:9f7cd9ee9b21d55e91689358978d2385ae65c22f307a63fb3269acf3f21e643d", size = 9345, upload-time = "2026-02-09T19:09:42.554Z" },
-]
-
 [[package]]
 name = "griffelib"
 version = "2.0.0"
@@ -1978,14 +1954,14 @@ name = "langchain-core"
 version = "1.2.13"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "jsonpatch" },
+    { name = "jsonpatch", marker = "python_full_version < '3.14'" },
-    { name = "langsmith" },
+    { name = "langsmith", marker = "python_full_version < '3.14'" },
-    { name = "packaging" },
+    { name = "packaging", marker = "python_full_version < '3.14'" },
-    { name = "pydantic" },
+    { name = "pydantic", marker = "python_full_version < '3.14'" },
-    { name = "pyyaml" },
+    { name = "pyyaml", marker = "python_full_version < '3.14'" },
-    { name = "tenacity" },
+    { name = "tenacity", marker = "python_full_version < '3.14'" },
-    { name = "typing-extensions" },
+    { name = "typing-extensions", marker = "python_full_version < '3.14'" },
-    { name = "uuid-utils" },
+    { name = "uuid-utils", marker = "python_full_version < '3.14'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/fb/bb/c501ca60556c11ac80d1454bdcac63cb33583ce4e64fc4535ad5a7d5c6ba/langchain_core-1.2.13.tar.gz", hash = "sha256:d2773d0d0130a356378db9a858cfeef64c3d64bc03722f1d4d6c40eb46fdf01b", size = 831612, upload-time = "2026-02-15T07:45:57.014Z" }
 wheels = [
@@ -1997,7 +1973,7 @@ name = "langchain-text-splitters"
 version = "1.1.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "langchain-core" },
+    { name = "langchain-core", marker = "python_full_version < '3.14'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/41/42/c178dcdc157b473330eb7cc30883ea69b8ec60078c7b85e2d521054c4831/langchain_text_splitters-1.1.0.tar.gz", hash = "sha256:75e58acb7585dc9508f3cd9d9809cb14751283226c2d6e21fb3a9ae57582ca22", size = 272230, upload-time = "2025-12-14T01:15:38.659Z" }
 wheels = [
@@ -2009,15 +1985,15 @@ name = "langsmith"
 version = "0.7.3"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "httpx" },
+    { name = "httpx", marker = "python_full_version < '3.14'" },
-    { name = "orjson", marker = "platform_python_implementation != 'PyPy'" },
+    { name = "orjson", marker = "python_full_version < '3.14' and platform_python_implementation != 'PyPy'" },
-    { name = "packaging" },
+    { name = "packaging", marker = "python_full_version < '3.14'" },
-    { name = "pydantic" },
+    { name = "pydantic", marker = "python_full_version < '3.14'" },
-    { name = "requests" },
+    { name = "requests", marker = "python_full_version < '3.14'" },
-    { name = "requests-toolbelt" },
+    { name = "requests-toolbelt", marker = "python_full_version < '3.14'" },
-    { name = "uuid-utils" },
+    { name = "uuid-utils", marker = "python_full_version < '3.14'" },
-    { name = "xxhash" },
+    { name = "xxhash", marker = "python_full_version < '3.14'" },
-    { name = "zstandard" },
+    { name = "zstandard", marker = "python_full_version < '3.14'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/8d/bc/8172fefad4f2da888a6d564a27d1fb7d4dbf3c640899c2b40c46235cbe98/langsmith-0.7.3.tar.gz", hash = "sha256:0223b97021af62d2cf53c8a378a27bd22e90a7327e45b353e0069ae60d5d6f9e", size = 988575, upload-time = "2026-02-13T23:25:32.916Z" }
 wheels = [
@@ -3626,20 +3602,20 @@ email = [
 
 [[package]]
 name = "pydantic-ai-slim"
-version = "1.60.0"
+version = "1.63.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "genai-prices" },
-    { name = "griffe" },
+    { name = "griffelib" },
     { name = "httpx" },
     { name = "opentelemetry-api" },
     { name = "pydantic" },
     { name = "pydantic-graph" },
     { name = "typing-inspection" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/16/97/f73f439f3d415d43f38250f76852121188d0ef6114ec702e84e7d69c301d/pydantic_ai_slim-1.60.0.tar.gz", hash = "sha256:12ba3e6ef933fcb9fc6a307dbdaa43ca15bbc1b8ec77521afd1b7a526d12330f", size = 418839, upload-time = "2026-02-17T00:33:29.672Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/da/6d/2b5c0c60b42e6af49830f6a09b5d38fecdb1f20d9659152691eba95613b4/pydantic_ai_slim-1.63.0.tar.gz", hash = "sha256:9377afecdfe4bc17f5c9ed72c758e460703ac5876931aa2f18ace8ac0e69312a", size = 426862, upload-time = "2026-02-23T17:56:36.215Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/e3/40/8cb494a4d2ba62b5f92ae8bc79a2abbbf8509cb692edd4bc841695187780/pydantic_ai_slim-1.60.0-py3-none-any.whl", hash = "sha256:6865188a225a2979c82bb022a299d438d805d258c6d3f9810f7fe4e3c86af80a", size = 546410, upload-time = "2026-02-17T00:33:21.901Z" },
+    { url = "https://files.pythonhosted.org/packages/f2/ca/c4e39eec1cff5a294b64313a8a959b38d326819e0f0a41f48e61ce019a22/pydantic_ai_slim-1.63.0-py3-none-any.whl", hash = "sha256:ed393b0f871b748171f65bec5191c3025b5abb8a4fc616afee17eb9dc2dfa15d", size = 554190, upload-time = "2026-02-23T17:56:29.533Z" },
 ]
 
 [package.optional-dependencies]
@@ -3683,7 +3659,7 @@ vertexai = [
     { name = "requests" },
 ]
 voyageai = [
-    { name = "voyageai" },
+    { name = "voyageai", marker = "python_full_version < '3.14'" },
 ]
 
 [[package]]
@@ -3759,7 +3735,7 @@ wheels = [
 
 [[package]]
 name = "pydantic-evals"
-version = "1.60.0"
+version = "1.63.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "anyio" },
@@ -3769,14 +3745,14 @@ dependencies = [
     { name = "pyyaml" },
     { name = "rich" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/ec/2c/bed606a726b09adc9ee414bdb919ffe499edf7f8c631ba03b0ff3aa34435/pydantic_evals-1.60.0.tar.gz", hash = "sha256:ae3edd6667075acd8ef04c0d6fffb1ebe72c37ff077295fdbd6319e59284580b", size = 54214, upload-time = "2026-02-17T00:33:31.697Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/99/43/21b6ddf65b56f7401c344f98e4e6258a02d2868c8a52a8b79c0e0e701029/pydantic_evals-1.63.0.tar.gz", hash = "sha256:eed56a7192e07c8be8cf16e53bb2ef652b4f7f7b8527650ac45fde865a4ecf9d", size = 56365, upload-time = "2026-02-23T17:56:37.71Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/fe/a2/60790b2c971f6ce78fea2db2b65800489982c341ad9aa6db750070a93dcd/pydantic_evals-1.60.0-py3-none-any.whl", hash = "sha256:7a7414535002cae63ba0d0d9b15c6e72c28252cf120290b18372c9852c91fcfe", size = 65278, upload-time = "2026-02-17T00:33:23.411Z" },
+    { url = "https://files.pythonhosted.org/packages/9b/f2/7174ad6abca2457e35a1b902ca4fa78aa8ee72e4ec2e9cd5dc8904014ec9/pydantic_evals-1.63.0-py3-none-any.whl", hash = "sha256:2e92a3af579a5670b2babf2044081d0ef99ab5a9ef141972616d71fd7e5bfd0e", size = 67279, upload-time = "2026-02-23T17:56:31.008Z" },
 ]
 
 [[package]]
 name = "pydantic-graph"
-version = "1.60.0"
+version = "1.63.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "httpx" },
@@ -3784,9 +3760,9 @@ dependencies = [
     { name = "pydantic" },
     { name = "typing-inspection" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/f9/e6/1cae7cd39ab29f2eebc87c0a82c7ffcdfbe88492d2fb4afaaad91534e1ff/pydantic_graph-1.60.0.tar.gz", hash = "sha256:9710e457c2f8c113fd63629f05174e45bdca917d90c69ec8cf558649f995505f", size = 58492, upload-time = "2026-02-17T00:33:32.66Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/7a/c8/aa3cb56552562b799f31e9de291c8bd88306308cfc9647d220dfff2bea18/pydantic_graph-1.63.0.tar.gz", hash = "sha256:5fd98bb22fa6181f0357a6ffad38a3214af12868bd46492d6456c5db434466b4", size = 58528, upload-time = "2026-02-23T17:56:39.118Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/d8/6c/ce1c0eca77c6efbf7c7168c05a244c2756bde876a52575f750471d522024/pydantic_graph-1.60.0-py3-none-any.whl", hash = "sha256:741fa1e48424b0def86079a01100ad0652e75882f0352cd157232b75ace468a5", size = 72345, upload-time = "2026-02-17T00:33:25.077Z" },
+    { url = "https://files.pythonhosted.org/packages/a4/1c/8dcae24c824dd2690fbe7375083b369b10ed1ad773e2b9d1122bb6c0fcdc/pydantic_graph-1.63.0-py3-none-any.whl", hash = "sha256:d9b7a387116f358d470c042b07aa08125cadfcfa8c08ef01769746a489aef0d5", size = 72353, upload-time = "2026-02-23T17:56:32.304Z" },
 ]
 
 [[package]]
@@ -4398,7 +4374,7 @@ name = "requests-toolbelt"
 version = "1.0.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "requests" },
+    { name = "requests", marker = "python_full_version < '3.14'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/f3/61/d7545dafb7ac2230c70d38d31cbfe4cc64f7144dc41f6e4e4b78ecd9f5bb/requests-toolbelt-1.0.0.tar.gz", hash = "sha256:7681a0a3d047012b5bdc0ee37d7f8f07ebe76ab08caeccfc3921ce23c88d5bc6", size = 206888, upload-time = "2023-05-01T04:11:33.229Z" }
 wheels = [
@@ -5371,16 +5347,16 @@ name = "voyageai"
 version = "0.3.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "aiohttp" },
+    { name = "aiohttp", marker = "python_full_version < '3.14'" },
-    { name = "aiolimiter" },
+    { name = "aiolimiter", marker = "python_full_version < '3.14'" },
-    { name = "ffmpeg-python" },
+    { name = "ffmpeg-python", marker = "python_full_version < '3.14'" },
-    { name = "langchain-text-splitters" },
+    { name = "langchain-text-splitters", marker = "python_full_version < '3.14'" },
     { name = "numpy", marker = "python_full_version < '3.14'" },
-    { name = "pillow" },
+    { name = "pillow", marker = "python_full_version < '3.14'" },
-    { name = "pydantic" },
+    { name = "pydantic", marker = "python_full_version < '3.14'" },
-    { name = "requests" },
+    { name = "requests", marker = "python_full_version < '3.14'" },
-    { name = "tenacity" },
+    { name = "tenacity", marker = "python_full_version < '3.14'" },
-    { name = "tokenizers" },
+    { name = "tokenizers", marker = "python_full_version < '3.14'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/94/16/1b46b3cd401e1717a68197c1fe336d7bb4e0a1833f8105e1738f5b1add05/voyageai-0.3.7.tar.gz", hash = "sha256:826cd97f97223f42b5babc5c459c9c80f3a8215ce5c0e007b0b276550f790d24", size = 26485, upload-time = "2025-12-16T18:43:05.26Z" }
 wheels = [