414 lines
13 KiB
Python
414 lines
13 KiB
Python
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
|
|
|
|
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": type,
|
|
"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(),
|
|
"ask": self._make_ask(),
|
|
}
|
|
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
|
|
docs = await self.client.list_documents(
|
|
filter=f"title = '{id_or_title}'"
|
|
)
|
|
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 = '{id_or_title}'")
|
|
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()
|
|
docs = await self.client.list_documents(
|
|
filter=f"title = '{id_or_title}'"
|
|
)
|
|
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 = '{id_or_title}'")
|
|
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_ask(self):
|
|
"""Create sync ask function that uses QA agent."""
|
|
|
|
def ask(question: str) -> str:
|
|
async def _ask():
|
|
answer, citations = await self.client.ask(
|
|
question, filter=self.context.filter
|
|
)
|
|
for c in citations:
|
|
for sr in self.context.search_results:
|
|
if sr.chunk_id == c.chunk_id:
|
|
break
|
|
else:
|
|
from haiku.rag.store.models import SearchResult
|
|
|
|
self.context.search_results.append(
|
|
SearchResult(
|
|
chunk_id=c.chunk_id,
|
|
document_id=c.document_id,
|
|
document_title=c.document_title or "",
|
|
document_uri=c.document_uri,
|
|
content=c.content,
|
|
score=1.0,
|
|
page_numbers=c.page_numbers,
|
|
headings=c.headings or [],
|
|
)
|
|
)
|
|
return answer
|
|
|
|
return self._run_async_from_thread(_ask())
|
|
|
|
return ask
|
|
|
|
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"
|
|
)
|
|
|
|
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
|