Merge pull request #295 from ggozad/fix/tool-calls

Use ToolOutput for structured output

haiku_rag_slim/haiku/rag/agents/qa/agent.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 
 from pydantic_ai import Agent
+from pydantic_ai.output import ToolOutput
 
 from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
 from haiku.rag.agents.research.models import (
@@ -59,8 +60,7 @@ class QuestionAnswerAgent:
         agent: Agent[_QARunDeps, RawSearchAnswer] = Agent(  # ty: ignore[invalid-assignment]
             model=get_model(self._model_config, self._config),
             deps_type=_QARunDeps,
-            output_type=RawSearchAnswer,
+            output_type=ToolOutput(RawSearchAnswer, max_retries=3),
-            output_retries=3,
             instructions=self._system_prompt,
             toolsets=[search_toolset],
             retries=3,

haiku_rag_slim/haiku/rag/agents/research/graph.py

@@ -68,10 +68,9 @@ async def _iterative_plan_logic(
 
     plan_agent: Agent[ResearchDependencies, IterativePlanResult] = Agent(  # type: ignore[assignment]
         model=get_model(model_config, config),
-        output_type=IterativePlanResult,
+        output_type=ToolOutput(IterativePlanResult, max_retries=3),
         instructions=effective_prompt,
         retries=3,
-        output_retries=3,
         deps_type=ResearchDependencies,
     )
 
@@ -219,10 +218,9 @@ def build_research_graph(
 
         agent: Agent[ResearchDependencies, ResearchReport] = Agent(  # type: ignore[assignment]
             model=get_model(model_config, config),
-            output_type=ResearchReport,
+            output_type=ToolOutput(ResearchReport, max_retries=3),
             instructions=synthesis_prompt,
             retries=3,
-            output_retries=3,
             deps_type=ResearchDependencies,
         )
 

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

@@ -1,4 +1,5 @@
 from pydantic_ai import Agent, RunContext
+from pydantic_ai.output import ToolOutput
 
 from haiku.rag.agents.rlm.dependencies import RLMDeps
 from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
@@ -25,7 +26,7 @@ def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
     agent: Agent[RLMDeps, RLMResult] = Agent(  # type: ignore[invalid-assignment]
         model,
         deps_type=RLMDeps,
-        output_type=RLMResult,
+        output_type=ToolOutput(RLMResult, max_retries=3),
         instructions=RLM_SYSTEM_PROMPT,
         retries=3,
     )

tests/agents/rlm/test_agent.py

@@ -2,6 +2,7 @@ from pathlib import Path
 
 import pytest
 from pydantic_ai import Agent
+from pydantic_ai.output import ToolOutput
 
 from haiku.rag.agents.rlm.agent import create_rlm_agent
 from haiku.rag.agents.rlm.dependencies import RLMDeps
@@ -19,7 +20,8 @@ class TestCreateRLMAgent:
         agent = create_rlm_agent(Config)
         assert isinstance(agent, Agent)
         assert agent.deps_type is RLMDeps
-        assert agent.output_type is RLMResult
+        assert isinstance(agent.output_type, ToolOutput)
+        assert agent.output_type.output is RLMResult
 
     def test_agent_has_execute_code_tool(self):
         agent = create_rlm_agent(Config)

tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_count_documents.yaml

@@ -128,7 +128,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '7325'
+      - '7704'
       content-type:
       - application/json
       host:
@@ -139,15 +139,13 @@ interactions:
       - content: |-
           You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
 
-          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.
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. 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 call them with `await`:
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
           - 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 with `await`, no imports needed):
-
           ## Available Functions
 
           ### await search(query, limit=10) -> list[dict]
@@ -167,6 +165,26 @@ interactions:
           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.
 
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
           ### 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
@@ -180,7 +198,7 @@ interactions:
           for doc in documents:
               print(doc['title'], len(doc['content']))
           ```
-          Check if it exists with: `if 'documents' in dir(): ...`
+          Check if it exists with: `try: documents ... except NameError: ...`
 
           ## Available Python Features
 
@@ -188,17 +206,15 @@ interactions:
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
-          For pattern matching or text extraction, use string methods (`str.split`, `str.find`, `str.startswith`, `in` operator) or the `llm()` function.
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
 
           ## 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).
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
-          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), 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.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
-          5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
+          5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
-          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
 
@@ -214,44 +230,37 @@ interactions:
           print(f"Total: {count}")
           ```
 
-          ### Extracting data with llm()
+          ### Extracting data with regex
           ```python
           numbers = []
           results = await search("financial data", limit=20)
           for r in results:
-              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
-              for part in extracted.split(','):
+              for a in amounts:
-                  part = part.strip().replace(',', '')
+                  numbers.append(int(a.replace(',', '')))
-                  if part.isdigit():
-                      numbers.append(int(part))
           if numbers:
               print(f"Average: {sum(numbers) / len(numbers)}")
           ```
 
-          ### Using search results with get_chunk for citations
+          ### Extracting tables from a document
           ```python
-          results = await search("safety requirements", limit=5)
+          docs = await list_documents(limit=10)
-          for r in results:
+          for d in docs:
-              chunk = await get_chunk(r['chunk_id'])
+              doc = await get_docling_document(d['id'])
-              print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
           ```
 
-          ### Using llm() for classification
-          ```python
-          content = await get_document("Q1 Report")
-          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
-          print(sentiment)
-          ```
-
-          ## Workflow
-
-          1. **ALWAYS start by using execute_code** to explore the knowledge base
-          2. Run multiple code blocks as needed to gather information
-          3. After collecting data, provide your final answer
-
           ## Output Format
 
-          CRITICAL: Your final response MUST be valid JSON matching this exact schema:
+          Your final response MUST be valid JSON matching this exact schema:
           ```json
           {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
           ```
@@ -261,7 +270,7 @@ interactions:
 
           Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
 
-          CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
         role: system
       - content: How many documents are in the database?
         role: user
@@ -315,7 +324,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '514'
+      - '516'
       content-type:
       - application/json
     parsed_body:
@@ -324,24 +333,24 @@ interactions:
         index: 0
         message:
           content: ''
-          reasoning: Need to list docs.
+          reasoning: Need list_documents.
           role: assistant
           tool_calls:
           - function:
               arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
               name: execute_code
-            id: call_stp0fimx
+            id: call_oyaoz18v
             index: 0
             type: function
-      created: 1771924497
+      created: 1772549310
-      id: chatcmpl-750
+      id: chatcmpl-325
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 44
+        completion_tokens: 43
-        prompt_tokens: 1623
+        prompt_tokens: 1730
-        total_tokens: 1667
+        total_tokens: 1773
     status:
       code: 200
       message: OK
@@ -354,7 +363,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '7759'
+      - '8140'
       content-type:
       - application/json
       host:
@@ -365,15 +374,13 @@ interactions:
       - content: |-
           You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
 
-          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.
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. 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 call them with `await`:
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
           - 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 with `await`, no imports needed):
-
           ## Available Functions
 
           ### await search(query, limit=10) -> list[dict]
@@ -393,6 +400,26 @@ interactions:
           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.
 
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
           ### 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
@@ -406,7 +433,7 @@ interactions:
           for doc in documents:
               print(doc['title'], len(doc['content']))
           ```
-          Check if it exists with: `if 'documents' in dir(): ...`
+          Check if it exists with: `try: documents ... except NameError: ...`
 
           ## Available Python Features
 
@@ -414,17 +441,15 @@ interactions:
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
-          For pattern matching or text extraction, use string methods (`str.split`, `str.find`, `str.startswith`, `in` operator) or the `llm()` function.
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
 
           ## 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).
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
-          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), 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.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
-          5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
+          5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
-          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
 
@@ -440,44 +465,37 @@ interactions:
           print(f"Total: {count}")
           ```
 
-          ### Extracting data with llm()
+          ### Extracting data with regex
           ```python
           numbers = []
           results = await search("financial data", limit=20)
           for r in results:
-              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
-              for part in extracted.split(','):
+              for a in amounts:
-                  part = part.strip().replace(',', '')
+                  numbers.append(int(a.replace(',', '')))
-                  if part.isdigit():
-                      numbers.append(int(part))
           if numbers:
               print(f"Average: {sum(numbers) / len(numbers)}")
           ```
 
-          ### Using search results with get_chunk for citations
+          ### Extracting tables from a document
           ```python
-          results = await search("safety requirements", limit=5)
+          docs = await list_documents(limit=10)
-          for r in results:
+          for d in docs:
-              chunk = await get_chunk(r['chunk_id'])
+              doc = await get_docling_document(d['id'])
-              print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
           ```
 
-          ### Using llm() for classification
-          ```python
-          content = await get_document("Q1 Report")
-          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
-          print(sentiment)
-          ```
-
-          ## Workflow
-
-          1. **ALWAYS start by using execute_code** to explore the knowledge base
-          2. Run multiple code blocks as needed to gather information
-          3. After collecting data, provide your final answer
-
           ## Output Format
 
-          CRITICAL: Your final response MUST be valid JSON matching this exact schema:
+          Your final response MUST be valid JSON matching this exact schema:
           ```json
           {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
           ```
@@ -487,22 +505,22 @@ interactions:
 
           Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
 
-          CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
         role: system
       - content: How many documents are in the database?
         role: user
       - content: null
-        reasoning: Need to list docs.
+        reasoning: Need list_documents.
         role: assistant
         tool_calls:
         - function:
             arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
             name: execute_code
-          id: call_stp0fimx
+          id: call_oyaoz18v
           type: function
       - content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"3\n","stderr":"","success":true}'
         role: tool
-        tool_call_id: call_stp0fimx
+        tool_call_id: call_oyaoz18v
       model: gpt-oss
       reasoning_effort: low
       stream: false
@@ -563,15 +581,270 @@ interactions:
         message:
           content: '{"answer":"There are 3 documents in the database.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
           role: assistant
-      created: 1771924498
+      created: 1772549311
-      id: chatcmpl-945
+      id: chatcmpl-670
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
         completion_tokens: 38
-        prompt_tokens: 1709
+        prompt_tokens: 1815
-        total_tokens: 1747
+        total_tokens: 1853
+    status:
+      code: 200
+      message: OK
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      content-length:
+      - '8432'
+      content-type:
+      - application/json
+      host:
+      - localhost:11434
+    method: POST
+    parsed_body:
+      messages:
+      - content: |-
+          You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
+
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do.
+
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
+          - results = await search("query")  ✓ CORRECT
+          - import search  ✗ WRONG - will fail
+          - results = search("query")  ✗ WRONG - must use await
+
+          ## Available Functions
+
+          ### 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
+
+          ### 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
+
+          ### 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.
+
+          ### 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.
+
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
+          ### 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.
+
+          ## 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: `try: documents ... except NameError: ...`
+
+          ## Available Python Features
+
+          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.
+
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
+
+          ## Strategy Guide
+
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), or `await search()` to find relevant content.
+          3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
+          5. **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 = await list_documents(limit=100)
+          count = 0
+          for doc in docs:
+              content = await get_document(doc['id'])
+              if content and 'keyword' in content.lower():
+                  count += 1
+                  print(f"Found in: {doc['title']}")
+          print(f"Total: {count}")
+          ```
+
+          ### Extracting data with regex
+          ```python
+          numbers = []
+          results = await search("financial data", limit=20)
+          for r in results:
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
+              for a in amounts:
+                  numbers.append(int(a.replace(',', '')))
+          if numbers:
+              print(f"Average: {sum(numbers) / len(numbers)}")
+          ```
+
+          ### Extracting tables from a document
+          ```python
+          docs = await list_documents(limit=10)
+          for d in docs:
+              doc = await get_docling_document(d['id'])
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
+          ```
+
+          ## Output Format
+
+          Your final response MUST be valid JSON matching this exact schema:
+          ```json
+          {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
+          ```
+
+          - `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
+          - `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
+
+          Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
+
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+        role: system
+      - content: How many documents are in the database?
+        role: user
+      - content: null
+        reasoning: Need list_documents.
+        role: assistant
+        tool_calls:
+        - function:
+            arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
+            name: execute_code
+          id: call_oyaoz18v
+          type: function
+      - content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"3\n","stderr":"","success":true}'
+        role: tool
+        tool_call_id: call_oyaoz18v
+      - content: '{"answer":"There are 3 documents in the database.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
+        role: assistant
+      - content: |-
+          Validation feedback:
+          Please include your response in a tool call.
+
+          Fix the errors and try again.
+        role: user
+      model: gpt-oss
+      reasoning_effort: low
+      stream: false
+      tool_choice: auto
+      tools:
+      - function:
+          description: |-
+            <summary>Execute Python code in a sandboxed interpreter.
+
+            The code has access to haiku.rag functions (search, list_documents,
+            get_document, get_chunk, llm).
+
+            Use print() to output results.</summary>
+            <returns>
+            <description>Structured result with success status, stdout, and stderr.</description>
+            </returns>
+          name: execute_code
+          parameters:
+            additionalProperties: false
+            properties:
+              code:
+                description: Python code to execute.
+                type: string
+            required:
+            - code
+            type: object
+          strict: true
+        type: function
+      - function:
+          description: Result from RLM agent execution.
+          name: final_result
+          parameters:
+            additionalProperties: false
+            properties:
+              answer:
+                description: The answer to the user's question
+                type: string
+              program:
+                description: The final consolidated program
+                type: string
+            required:
+            - answer
+            - program
+            type: object
+          strict: true
+        type: function
+    uri: http://localhost:11434/v1/chat/completions
+  response:
+    headers:
+      content-length:
+      - '610'
+      content-type:
+      - application/json
+    parsed_body:
+      choices:
+      - finish_reason: tool_calls
+        index: 0
+        message:
+          content: ''
+          reasoning: Need to return via tool call? We should use final_result.
+          role: assistant
+          tool_calls:
+          - function:
+              arguments: '{"answer":"There are 3 documents in the database.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
+              name: final_result
+            id: call_46d25765
+            index: 0
+            type: function
+      created: 1772549312
+      id: chatcmpl-412
+      model: gpt-oss
+      object: chat.completion
+      system_fingerprint: fp_ollama
+      usage:
+        completion_tokens: 63
+        prompt_tokens: 1865
+        total_tokens: 1928
     status:
       code: 200
       message: OK

tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_search_and_get_chunk.yaml

@@ -48,7 +48,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '7359'
+      - '7738'
       content-type:
       - application/json
       host:
@@ -59,15 +59,13 @@ interactions:
       - content: |-
           You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
 
-          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.
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. 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 call them with `await`:
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
           - 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 with `await`, no imports needed):
-
           ## Available Functions
 
           ### await search(query, limit=10) -> list[dict]
@@ -87,6 +85,26 @@ interactions:
           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.
 
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
           ### 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
@@ -100,7 +118,7 @@ interactions:
           for doc in documents:
               print(doc['title'], len(doc['content']))
           ```
-          Check if it exists with: `if 'documents' in dir(): ...`
+          Check if it exists with: `try: documents ... except NameError: ...`
 
           ## Available Python Features
 
@@ -108,17 +126,15 @@ interactions:
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
-          For pattern matching or text extraction, use string methods (`str.split`, `str.find`, `str.startswith`, `in` operator) or the `llm()` function.
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
 
           ## 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).
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
-          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), 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.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
-          5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
+          5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
-          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
 
@@ -134,44 +150,37 @@ interactions:
           print(f"Total: {count}")
           ```
 
-          ### Extracting data with llm()
+          ### Extracting data with regex
           ```python
           numbers = []
           results = await search("financial data", limit=20)
           for r in results:
-              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
-              for part in extracted.split(','):
+              for a in amounts:
-                  part = part.strip().replace(',', '')
+                  numbers.append(int(a.replace(',', '')))
-                  if part.isdigit():
-                      numbers.append(int(part))
           if numbers:
               print(f"Average: {sum(numbers) / len(numbers)}")
           ```
 
-          ### Using search results with get_chunk for citations
+          ### Extracting tables from a document
           ```python
-          results = await search("safety requirements", limit=5)
+          docs = await list_documents(limit=10)
-          for r in results:
+          for d in docs:
-              chunk = await get_chunk(r['chunk_id'])
+              doc = await get_docling_document(d['id'])
-              print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
           ```
 
-          ### Using llm() for classification
-          ```python
-          content = await get_document("Q1 Report")
-          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
-          print(sentiment)
-          ```
-
-          ## Workflow
-
-          1. **ALWAYS start by using execute_code** to explore the knowledge base
-          2. Run multiple code blocks as needed to gather information
-          3. After collecting data, provide your final answer
-
           ## Output Format
 
-          CRITICAL: Your final response MUST be valid JSON matching this exact schema:
+          Your final response MUST be valid JSON matching this exact schema:
           ```json
           {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
           ```
@@ -181,7 +190,7 @@ interactions:
 
           Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
 
-          CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
         role: system
       - content: Search for content about animals and tell me which document it came from.
         role: user
@@ -235,7 +244,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '537'
+      - '625'
       content-type:
       - application/json
     parsed_body:
@@ -244,24 +253,25 @@ interactions:
         index: 0
         message:
           content: ''
-          reasoning: We need to search for animals.
+          reasoning: Need to search for "animals".
           role: assistant
           tool_calls:
           - function:
-              arguments: '{"code":"results = await search(\"animals\", limit=10)\nprint(results)\n"}'
+              arguments: '{"code":"results = await search(\"animals\", limit=5)\nprint(len(results))\nfor r in results:\n    print(r[''document_title''],
+                r[''score''])\nprint(results[:5])"}'
               name: execute_code
-            id: call_pvjujvr9
+            id: call_4vtaz637
             index: 0
             type: function
-      created: 1771924521
+      created: 1772549356
-      id: chatcmpl-217
+      id: chatcmpl-763
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 46
+        completion_tokens: 72
-        prompt_tokens: 1629
+        prompt_tokens: 1736
-        total_tokens: 1675
+        total_tokens: 1808
     status:
       code: 200
       message: OK
@@ -314,7 +324,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '8119'
+      - '8713'
       content-type:
       - application/json
       host:
@@ -325,15 +335,13 @@ interactions:
       - content: |-
           You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
 
-          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.
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. 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 call them with `await`:
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
           - 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 with `await`, no imports needed):
-
           ## Available Functions
 
           ### await search(query, limit=10) -> list[dict]
@@ -353,6 +361,26 @@ interactions:
           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.
 
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
           ### 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
@@ -366,7 +394,7 @@ interactions:
           for doc in documents:
               print(doc['title'], len(doc['content']))
           ```
-          Check if it exists with: `if 'documents' in dir(): ...`
+          Check if it exists with: `try: documents ... except NameError: ...`
 
           ## Available Python Features
 
@@ -374,17 +402,15 @@ interactions:
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
-          For pattern matching or text extraction, use string methods (`str.split`, `str.find`, `str.startswith`, `in` operator) or the `llm()` function.
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
 
           ## 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).
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
-          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), 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.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
-          5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
+          5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
-          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
 
@@ -400,44 +426,37 @@ interactions:
           print(f"Total: {count}")
           ```
 
-          ### Extracting data with llm()
+          ### Extracting data with regex
           ```python
           numbers = []
           results = await search("financial data", limit=20)
           for r in results:
-              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
-              for part in extracted.split(','):
+              for a in amounts:
-                  part = part.strip().replace(',', '')
+                  numbers.append(int(a.replace(',', '')))
-                  if part.isdigit():
-                      numbers.append(int(part))
           if numbers:
               print(f"Average: {sum(numbers) / len(numbers)}")
           ```
 
-          ### Using search results with get_chunk for citations
+          ### Extracting tables from a document
           ```python
-          results = await search("safety requirements", limit=5)
+          docs = await list_documents(limit=10)
-          for r in results:
+          for d in docs:
-              chunk = await get_chunk(r['chunk_id'])
+              doc = await get_docling_document(d['id'])
-              print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
           ```
 
-          ### Using llm() for classification
-          ```python
-          content = await get_document("Q1 Report")
-          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
-          print(sentiment)
-          ```
-
-          ## Workflow
-
-          1. **ALWAYS start by using execute_code** to explore the knowledge base
-          2. Run multiple code blocks as needed to gather information
-          3. After collecting data, provide your final answer
-
           ## Output Format
 
-          CRITICAL: Your final response MUST be valid JSON matching this exact schema:
+          Your final response MUST be valid JSON matching this exact schema:
           ```json
           {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
           ```
@@ -447,25 +466,27 @@ interactions:
 
           Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
 
-          CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
         role: system
       - content: Search for content about animals and tell me which document it came from.
         role: user
       - content: null
-        reasoning: We need to search for animals.
+        reasoning: Need to search for "animals".
         role: assistant
         tool_calls:
         - function:
-            arguments: '{"code":"results = await search(\"animals\", limit=10)\nprint(results)\n"}'
+            arguments: '{"code":"results = await search(\"animals\", limit=5)\nprint(len(results))\nfor r in results:\n    print(r[''document_title''],
+              r[''score''])\nprint(results[:5])"}'
             name: execute_code
-          id: call_pvjujvr9
+          id: call_4vtaz637
           type: function
-      - content: '{"code":"results = await search(\"animals\", limit=10)\nprint(results)\n","stdout":"[{''chunk_id'': ''503004ae-51ea-4953-93c3-48acac4a929c'',
+      - content: '{"code":"results = await search(\"animals\", limit=5)\nprint(len(results))\nfor r in results:\n    print(r[''document_title''],
-          ''content'': ''The quick brown fox jumps over the lazy dog.'', ''document_id'': ''ee57cfe6-fe82-4162-afd9-f206002ae08e'',
+          r[''score''])\nprint(results[:5])","stdout":"1\nAnimal Facts 0.016393441706895828\n[{''chunk_id'': ''69018e55-d141-4f97-89bb-7d7a19ffc273'',
+          ''content'': ''The quick brown fox jumps over the lazy dog.'', ''document_id'': ''7a3553dd-19ff-4660-8ff1-23c7e9b3aa89'',
           ''document_title'': ''Animal Facts'', ''document_uri'': None, ''score'': 0.016393441706895828, ''page_numbers'':
           [], ''headings'': None}]\n","stderr":"","success":true}'
         role: tool
-        tool_call_id: call_pvjujvr9
+        tool_call_id: call_4vtaz637
       model: gpt-oss
       reasoning_effort: low
       stream: false
@@ -516,7 +537,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '888'
+      - '1314'
       content-type:
       - application/json
     parsed_body:
@@ -524,21 +545,298 @@ interactions:
       - finish_reason: stop
         index: 0
         message:
-          content: '{"answer":"The document containing content about animals is \"Animal Facts\" (document ID: ee57cfe6-fe82-4162-afd9-f206002ae08e).","program":"import
+          content: '{"answer":"The search result for the query “animals” came from the document titled **\"Animal Facts\"**.
-            asyncio\n\nasync def main():\n    results = await search(\"animals\", limit=10)\n    if results:\n        chunk
+            The relevant chunk excerpted is: “The quick brown fox jumps over the lazy dog.”.","program":"import asyncio\n\nasync
-            = results[0]\n        print(f\"Document: {chunk[''document_title'']} (ID: {chunk[''document_id'']})\")\n    else:\n        print(\"No
+            def main():\n    # Search for the term \"animals\" in the knowledge base\n    results = await search(\"animals\",
-            animal-related content found.\")\n\nasyncio.run(main())"}'
+            limit=5)\n    # Print the number of results found\n    print(f\"Found {len(results)} result(s).\")\n    # Output
-          reasoning: It found one chunk. We need to give answer with source. Provide program that searches and prints answer.
+            the source document for each result\n    for r in results:\n        print(f\"Chunk ID: {r[''chunk_id'']}\")\n        print(f\"Document
+            ID: {r[''document_id'']}\")\n        print(f\"Document Title: {r[''document_title'']}\")\n        print(f\"Score:
+            {r[''score'']:.6f}\")\n        print(f\"Excerpt: {r[''content''][:200]}\\n\")\n\nif __name__ == \"__main__\":\n    asyncio.run(main())"}'
+          reasoning: The only result is chunk about fox and dog from document Animal Facts. Provide answer and program that
+            runs search.
           role: assistant
-      created: 1771924525
+      created: 1772549361
-      id: chatcmpl-109
+      id: chatcmpl-58
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 168
+        completion_tokens: 265
-        prompt_tokens: 1822
+        prompt_tokens: 2001
-        total_tokens: 1990
+        total_tokens: 2266
+    status:
+      code: 200
+      message: OK
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      content-length:
+      - '9903'
+      content-type:
+      - application/json
+      host:
+      - localhost:11434
+    method: POST
+    parsed_body:
+      messages:
+      - content: |-
+          You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
+
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do.
+
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
+          - results = await search("query")  ✓ CORRECT
+          - import search  ✗ WRONG - will fail
+          - results = search("query")  ✗ WRONG - must use await
+
+          ## Available Functions
+
+          ### 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
+
+          ### 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
+
+          ### 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.
+
+          ### 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.
+
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
+          ### 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.
+
+          ## 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: `try: documents ... except NameError: ...`
+
+          ## Available Python Features
+
+          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.
+
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
+
+          ## Strategy Guide
+
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), or `await search()` to find relevant content.
+          3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
+          5. **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 = await list_documents(limit=100)
+          count = 0
+          for doc in docs:
+              content = await get_document(doc['id'])
+              if content and 'keyword' in content.lower():
+                  count += 1
+                  print(f"Found in: {doc['title']}")
+          print(f"Total: {count}")
+          ```
+
+          ### Extracting data with regex
+          ```python
+          numbers = []
+          results = await search("financial data", limit=20)
+          for r in results:
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
+              for a in amounts:
+                  numbers.append(int(a.replace(',', '')))
+          if numbers:
+              print(f"Average: {sum(numbers) / len(numbers)}")
+          ```
+
+          ### Extracting tables from a document
+          ```python
+          docs = await list_documents(limit=10)
+          for d in docs:
+              doc = await get_docling_document(d['id'])
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
+          ```
+
+          ## Output Format
+
+          Your final response MUST be valid JSON matching this exact schema:
+          ```json
+          {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
+          ```
+
+          - `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
+          - `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
+
+          Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
+
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+        role: system
+      - content: Search for content about animals and tell me which document it came from.
+        role: user
+      - content: null
+        reasoning: Need to search for "animals".
+        role: assistant
+        tool_calls:
+        - function:
+            arguments: '{"code":"results = await search(\"animals\", limit=5)\nprint(len(results))\nfor r in results:\n    print(r[''document_title''],
+              r[''score''])\nprint(results[:5])"}'
+            name: execute_code
+          id: call_4vtaz637
+          type: function
+      - content: '{"code":"results = await search(\"animals\", limit=5)\nprint(len(results))\nfor r in results:\n    print(r[''document_title''],
+          r[''score''])\nprint(results[:5])","stdout":"1\nAnimal Facts 0.016393441706895828\n[{''chunk_id'': ''69018e55-d141-4f97-89bb-7d7a19ffc273'',
+          ''content'': ''The quick brown fox jumps over the lazy dog.'', ''document_id'': ''7a3553dd-19ff-4660-8ff1-23c7e9b3aa89'',
+          ''document_title'': ''Animal Facts'', ''document_uri'': None, ''score'': 0.016393441706895828, ''page_numbers'':
+          [], ''headings'': None}]\n","stderr":"","success":true}'
+        role: tool
+        tool_call_id: call_4vtaz637
+      - content: '{"answer":"The search result for the query “animals” came from the document titled **\"Animal Facts\"**.
+          The relevant chunk excerpted is: “The quick brown fox jumps over the lazy dog.”.","program":"import asyncio\n\nasync
+          def main():\n    # Search for the term \"animals\" in the knowledge base\n    results = await search(\"animals\",
+          limit=5)\n    # Print the number of results found\n    print(f\"Found {len(results)} result(s).\")\n    # Output
+          the source document for each result\n    for r in results:\n        print(f\"Chunk ID: {r[''chunk_id'']}\")\n        print(f\"Document
+          ID: {r[''document_id'']}\")\n        print(f\"Document Title: {r[''document_title'']}\")\n        print(f\"Score:
+          {r[''score'']:.6f}\")\n        print(f\"Excerpt: {r[''content''][:200]}\\n\")\n\nif __name__ == \"__main__\":\n    asyncio.run(main())"}'
+        reasoning: The only result is chunk about fox and dog from document Animal Facts. Provide answer and program that
+          runs search.
+        role: assistant
+      - content: |-
+          Validation feedback:
+          Please include your response in a tool call.
+
+          Fix the errors and try again.
+        role: user
+      model: gpt-oss
+      reasoning_effort: low
+      stream: false
+      tool_choice: auto
+      tools:
+      - function:
+          description: |-
+            <summary>Execute Python code in a sandboxed interpreter.
+
+            The code has access to haiku.rag functions (search, list_documents,
+            get_document, get_chunk, llm).
+
+            Use print() to output results.</summary>
+            <returns>
+            <description>Structured result with success status, stdout, and stderr.</description>
+            </returns>
+          name: execute_code
+          parameters:
+            additionalProperties: false
+            properties:
+              code:
+                description: Python code to execute.
+                type: string
+            required:
+            - code
+            type: object
+          strict: true
+        type: function
+      - function:
+          description: Result from RLM agent execution.
+          name: final_result
+          parameters:
+            additionalProperties: false
+            properties:
+              answer:
+                description: The answer to the user's question
+                type: string
+              program:
+                description: The final consolidated program
+                type: string
+            required:
+            - answer
+            - program
+            type: object
+          strict: true
+        type: function
+    uri: http://localhost:11434/v1/chat/completions
+  response:
+    headers:
+      content-length:
+      - '1233'
+      content-type:
+      - application/json
+    parsed_body:
+      choices:
+      - finish_reason: tool_calls
+        index: 0
+        message:
+          content: ''
+          reasoning: Need to output JSON via tool call, not plain. Use function final_result.
+          role: assistant
+          tool_calls:
+          - function:
+              arguments: '{"answer":"The search result for the query “animals” came from the document titled *\"Animal Facts\"*.
+                The relevant chunk excerpted is: “The quick brown fox jumps over the lazy dog.”.","program":"import asyncio\n\nasync
+                def main():\n    results = await search(\"animals\", limit=5)\n    print(f\"Found {len(results)} result(s).\")\n    for
+                r in results:\n        print(f\"Chunk ID: {r[''chunk_id'']}\")\n        print(f\"Document ID: {r[''document_id'']}\")\n        print(f\"Document
+                Title: {r[''document_title'']}\")\n        print(f\"Score: {r[''score'']:.6f}\")\n        print(f\"Excerpt:
+                {r[''content''][:200]}\\n\")\n\nif __name__ == \"__main__\":\n    asyncio.run(main())"}'
+              name: final_result
+            id: call_y5fedbyn
+            index: 0
+            type: function
+      created: 1772549367
+      id: chatcmpl-895
+      model: gpt-oss
+      object: chat.completion
+      system_fingerprint: fp_ollama
+      usage:
+        completion_tokens: 231
+        prompt_tokens: 2247
+        total_tokens: 2478
     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:
-      - '7319'
+      - '7698'
       content-type:
       - application/json
       host:
@@ -139,15 +139,13 @@ interactions:
       - content: |-
           You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
 
-          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.
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. 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 call them with `await`:
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
           - 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 with `await`, no imports needed):
-
           ## Available Functions
 
           ### await search(query, limit=10) -> list[dict]
@@ -167,6 +165,26 @@ interactions:
           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.
 
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
           ### 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
@@ -180,7 +198,7 @@ interactions:
           for doc in documents:
               print(doc['title'], len(doc['content']))
           ```
-          Check if it exists with: `if 'documents' in dir(): ...`
+          Check if it exists with: `try: documents ... except NameError: ...`
 
           ## Available Python Features
 
@@ -188,17 +206,15 @@ interactions:
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
-          For pattern matching or text extraction, use string methods (`str.split`, `str.find`, `str.startswith`, `in` operator) or the `llm()` function.
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
 
           ## 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).
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
-          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), 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.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
-          5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
+          5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
-          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
 
@@ -214,44 +230,37 @@ interactions:
           print(f"Total: {count}")
           ```
 
-          ### Extracting data with llm()
+          ### Extracting data with regex
           ```python
           numbers = []
           results = await search("financial data", limit=20)
           for r in results:
-              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
-              for part in extracted.split(','):
+              for a in amounts:
-                  part = part.strip().replace(',', '')
+                  numbers.append(int(a.replace(',', '')))
-                  if part.isdigit():
-                      numbers.append(int(part))
           if numbers:
               print(f"Average: {sum(numbers) / len(numbers)}")
           ```
 
-          ### Using search results with get_chunk for citations
+          ### Extracting tables from a document
           ```python
-          results = await search("safety requirements", limit=5)
+          docs = await list_documents(limit=10)
-          for r in results:
+          for d in docs:
-              chunk = await get_chunk(r['chunk_id'])
+              doc = await get_docling_document(d['id'])
-              print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
           ```
 
-          ### Using llm() for classification
-          ```python
-          content = await get_document("Q1 Report")
-          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
-          print(sentiment)
-          ```
-
-          ## Workflow
-
-          1. **ALWAYS start by using execute_code** to explore the knowledge base
-          2. Run multiple code blocks as needed to gather information
-          3. After collecting data, provide your final answer
-
           ## Output Format
 
-          CRITICAL: Your final response MUST be valid JSON matching this exact schema:
+          Your final response MUST be valid JSON matching this exact schema:
           ```json
           {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
           ```
@@ -261,7 +270,7 @@ interactions:
 
           Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
 
-          CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
         role: system
       - content: How many documents are available?
         role: user
@@ -315,7 +324,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '547'
+      - '519'
       content-type:
       - application/json
     parsed_body:
@@ -328,20 +337,20 @@ interactions:
           role: assistant
           tool_calls:
           - function:
-              arguments: '{"code":"docs = await list_documents(limit=1000, offset=0)\nprint(len(docs))\nprint(docs[:3])"}'
+              arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
               name: execute_code
-            id: call_pu34e9fx
+            id: call_qqoyb2of
             index: 0
             type: function
-      created: 1771924517
+      created: 1772548188
-      id: chatcmpl-236
+      id: chatcmpl-356
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 56
+        completion_tokens: 44
-        prompt_tokens: 1621
+        prompt_tokens: 1728
-        total_tokens: 1677
+        total_tokens: 1772
     status:
       code: 200
       message: OK
@@ -354,7 +363,7 @@ interactions:
       connection:
       - keep-alive
       content-length:
-      - '7939'
+      - '8137'
       content-type:
       - application/json
       host:
@@ -365,15 +374,13 @@ interactions:
       - content: |-
           You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
 
-          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.
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. 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 call them with `await`:
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
           - 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 with `await`, no imports needed):
-
           ## Available Functions
 
           ### await search(query, limit=10) -> list[dict]
@@ -393,6 +400,26 @@ interactions:
           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.
 
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
           ### 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
@@ -406,7 +433,7 @@ interactions:
           for doc in documents:
               print(doc['title'], len(doc['content']))
           ```
-          Check if it exists with: `if 'documents' in dir(): ...`
+          Check if it exists with: `try: documents ... except NameError: ...`
 
           ## Available Python Features
 
@@ -414,17 +441,15 @@ interactions:
 
           Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
 
-          For pattern matching or text extraction, use string methods (`str.split`, `str.find`, `str.startswith`, `in` operator) or the `llm()` function.
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
 
           ## 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).
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
-          2. **If get_document returns None**: Use `await list_documents()` to see actual document titles, or `await search()` to find relevant content.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), 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.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
-          5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and data structures.
+          5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
-          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
 
@@ -440,44 +465,37 @@ interactions:
           print(f"Total: {count}")
           ```
 
-          ### Extracting data with llm()
+          ### Extracting data with regex
           ```python
           numbers = []
           results = await search("financial data", limit=20)
           for r in results:
-              extracted = await llm(f"Extract all dollar amounts from this text as a comma-separated list of numbers (no $ signs): {r['content']}")
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
-              for part in extracted.split(','):
+              for a in amounts:
-                  part = part.strip().replace(',', '')
+                  numbers.append(int(a.replace(',', '')))
-                  if part.isdigit():
-                      numbers.append(int(part))
           if numbers:
               print(f"Average: {sum(numbers) / len(numbers)}")
           ```
 
-          ### Using search results with get_chunk for citations
+          ### Extracting tables from a document
           ```python
-          results = await search("safety requirements", limit=5)
+          docs = await list_documents(limit=10)
-          for r in results:
+          for d in docs:
-              chunk = await get_chunk(r['chunk_id'])
+              doc = await get_docling_document(d['id'])
-              print(f"From '{chunk['document_title']}', page {chunk['page_numbers']}: {chunk['content'][:100]}")
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
           ```
 
-          ### Using llm() for classification
-          ```python
-          content = await get_document("Q1 Report")
-          sentiment = await llm(f"Classify the sentiment as positive, negative, or mixed: {content}")
-          print(sentiment)
-          ```
-
-          ## Workflow
-
-          1. **ALWAYS start by using execute_code** to explore the knowledge base
-          2. Run multiple code blocks as needed to gather information
-          3. After collecting data, provide your final answer
-
           ## Output Format
 
-          CRITICAL: Your final response MUST be valid JSON matching this exact schema:
+          Your final response MUST be valid JSON matching this exact schema:
           ```json
           {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
           ```
@@ -487,7 +505,7 @@ interactions:
 
           Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
 
-          CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
         role: system
       - content: How many documents are available?
         role: user
@@ -496,14 +514,13 @@ interactions:
         role: assistant
         tool_calls:
         - function:
-            arguments: '{"code":"docs = await list_documents(limit=1000, offset=0)\nprint(len(docs))\nprint(docs[:3])"}'
+            arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
             name: execute_code
-          id: call_pu34e9fx
+          id: call_qqoyb2of
           type: function
-      - content: '{"code":"docs = await list_documents(limit=1000, offset=0)\nprint(len(docs))\nprint(docs[:3])","stdout":"1\n[{''id'':
+      - content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"1\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_pu34e9fx
+        tool_call_id: call_qqoyb2of
       model: gpt-oss
       reasoning_effort: low
       stream: false
@@ -554,7 +571,7 @@ interactions:
   response:
     headers:
       content-length:
-      - '704'
+      - '430'
       content-type:
       - application/json
     parsed_body:
@@ -562,20 +579,272 @@ interactions:
       - finish_reason: stop
         index: 0
         message:
-          content: '{"answer":"There is 1 document available in the knowledge base.","program":"# List all documents\nimport
+          content: '{"answer":"There is 1 document available in the knowledge base.","program":"docs = await list_documents(limit=1000)\nprint(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: 1771924519
+      created: 1772548189
-      id: chatcmpl-487
+      id: chatcmpl-173
       model: gpt-oss
       object: chat.completion
       system_fingerprint: fp_ollama
       usage:
-        completion_tokens: 117
+        completion_tokens: 40
-        prompt_tokens: 1793
+        prompt_tokens: 1814
-        total_tokens: 1910
+        total_tokens: 1854
+    status:
+      code: 200
+      message: OK
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      content-length:
+      - '8443'
+      content-type:
+      - application/json
+      host:
+      - localhost:11434
+    method: POST
+    parsed_body:
+      messages:
+      - content: |-
+          You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
+
+          You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do.
+
+          Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
+          - results = await search("query")  ✓ CORRECT
+          - import search  ✗ WRONG - will fail
+          - results = search("query")  ✗ WRONG - must use await
+
+          ## Available Functions
+
+          ### 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
+
+          ### 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
+
+          ### 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.
+
+          ### 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.
+
+          ### await get_docling_document(document_id) -> dict | None
+          Get the full document structure as a dict (DoclingDocument format).
+          Use `list_documents()` or search results to get document IDs first.
+          - `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
+          - `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
+          - `pictures`: list of figures/images with metadata
+          - `pages`: page dimensions and metadata
+
+          ### await regex_findall(pattern, text) -> list[str]
+          Find all non-overlapping matches of a regular expression pattern in text.
+
+          ### await regex_sub(pattern, repl, text) -> str
+          Replace all occurrences of a regular expression pattern with a replacement string.
+
+          ### await regex_search(pattern, text) -> dict | None
+          Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
+
+          ### await regex_split(pattern, text) -> list[str]
+          Split text by a regular expression pattern.
+
+          ### 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.
+
+          ## 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: `try: documents ... except NameError: ...`
+
+          ## Available Python Features
+
+          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.
+
+          For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
+
+          ## Strategy Guide
+
+          1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
+          2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), or `await search()` to find relevant content.
+          3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
+          4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
+          5. **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 = await list_documents(limit=100)
+          count = 0
+          for doc in docs:
+              content = await get_document(doc['id'])
+              if content and 'keyword' in content.lower():
+                  count += 1
+                  print(f"Found in: {doc['title']}")
+          print(f"Total: {count}")
+          ```
+
+          ### Extracting data with regex
+          ```python
+          numbers = []
+          results = await search("financial data", limit=20)
+          for r in results:
+              amounts = await regex_findall(r'\$([\d,]+)', r['content'])
+              for a in amounts:
+                  numbers.append(int(a.replace(',', '')))
+          if numbers:
+              print(f"Average: {sum(numbers) / len(numbers)}")
+          ```
+
+          ### Extracting tables from a document
+          ```python
+          docs = await list_documents(limit=10)
+          for d in docs:
+              doc = await get_docling_document(d['id'])
+              if doc:
+                  tables = doc.get('tables', [])
+                  if tables:
+                      print(f"{d['title']}: {len(tables)} table(s)")
+                      for i, table in enumerate(tables):
+                          grid = table.get('data', {}).get('grid', [])
+                          for row in grid:
+                              cells = [cell.get('text', '') for cell in row]
+                              print(f"  Table {i}: {cells}")
+          ```
+
+          ## Output Format
+
+          Your final response MUST be valid JSON matching this exact schema:
+          ```json
+          {"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
+          ```
+
+          - `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
+          - `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
+
+          Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
+
+          You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
+        role: system
+      - content: How many documents are available?
+        role: user
+      - content: null
+        reasoning: Need to list documents.
+        role: assistant
+        tool_calls:
+        - function:
+            arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
+            name: execute_code
+          id: call_qqoyb2of
+          type: function
+      - content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"1\n","stderr":"","success":true}'
+        role: tool
+        tool_call_id: call_qqoyb2of
+      - content: '{"answer":"There is 1 document available in the knowledge base.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
+        role: assistant
+      - content: |-
+          Validation feedback:
+          Please include your response in a tool call.
+
+          Fix the errors and try again.
+        role: user
+      model: gpt-oss
+      reasoning_effort: low
+      stream: false
+      tool_choice: auto
+      tools:
+      - function:
+          description: |-
+            <summary>Execute Python code in a sandboxed interpreter.
+
+            The code has access to haiku.rag functions (search, list_documents,
+            get_document, get_chunk, llm).
+
+            Use print() to output results.</summary>
+            <returns>
+            <description>Structured result with success status, stdout, and stderr.</description>
+            </returns>
+          name: execute_code
+          parameters:
+            additionalProperties: false
+            properties:
+              code:
+                description: Python code to execute.
+                type: string
+            required:
+            - code
+            type: object
+          strict: true
+        type: function
+      - function:
+          description: Result from RLM agent execution.
+          name: final_result
+          parameters:
+            additionalProperties: false
+            properties:
+              answer:
+                description: The answer to the user's question
+                type: string
+              program:
+                description: The final consolidated program
+                type: string
+            required:
+            - answer
+            - program
+            type: object
+          strict: true
+        type: function
+    uri: http://localhost:11434/v1/chat/completions
+  response:
+    headers:
+      content-length:
+      - '620'
+      content-type:
+      - application/json
+    parsed_body:
+      choices:
+      - finish_reason: tool_calls
+        index: 0
+        message:
+          content: ''
+          reasoning: Need to output JSON in a tool call. Use final_result.
+          role: assistant
+          tool_calls:
+          - function:
+              arguments: '{"answer":"There is 1 document available in the knowledge base.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
+              name: final_result
+            id: call_hli4bq9m
+            index: 0
+            type: function
+      created: 1772548191
+      id: chatcmpl-603
+      model: gpt-oss
+      object: chat.completion
+      system_fingerprint: fp_ollama
+      usage:
+        completion_tokens: 65
+        prompt_tokens: 1865
+        total_tokens: 1930
     status:
       code: 200
       message: OK