84 lines
2.2 KiB
Python
84 lines
2.2 KiB
Python
import asyncio
|
|
import openai
|
|
import websocket
|
|
import json
|
|
import traceback
|
|
import pprint
|
|
import os
|
|
import base64
|
|
from urllib.parse import urlparse
|
|
from typing import List, Dict, Any
|
|
from openai import utils
|
|
from openai.errors import AuthenticationError, NotFoundError, RateLimitError
|
|
|
|
# Constants
|
|
SHOULD_MOCK_AI_RESPONSE = os.getenv("MOCK_AI_RESPONSE", "false") == "true"
|
|
IS_PROD = os.getenv("PRODUCTION", "false") == "true"
|
|
openai_api_key = os.getenv("OPENAI_API_KEY")
|
|
openai_base_url = os.getenv("OPENAI_BASE_URL", "https://api.openai.com")
|
|
|
|
# Utility functions
|
|
def create_alt_url_mapping(html: str) -> Dict[str, str]:
|
|
# ...
|
|
|
|
def pprint_prompt(prompt_messages: List[Dict[str, Any]]) -> None:
|
|
# ...
|
|
|
|
def write_logs(prompt_messages: List[Dict[str, Any]], completion: str) -> None:
|
|
# ...
|
|
|
|
async def mock_completion(process_chunk: Callable[[str], None]) -> str:
|
|
# ...
|
|
|
|
async def generate_images(
|
|
completion: str,
|
|
api_key: str,
|
|
base_url: str,
|
|
image_cache: Dict[str, str],
|
|
) -> str:
|
|
# ...
|
|
|
|
async def throw_error(error_message: str) -> None:
|
|
# ...
|
|
|
|
async def handle_websocket(websocket: websocket.WebSocketClientProtocol) -> None:
|
|
# ...
|
|
|
|
# Parse the initial message
|
|
initial_message = await websocket.recv()
|
|
params = json.loads(initial_message)
|
|
|
|
# Check if we should generate images
|
|
should_generate_images = params.get("shouldGenerateImages", False)
|
|
|
|
# Generate the prompt messages
|
|
prompt_messages = []
|
|
|
|
# ...
|
|
|
|
# Process the response
|
|
async def process_chunk(chunk: str) -> None:
|
|
# ...
|
|
|
|
# Generate the code
|
|
if params["generationType"] == "update":
|
|
# ...
|
|
|
|
# Write the messages dict into a log so that we can debug later
|
|
write_logs(prompt_messages, completion)
|
|
|
|
# Send the updated HTML to the frontend
|
|
await websocket.send_json({"type": "setCode", "value": updated_html})
|
|
await websocket.send_json({"type": "status", "value": "Code generation complete."})
|
|
|
|
# Close the websocket
|
|
await websocket.close()
|
|
|
|
# Start the server
|
|
async def main() -> None:
|
|
async with websocket.connect("ws://localhost:8080") as websocket:
|
|
await handle_websocket(websocket)
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|