fix up pretty printing

abi · abi · commit 0301f24fd89a · 2023-12-07T11:26:48.000-05:00
diff --git a/backend/main.py b/backend/main.py
@@ -1,6 +1,7 @@
 # Load environment variables first
 from dotenv import load_dotenv
 
+
 load_dotenv()
 
 
@@ -14,6 +15,7 @@
 import openai
 from llm import stream_openai_response
 from mock import mock_completion
+from utils import pprint_prompt
 from image_generation import create_alt_url_mapping, generate_images
 from prompts import assemble_prompt
 from routes import screenshot
@@ -183,7 +185,6 @@ async def process_chunk(content):
             prompt_messages += [
                 {"role": "assistant" if index % 2 == 0 else "user", "content": text}
             ]
-
         image_cache = create_alt_url_mapping(params["history"][-2])
 
     if SHOULD_MOCK_AI_RESPONSE:
diff --git a/backend/utils.py b/backend/utils.py
@@ -1,4 +1,9 @@
 import copy
+import json
+
+
+def pprint_prompt(prompt_messages):
+    print(json.dumps(truncate_data_strings(prompt_messages), indent=4))
 
 
 def truncate_data_strings(data):
@@ -10,9 +15,12 @@ def truncate_data_strings(data):
             # Recursively call the function if the value is a dictionary or a list
             if isinstance(value, (dict, list)):
                 cloned_data[key] = truncate_data_strings(value)
-            # Truncate the string if it starts with 'data:'
-            elif isinstance(value, str) and value.startswith("data:"):
-                cloned_data[key] = value[:20]
+            # Truncate the string if it it's long and add ellipsis and length
+            elif isinstance(value, str):
+                cloned_data[key] = value[:40]
+                if len(value) > 40:
+                    cloned_data[key] += "..." + f" ({len(value)} chars)"
+
     elif isinstance(cloned_data, list):
         # Process each item in the list
         cloned_data = [truncate_data_strings(item) for item in cloned_data]
