From 167c837514f65749d81a720e5211e099f4422ce1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 20 Apr 2026 11:57:39 +0300 Subject: [PATCH] repository methods, module-level executor, read-only VFS --- .../haiku/rag/agents/analysis/sandbox.py | 48 +++++++++---------- .../haiku/rag/store/repositories/document.py | 14 ++++++ .../rag/store/repositories/document_item.py | 12 +++++ 3 files changed, 49 insertions(+), 25 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py index 635e9caa..bf44dd04 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py @@ -26,10 +26,12 @@ class SandboxResult: success: bool +_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + + def _run_async(coro: Any) -> Any: """Run an async coroutine from a sync context (CallbackFile read).""" - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(asyncio.run, coro).result() + return _executor.submit(asyncio.run, coro).result() class Sandbox: @@ -168,20 +170,12 @@ class Sandbox: def read_content(_path: "PurePosixPath") -> str: async def _fetch() -> str: from haiku.rag.client import HaikuRAG - from haiku.rag.utils import escape_sql_string async with HaikuRAG( db_path, config=config, read_only=True ) as rag: - safe_id = escape_sql_string(did) - rows = list( - rag.store.documents_table.search() - .select(["content"]) - .where(f"id = '{safe_id}'") - .limit(1) - .to_list() - ) - return rows[0]["content"] if rows else "" + content = await rag.document_repository.get_content(did) + return content or "" return _run_async(_fetch()) @@ -197,10 +191,8 @@ class Sandbox: async with HaikuRAG( db_path, config=config, read_only=True ) as rag: - items = ( - await rag.document_item_repository.get_items_in_range( - did, 0, 999999 - ) + items = await rag.document_item_repository.get_all_items( + did ) lines = [] for item in items: @@ -222,24 +214,27 @@ class Sandbox: return read_items + def _deny_write(_path: "PurePosixPath", _content: str | bytes) -> None: + raise PermissionError(f"Document files are read-only: {_path}") + files.append( CallbackFile( f"{doc_dir}/content.txt", read=_make_content_reader(doc_id), - write=lambda _p, _c: None, + write=_deny_write, ) ) files.append( CallbackFile( f"{doc_dir}/items.jsonl", read=_make_items_reader(doc_id), - write=lambda _p, _c: None, + write=_deny_write, ) ) return OSAccess(files) - async def _ensure_initialized(self) -> None: + async def _ensure_initialized(self) -> tuple[MontyRepl, OSAccess]: """Initialize the REPL session and VFS on first use.""" if self._repl is None: self._vfs = await self._build_vfs() @@ -266,16 +261,19 @@ class Sandbox: external_functions=self._build_external_functions(), os=self._vfs, ) + # Both are guaranteed non-None after initialization + repl = self._repl + vfs = self._vfs + if repl is None or vfs is None: + raise RuntimeError("Sandbox initialization failed") + return repl, vfs async def execute(self, code: str) -> SandboxResult: """Execute Python code in the Monty REPL. Variables persist across calls within the same Sandbox instance. """ - await self._ensure_initialized() - assert self._repl is not None - assert self._vfs is not None - + repl, vfs = await self._ensure_initialized() external_fns = self._build_external_functions() stdout_lines: list[str] = [] @@ -287,11 +285,11 @@ class Sandbox: try: output = await pydantic_monty.run_repl_async( - self._repl, + repl, code, external_functions=external_fns, print_callback=print_callback, - os=self._vfs, + os=vfs, ) except ( pydantic_monty.MontySyntaxError, diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 0c9de830..4f5e89e1 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -100,6 +100,20 @@ class DocumentRepository: return self._record_to_document(results[0]) + async def get_content(self, entity_id: str) -> str | None: + """Get only the text content of a document (skips docling blobs).""" + safe_id = escape_sql_string(entity_id) + results = list( + self.store.documents_table.search() + .select(["content"]) + .where(f"id = '{safe_id}'") + .limit(1) + .to_list() + ) + if not results: + return None + return results[0]["content"] + _DOCLING_COLUMNS = ["id", "docling_document", "docling_version"] async def get_docling_data(self, entity_id: str) -> Document | None: diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index 9b30a271..eeb5523d 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -40,6 +40,18 @@ class DocumentItemRepository: ] self.store.document_items_table.add(records) + async def get_all_items(self, document_id: str) -> list[DocumentItem]: + """Get all items for a document, sorted by position.""" + safe_id = escape_sql_string(document_id) + rows = ( + self.store.document_items_table.search() + .where(f"document_id = '{safe_id}'") + .to_list() + ) + items = [self._record_to_item(row) for row in rows] + items.sort(key=lambda x: x.position) + return items + async def get_items_in_range( self, document_id: str, start: int, end: int ) -> list[DocumentItem]: