From d70dd412527f657aff0510cffbad5a1c0d79b9bf Mon Sep 17 00:00:00 2001 From: DaEpic <107145976+DaEpicR@users.noreply.github.com> Date: Fri, 17 Nov 2023 19:43:05 +0000 Subject: [PATCH] Update main.py ## Improving Code Structure: 1-Group related imports together . 2-separate the code into functions to enhance readability and maintainability . --- backend/main.py | 155 ++++++++++++++++++++---------------------------- 1 file changed, 63 insertions(+), 92 deletions(-) diff --git a/backend/main.py b/backend/main.py index c6e5556..7a65fce 100644 --- a/backend/main.py +++ b/backend/main.py @@ -1,64 +1,95 @@ -# Load environment variables first -from dotenv import load_dotenv - -load_dotenv() - - +from datetime import datetime import json import os import traceback -from datetime import datetime from fastapi import FastAPI, WebSocket +from dotenv import load_dotenv from llm import stream_openai_response from mock import mock_completion from image_generation import create_alt_url_mapping, generate_images from prompts import assemble_prompt + app = FastAPI() - -# Useful for debugging purposes when you don't want to waste GPT4-Vision credits -# Setting to True will stream a mock response instead of calling the OpenAI API -SHOULD_MOCK_AI_RESPONSE = False +load_dotenv() -def write_logs(prompt_messages, completion): - # Get the logs path from environment, default to the current working directory - logs_path = os.environ.get("LOGS_PATH", os.getcwd()) +def get_openai_api_key(params): + return params.get("openAiApiKey") or os.environ.get("OPENAI_API_KEY") - # Create run_logs directory if it doesn't exist within the specified logs path + +def create_logs_directory(logs_path): logs_directory = os.path.join(logs_path, "run_logs") if not os.path.exists(logs_directory): os.makedirs(logs_directory) + return logs_directory - print("Writing to logs directory:", logs_directory) - # Generate a unique filename using the current timestamp within the logs directory +def write_logs(logs_directory, prompt_messages, completion): 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})) +async def send_status_message(websocket, message): + await websocket.send_json({"type": "status", "value": message}) + + +async def process_chunk(websocket, content): + await websocket.send_json({"type": "chunk", "value": content}) + + +async def generate_code(websocket, params, openai_api_key): + should_generate_images = params.get("isImageGenerationEnabled", True) + await send_status_message(websocket, "Generating code...") + + prompt_messages = assemble_prompt(params["image"]) + image_cache = {} + + if params["generationType"] == "update": + for index, text in enumerate(params["history"]): + 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: + completion = await mock_completion(lambda x: process_chunk(websocket, x)) + else: + completion = await stream_openai_response( + prompt_messages, api_key=openai_api_key, callback=lambda x: process_chunk(websocket, x) + ) + + logs_directory = create_logs_directory(os.environ.get("LOGS_PATH", os.getcwd())) + write_logs(logs_directory, prompt_messages, completion) + + try: + if should_generate_images: + await send_status_message(websocket, "Generating images...") + updated_html = await generate_images(completion, api_key=openai_api_key, image_cache=image_cache) + else: + updated_html = completion + + await websocket.send_json({"type": "setCode", "value": updated_html}) + await send_status_message(websocket, "Code generation complete.") + + except Exception as e: + traceback.print_exc() + print("Image generation failed", e) + await send_status_message(websocket, "Image generation failed but code is complete.") + + finally: + await websocket.close() + + @app.websocket("/generate-code") async def stream_code_test(websocket: WebSocket): await websocket.accept() - params = await websocket.receive_json() - # Get the OpenAI API key from the request. Fall back to environment variable if not provided. - # If neither is provided, we throw an error. - if params["openAiApiKey"]: - openai_api_key = params["openAiApiKey"] - print("Using OpenAI API key from client-side settings dialog") - else: - openai_api_key = os.environ.get("OPENAI_API_KEY") - if openai_api_key: - print("Using OpenAI API key from environment variable") - + openai_api_key = get_openai_api_key(params) if not openai_api_key: - print("OpenAI API key not found") await websocket.send_json( { "type": "error", @@ -67,64 +98,4 @@ async def stream_code_test(websocket: WebSocket): ) return - should_generate_images = ( - params["isImageGenerationEnabled"] - if "isImageGenerationEnabled" in params - else True - ) - - print("generating code...") - await websocket.send_json({"type": "status", "value": "Generating code..."}) - - async def process_chunk(content): - await websocket.send_json({"type": "chunk", "value": content}) - - prompt_messages = assemble_prompt(params["image"]) - - # Image cache for updates so that we don't have to regenerate images - image_cache = {} - - if params["generationType"] == "update": - # Transform into message format - # TODO: Move this to frontend - for index, text in enumerate(params["history"]): - 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: - completion = await mock_completion(process_chunk) - else: - completion = await stream_openai_response( - prompt_messages, - api_key=openai_api_key, - callback=lambda x: process_chunk(x), - ) - - # Write the messages dict into a log so that we can debug later - write_logs(prompt_messages, completion) - - try: - if should_generate_images: - await websocket.send_json( - {"type": "status", "value": "Generating images..."} - ) - updated_html = await generate_images( - completion, api_key=openai_api_key, image_cache=image_cache - ) - else: - updated_html = completion - await websocket.send_json({"type": "setCode", "value": updated_html}) - await websocket.send_json( - {"type": "status", "value": "Code generation complete."} - ) - except Exception as e: - traceback.print_exc() - print("Image generation failed", e) - await websocket.send_json( - {"type": "status", "value": "Image generation failed but code is complete."} - ) - finally: - await websocket.close() + await generate_code(websocket, params, openai_api_key)