Some printing fixes

plasma-umass · Feb 24, 2024 · 13db5bf · 13db5bf
1 parent eaff622
commit 13db5bf
Showing 1 changed file with 20 additions and 14 deletions.
diff --git a/src/chatdbg/assistant/lite_assistant.py b/src/chatdbg/assistant/lite_assistant.py
@@ -51,7 +51,7 @@ def _print_to_file(file, indent):
             elif hasattr(message, "content"):
                 content = message.content
 
-            assert bool(tool_calls) != bool(content)
+            assert content or tool_calls
 
             # The longest role string is 'assistant'.
             max_role_length = 9
@@ -61,26 +61,32 @@ def _print_to_file(file, indent):
             role = message["role"].upper()
             role_indent = max_role_length - len(role)
 
-            if tool_calls:
-                print(
-                    f"{' ' * indent}[{role}]{' ' * role_indent} Function calls:",
-                    file=file,
-                )
-                for tool_call in tool_calls:
-                    arguments = json.loads(tool_call.function.arguments)
-                    print(
-                        f"{' ' * (subindent + 4)}{tool_call.function.name}({', '.join([f'{k}={v}' for k, v in arguments.items()])})",
-                        file=file,
-                    )
-            else:
+            if content:
                 content = llm_utils.word_wrap_except_code_blocks(
                     content, wrap - len(role) - indent - 3
                 )
                 first, *rest = content.split("\n")
                 print(f"{' ' * indent}[{role}]{' ' * role_indent} {first}", file=file)
                 for line in rest:
                     print(f"{' ' * subindent}{line}", file=file)
-            print("\n\n", file=file)
+                print()
+
+            if tool_calls:
+                if content:
+                    print(f"{' ' * subindent} Function calls:", file=file)
+                else:
+                    print(
+                        f"{' ' * indent}[{role}]{' ' * role_indent} Function calls:",
+                        file=file,
+                    )
+                for tool_call in tool_calls:
+                    arguments = json.loads(tool_call.function.arguments)
+                    print(
+                        f"{' ' * (subindent + 4)}{tool_call.function.name}({', '.join([f'{k}={v}' for k, v in arguments.items()])})",
+                        file=file,
+                    )
+                print()
+            print("\n", file=file)
 
         # None is the default file value for print().
         _print_to_file(None, indent)