fix up pretty printing

This commit is contained in:
Abi Raja 2023-12-07 11:26:48 -05:00
parent a559cd3120
commit 0301f24fd8
2 changed files with 13 additions and 4 deletions

View File

@ -1,6 +1,7 @@
# Load environment variables first # Load environment variables first
from dotenv import load_dotenv from dotenv import load_dotenv
load_dotenv() load_dotenv()
@ -14,6 +15,7 @@ from fastapi.responses import HTMLResponse
import openai import openai
from llm import stream_openai_response from llm import stream_openai_response
from mock import mock_completion from mock import mock_completion
from utils import pprint_prompt
from image_generation import create_alt_url_mapping, generate_images from image_generation import create_alt_url_mapping, generate_images
from prompts import assemble_prompt from prompts import assemble_prompt
from routes import screenshot from routes import screenshot
@ -183,7 +185,6 @@ async def stream_code(websocket: WebSocket):
prompt_messages += [ prompt_messages += [
{"role": "assistant" if index % 2 == 0 else "user", "content": text} {"role": "assistant" if index % 2 == 0 else "user", "content": text}
] ]
image_cache = create_alt_url_mapping(params["history"][-2]) image_cache = create_alt_url_mapping(params["history"][-2])
if SHOULD_MOCK_AI_RESPONSE: if SHOULD_MOCK_AI_RESPONSE:

View File

@ -1,4 +1,9 @@
import copy import copy
import json
def pprint_prompt(prompt_messages):
print(json.dumps(truncate_data_strings(prompt_messages), indent=4))
def truncate_data_strings(data): 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 # Recursively call the function if the value is a dictionary or a list
if isinstance(value, (dict, list)): if isinstance(value, (dict, list)):
cloned_data[key] = truncate_data_strings(value) cloned_data[key] = truncate_data_strings(value)
# Truncate the string if it starts with 'data:' # Truncate the string if it it's long and add ellipsis and length
elif isinstance(value, str) and value.startswith("data:"): elif isinstance(value, str):
cloned_data[key] = value[:20] cloned_data[key] = value[:40]
if len(value) > 40:
cloned_data[key] += "..." + f" ({len(value)} chars)"
elif isinstance(cloned_data, list): elif isinstance(cloned_data, list):
# Process each item in the list # Process each item in the list
cloned_data = [truncate_data_strings(item) for item in cloned_data] cloned_data = [truncate_data_strings(item) for item in cloned_data]