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
from dotenv import load_dotenv
load_dotenv()
@ -14,6 +15,7 @@ from fastapi.responses import HTMLResponse
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 stream_code(websocket: WebSocket):
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:

View File

@ -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]