diff --git a/backend/.gitignore b/backend/.gitignore index d9005f2..128eab6 100644 --- a/backend/.gitignore +++ b/backend/.gitignore @@ -150,3 +150,7 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + + +# Temporary eval output +evals diff --git a/backend/eval.py b/backend/eval.py new file mode 100644 index 0000000..e1c7a81 --- /dev/null +++ b/backend/eval.py @@ -0,0 +1,72 @@ +# Load environment variables first +from typing import Any, Coroutine +from dotenv import load_dotenv + +load_dotenv() + +import base64 +import os +from llm import stream_openai_response +from prompts import assemble_prompt +import asyncio + +from utils import pprint_prompt + + +async def generate_code_core(image_url: str, stack: str) -> str: + prompt_messages = assemble_prompt(image_url, stack) + openai_api_key = os.environ.get("OPENAI_API_KEY") + openai_base_url = None + + pprint_prompt(prompt_messages) + + async def process_chunk(content: str): + pass + + if not openai_api_key: + raise Exception("OpenAI API key not found") + + completion = await stream_openai_response( + prompt_messages, + api_key=openai_api_key, + base_url=openai_base_url, + callback=lambda x: process_chunk(x), + ) + + return completion + + +async def image_to_data_url(filepath: str): + with open(filepath, "rb") as image_file: + encoded_string = base64.b64encode(image_file.read()).decode() + return f"data:image/png;base64,{encoded_string}" + + +async def main(): + EVALS_DIR = "./evals" + INPUT_DIR = EVALS_DIR + "/inputs" + OUTPUT_DIR = EVALS_DIR + "/outputs" + + # Get all the files in the directory (only grab pngs) + evals = [f for f in os.listdir(INPUT_DIR) if f.endswith(".png")] + + tasks: list[Coroutine[Any, Any, str]] = [] + for filename in evals: + filepath = os.path.join(INPUT_DIR, filename) + data_url = await image_to_data_url(filepath) + task = generate_code_core(data_url, "html_tailwind") + tasks.append(task) + + results = await asyncio.gather(*tasks) + + os.makedirs(OUTPUT_DIR, exist_ok=True) + + for filename, content in zip(evals, results): + # File name is derived from the original filename in evals + output_filename = f"{os.path.splitext(filename)[0]}.html" + output_filepath = os.path.join(OUTPUT_DIR, output_filename) + with open(output_filepath, "w") as file: + file.write(content) + + +asyncio.run(main())