danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Docker-based sandbox. Reused within a single rlm() call for latency

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-05 19:59:00 +01:00
parent 4241a4b09e
commit b426827fc2
No known key found for this signature in database
21 changed files with 974 additions and 1328 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
.github/workflows/test.yml vendored
View file

@ -74,3 +74,22 @@ jobs:
token: ${{ secrets.CODECOV_TOKEN }}
files: ./coverage.xml
fail_ci_if_error: false
test-docker-sandbox:
needs: [lint]
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: astral-sh/setup-uv@v4
with:
enable-cache: true
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: "pyproject.toml"
- name: Install dependencies
run: uv sync --all-extras
- name: Build Docker image
run: docker build -t haiku-rag-slim:test -f docker/Dockerfile.slim .
- name: Run Docker integration tests
run: uv run pytest tests/agents/rlm/test_sandbox.py -v

5
CHANGELOG.md
View file

@ -8,13 +8,16 @@
- Allows disabling OCR via config when running docling-serve in read-only containers
- **RLM Agent (Recursive Language Model)**: New agent for complex analytical tasks via sandboxed Python code execution
- Solves problems traditional RAG can't handle: aggregation, computation, multi-document analysis
- Sandboxed execution with safe builtins and allowed imports (json, re, math, statistics, etc.)
- Docker-based sandbox with full Python environment (no import restrictions)
- Container reuse within a single `rlm()` call for reduced latency
- Available functions: `search()`, `list_documents()`, `get_document()`, `get_docling_document()`, `llm()`
- Pre-loaded documents support via `documents` variable
- Context filter for scoping searches without LLM control
- New `client.rlm(question)` method on HaikuRAG client
- New `haiku-rag rlm` CLI command
- New `rlm_question` MCP tool
- New config options: `docker_image`, `docker_memory_limit`
- **CI**: Docker sandbox integration tests run in GitHub Actions
### Fixed

50
docs/rlm.md
View file

@ -132,19 +132,9 @@ for doc in documents:
Each document dict has keys: `id`, `title`, `uri`, `content`
## Allowed Imports
## Imports
The following standard library modules can be imported:
- `json` - JSON encoding/decoding
- `re` - Regular expressions
- `math` - Mathematical functions
- `statistics` - Statistical functions
- `collections` - Specialized containers
- `itertools` - Iterator utilities
- `functools` - Higher-order functions
- `datetime` - Date and time handling
- `typing` - Type hints
The sandbox runs in a Docker container with full Python available. Any module installed in the container image can be imported:
```python
import re
@ -161,15 +151,17 @@ for r in results:
print(Counter(error_types).most_common(10))
```
## Security
The default image (`ghcr.io/ggozad/haiku.rag-slim`) includes the Python standard library. Custom images can add additional packages like `pandas` or `numpy`.
The sandbox enforces several security measures:
## Docker Sandbox
- **Blocked builtins**: `eval`, `exec`, `compile`, `open`, `input`, `__import__`, `globals`, `locals`, `getattr`, `setattr`, `delattr`
- **Blocked imports**: `os`, `sys`, `subprocess`, `shutil`, `socket`, `requests`, `builtins`
- **Private attribute access blocked**: Cannot access `__dunder__` attributes (except common ones like `__init__`, `__str__`)
- **Execution timeout**: Code execution times out after configurable limit (default 60s)
Code executes in an isolated Docker container with:
- **Read-only database**: The LanceDB database is mounted read-only
- **Memory limits**: Configurable memory limit (default 512MB)
- **Execution timeout**: Code times out after configurable limit (default 60s)
- **Output truncation**: Large outputs are truncated to prevent memory issues
- **Container reuse**: Within a single `rlm()` call, the container stays warm for multiple code executions
## Context Filter
@ -201,4 +193,26 @@ rlm:
code_timeout: 60.0 # Max seconds for code execution
max_tool_calls: 20 # Max execute_code calls per question
max_output_chars: 50000 # Truncate output after this many chars
docker_image: "ghcr.io/ggozad/haiku.rag-slim:latest" # Container image
docker_memory_limit: "512m" # Container memory limit
```
### Custom Docker Image
To add additional Python packages, create a custom Dockerfile:
```dockerfile
FROM ghcr.io/ggozad/haiku.rag-slim:latest
RUN pip install pandas numpy
```
Build and configure:
```bash
docker build -t my-rlm-image .
```
```yaml
rlm:
docker_image: "my-rlm-image"
```

6
haiku_rag_slim/haiku/rag/agents/rlm/__init__.py
View file

@ -1,16 +1,16 @@
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
from haiku.rag.agents.rlm.docker_sandbox import DockerSandbox, SandboxResult
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.agents.rlm.prompts import RLM_SYSTEM_PROMPT
from haiku.rag.agents.rlm.sandbox import REPLEnvironment, REPLResult
__all__ = [
"CodeExecution",
"DockerSandbox",
"RLMContext",
"RLMDeps",
"RLMResult",
"RLM_SYSTEM_PROMPT",
"REPLEnvironment",
"REPLResult",
"SandboxResult",
"create_rlm_agent",
]

21
haiku_rag_slim/haiku/rag/agents/rlm/agent.py
View file

@ -3,24 +3,9 @@ from pydantic_ai import Agent, RunContext
from haiku.rag.agents.rlm.dependencies import RLMDeps
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.agents.rlm.prompts import RLM_SYSTEM_PROMPT
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import get_model
_repl_cache: dict[int, REPLEnvironment] = {}
def _get_or_create_repl(ctx) -> REPLEnvironment:
"""Get or create a REPL environment for this context."""
key = id(ctx.deps)
if key not in _repl_cache:
_repl_cache[key] = REPLEnvironment(
client=ctx.deps.client,
config=ctx.deps.config.rlm,
context=ctx.deps.context,
)
return _repl_cache[key]
def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
"""Create an RLM agent with code execution capability.
@ -54,7 +39,7 @@ def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
modules (json, re, collections, math, statistics, itertools,
functools, datetime, typing).
Use print() to output results. Variables persist between executions.
Use print() to output results.
Args:
code: Python code to execute.
@ -62,9 +47,7 @@ def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
Returns:
Structured result with success status, stdout, and stderr.
"""
repl = _get_or_create_repl(ctx)
result = await repl.execute_async(code)
result = await ctx.deps.sandbox.execute(code)
execution = CodeExecution(
code=code,

2
haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py
View file

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING
from haiku.rag.store.models import Document, SearchResult
if TYPE_CHECKING:
from haiku.rag.agents.rlm.docker_sandbox import DockerSandbox
from haiku.rag.agents.rlm.models import CodeExecution
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
@ -25,4 +26,5 @@ class RLMDeps:
client: "HaikuRAG"
config: "AppConfig"
sandbox: "DockerSandbox"
context: RLMContext = field(default_factory=RLMContext)

207
haiku_rag_slim/haiku/rag/agents/rlm/docker_sandbox.py Normal file
View file

@ -0,0 +1,207 @@
"""Docker-based sandboxed execution."""
import asyncio
import json
import os
import subprocess
import sys
from dataclasses import dataclass
from typing import TYPE_CHECKING
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.config.models import RLMConfig
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
@dataclass
class SandboxResult:
"""Result of executing code in the sandbox."""
stdout: str
stderr: str
success: bool
class DockerSandbox:
"""Execute code in a persistent Docker container.
Use as an async context manager to manage container lifecycle:
async with DockerSandbox(client, config, context) as sandbox:
result = await sandbox.execute("print('hello')")
result = await sandbox.execute("print('world')")
"""
DEFAULT_IMAGE = "ghcr.io/ggozad/haiku.rag-slim:latest"
def __init__(
self,
client: "HaikuRAG",
config: RLMConfig,
context: RLMContext,
image: str | None = None,
):
self.haiku_client = client
self.config = config
self.context = context
self.image = image or self.DEFAULT_IMAGE
self._process: subprocess.Popen[bytes] | None = None
def _build_docker_cmd(self) -> list[str]:
"""Build the docker run command."""
db_path = str(self.haiku_client.store.db_path)
env_list = ["-e", "HAIKU_DB_PATH=/data/db.lancedb"]
if self.context.filter:
env_list.extend(["-e", f"HAIKU_FILTER={self.context.filter}"])
ollama_host = os.environ.get("OLLAMA_HOST", "")
ollama_base_url = os.environ.get("OLLAMA_BASE_URL", "")
if sys.platform == "darwin":
if not ollama_host or "localhost" in ollama_host:
ollama_host = "http://host.docker.internal:11434"
if not ollama_base_url or "localhost" in ollama_base_url:
ollama_base_url = "http://host.docker.internal:11434"
if ollama_host:
env_list.extend(["-e", f"OLLAMA_HOST={ollama_host}"])
if ollama_base_url:
env_list.extend(["-e", f"OLLAMA_BASE_URL={ollama_base_url}"])
for key in [
"ANTHROPIC_API_KEY",
"OPENAI_API_KEY",
"VOYAGE_API_KEY",
"COHERE_API_KEY",
]:
if value := os.environ.get(key):
env_list.extend(["-e", f"{key}={value}"])
return [
"docker",
"run",
"--rm",
"-i",
"-v",
f"{db_path}:/data/db.lancedb:ro",
f"--memory={self.config.docker_memory_limit}",
"--network=host",
*env_list,
self.image,
"python",
"-m",
"haiku.rag.agents.rlm.runner",
]
async def __aenter__(self) -> "DockerSandbox":
"""Start the container."""
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self._start_container)
return self
async def __aexit__(
self, exc_type: object, exc_val: object, exc_tb: object
) -> None:
"""Stop the container."""
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self._stop_container)
def _start_container(self) -> None:
"""Start the persistent container process."""
if self._process is not None:
return
cmd = self._build_docker_cmd()
self._process = subprocess.Popen(
cmd,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
def _stop_container(self) -> None:
"""Stop the container process."""
if self._process is None:
return
try:
if self._process.stdin:
self._process.stdin.close()
self._process.terminate()
self._process.wait(timeout=5)
except subprocess.TimeoutExpired:
self._process.kill()
self._process.wait()
finally:
self._process = None
async def execute(self, code: str) -> SandboxResult:
"""Execute code in the container."""
if self._process is None:
return SandboxResult(
stdout="",
stderr="Container not started. Use 'async with' context manager.",
success=False,
)
loop = asyncio.get_running_loop()
return await loop.run_in_executor(None, self._execute_sync, code)
def _execute_sync(self, code: str) -> SandboxResult:
"""Send code to container and read result."""
assert self._process is not None and self._process.stdin is not None
try:
message = json.dumps({"code": code})
length_line = f"{len(message)}\n".encode()
self._process.stdin.write(length_line)
self._process.stdin.write(message.encode())
self._process.stdin.flush()
if self._process.stdout is None:
return SandboxResult(
stdout="", stderr="No stdout from container.", success=False
)
length_line = self._process.stdout.readline()
if not length_line:
stderr = ""
if self._process.stderr:
stderr = self._process.stderr.read().decode()
return SandboxResult(
stdout="",
stderr=stderr or "Container closed unexpectedly.",
success=False,
)
length = int(length_line.strip())
response = self._process.stdout.read(length).decode()
result_data = json.loads(response)
return SandboxResult(
stdout=result_data.get("stdout", ""),
stderr=result_data.get("stderr", ""),
success=result_data.get("success", False),
)
except subprocess.TimeoutExpired:
return SandboxResult(
stdout="",
stderr=f"Execution timed out after {self.config.code_timeout} seconds",
success=False,
)
except json.JSONDecodeError as e:
return SandboxResult(
stdout="",
stderr=f"Invalid response from container: {e}",
success=False,
)
except Exception as e:
return SandboxResult(
stdout="",
stderr=f"Execution error: {e}",
success=False,
)

214
haiku_rag_slim/haiku/rag/agents/rlm/runner.py Normal file
View file

@ -0,0 +1,214 @@
"""Entry point for sandboxed code execution in Docker container."""
import asyncio
import json
import sys
import traceback
from io import StringIO
from typing import Any
def build_namespace(
client: Any, config: Any, context: Any, loop: asyncio.AbstractEventLoop
) -> dict[str, Any]:
"""Build execution namespace with haiku.rag functions injected."""
from haiku.rag.store.repositories.document import _escape_sql_string
def run_async(coro: Any) -> Any:
"""Run async coroutine from sync context using thread-safe scheduling."""
future = asyncio.run_coroutine_threadsafe(coro, loop)
return future.result(timeout=config.rlm.code_timeout)
def search(query: str, limit: int = 10) -> list[dict]:
async def _search() -> Any:
return await client.search(query, limit=limit, filter=context.filter)
results = run_async(_search())
context.search_results.extend(results)
return [
{
"chunk_id": r.chunk_id,
"content": r.content,
"document_id": r.document_id,
"document_title": r.document_title,
"document_uri": r.document_uri,
"score": r.score,
"page_numbers": r.page_numbers,
"headings": r.headings,
}
for r in results
]
def list_documents(limit: int = 10, offset: int = 0) -> list[dict]:
async def _list() -> Any:
return await client.list_documents(
limit=limit, offset=offset, filter=context.filter
)
docs = run_async(_list())
return [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"created_at": str(d.created_at),
}
for d in docs
]
def get_document(id_or_title: str) -> str | None:
async def _get() -> str | None:
doc = await client.get_document_by_id(id_or_title)
if doc:
return doc.content
safe_input = _escape_sql_string(id_or_title)
docs = await client.list_documents(filter=f"title = '{safe_input}'")
if docs and docs[0].id:
full_doc = await client.get_document_by_id(docs[0].id)
return full_doc.content if full_doc else None
docs = await client.list_documents(filter=f"uri = '{safe_input}'")
if docs and docs[0].id:
full_doc = await client.get_document_by_id(docs[0].id)
return full_doc.content if full_doc else None
return None
return run_async(_get())
def get_docling_document(id_or_title: str) -> Any:
async def _get() -> Any:
doc = await client.get_document_by_id(id_or_title)
if doc:
return doc.get_docling_document()
safe_input = _escape_sql_string(id_or_title)
docs = await client.list_documents(filter=f"title = '{safe_input}'")
if docs and docs[0].id:
full_doc = await client.get_document_by_id(docs[0].id)
return full_doc.get_docling_document() if full_doc else None
docs = await client.list_documents(filter=f"uri = '{safe_input}'")
if docs and docs[0].id:
full_doc = await client.get_document_by_id(docs[0].id)
return full_doc.get_docling_document() if full_doc else None
return None
return run_async(_get())
def llm(prompt: str) -> str:
async def _llm() -> str:
from pydantic_ai import Agent
from haiku.rag.utils import get_model
model = get_model(config.rlm.model, config)
agent: Agent[None, str] = Agent(model, output_type=str)
result = await agent.run(prompt)
return result.output
return run_async(_llm())
namespace: dict[str, Any] = {
"search": search,
"list_documents": list_documents,
"get_document": get_document,
"get_docling_document": get_docling_document,
"llm": llm,
}
if context.documents:
namespace["documents"] = [
{"id": d.id, "title": d.title, "uri": d.uri, "content": d.content}
for d in context.documents
]
return namespace
def execute_code(
code: str, namespace: dict[str, Any], max_output_chars: int
) -> dict[str, Any]:
"""Execute code and capture output."""
stdout_capture = StringIO()
original_stdout = sys.stdout
try:
sys.stdout = stdout_capture
exec(code, namespace)
stdout = stdout_capture.getvalue()
if len(stdout) > max_output_chars:
stdout = stdout[:max_output_chars] + "\n... (output truncated)"
return {
"success": True,
"stdout": stdout,
"stderr": "",
}
except Exception:
return {
"success": False,
"stdout": stdout_capture.getvalue(),
"stderr": traceback.format_exc(),
}
finally:
sys.stdout = original_stdout
def send_response(result: dict[str, Any]) -> None:
"""Send length-prefixed JSON response."""
response = json.dumps(result)
sys.stdout.write(f"{len(response)}\n")
sys.stdout.write(response)
sys.stdout.flush()
async def main() -> None:
"""Main entry point for container execution.
Runs a loop reading length-prefixed JSON messages and executing code.
"""
import concurrent.futures
import os
from pathlib import Path
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
config = get_config()
db_path = Path(os.environ.get("HAIKU_DB_PATH", "/data/db.lancedb"))
filter_expr = os.environ.get("HAIKU_FILTER")
context = RLMContext(filter=filter_expr)
max_output_chars = config.rlm.max_output_chars
loop = asyncio.get_running_loop()
async with HaikuRAG(db_path, config=config, read_only=True) as client:
namespace = build_namespace(client, config, context, loop)
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
while True:
# Read length-prefixed message
length_line = sys.stdin.readline()
if not length_line:
break
try:
length = int(length_line.strip())
message = sys.stdin.read(length)
request = json.loads(message)
code = request.get("code", "")
result = await loop.run_in_executor(
executor, execute_code, code, namespace, max_output_chars
)
send_response(result)
except (ValueError, json.JSONDecodeError) as e:
send_response(
{
"success": False,
"stdout": "",
"stderr": f"Invalid request: {e}",
}
)
if __name__ == "__main__":
asyncio.run(main())

414
haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py
View file

@ -1,414 +0,0 @@
import ast
import asyncio
import concurrent.futures
import sys
import traceback
from io import StringIO
from typing import TYPE_CHECKING, Any
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.config.models import RLMConfig
from haiku.rag.store.repositories.document import _escape_sql_string
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
class REPLResult:
"""Result of executing code in the REPL environment."""
def __init__(
self, stdout: str, stderr: str, success: bool, locals_: dict | None = None
):
self.stdout = stdout
self.stderr = stderr
self.success = success
self.locals = locals_ or {}
def __repr__(self) -> str:
return f"REPLResult(success={self.success}, stdout={self.stdout!r}, stderr={self.stderr!r})"
class REPLEnvironment:
"""Sandboxed Python execution environment with haiku.rag access."""
SAFE_BUILTINS: dict[str, Any] = {
"True": True,
"False": False,
"None": None,
"__build_class__": __builtins__["__build_class__"]
if isinstance(__builtins__, dict)
else getattr(__builtins__, "__build_class__"),
"abs": abs,
"all": all,
"any": any,
"ascii": ascii,
"bin": bin,
"bool": bool,
"bytearray": bytearray,
"bytes": bytes,
"callable": callable,
"chr": chr,
"complex": complex,
"dict": dict,
"divmod": divmod,
"enumerate": enumerate,
"filter": filter,
"float": float,
"format": format,
"frozenset": frozenset,
"hash": hash,
"hex": hex,
"id": id,
"int": int,
"isinstance": isinstance,
"issubclass": issubclass,
"iter": iter,
"len": len,
"list": list,
"map": map,
"max": max,
"min": min,
"next": next,
"object": object,
"oct": oct,
"ord": ord,
"pow": pow,
"print": print,
"range": range,
"repr": repr,
"reversed": reversed,
"round": round,
"set": set,
"slice": slice,
"sorted": sorted,
"str": str,
"sum": sum,
"tuple": tuple,
"type": (
lambda obj: type(obj)
), # Single-arg only, blocks type(name, bases, dict)
"zip": zip,
"Exception": Exception,
"ValueError": ValueError,
"TypeError": TypeError,
"KeyError": KeyError,
"IndexError": IndexError,
"AttributeError": AttributeError,
"RuntimeError": RuntimeError,
"StopIteration": StopIteration,
"ZeroDivisionError": ZeroDivisionError,
"AssertionError": AssertionError,
}
ALLOWED_IMPORTS = {
"json",
"re",
"collections",
"math",
"statistics",
"itertools",
"functools",
"datetime",
"typing",
}
def __init__(
self,
client: "HaikuRAG",
config: RLMConfig,
context: RLMContext,
event_loop: asyncio.AbstractEventLoop | None = None,
):
self.client = client
self.config = config
self.context = context
self._event_loop = event_loop
self._setup_namespace()
def _run_async_from_thread(self, coro):
"""Run async coroutine from a worker thread using run_coroutine_threadsafe."""
if self._event_loop is None:
raise RuntimeError("Event loop not set. Cannot call async functions.")
future = asyncio.run_coroutine_threadsafe(coro, self._event_loop)
return future.result(timeout=self.config.code_timeout)
def _setup_namespace(self) -> None:
"""Build execution namespace with haiku.rag functions."""
self.globals: dict[str, Any] = {
"__builtins__": dict(self.SAFE_BUILTINS),
"__name__": "__sandbox__",
"search": self._make_search(),
"list_documents": self._make_list_documents(),
"get_document": self._make_get_document(),
"get_docling_document": self._make_get_docling_document(),
"llm": self._make_llm(),
}
self.locals: dict[str, Any] = {}
if self.context.documents:
self.globals["documents"] = [
{"id": d.id, "title": d.title, "uri": d.uri, "content": d.content}
for d in self.context.documents
]
def _make_search(self):
"""Create sync search function that bridges to async client."""
def search(query: str, limit: int = 10) -> list[dict]:
async def _search():
return await self.client.search(
query, limit=limit, filter=self.context.filter
)
results = self._run_async_from_thread(_search())
self.context.search_results.extend(results)
return [
{
"chunk_id": r.chunk_id,
"content": r.content,
"document_id": r.document_id,
"document_title": r.document_title,
"document_uri": r.document_uri,
"score": r.score,
"page_numbers": r.page_numbers,
"headings": r.headings,
}
for r in results
]
return search
def _make_list_documents(self):
"""Create sync list_documents function."""
def list_documents(limit: int = 10, offset: int = 0) -> list[dict]:
async def _list():
return await self.client.list_documents(
limit=limit, offset=offset, filter=self.context.filter
)
docs = self._run_async_from_thread(_list())
return [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"created_at": str(d.created_at),
}
for d in docs
]
return list_documents
def _make_get_document(self):
"""Create sync get_document function that returns text content."""
def get_document(id_or_title: str) -> str | None:
async def _get():
doc = await self.client.get_document_by_id(id_or_title)
if doc:
return doc.content
safe_input = _escape_sql_string(id_or_title)
docs = await self.client.list_documents(
filter=f"title = '{safe_input}'"
)
if docs and docs[0].id:
full_doc = await self.client.get_document_by_id(docs[0].id)
return full_doc.content if full_doc else None
docs = await self.client.list_documents(filter=f"uri = '{safe_input}'")
if docs and docs[0].id:
full_doc = await self.client.get_document_by_id(docs[0].id)
return full_doc.content if full_doc else None
return None
return self._run_async_from_thread(_get())
return get_document
def _make_get_docling_document(self):
"""Create sync get_docling_document function that returns DoclingDocument."""
def get_docling_document(id_or_title: str):
async def _get():
doc = await self.client.get_document_by_id(id_or_title)
if doc:
return doc.get_docling_document()
safe_input = _escape_sql_string(id_or_title)
docs = await self.client.list_documents(
filter=f"title = '{safe_input}'"
)
if docs and docs[0].id:
full_doc = await self.client.get_document_by_id(docs[0].id)
return full_doc.get_docling_document() if full_doc else None
docs = await self.client.list_documents(filter=f"uri = '{safe_input}'")
if docs and docs[0].id:
full_doc = await self.client.get_document_by_id(docs[0].id)
return full_doc.get_docling_document() if full_doc else None
return None
return self._run_async_from_thread(_get())
return get_docling_document
def _make_llm(self):
"""Create sync llm function for plain LLM calls without RAG."""
def llm(prompt: str) -> str:
async def _llm():
from pydantic_ai import Agent
from haiku.rag.utils import get_model
model = get_model(self.config.model)
agent: Agent[None, str] = Agent(model, output_type=str)
result = await agent.run(prompt)
return result.output
return self._run_async_from_thread(_llm())
return llm
def _safe_import(
self,
name: str,
globals: dict | None = None,
locals: dict | None = None,
fromlist: tuple = (),
level: int = 0,
):
"""Import hook that only allows safe modules."""
base_module = name.split(".")[0]
if base_module not in self.ALLOWED_IMPORTS:
raise ImportError(f"Import of '{name}' is not allowed in sandbox")
import importlib
module = importlib.import_module(name)
if fromlist:
for attr in fromlist:
if not hasattr(module, attr):
raise ImportError(f"cannot import name '{attr}' from '{name}'")
return module
return module
def _validate_code(self, code: str) -> None:
"""Validate code AST for security issues."""
tree = ast.parse(code)
for node in ast.walk(tree):
if isinstance(node, ast.Attribute):
if node.attr.startswith("_") and node.attr not in (
"__init__",
"__str__",
"__repr__",
"__class__",
"__name__",
"__doc__",
"__dict__",
):
raise SecurityError(
f"Access to private/dunder attribute '{node.attr}' is not allowed"
)
# Block dictionary key access to dunder/private strings
# This prevents type.__dict__['__subclasses__'] attacks
if isinstance(node, ast.Subscript):
if isinstance(node.slice, ast.Constant):
if isinstance(
node.slice.value, str
) and node.slice.value.startswith("_"):
raise SecurityError(
f"Dictionary access to '{node.slice.value}' is not allowed"
)
def _execute_sync(self, code: str) -> REPLResult:
"""Internal synchronous execution - must be called from executor thread."""
stdout_capture = StringIO()
stderr_capture = StringIO()
original_stdout = sys.stdout
original_stderr = sys.stderr
try:
self._validate_code(code)
except SyntaxError as e:
return REPLResult(
stdout="",
stderr=f"SyntaxError: {e}",
success=False,
)
except SecurityError as e:
return REPLResult(
stdout="",
stderr=str(e),
success=False,
)
exec_globals = dict(self.globals)
exec_globals["__builtins__"] = dict(self.SAFE_BUILTINS)
exec_globals["__builtins__"]["__import__"] = self._safe_import
try:
sys.stdout = stdout_capture
sys.stderr = stderr_capture
exec(code, exec_globals, self.locals)
for key, value in self.locals.items():
if not key.startswith("_"):
self.globals[key] = value
stdout = stdout_capture.getvalue()
if len(stdout) > self.config.max_output_chars:
stdout = (
stdout[: self.config.max_output_chars] + "\n... (output truncated)"
)
return REPLResult(
stdout=stdout,
stderr=stderr_capture.getvalue(),
success=True,
locals_=dict(self.locals),
)
except Exception:
tb = traceback.format_exc()
return REPLResult(
stdout=stdout_capture.getvalue(),
stderr=tb,
success=False,
)
finally:
sys.stdout = original_stdout
sys.stderr = original_stderr
def execute(self, code: str) -> REPLResult:
"""Execute code in sandbox synchronously.
This method runs code directly in the current thread.
For async contexts, use execute_async() instead.
"""
return self._execute_sync(code)
async def execute_async(self, code: str) -> REPLResult:
"""Execute code in sandbox from async context.
Runs the synchronous code in a thread executor, allowing
sandbox functions to call back to async client methods.
"""
loop = asyncio.get_running_loop()
self._event_loop = loop
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
result = await asyncio.wait_for(
loop.run_in_executor(executor, self._execute_sync, code),
timeout=self.config.code_timeout,
)
return result
class SecurityError(Exception):
"""Raised when sandbox security is violated."""
pass

26
haiku_rag_slim/haiku/rag/client.py
View file

@ -1316,7 +1316,12 @@ class HaikuRAG:
Returns:
The answer as a string.
"""
from haiku.rag.agents.rlm import RLMContext, RLMDeps, create_rlm_agent
from haiku.rag.agents.rlm import (
DockerSandbox,
RLMContext,
RLMDeps,
create_rlm_agent,
)
context = RLMContext(filter=filter)
@ -1333,16 +1338,23 @@ class HaikuRAG:
loaded_docs.append(doc)
context.documents = loaded_docs if loaded_docs else None
deps = RLMDeps(
async with DockerSandbox(
client=self,
config=self._config,
config=self._config.rlm,
context=context,
)
image=self._config.rlm.docker_image,
) as sandbox:
deps = RLMDeps(
client=self,
config=self._config,
sandbox=sandbox,
context=context,
)
agent = create_rlm_agent(self._config)
result = await agent.run(question, deps=deps)
agent = create_rlm_agent(self._config)
result = await agent.run(question, deps=deps)
return result.output.answer
return result.output.answer
async def visualize_chunk(self, chunk: Chunk) -> list:
"""Render page images with bounding box highlights for a chunk.

2
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -105,6 +105,8 @@ class RLMConfig(BaseModel):
code_timeout: float = 60.0
max_output_chars: int = 50_000
max_tool_calls: int = 20
docker_image: str = "ghcr.io/ggozad/haiku.rag-slim:latest"
docker_memory_limit: str = "512m"
class PictureDescriptionConfig(BaseModel):

1
haiku_rag_slim/pyproject.toml
View file

@ -22,6 +22,7 @@ classifiers = [
]
dependencies = [
"docker>=7.1.0",
"docling-core==2.60.1",
"httpx>=0.28.1",
"jsonpatch>=1.33",

43
tests/agents/rlm/conftest.py
View file

@ -1,10 +1,40 @@
import os
import subprocess
from pathlib import Path
import pytest
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.agents.rlm.docker_sandbox import DockerSandbox
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
TEST_DOCKER_IMAGE = os.environ.get("HAIKU_TEST_DOCKER_IMAGE", "haiku-rag-slim:test")
@pytest.fixture(scope="session")
def test_docker_image():
"""Build and return the Docker image for testing."""
if os.environ.get("CI"):
return TEST_DOCKER_IMAGE
project_root = Path(__file__).parent.parent.parent.parent
dockerfile = project_root / "docker" / "Dockerfile.slim"
if not dockerfile.exists():
pytest.skip(f"Dockerfile.slim not found at {dockerfile}")
result = subprocess.run(
["docker", "build", "-t", TEST_DOCKER_IMAGE, "-f", str(dockerfile), "."],
cwd=project_root,
capture_output=True,
text=True,
)
if result.returncode != 0:
pytest.fail(f"Failed to build Docker image:\n{result.stderr}")
return TEST_DOCKER_IMAGE
@pytest.fixture
async def empty_client(temp_db_path):
@ -14,8 +44,11 @@ async def empty_client(temp_db_path):
@pytest.fixture
async def repl_env_empty(empty_client):
"""Create a REPL environment without documents."""
config = RLMConfig()
async def docker_sandbox(empty_client, test_docker_image):
"""Create a Docker sandbox for testing."""
config = RLMConfig(docker_image=test_docker_image)
context = RLMContext()
return REPLEnvironment(client=empty_client, config=config, context=context)
async with DockerSandbox(
client=empty_client, config=config, context=context, image=test_docker_image
) as sandbox:
yield sandbox

77
tests/agents/rlm/test_agent.py
View file

@ -4,7 +4,7 @@ import pytest
from pydantic_ai import Agent
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
from haiku.rag.agents.rlm.dependencies import RLMDeps
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.config import Config
@ -27,56 +27,8 @@ class TestCreateRLMAgent:
assert "execute_code" in tool_names
class TestExecuteCodeTool:
@pytest.mark.asyncio
async def test_execute_code_returns_structured_result(self, empty_client):
"""Test that execute_code tool produces structured CodeExecution output."""
from haiku.rag.agents.rlm.agent import _get_or_create_repl
context = RLMContext()
deps = RLMDeps(
client=empty_client,
config=Config,
context=context,
)
class MockCtx:
def __init__(self, deps):
self.deps = deps
ctx = MockCtx(deps)
repl = _get_or_create_repl(ctx)
result = await repl.execute_async("print(1 + 1)")
assert result.success
assert "2" in result.stdout
@pytest.mark.asyncio
async def test_execute_code_tracks_executions_in_context(self, empty_client):
"""Test that code executions are tracked as CodeExecution objects in RLMContext."""
from haiku.rag.agents.rlm.agent import _get_or_create_repl
context = RLMContext()
deps = RLMDeps(
client=empty_client,
config=Config,
context=context,
)
class MockCtx:
def __init__(self, deps):
self.deps = deps
ctx = MockCtx(deps)
repl = _get_or_create_repl(ctx)
assert len(context.code_executions) == 0
result = await repl.execute_async("x = 42")
assert result.success
@pytest.mark.asyncio
async def test_code_execution_has_correct_fields(self, empty_client):
class TestCodeExecutionModel:
def test_code_execution_has_correct_fields(self):
"""Test that CodeExecution has all expected fields."""
execution = CodeExecution(
code="print('hello')",
@ -89,29 +41,6 @@ class TestExecuteCodeTool:
assert execution.stderr == ""
assert execution.success is True
@pytest.mark.asyncio
async def test_code_execution_captures_errors(self, empty_client):
"""Test that failed executions are properly captured."""
from haiku.rag.agents.rlm.agent import _get_or_create_repl
context = RLMContext()
deps = RLMDeps(
client=empty_client,
config=Config,
context=context,
)
class MockCtx:
def __init__(self, deps):
self.deps = deps
ctx = MockCtx(deps)
repl = _get_or_create_repl(ctx)
result = await repl.execute_async("1/0")
assert result.success is False
assert "ZeroDivisionError" in result.stderr
class TestClientRLMIntegration:
"""Integration tests for client.rlm() method."""

909
tests/agents/rlm/test_sandbox.py
View file

@ -2,747 +2,246 @@ from pathlib import Path
import pytest
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.docker_sandbox import DockerSandbox, SandboxResult
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_sandbox")
class TestSafeBuiltins:
"""Test that safe builtins are available."""
def is_docker_available() -> bool:
"""Check if Docker daemon is available."""
try:
import docker
client = docker.from_env()
client.ping()
return True
except Exception:
return False
docker_required = pytest.mark.skipif(
not is_docker_available(),
reason="Docker daemon not available",
)
@pytest.mark.integration
class TestDockerSandboxBasics:
"""Test basic Docker sandbox functionality."""
@docker_required
@pytest.mark.asyncio
async def test_print_available(self, repl_env_empty):
result = await repl_env_empty.execute_async("print('hello')")
async def test_execute_simple_code(self, docker_sandbox):
"""Test executing simple code in the sandbox."""
result = await docker_sandbox.execute("print('hello world')")
assert isinstance(result, SandboxResult)
assert result.success
assert "hello" in result.stdout
assert "hello world" in result.stdout
assert result.stderr == ""
@pytest.mark.integration
class TestDockerSandboxErrors:
"""Test error handling in Docker sandbox."""
@docker_required
@pytest.mark.asyncio
async def test_len_available(self, repl_env_empty):
result = await repl_env_empty.execute_async("print(len([1, 2, 3]))")
assert result.success
assert "3" in result.stdout
@pytest.mark.asyncio
async def test_range_available(self, repl_env_empty):
result = await repl_env_empty.execute_async("print(list(range(3)))")
assert result.success
assert "[0, 1, 2]" in result.stdout
@pytest.mark.asyncio
async def test_enumerate_available(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"print(list(enumerate(['a', 'b'])))"
)
assert result.success
assert "[(0, 'a'), (1, 'b')]" in result.stdout
@pytest.mark.asyncio
async def test_sorted_available(self, repl_env_empty):
result = await repl_env_empty.execute_async("print(sorted([3, 1, 2]))")
assert result.success
assert "[1, 2, 3]" in result.stdout
@pytest.mark.asyncio
async def test_sum_available(self, repl_env_empty):
result = await repl_env_empty.execute_async("print(sum([1, 2, 3]))")
assert result.success
assert "6" in result.stdout
@pytest.mark.asyncio
async def test_min_max_available(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"print(min([3, 1, 2]), max([3, 1, 2]))"
)
assert result.success
assert "1 3" in result.stdout
@pytest.mark.asyncio
async def test_all_any_available(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"print(all([True, True]), any([False, True]))"
)
assert result.success
assert "True True" in result.stdout
@pytest.mark.asyncio
async def test_dict_list_set_tuple_available(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"print(dict(a=1), list((1,2)), set([1,2,1]), tuple([1,2]))"
)
assert result.success
assert "{'a': 1}" in result.stdout
@pytest.mark.asyncio
async def test_str_int_float_bool_available(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"print(str(1), int('2'), float('3.0'), bool(1))"
)
assert result.success
assert "1 2 3.0 True" in result.stdout
@pytest.mark.asyncio
async def test_zip_map_filter_available(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"print(list(zip([1,2], ['a','b'])), "
"list(map(str, [1,2])), "
"list(filter(lambda x: x > 1, [1,2,3])))"
)
assert result.success
assert "[(1, 'a'), (2, 'b')]" in result.stdout
@pytest.mark.asyncio
async def test_isinstance_type_available(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"print(isinstance(1, int), type([]))"
)
assert result.success
assert "True" in result.stdout
class TestDangerousBuiltinsBlocked:
"""Test that dangerous builtins are blocked."""
@pytest.mark.asyncio
async def test_eval_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("eval('1+1')")
async def test_syntax_error(self, docker_sandbox):
"""Test that syntax errors are reported."""
result = await docker_sandbox.execute("def foo(")
assert not result.success
assert "eval" in result.stderr.lower() or "not defined" in result.stderr.lower()
assert "SyntaxError" in result.stderr
@docker_required
@pytest.mark.asyncio
async def test_exec_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("exec('x = 1')")
async def test_runtime_error(self, docker_sandbox):
"""Test that runtime errors are reported."""
result = await docker_sandbox.execute("x = 1/0")
assert not result.success
assert "exec" in result.stderr.lower() or "not defined" in result.stderr.lower()
assert "ZeroDivisionError" in result.stderr
@docker_required
@pytest.mark.asyncio
async def test_compile_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"compile('1+1', '<string>', 'eval')"
)
async def test_name_error(self, docker_sandbox):
"""Test that name errors are reported."""
result = await docker_sandbox.execute("print(undefined_variable)")
assert not result.success
assert (
"compile" in result.stderr.lower() or "not defined" in result.stderr.lower()
)
assert "NameError" in result.stderr
@docker_required
@pytest.mark.asyncio
async def test_open_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("open('/etc/passwd')")
assert not result.success
assert "open" in result.stderr.lower() or "not defined" in result.stderr.lower()
async def test_missing_image(self, temp_db_path):
"""Test error when Docker image is not found."""
async with HaikuRAG(temp_db_path, create=True) as client:
config = RLMConfig(docker_image="nonexistent-image:v999.999.999")
context = RLMContext()
async with DockerSandbox(
client=client, config=config, context=context, image=config.docker_image
) as sandbox:
result = await sandbox.execute("print('hello')")
assert not result.success
assert (
"not found" in result.stderr.lower()
or "error" in result.stderr.lower()
)
@pytest.mark.integration
class TestDockerSandboxHaikuRAG:
"""Test haiku.rag functions in Docker sandbox."""
@docker_required
@pytest.mark.asyncio
async def test_input_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("input('Enter: ')")
assert not result.success
assert (
"input" in result.stderr.lower() or "not defined" in result.stderr.lower()
)
@pytest.mark.asyncio
async def test___import___blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("__import__('os')")
assert not result.success
@pytest.mark.asyncio
async def test_globals_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("globals()")
assert not result.success
@pytest.mark.asyncio
async def test_locals_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("locals()")
assert not result.success
@pytest.mark.asyncio
async def test_breakpoint_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("breakpoint()")
assert not result.success
@pytest.mark.asyncio
async def test_getattr_setattr_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("getattr(object, '__class__')")
assert not result.success
@pytest.mark.asyncio
async def test_delattr_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("delattr(object, 'x')")
assert not result.success
class TestAllowedImports:
"""Test that allowed imports work."""
@pytest.mark.asyncio
async def test_json_import(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"import json\nprint(json.dumps({'a': 1}))"
)
assert result.success
assert '{"a": 1}' in result.stdout
@pytest.mark.asyncio
async def test_re_import(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"import re\nprint(re.match(r'\\d+', '123').group())"
)
assert result.success
assert "123" in result.stdout
@pytest.mark.asyncio
async def test_math_import(self, repl_env_empty):
result = await repl_env_empty.execute_async("import math\nprint(math.sqrt(4))")
assert result.success
assert "2.0" in result.stdout
@pytest.mark.asyncio
async def test_statistics_import(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"import statistics\nprint(statistics.mean([1, 2, 3]))"
)
assert result.success
assert "2" in result.stdout
@pytest.mark.asyncio
async def test_collections_import(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"from collections import Counter\nprint(Counter(['a', 'b', 'a']))"
)
assert result.success
assert "'a': 2" in result.stdout
@pytest.mark.asyncio
async def test_itertools_import(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"from itertools import chain\nprint(list(chain([1], [2])))"
)
assert result.success
assert "[1, 2]" in result.stdout
@pytest.mark.asyncio
async def test_functools_import(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"from functools import reduce\nprint(reduce(lambda a, b: a+b, [1,2,3]))"
)
assert result.success
assert "6" in result.stdout
@pytest.mark.asyncio
async def test_datetime_import(self, repl_env_empty):
result = await repl_env_empty.execute_async(
"from datetime import date\nprint(date(2025, 1, 1))"
)
assert result.success
assert "2025-01-01" in result.stdout
class TestDangerousImportsBlocked:
"""Test that dangerous imports are blocked."""
@pytest.mark.asyncio
async def test_os_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("import os")
assert not result.success
assert (
"not allowed" in result.stderr.lower() or "error" in result.stderr.lower()
)
@pytest.mark.asyncio
async def test_sys_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("import sys")
assert not result.success
@pytest.mark.asyncio
async def test_subprocess_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("import subprocess")
assert not result.success
@pytest.mark.asyncio
async def test_shutil_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("import shutil")
assert not result.success
@pytest.mark.asyncio
async def test_socket_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("import socket")
assert not result.success
@pytest.mark.asyncio
async def test_requests_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("import requests")
assert not result.success
@pytest.mark.asyncio
async def test_builtins_module_blocked(self, repl_env_empty):
result = await repl_env_empty.execute_async("import builtins")
assert not result.success
class TestHaikuRAGBridgeFunctions:
"""Test haiku.rag bridge functions in sandbox."""
@pytest.mark.asyncio
async def test_search(self, repl_env_empty):
"""Test search function calls client with correct args."""
from unittest.mock import AsyncMock
from haiku.rag.store.models import SearchResult
mock_results = [
SearchResult(
chunk_id="chunk-1",
document_id="doc-1",
document_title="Test Doc",
document_uri="test://doc",
content="Test content about foxes",
score=0.9,
page_numbers=[1],
headings=["Heading"],
)
]
repl_env_empty.client.search = AsyncMock(return_value=mock_results)
result = await repl_env_empty.execute_async(
"results = search('fox', limit=5)\n"
"print(len(results), results[0]['chunk_id'], 'fox' in results[0]['content'].lower())"
)
assert result.success
assert "1 chunk-1 True" in result.stdout
repl_env_empty.client.search.assert_called_once_with(
"fox", limit=5, filter=None
)
@pytest.mark.asyncio
async def test_list_documents(self, repl_env_empty):
"""Test list_documents returns list structure."""
result = await repl_env_empty.execute_async(
async def test_list_documents_empty(self, docker_sandbox):
"""Test list_documents returns empty list for empty database."""
result = await docker_sandbox.execute(
"docs = list_documents()\nprint(type(docs).__name__, len(docs))"
)
assert result.success
assert "list 0" in result.stdout
@pytest.mark.asyncio
async def test_get_document(self, repl_env_empty):
"""Test get_document calls client correctly."""
from unittest.mock import AsyncMock
from haiku.rag.store.models import Document
mock_doc = Document(
id="doc-1",
uri="test://doc",
title="Test Doc",
content="The quick brown fox",
)
repl_env_empty.client.get_document_by_id = AsyncMock(return_value=mock_doc)
result = await repl_env_empty.execute_async(
"doc = get_document('doc-1')\nprint('fox' in doc.lower())"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.asyncio
async def test_get_document_missing(self, repl_env_empty):
"""Test get_document returns None for missing document."""
result = await repl_env_empty.execute_async(
"doc = get_document('Nonexistent')\nprint(doc is None)"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.asyncio
async def test_llm(self, repl_env_empty):
"""Test llm function is available in sandbox."""
result = await repl_env_empty.execute_async("print(callable(llm))")
assert result.success
assert "True" in result.stdout
class TestSandboxExecution:
"""Test general sandbox execution behavior."""
@pytest.mark.asyncio
async def test_variable_persistence(self, repl_env_empty):
"""Variables persist across executions."""
await repl_env_empty.execute_async("x = 42")
result = await repl_env_empty.execute_async("print(x)")
assert result.success
assert "42" in result.stdout
@pytest.mark.asyncio
async def test_function_definition(self, repl_env_empty):
"""Can define and call functions."""
result = await repl_env_empty.execute_async(
"def add(a, b):\n return a + b\nprint(add(1, 2))"
)
assert result.success
assert "3" in result.stdout
@pytest.mark.asyncio
async def test_class_definition(self, repl_env_empty):
"""Can define and use classes."""
result = await repl_env_empty.execute_async(
"class Point:\n"
" def __init__(self, x, y):\n"
" self.x = x\n"
" self.y = y\n"
"p = Point(1, 2)\n"
"print(p.x, p.y)"
)
assert result.success
assert "1 2" in result.stdout
@pytest.mark.asyncio
async def test_list_comprehension(self, repl_env_empty):
"""List comprehensions work."""
result = await repl_env_empty.execute_async("print([x**2 for x in range(5)])")
assert result.success
assert "[0, 1, 4, 9, 16]" in result.stdout
@pytest.mark.asyncio
async def test_dict_comprehension(self, repl_env_empty):
"""Dict comprehensions work."""
result = await repl_env_empty.execute_async(
"print({x: x**2 for x in range(3)})"
)
assert result.success
assert "{0: 0, 1: 1, 2: 4}" in result.stdout
@pytest.mark.asyncio
async def test_exception_handling(self, repl_env_empty):
"""Can catch and handle exceptions."""
result = await repl_env_empty.execute_async(
"try:\n x = 1/0\nexcept ZeroDivisionError:\n print('caught')"
)
assert result.success
assert "caught" in result.stdout
@pytest.mark.asyncio
async def test_uncaught_exception_reports_error(self, repl_env_empty):
"""Uncaught exceptions are reported."""
result = await repl_env_empty.execute_async("x = 1/0")
assert not result.success
assert "ZeroDivisionError" in result.stderr
@pytest.mark.asyncio
async def test_syntax_error_reports_error(self, repl_env_empty):
"""Syntax errors are reported."""
result = await repl_env_empty.execute_async("def foo(")
assert not result.success
assert "SyntaxError" in result.stderr
@pytest.mark.asyncio
async def test_output_truncation(self, repl_env_empty):
"""Output is truncated if too long."""
repl_env_empty.config.max_output_chars = 100
result = await repl_env_empty.execute_async("print('x' * 1000)")
assert result.success
assert (
len(result.stdout) <= 100 + 50
) # Allow some margin for truncation message
class TestContextFilter:
"""Test that context filter is applied to all searches."""
@pytest.mark.asyncio
async def test_context_filter_applied_to_search(self, temp_db_path):
"""Search applies context filter automatically."""
from unittest.mock import AsyncMock
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
async with HaikuRAG(temp_db_path, create=True) as client:
context = RLMContext(filter="uri LIKE '%medical%'")
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
client.search = AsyncMock(return_value=[])
await repl.execute_async("search('test query')")
client.search.assert_called_once_with(
"test query", limit=10, filter="uri LIKE '%medical%'"
)
@pytest.mark.asyncio
async def test_context_filter_applied_to_list_documents(self, temp_db_path):
"""list_documents applies context filter automatically."""
from unittest.mock import AsyncMock
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
async with HaikuRAG(temp_db_path, create=True) as client:
context = RLMContext(filter="title = 'Report'")
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
client.list_documents = AsyncMock(return_value=[])
await repl.execute_async("list_documents()")
client.list_documents.assert_called_once_with(
limit=10, offset=0, filter="title = 'Report'"
)
class TestPreloadedDocuments:
"""Test pre-loaded documents context variable."""
@pytest.mark.asyncio
async def test_documents_variable_available_when_preloaded(self, temp_db_path):
"""documents variable is available when context.documents is set."""
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
from haiku.rag.store.models import Document
async with HaikuRAG(temp_db_path, create=True) as client:
preloaded = [
Document(
id="doc-1",
title="First Doc",
uri="test://first",
content="Content of first document about cats.",
),
Document(
id="doc-2",
title="Second Doc",
uri="test://second",
content="Content of second document about dogs.",
),
]
context = RLMContext(documents=preloaded)
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
result = await repl.execute_async(
"print(len(documents))\n"
"print([d['title'] for d in documents])\n"
"print('cats' in documents[0]['content'])"
)
assert result.success
assert "2" in result.stdout
assert "First Doc" in result.stdout
assert "Second Doc" in result.stdout
assert "True" in result.stdout
@pytest.mark.asyncio
async def test_documents_variable_not_available_without_preload(self, temp_db_path):
"""documents variable is not available when context.documents is None."""
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
async with HaikuRAG(temp_db_path, create=True) as client:
context = RLMContext()
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
result = await repl.execute_async("print(documents)")
assert not result.success
assert "NameError" in result.stderr
@pytest.mark.asyncio
async def test_documents_has_expected_fields(self, temp_db_path):
"""documents variable contains expected dict fields."""
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
from haiku.rag.store.models import Document
async with HaikuRAG(temp_db_path, create=True) as client:
preloaded = [
Document(
id="doc-1",
title="Test Doc",
uri="test://doc",
content="Test content",
),
]
context = RLMContext(documents=preloaded)
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
result = await repl.execute_async(
"d = documents[0]\n"
"print(sorted(d.keys()))\n"
"print(d['id'], d['title'], d['uri'])"
)
assert result.success
assert "['content', 'id', 'title', 'uri']" in result.stdout
assert "doc-1" in result.stdout
assert "Test Doc" in result.stdout
assert "test://doc" in result.stdout
class TestSandboxEscapeVectors:
"""Test that known sandbox escape techniques are blocked.
Each test contains actual exploit code that would work without the fix.
"""
@pytest.mark.asyncio
async def test_type_dict_subclasses_escape_blocked(self, repl_env_empty):
"""Cannot escape via type.__dict__['__subclasses__'].
Without fix: This would enumerate all loaded classes and find
subprocess.Popen to execute arbitrary shell commands.
"""
result = await repl_env_empty.execute_async("""
# EXPLOIT: Access __subclasses__ via dict to bypass AST check
subclasses_method = type.__dict__['__subclasses__']
all_classes = subclasses_method(object)
print(f"Found {len(all_classes)} classes")
""")
assert not result.success
assert "not allowed" in result.stderr.lower()
@pytest.mark.asyncio
async def test_popen_shell_execution_blocked(self, repl_env_empty):
"""Cannot execute shell commands via Popen.
Without fix: This would execute 'whoami' and return the username.
"""
result = await repl_env_empty.execute_async("""
# EXPLOIT: Find subprocess.Popen and execute shell commands
subclasses_method = type.__dict__['__subclasses__']
all_classes = subclasses_method(object)
popen = [c for c in all_classes if c.__name__ == 'Popen'][0]
proc = popen('whoami', shell=True, stdout=-1)
print(proc.stdout.read())
""")
assert not result.success
@pytest.mark.asyncio
async def test_socket_creation_blocked(self, repl_env_empty):
"""Cannot create network sockets for data exfiltration.
Without fix: This would create a socket that could connect to external servers.
"""
result = await repl_env_empty.execute_async("""
# EXPLOIT: Find socket class and create network connection
subclasses_method = type.__dict__['__subclasses__']
all_classes = subclasses_method(object)
socket_cls = [c for c in all_classes if c.__name__ == 'socket'][0]
s = socket_cls(2, 1) # AF_INET, SOCK_STREAM
print(f"Created socket: {s}")
""")
assert not result.success
@pytest.mark.asyncio
async def test_type_three_arg_class_creation_blocked(self, repl_env_empty):
"""Cannot use type() with 3 arguments to create classes dynamically."""
result = await repl_env_empty.execute_async(
"EvilClass = type('EvilClass', (object,), {'x': 1})"
)
assert not result.success
@pytest.mark.asyncio
async def test_dict_key_dunder_access_blocked(self, repl_env_empty):
"""Cannot access dunder methods via dictionary key access."""
result = await repl_env_empty.execute_async(
"method = str.__dict__['__add__']\nprint(method)"
)
assert not result.success
assert "not allowed" in result.stderr.lower()
@pytest.mark.asyncio
async def test_dict_key_private_access_blocked(self, repl_env_empty):
"""Cannot access private attributes via dictionary key access."""
result = await repl_env_empty.execute_async(
"method = object.__dict__['_private']\nprint(method)"
)
assert not result.success
assert "not allowed" in result.stderr.lower()
@docker_required
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_sql_injection_in_get_document_blocked(self, temp_db_path):
"""SQL injection in get_document cannot bypass context filter.
Without fix: Injecting quotes would leak documents that should be
protected by the context filter.
"""
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
async def test_list_documents_with_data(self, temp_db_path, test_docker_image):
"""Test list_documents returns documents when populated."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Create documents: one secret, one public
await client.create_document(
content="TOP SECRET: Launch codes 1234",
uri="secret://classified",
title="Classified Intel",
)
await client.create_document(
content="Public weather report",
uri="public://weather",
title="Weather",
content="Test content",
uri="test://doc1",
title="Test Document",
)
# Sandbox restricted to public:// only
config = RLMConfig(docker_image=test_docker_image)
context = RLMContext()
async with DockerSandbox(
client=client, config=config, context=context, image=test_docker_image
) as sandbox:
result = await sandbox.execute(
"docs = list_documents()\nprint(len(docs))\nprint(docs[0]['title'])"
)
assert result.success
assert "1" in result.stdout
assert "Test Document" in result.stdout
@docker_required
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_search_with_data(self, temp_db_path, test_docker_image):
"""Test search function works."""
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content="The quick brown fox jumps over the lazy dog.",
uri="test://animals",
title="Animals",
)
config = RLMConfig(docker_image=test_docker_image)
context = RLMContext()
async with DockerSandbox(
client=client, config=config, context=context, image=test_docker_image
) as sandbox:
result = await sandbox.execute(
"results = search('fox', limit=5)\n"
"print(len(results))\n"
"if results:\n"
" print('fox' in results[0]['content'].lower())"
)
assert result.success
# Search should return at least one result
assert "True" in result.stdout or "1" in result.stdout
@docker_required
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_get_document(self, temp_db_path, test_docker_image):
"""Test get_document function."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Content about foxes and dogs.",
uri="test://doc",
title="Fox Document",
)
config = RLMConfig(docker_image=test_docker_image)
context = RLMContext()
async with DockerSandbox(
client=client, config=config, context=context, image=test_docker_image
) as sandbox:
result = await sandbox.execute(
f"content = get_document('{doc.id}')\n"
"print('foxes' in content.lower() if content else 'None')"
)
assert result.success
assert "True" in result.stdout
@docker_required
@pytest.mark.asyncio
async def test_get_document_not_found(self, docker_sandbox):
"""Test get_document returns None for missing document."""
result = await docker_sandbox.execute(
"content = get_document('nonexistent-id')\nprint(content is None)"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.integration
class TestDockerSandboxContextFilter:
"""Test context filter is applied."""
@docker_required
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_filter_applied_to_list_documents(
self, temp_db_path, test_docker_image
):
"""Test that context filter is passed to list_documents."""
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content="Public content",
uri="public://doc1",
title="Public Doc",
)
await client.create_document(
content="Private content",
uri="private://doc2",
title="Private Doc",
)
config = RLMConfig(docker_image=test_docker_image)
context = RLMContext(filter="uri LIKE 'public://%'")
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
# EXPLOIT: SQL injection to access secret document
result = await repl.execute_async("""
# Injection payload breaks out of quotes and adds OR clause
content = get_document("x' OR uri LIKE 'secret://%")
if content:
print(f"LEAKED: {content}")
else:
print("NO LEAK")
""")
assert result.success
assert "TOP SECRET" not in result.stdout
assert "Launch codes" not in result.stdout
async with DockerSandbox(
client=client, config=config, context=context, image=test_docker_image
) as sandbox:
result = await sandbox.execute(
"docs = list_documents()\n"
"print(len(docs))\n"
"if docs:\n"
" print(docs[0]['title'])"
)
assert result.success
assert "1" in result.stdout
assert "Public Doc" in result.stdout
assert "Private Doc" not in result.stdout
class TestSecurityEscapes:
"""Test that common security escape attempts are blocked."""
@pytest.mark.integration
class TestDockerSandboxPreloadedDocuments:
"""Test pre-loaded documents context variable."""
@docker_required
@pytest.mark.asyncio
async def test_eval_via_builtins_dict(self, repl_env_empty):
"""Cannot access eval through __builtins__."""
result = await repl_env_empty.execute_async("__builtins__['eval']('1+1')")
assert not result.success
@pytest.mark.asyncio
async def test_import_via_builtins(self, repl_env_empty):
"""Cannot import os through builtins trickery."""
result = await repl_env_empty.execute_async("__builtins__.__import__('os')")
assert not result.success
@pytest.mark.asyncio
async def test_class_bases_escape(self, repl_env_empty):
"""Cannot escape through __class__.__bases__."""
result = await repl_env_empty.execute_async(
"().__class__.__bases__[0].__subclasses__()"
)
assert not result.success
@pytest.mark.asyncio
async def test_code_object_escape(self, repl_env_empty):
"""Cannot create code objects."""
result = await repl_env_empty.execute_async(
"def f(): pass\n"
"type(f.__code__)(0, 0, 0, 0, 0, 0, b'', (), (), (), '', '', 0, b'')"
)
assert not result.success
@pytest.mark.asyncio
async def test_import_system_escape(self, repl_env_empty):
"""Cannot escape through importlib."""
result = await repl_env_empty.execute_async("import importlib")
assert not result.success
@pytest.mark.asyncio
async def test_pickle_escape(self, repl_env_empty):
"""Cannot use pickle for code execution."""
result = await repl_env_empty.execute_async("import pickle")
async def test_documents_variable_not_available_without_preload(
self, docker_sandbox
):
"""documents variable is not available when context.documents is None."""
result = await docker_sandbox.execute("print(documents)")
assert not result.success
assert "NameError" in result.stderr

82
tests/cassettes/test_sandbox/TestDockerSandboxContextFilter.test_filter_applied_to_list_documents.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_sandbox/TestDockerSandboxHaikuRAG.test_get_document.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_sandbox/TestDockerSandboxHaikuRAG.test_list_documents_with_data.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_sandbox/TestDockerSandboxHaikuRAG.test_search_with_data.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_sandbox/TestSandboxEscapeVectors.test_sql_injection_in_get_document_blocked.yaml
View file

File diff suppressed because one or more lines are too long

16
uv.lock
View file

@ -739,6 +739,20 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ba/5a/18ad964b0086c6e62e2e7500f7edc89e3faa45033c71c1893d34eed2b2de/dnspython-2.8.0-py3-none-any.whl", hash = "sha256:01d9bbc4a2d76bf0db7c1f729812ded6d912bd318d3b1cf81d30c0f845dbf3af", size = 331094, upload-time = "2025-09-07T18:57:58.071Z" },
]
[[package]]
name = "docker"
version = "7.1.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pywin32", marker = "sys_platform == 'win32'" },
{ name = "requests" },
{ name = "urllib3" },
]
sdist = { url = "https://files.pythonhosted.org/packages/91/9b/4a2ea29aeba62471211598dac5d96825bb49348fa07e906ea930394a83ce/docker-7.1.0.tar.gz", hash = "sha256:ad8c70e6e3f8926cb8a92619b832b4ea5299e2831c14284663184e200546fa6c", size = 117834, upload-time = "2024-05-23T11:13:57.216Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e3/26/57c6fb270950d476074c087527a558ccb6f4436657314bfb6cdf484114c4/docker-7.1.0-py3-none-any.whl", hash = "sha256:c96b93b7f0a746f9e77d325bcfb87422a3d8bd4f03136ae8a85b37f1898d5fc0", size = 147774, upload-time = "2024-05-23T11:13:55.01Z" },
]
[[package]]
name = "docling"
version = "2.69.1"
@ -1366,6 +1380,7 @@ name = "haiku-rag-slim"
version = "0.28.0"
source = { editable = "haiku_rag_slim" }
dependencies = [
{ name = "docker" },
{ name = "docling-core" },
{ name = "httpx" },
{ name = "jsonpatch" },
@ -1427,6 +1442,7 @@ zeroentropy = [
[package.metadata]
requires-dist = [
{ name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.20.1" },
{ name = "docker", specifier = ">=7.1.0" },
{ name = "docling", marker = "extra == 'docling'", specifier = "==2.69.1" },
{ name = "docling-core", specifier = "==2.60.1" },
{ name = "httpx", specifier = ">=0.28.1" },