diff --git a/src/google/adk/environment/_base_environment.py b/src/google/adk/environment/_base_environment.py index 1217d6115ad..8779bf8d7db 100644 --- a/src/google/adk/environment/_base_environment.py +++ b/src/google/adk/environment/_base_environment.py @@ -121,6 +121,37 @@ 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]: + """Reads lines from a file, avoiding loading the entire file if possible. + + Args: + path: Absolute or working-dir-relative path to the file. + start_line: First line to return (1-based, inclusive). Defaults to 1. + end_line: Last line to return (1-based, inclusive). Defaults to end of + file. + + Returns: + A tuple containing: + - A list of lines as bytes. + - The total number of lines in the file. + + Raises: + FileNotFoundError: If the file does not exist. + """ + data_bytes = await self.read_file(Path(path)) + lines_bytes = data_bytes.splitlines(keepends=True) + total = len(lines_bytes) + start = max(1, start_line) + end = total if end_line is None else min(total, end_line) + if start > total: + return [], total + return lines_bytes[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..bf98cbf7382 100644 --- a/src/google/adk/environment/_local_environment.py +++ b/src/google/adk/environment/_local_environment.py @@ -207,6 +207,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 +247,21 @@ 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 = None + ) -> tuple[list[bytes], int]: + selected_lines: list[bytes] = [] + total_lines = 0 + with open(path, 'rb') as f: + for line in f: + total_lines += 1 + if start_line <= total_lines and ( + end_line is None or total_lines <= end_line + ): + selected_lines.append(line) + return selected_lines, total_lines + @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..f63ce3891c5 100644 --- a/src/google/adk/tools/environment/_read_file_tool.py +++ b/src/google/adk/tools/environment/_read_file_tool.py @@ -114,13 +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) - end = min(total, end_line or total) + selected_bytes, total = await self._environment.read_file_lines( + path, start_line=start, end_line=end_line + ) if start > total: return { 'status': 'error', @@ -129,13 +126,13 @@ async def run_async( ), 'total_lines': total, } + end = min(total, end_line or total) if start > end: return { 'status': 'error', '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/tools/environment/test_read_file_tool.py b/tests/unittests/tools/environment/test_read_file_tool.py index 8cc66d52785..bd16322acc6 100644 --- a/tests/unittests/tools/environment/test_read_file_tool.py +++ b/tests/unittests/tools/environment/test_read_file_tool.py @@ -194,3 +194,67 @@ async def test_read_file_rejects_boolean_line_numbers( 'status': 'error', 'error': '`end_line` must be an integer if provided.', } + + @pytest.mark.asyncio + async def test_read_file_lines_empty_file(self, env: LocalEnvironment): + """Test reading from an empty file.""" + await env.write_file('empty.txt', '') + + tool = ReadFileTool(env) + result = await tool.run_async( + args={'path': 'empty.txt', 'start_line': 1}, + tool_context=None, + ) + + assert result == { + 'status': 'error', + 'error': '`start_line` 1 exceeds file length (0 lines).', + 'total_lines': 0, + } + + @pytest.mark.asyncio + async def test_read_file_lines_exceeds_total(self, env: LocalEnvironment): + """Test start_line exceeding total lines of file.""" + await env.write_file('sample.txt', 'line1\nline2\n') + + tool = ReadFileTool(env) + result = await tool.run_async( + args={'path': 'sample.txt', 'start_line': 5}, + tool_context=None, + ) + + assert result == { + 'status': 'error', + 'error': '`start_line` 5 exceeds file length (2 lines).', + 'total_lines': 2, + } + + @pytest.mark.asyncio + async def test_read_file_lines_order_violation(self, env: LocalEnvironment): + """Test start_line after end_line error handling.""" + await env.write_file('sample.txt', 'line1\nline2\n') + + tool = ReadFileTool(env) + result = await tool.run_async( + args={'path': 'sample.txt', 'start_line': 3, 'end_line': 2}, + tool_context=None, + ) + + # Note that start_line is 3, which is also exceeding the file length. + # Therefore, the start_line exceeding check is fired first. + assert result == { + 'status': 'error', + 'error': '`start_line` 3 exceeds file length (2 lines).', + 'total_lines': 2, + } + + # Now let's try with start_line = 2, end_line = 1 (valid lines but out of order) + result_out_of_order = await tool.run_async( + args={'path': 'sample.txt', 'start_line': 2, 'end_line': 1}, + tool_context=None, + ) + assert result_out_of_order == { + 'status': 'error', + 'error': '`start_line` (2) is after `end_line` (1).', + 'total_lines': 2, + }