From 9a43e15bf4847b965fead59d26e9ca75beefbdda Mon Sep 17 00:00:00 2001
From: you-n-g <you-n-g@users.noreply.github.com>
Date: Fri, 21 Mar 2025 00:41:52 +0800
Subject: [PATCH] feat: Add line length limit to shrink_text function and
 settings (#715)

---
 rdagent/core/conf.py       |  1 +
 rdagent/core/experiment.py |  1 +
 rdagent/utils/fmt.py       | 24 +++++++++++++++++++-----
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py
index 7eaff5cc..c51b92c1 100644
--- a/rdagent/core/conf.py
+++ b/rdagent/core/conf.py
@@ -78,6 +78,7 @@ class RDAgentSettings(ExtendedBaseSettings):
     # misc
     """The limitation of context stdout"""
     stdout_context_len: int = 400
+    stdout_line_len: int = 10000
 
 
 RD_AGENT_SETTINGS = RDAgentSettings()
diff --git a/rdagent/core/experiment.py b/rdagent/core/experiment.py
index b5bd2154..527937ab 100644
--- a/rdagent/core/experiment.py
+++ b/rdagent/core/experiment.py
@@ -266,6 +266,7 @@ class FBWorkspace(Workspace):
             shrink_text(
                 filter_redundant_text(stdout),
                 context_lines=RD_AGENT_SETTINGS.stdout_context_len,
+                line_len=RD_AGENT_SETTINGS.stdout_line_len,
             ),
             return_code,
         )
diff --git a/rdagent/utils/fmt.py b/rdagent/utils/fmt.py
index 3d382fb5..79af8657 100644
--- a/rdagent/utils/fmt.py
+++ b/rdagent/utils/fmt.py
@@ -3,19 +3,33 @@ Tools that support generating better formats.
 """
 
 
-def shrink_text(text: str, context_lines: int = 200) -> str:
+def shrink_text(text: str, context_lines: int = 200, line_len: int = 5000) -> str:
     """
     When the context is too long, hide the part in the middle.
 
-        text before
-        ... (XXXXX lines are hidden) ...
-        text after
+    >>> shrink_text("line1\\nline2\\nline3", context_lines=2, line_len=5)
+    'line1\\n... (1 lines are hidden) ...\\nline3'
+
+    >>> shrink_text("short line", context_lines=2, line_len=5)
+    'sh... (5 chars are hidden) ...ine'
+
+    >>> shrink_text("a" * 5010, context_lines=2, line_len=10)
+    'aaaaa... (5000 chars are hidden) ...aaaaa'
     """
+
     lines = text.splitlines()
     total_lines = len(lines)
 
+    new_lines = []
+    for line in lines:
+        if len(line) > line_len:
+            # If any line is longer than line_len, we can't shrink it
+            line = f"{line[:line_len // 2]}... ({len(line) - line_len} chars are hidden) ...{line[- line_len + line_len // 2:]}"
+        new_lines.append(line)
+    lines = new_lines
+
     if total_lines <= context_lines:
-        return text
+        return "\n".join(lines)
 
     # Calculate how many lines to show from start and end
     half_lines = context_lines // 2