Improve file writing and activity replay

src/agent_tools.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import base64
+import csv
 import hashlib
+import io
 import json
 import os
 import re
@@ -1872,31 +1874,69 @@ def _write_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str
 def _coerce_write_file_content(arguments: dict[str, Any]) -> str:
     provided = [
         key
-        for key in ('content', 'content_lines', 'content_base64')
+        for key in (
+            'content_lines',
+            'content',
+            'content_base64',
+            'json_content',
+            'jsonl_records',
+            'csv_rows',
+        )
         if key in arguments and arguments[key] is not None
     ]
     if len(provided) != 1:
         raise ToolExecutionError(
-            'Provide exactly one of content, content_lines, or content_base64'
+            'Provide exactly one of content_lines, content, content_base64, json_content, jsonl_records, or csv_rows'
         )
     source = provided[0]
-    if source == 'content':
-        content = arguments['content']
-        if not isinstance(content, str):
-            raise ToolExecutionError('content must be a string')
-        return content
     if source == 'content_lines':
         content_lines = arguments['content_lines']
         if not isinstance(content_lines, list) or not all(isinstance(line, str) for line in content_lines):
             raise ToolExecutionError('content_lines must be an array of strings')
         return '\n'.join(content_lines)
-    encoded = arguments['content_base64']
-    if not isinstance(encoded, str):
-        raise ToolExecutionError('content_base64 must be a string')
-    try:
-        return base64.b64decode(encoded, validate=True).decode('utf-8')
-    except (ValueError, UnicodeDecodeError) as exc:
-        raise ToolExecutionError('content_base64 must be valid base64-encoded UTF-8 text') from exc
+    if source == 'content':
+        content = arguments['content']
+        if not isinstance(content, str):
+            raise ToolExecutionError('content must be a string')
+        return content
+    if source == 'content_base64':
+        encoded = arguments['content_base64']
+        if not isinstance(encoded, str):
+            raise ToolExecutionError('content_base64 must be a string')
+        try:
+            return base64.b64decode(encoded, validate=True).decode('utf-8')
+        except (ValueError, UnicodeDecodeError) as exc:
+            raise ToolExecutionError('content_base64 must be valid base64-encoded UTF-8 text') from exc
+    if source == 'json_content':
+        return json.dumps(arguments['json_content'], ensure_ascii=False, indent=2)
+    if source == 'jsonl_records':
+        records = arguments['jsonl_records']
+        if not isinstance(records, list) or not all(isinstance(record, dict) for record in records):
+            raise ToolExecutionError('jsonl_records must be an array of objects')
+        return '\n'.join(json.dumps(record, ensure_ascii=False, separators=(',', ':')) for record in records)
+    csv_rows = arguments['csv_rows']
+    if not isinstance(csv_rows, list) or not all(isinstance(row, list) for row in csv_rows):
+        raise ToolExecutionError('csv_rows must be an array of arrays')
+    csv_headers = arguments.get('csv_headers')
+    if csv_headers is not None and (
+        not isinstance(csv_headers, list)
+        or not all(isinstance(header, str) for header in csv_headers)
+    ):
+        raise ToolExecutionError('csv_headers must be an array of strings')
+    output = io.StringIO()
+    writer = csv.writer(output, lineterminator='\n')
+    if csv_headers:
+        writer.writerow(csv_headers)
+    writer.writerows([[_stringify_csv_cell(cell) for cell in row] for row in csv_rows])
+    return output.getvalue().rstrip('\n')
+
+
+def _stringify_csv_cell(value: Any) -> str:
+    if value is None:
+        return ''
+    if isinstance(value, (str, int, float, bool)):
+        return str(value)
+    return json.dumps(value, ensure_ascii=False, separators=(',', ':'))
 
 
 def _edit_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str: