Outputter class does line wrapping incrementally

trustgraph-cli/trustgraph/cli/invoke_agent.py

@@ -14,6 +14,75 @@ default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
 default_user = 'trustgraph'
 default_collection = 'default'
 
+class Outputter:
+    def __init__(self, width=75, prefix="> "):
+        self.width = width
+        self.prefix = prefix
+        self.column = 0
+        self.word_buffer = ""
+        self.just_wrapped = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # Flush remaining word buffer
+        if self.word_buffer:
+            print(self.word_buffer, end="", flush=True)
+            self.column += len(self.word_buffer)
+            self.word_buffer = ""
+
+        # Add final newline if not at line start
+        if self.column > 0:
+            print(flush=True)
+            self.column = 0
+
+    def output(self, text):
+        for char in text:
+            # Handle whitespace (space/tab)
+            if char in (' ', '\t'):
+                # Flush word buffer if present
+                if self.word_buffer:
+                    # Check if word + space would exceed width
+                    if self.column + len(self.word_buffer) + 1 > self.width:
+                        # Wrap: newline + prefix
+                        print(flush=True)
+                        print(self.prefix, end="", flush=True)
+                        self.column = len(self.prefix)
+                        self.just_wrapped = True
+
+                    # Output word buffer
+                    print(self.word_buffer, end="", flush=True)
+                    self.column += len(self.word_buffer)
+                    self.word_buffer = ""
+
+                # Output the space
+                print(char, end="", flush=True)
+                self.column += 1
+                self.just_wrapped = False
+
+            # Handle newline
+            elif char == '\n':
+                if self.just_wrapped:
+                    # Skip this newline (already wrapped)
+                    self.just_wrapped = False
+                else:
+                    # Flush word buffer if any
+                    if self.word_buffer:
+                        print(self.word_buffer, end="", flush=True)
+                        self.word_buffer = ""
+
+                    # Output newline + prefix
+                    print(flush=True)
+                    print(self.prefix, end="", flush=True)
+                    self.column = len(self.prefix)
+                    self.just_wrapped = False
+
+            # Regular character - add to word buffer
+            else:
+                self.word_buffer += char
+                self.just_wrapped = False
+
 def wrap(text, width=75):
     if text is None: text = "n/a"
     out = textwrap.wrap(
@@ -41,30 +110,9 @@ async def question(
         output(wrap(question), "\U00002753 ")
         print()
 
-    # Track last chunk type and line buffer for incremental wrapping
+    # Track last chunk type and accumulated text for current message
     last_chunk_type = None
-    line_buffer = ""  # Buffer for current message being accumulated
+    current_message = ""
-
-    def flush_complete_lines(text, prefix, keep_last=True):
-        """Wrap text and output complete lines, optionally keeping last partial line in buffer."""
-        if not text:
-            return ""
-
-        lines = textwrap.wrap(text, width=75)
-        if not lines:
-            return text  # Nothing to wrap, keep in buffer
-
-        if keep_last and len(lines) > 0:
-            # Output all but last line (last might be incomplete)
-            for line in lines[:-1]:
-                print(f"{prefix} {line}")
-            # Return last line to keep in buffer
-            return lines[-1] if len(lines) > 0 else ""
-        else:
-            # Output all lines
-            for line in lines:
-                print(f"{prefix} {line}")
-            return ""
 
     def think(x):
         if verbose:
@@ -124,31 +172,22 @@ async def question(
 
                 # Check if we're switching to a new message type
                 if last_chunk_type != chunk_type:
-                    # Flush any remaining buffer from previous message type
+                    # When switching message types, flush accumulated message
-                    if line_buffer:
+                    if current_message:
-                        if last_chunk_type == "thought" and verbose:
+                        if last_chunk_type == "thought":
-                            flush_complete_lines(line_buffer, "\U0001f914", keep_last=False)
+                            think(current_message)
-                        elif last_chunk_type == "observation" and verbose:
+                        elif last_chunk_type == "observation":
-                            flush_complete_lines(line_buffer, "\U0001f4a1", keep_last=False)
+                            observe(current_message)
                         elif last_chunk_type == "answer":
-                            print(line_buffer)
+                            print(current_message)
-                        line_buffer = ""
+                            if not current_message.endswith('\n'):
-                        print()  # Blank line between message types
+                                print()
+                        current_message = ""
 
                     last_chunk_type = chunk_type
 
-                # Add chunk to buffer
+                # Accumulate content for current message type
-                line_buffer += content
+                current_message += content
-
-                # Flush complete lines while keeping last partial line
-                if chunk_type == "thought" and verbose:
-                    line_buffer = flush_complete_lines(line_buffer, "\U0001f914", keep_last=True)
-                elif chunk_type == "observation" and verbose:
-                    line_buffer = flush_complete_lines(line_buffer, "\U0001f4a1", keep_last=True)
-                elif chunk_type == "answer":
-                    # For answers, just print as-is without wrapping
-                    print(content, end="")
-                    line_buffer = ""
             else:
                 # Handle legacy format (backward compatibility)
                 if "thought" in response:
@@ -164,15 +203,16 @@ async def question(
                     raise RuntimeError(response["error"])
 
             if obj["complete"]:
-                # Flush any remaining buffer
+                # Flush any remaining message
-                if line_buffer:
+                if current_message:
-                    if last_chunk_type == "thought" and verbose:
+                    if last_chunk_type == "thought":
-                        flush_complete_lines(line_buffer, "\U0001f914", keep_last=False)
+                        think(current_message)
-                    elif last_chunk_type == "observation" and verbose:
+                    elif last_chunk_type == "observation":
-                        flush_complete_lines(line_buffer, "\U0001f4a1", keep_last=False)
+                        observe(current_message)
                     elif last_chunk_type == "answer":
-                        print(line_buffer)
+                        print(current_message)
-                    print()  # Final newline
+                        if not current_message.endswith('\n'):
+                            print()
                 break
 
         await ws.close()
@@ -266,4 +306,4 @@ def main():
         print("Exception:", e, flush=True)
 
 if __name__ == "__main__":
-    main()
+    main()