From 6759a2e0a5c1a42ac37ad7c4f8689ffb8e3319ca Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 24 Feb 2026 13:12:56 +0200 Subject: [PATCH] Give regex to monty as externals --- CHANGELOG.md | 1 + docs/agents/rlm.md | 3 +- .../haiku/rag/agents/rlm/prompts.py | 24 +++++-- .../haiku/rag/agents/rlm/sandbox.py | 25 ++++++++ tests/agents/rlm/test_sandbox.py | 64 +++++++++++++++++++ 5 files changed, 109 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index badb9a92..3acba6af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ### Added +- **RLM sandbox regex functions**: `regex_findall`, `regex_sub`, `regex_search`, `regex_split` for pattern matching without LLM calls - **`HaikuRAG.get_chunk_by_id()`**: Public method for chunk lookup by ID ### Removed diff --git a/docs/agents/rlm.md b/docs/agents/rlm.md index 30a7346e..c339a7c0 100644 --- a/docs/agents/rlm.md +++ b/docs/agents/rlm.md @@ -64,6 +64,7 @@ The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https: | `get_chunk(chunk_id)` | Get a chunk with metadata (headings, page numbers, labels) for citations | | `get_docling_document(document_id)` | Get the full DoclingDocument structure as a dict (texts, tables, pictures, pages) | | `llm(prompt)` | Call an LLM for classification, summarization, or extraction | +| `regex_findall(pattern, text)`, `regex_sub(pattern, repl, text)`, `regex_search(pattern, text)`, `regex_split(pattern, text)` | Regular expression matching via Python's `re` module | When documents are pre-loaded via the `documents` parameter, they are injected as a `documents` variable accessible in the sandbox code. @@ -71,7 +72,7 @@ When documents are pre-loaded via the `documents` parameter, they are injected a The interpreter supports a subset of Python: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, try/except, and the `json` module. -Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, the agent can use string methods or the `llm()` function. +Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, the agent can use the `regex_*` functions, string methods, or the `llm()` function. ### Security diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py b/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py index 9b415d71..cc980ddc 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/rlm/prompts.py @@ -34,6 +34,18 @@ Use `list_documents()` or search results to get document IDs first. - `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 @@ -55,7 +67,7 @@ The interpreter supports: variables, arithmetic, strings, f-strings, lists, dict 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 @@ -79,16 +91,14 @@ for doc in docs: 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']}") - for part in extracted.split(','): - part = part.strip().replace(',', '') - if part.isdigit(): - numbers.append(int(part)) + 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)}") ``` diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py b/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py index f845a097..8f13f5b1 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py +++ b/haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py @@ -1,4 +1,5 @@ import json +import re from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal @@ -127,6 +128,26 @@ class Sandbox: result = await agent.run(prompt) return result.output + async def regex_findall(pattern: str, text: str) -> list[str]: + return re.findall(pattern, text) + + async def regex_sub(pattern: str, repl: str, text: str) -> str: + return re.sub(pattern, repl, text) + + async def regex_search(pattern: str, text: str) -> dict[str, Any] | None: + m = re.search(pattern, text) + if m is None: + return None + return { + "group": m.group(), + "groups": list(m.groups()), + "start": m.start(), + "end": m.end(), + } + + async def regex_split(pattern: str, text: str) -> list[str]: + return re.split(pattern, text) + return { "search": search, "list_documents": list_documents, @@ -134,6 +155,10 @@ class Sandbox: "get_chunk": get_chunk, "get_docling_document": get_docling_document, "llm": llm, + "regex_findall": regex_findall, + "regex_sub": regex_sub, + "regex_search": regex_search, + "regex_split": regex_split, } async def execute(self, code: str) -> SandboxResult: diff --git a/tests/agents/rlm/test_sandbox.py b/tests/agents/rlm/test_sandbox.py index adc34daa..8ab1c62d 100644 --- a/tests/agents/rlm/test_sandbox.py +++ b/tests/agents/rlm/test_sandbox.py @@ -381,6 +381,70 @@ class TestSandboxDoclingDocument: assert "True" in result.stdout +class TestSandboxRegex: + """Test regex external functions.""" + + @pytest.mark.asyncio + async def test_regex_findall(self, sandbox): + """regex_findall extracts all matches.""" + result = await sandbox.execute( + r"matches = await regex_findall(r'\d+', 'abc 123 def 456')" + "\nprint(matches)" + ) + assert result.success + assert "['123', '456']" in result.stdout + + @pytest.mark.asyncio + async def test_regex_sub(self, sandbox): + """regex_sub replaces matches.""" + result = await sandbox.execute( + r"out = await regex_sub(r'\d+', 'X', 'abc 123 def 456')" + "\nprint(out)" + ) + assert result.success + assert "abc X def X" in result.stdout + + @pytest.mark.asyncio + async def test_regex_search_match(self, sandbox): + """regex_search returns match dict when pattern matches.""" + result = await sandbox.execute( + r"m = await regex_search(r'(\d+)', 'abc 123')" + "\nprint(m['group'])" + "\nprint(m['start'])" + "\nprint(m['end'])" + ) + assert result.success + assert "123" in result.stdout + assert "4" in result.stdout + assert "7" in result.stdout + + @pytest.mark.asyncio + async def test_regex_search_no_match(self, sandbox): + """regex_search returns None when pattern doesn't match.""" + result = await sandbox.execute( + r"m = await regex_search(r'\d+', 'abc')" + "\nprint(m is None)" + ) + assert result.success + assert "True" in result.stdout + + @pytest.mark.asyncio + async def test_regex_split(self, sandbox): + """regex_split splits on pattern.""" + result = await sandbox.execute( + "out = await regex_split(',', 'a,b,,c')\nprint(out)" + ) + assert result.success + assert "['a', 'b', '', 'c']" in result.stdout + + @pytest.mark.asyncio + async def test_regex_invalid_pattern(self, sandbox): + """Invalid regex pattern surfaces as an error.""" + result = await sandbox.execute("await regex_findall('[invalid', 'text')") + assert not result.success + assert result.stderr != "" + + class TestSandboxLLM: """Test llm() external function."""