refactor

backend/logging/core.py

@@ -0,0 +1,23 @@
+from datetime import datetime
+import json
+import os
+from openai.types.chat import ChatCompletionMessageParam
+
+
+def write_logs(prompt_messages: list[ChatCompletionMessageParam], completion: str):
+    # Get the logs path from environment, default to the current working directory
+    logs_path = os.environ.get("LOGS_PATH", os.getcwd())
+
+    # Create run_logs directory if it doesn't exist within the specified logs path
+    logs_directory = os.path.join(logs_path, "run_logs")
+    if not os.path.exists(logs_directory):
+        os.makedirs(logs_directory)
+
+    print("Writing to logs directory:", logs_directory)
+
+    # Generate a unique filename using the current timestamp within the logs directory
+    filename = datetime.now().strftime(f"{logs_directory}/messages_%Y%m%d_%H%M%S.json")
+
+    # Write the messages dict into a new file for each run
+    with open(filename, "w") as f:
+        f.write(json.dumps({"prompt": prompt_messages, "completion": completion}))

backend/routes/generate_code.py

@@ -14,12 +14,11 @@ from llm import (
     stream_openai_response,
 )
 from openai.types.chat import ChatCompletionMessageParam
+from logging.core import write_logs
 from mock_llm import mock_completion
 from typing import Any, Callable, Coroutine, Dict, List, Literal, Union, cast, get_args
 from image_generation import create_alt_url_mapping, generate_images
 from prompts import assemble_imported_code_prompt, assemble_prompt
-from datetime import datetime
-import json
 from prompts.claude_prompts import VIDEO_PROMPT
 from prompts.types import Stack
 
@@ -31,25 +30,6 @@ from ws.constants import APP_ERROR_WEB_SOCKET_CODE  # type: ignore
 router = APIRouter()
 
 
-def write_logs(prompt_messages: List[ChatCompletionMessageParam], completion: str):
-    # Get the logs path from environment, default to the current working directory
-    logs_path = os.environ.get("LOGS_PATH", os.getcwd())
-
-    # Create run_logs directory if it doesn't exist within the specified logs path
-    logs_directory = os.path.join(logs_path, "run_logs")
-    if not os.path.exists(logs_directory):
-        os.makedirs(logs_directory)
-
-    print("Writing to logs directory:", logs_directory)
-
-    # Generate a unique filename using the current timestamp within the logs directory
-    filename = datetime.now().strftime(f"{logs_directory}/messages_%Y%m%d_%H%M%S.json")
-
-    # Write the messages dict into a new file for each run
-    with open(filename, "w") as f:
-        f.write(json.dumps({"prompt": prompt_messages, "completion": completion}))
-
-
 # Generate images and return updated completions
 async def process_completion(
     completion: str,
@@ -394,7 +374,7 @@ async def stream_code(websocket: WebSocket):
     completions = [extract_html_content(completion) for completion in completions]
 
     # Write the messages dict into a log so that we can debug later
-    write_logs(prompt_messages, completions[0])  # type: ignore
+    write_logs(prompt_messages, completions[0])
 
     try:
         image_generation_tasks = [