diff --git a/src/google/adk/environment/_base_environment.py b/src/google/adk/environment/_base_environment.py index 1217d6115ad..7975b9563d2 100644 --- a/src/google/adk/environment/_base_environment.py +++ b/src/google/adk/environment/_base_environment.py @@ -121,6 +121,41 @@ async def read_file(self, path: Path) -> bytes: FileNotFoundError: If the file does not exist. """ + async def read_file_lines( + self, + path: str | Path, + start_line: int = 1, + end_line: int | None = None, + ) -> tuple[list[bytes], int]: + """Read a range of lines from a file. + + The default implementation reads the whole file via ``read_file`` and + slices it in memory. Subclasses backed by a real filesystem should + override this to stream the file instead, so that requesting a small + range from a large file does not require buffering it entirely. + + Args: + path: Absolute or working-dir-relative path to the file. + start_line: First line to return (1-based, inclusive). + end_line: Last line to return (1-based, inclusive). ``None`` means + through the end of the file. + + Returns: + A tuple of (selected lines as bytes including line endings, total + number of lines in the file). + + Raises: + FileNotFoundError: If the file does not exist. + """ + data = await self.read_file(path) + lines = data.splitlines(keepends=True) + total = len(lines) + start = max(1, start_line) + end = total if end_line is None else min(total, end_line) + if start > end: + return [], total + return lines[start - 1 : end], total + @abstractmethod async def write_file(self, path: Path, content: str | bytes) -> None: """Write content to a file in the environment's filesystem. diff --git a/src/google/adk/environment/_local_environment.py b/src/google/adk/environment/_local_environment.py index 33b99c56965..cc00d0d1d4e 100644 --- a/src/google/adk/environment/_local_environment.py +++ b/src/google/adk/environment/_local_environment.py @@ -23,6 +23,8 @@ import shutil import signal import tempfile +from typing import BinaryIO +from typing import Iterator from typing_extensions import override @@ -37,6 +39,50 @@ # down cannot itself block forever. _TERMINATE_GRACE_SECONDS = 5 +# Chunk size for streaming a file while looking for line boundaries, so a +# single very long line cannot force reading the rest of the file at once. +_LINE_SCAN_CHUNK_SIZE = 65536 + + +def _iter_binary_lines(f: BinaryIO) -> Iterator[bytes]: + """Yields lines the same way `bytes.splitlines(keepends=True)` would. + + Unlike iterating a binary file object directly (which only splits on + `b'\\n'`), this also splits on a bare `b'\\r'` and on `b'\\r\\n'`, matching + `bytes.splitlines()` so streamed reads agree with the whole-file path other + environments use. + """ + pending = b'' + while True: + chunk = f.read(_LINE_SCAN_CHUNK_SIZE) + if not chunk: + break + buf = pending + chunk + pos = 0 + end = len(buf) + while pos < end: + idx_n = buf.find(b'\n', pos) + idx_r = buf.find(b'\r', pos) + if idx_n == -1 and idx_r == -1: + break + if idx_r != -1 and (idx_n == -1 or idx_r < idx_n): + if idx_r == end - 1: + # A lone `\r` at the end of this chunk might be the start of a + # `\r\n` pair split across the chunk boundary; defer it. + break + if buf[idx_r + 1 : idx_r + 2] == b'\n': + yield buf[pos : idx_r + 2] + pos = idx_r + 2 + else: + yield buf[pos : idx_r + 1] + pos = idx_r + 1 + else: + yield buf[pos : idx_n + 1] + pos = idx_n + 1 + pending = buf[pos:] + if pending: + yield pending + def _signal_group(group: int, sig: int) -> None: """Signals every process left in a group, tolerating an empty one.""" @@ -207,6 +253,21 @@ async def read_file(self, path: str | Path) -> bytes: resolved = self._resolve_path(path) return await asyncio.to_thread(self._sync_read, resolved) + @override + async def read_file_lines( + self, + path: str | Path, + start_line: int = 1, + end_line: int | None = None, + ) -> tuple[list[bytes], int]: + if self._working_dir is None: + raise RuntimeError('`working_dir` is not set. Call initialize() first.') + + resolved = self._resolve_path(path) + return await asyncio.to_thread( + self._sync_read_lines, resolved, start_line, end_line + ) + @override async def write_file(self, path: str | Path, content: str | bytes) -> None: if self._working_dir is None: @@ -232,6 +293,22 @@ def _sync_read(path: Path) -> bytes: with open(path, 'rb') as f: return f.read() + @staticmethod + def _sync_read_lines( + path: Path, start_line: int, end_line: int | None + ) -> tuple[list[bytes], int]: + """Streams *path* line-by-line, keeping only the requested range.""" + start = max(1, start_line) + selected: list[bytes] = [] + total = 0 + with open(path, 'rb') as f: + for line in _iter_binary_lines(f): + total += 1 + if total < start or (end_line is not None and total > end_line): + continue + selected.append(line) + return selected, total + @staticmethod def _sync_write(path: Path, content: str | bytes) -> None: os.makedirs(path.parent, exist_ok=True) diff --git a/src/google/adk/tools/environment/_read_file_tool.py b/src/google/adk/tools/environment/_read_file_tool.py index d49dd0d054a..c7440547a31 100644 --- a/src/google/adk/tools/environment/_read_file_tool.py +++ b/src/google/adk/tools/environment/_read_file_tool.py @@ -114,12 +114,10 @@ async def run_async( } try: - # TODO: Avoid loading the entire file into memory to prevent OOM on large files. - data_bytes = await self._environment.read_file(path) - # Slice data_bytes by line boundaries before decoding. - lines_bytes = data_bytes.splitlines(keepends=True) - total = len(lines_bytes) start = max(1, start_line or 1) + selected_bytes, total = await self._environment.read_file_lines( + path, start_line=start, end_line=end_line + ) end = min(total, end_line or total) if start > total: return { @@ -135,7 +133,6 @@ async def run_async( 'error': f'`start_line` ({start}) is after `end_line` ({end}).', 'total_lines': total, } - selected_bytes = lines_bytes[start - 1 : end] lines = [ line_bytes.decode('utf-8', errors='replace') for line_bytes in selected_bytes diff --git a/tests/unittests/environment/test_local_environment.py b/tests/unittests/environment/test_local_environment.py index 5cdafda5af2..cc4770c5bda 100644 --- a/tests/unittests/environment/test_local_environment.py +++ b/tests/unittests/environment/test_local_environment.py @@ -127,6 +127,67 @@ async def test_read_nonexistent_raises(self, env: LocalEnvironment): await env.read_file(Path("does_not_exist.txt")) +class TestReadFileLines: + """Verify read_file_lines streams a range without buffering the file.""" + + @pytest.mark.asyncio + async def test_returns_requested_range_and_total(self, env: LocalEnvironment): + """Selects the requested 1-based, inclusive line range.""" + await env.write_file("lines.txt", "one\ntwo\nthree\nfour\n") + + selected, total = await env.read_file_lines( + "lines.txt", start_line=2, end_line=3 + ) + + assert selected == [b"two\n", b"three\n"] + assert total == 4 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "content", + [ + pytest.param(b"one\rtwo\rthree\rfour\r", id="lone_cr"), + pytest.param(b"one\r\ntwo\rthree\nfour\r\n", id="mixed"), + ], + ) + async def test_line_splitting_matches_splitlines_semantics( + self, env: LocalEnvironment, content: bytes + ): + """Streamed splitting agrees with `bytes.splitlines()`, CR-only included.""" + await env.write_file("lines.bin", content) + + selected, total = await env.read_file_lines( + "lines.bin", start_line=2, end_line=3 + ) + + expected_lines = content.splitlines(keepends=True) + assert selected == expected_lines[1:3] + assert total == len(expected_lines) + + @pytest.mark.asyncio + async def test_does_not_read_the_whole_file_at_once( + self, env: LocalEnvironment, monkeypatch: pytest.MonkeyPatch + ): + """A ranged read must not go through the full-buffer `read_file` path.""" + await env.write_file("lines.txt", "one\ntwo\nthree\nfour\n") + + def _fail_full_read(path: Path) -> bytes: + raise AssertionError( + "read_file_lines must not buffer the whole file into memory" + ) + + monkeypatch.setattr( + LocalEnvironment, "_sync_read", staticmethod(_fail_full_read) + ) + + selected, total = await env.read_file_lines( + "lines.txt", start_line=2, end_line=2 + ) + + assert selected == [b"two\n"] + assert total == 4 + + class TestExecuteTimeout: """Timeout and cancellation must reach the whole process tree.""" diff --git a/tests/unittests/tools/environment/test_read_file_tool.py b/tests/unittests/tools/environment/test_read_file_tool.py index 8cc66d52785..977c1c022bc 100644 --- a/tests/unittests/tools/environment/test_read_file_tool.py +++ b/tests/unittests/tools/environment/test_read_file_tool.py @@ -31,6 +31,7 @@ class _StubEnvironment(BaseEnvironment): def __init__(self, files: dict[str, bytes]): self._files = files self.execute_calls: list[str] = [] + self.read_file_calls: list[str] = [] @property def working_dir(self) -> Path: @@ -48,10 +49,29 @@ async def execute( async def read_file(self, path: Path) -> bytes: key = str(path) + self.read_file_calls.append(key) if key not in self._files: raise FileNotFoundError(key) return self._files[key] + async def read_file_lines( + self, + path: Path, + start_line: int = 1, + end_line: Optional[int] = None, + ) -> tuple[list[bytes], int]: + """Slices the in-memory file without ever calling `read_file`.""" + key = str(path) + if key not in self._files: + raise FileNotFoundError(key) + lines = self._files[key].splitlines(keepends=True) + total = len(lines) + start = max(1, start_line) + end = total if end_line is None else min(total, end_line) + if start > end: + return [], total + return lines[start - 1 : end], total + async def write_file(self, path: Path, content: str | bytes) -> None: del path, content raise NotImplementedError @@ -75,6 +95,7 @@ async def test_read_file_with_line_range_uses_direct_file_read(): 'total_lines': 4, } assert environment.execute_calls == [] + assert environment.read_file_calls == [] @pytest.mark.asyncio