""" Generate images for the Screenshots section in lesson markdown files. """ import os import re import glob from typing import Dict, Any, List from slugify import slugify from utils.io import save_output, ensure_dir from chains.base import build_chain from utils.storage import write_bytes def extract_image_descriptions(markdown_content: str) -> List[str]: """ Extract image descriptions from the Screenshots section of a lesson markdown. Args: markdown_content: String containing the lesson markdown Returns: List of image descriptions (text inside square brackets) """ screenshot_pattern = r'\*\*Screenshots:\*\*(.*?)(?:\n\n|\Z)' screenshot_match = re.search(screenshot_pattern, markdown_content, re.DOTALL) if not screenshot_match: return [] screenshot_section = screenshot_match.group(1).strip() description_pattern = r'\[([^\]]+)\]' descriptions = re.findall(description_pattern, screenshot_section) return descriptions def extract_lesson_context(markdown_content: str, max_length: int = 1000) -> str: """ Extract key context from a lesson to help with image generation. Args: markdown_content: String containing the lesson markdown max_length: Maximum length of context to extract Returns: String containing the most relevant context for image generation """ title_match = re.search(r'## Lesson: (.*?)(?:\n|$)', markdown_content) title = title_match.group(1) if title_match else "Unknown Lesson" topic_match = re.search(r'\*\*Topic:\*\* (.*?)(?:\n|$)', markdown_content) topic = topic_match.group(1) if topic_match else "" intro_pattern = r'\*\*Intro:\*\*(.*?)(?:\*\*Content:|$)' intro_match = re.search(intro_pattern, markdown_content, re.DOTALL) intro = intro_match.group(1).strip() if intro_match else "" context = f"Lesson Title: {title}\nTopic: {topic}\n\nIntro: {intro}" if len(context) > max_length: context = context[:max_length] + "..." return context from openai import OpenAI import requests import base64 def generate_image_for_description( description: str, lesson_context: str, run_id: str ) -> Dict[str, Any]: """ Generate an image using OpenAI's GPT Image API (gpt-image-1 model). Args: description: Description of the image to generate lesson_context: Context from the lesson to inform image generation run_id: Run identifier Returns: Dictionary with image generation result """ prompt = f"{description}\nContext: {lesson_context}" try: client = OpenAI() response = client.images.generate( model="gpt-image-1", prompt=prompt, n=1, size="1024x1024" ) if not response.data: raise ValueError("No image data returned from image generation API") image_info = response.data[0] if image_info.url: image_url = image_info.url image_response = requests.get(image_url) image_response.raise_for_status() image_data = image_response.content elif image_info.b64_json: import base64 image_data = base64.b64decode(image_info.b64_json) else: raise ValueError("No image URL or base64 data returned from image generation API") return {"image_data": image_data} except Exception as e: print(f"Error generating image: {e}") raise def save_generated_image( image_data: bytes, description: str, lesson_slug: str, run_id: str, index: int ) -> str: """ Save a generated image to the appropriate directory. Args: image_data: Binary image data description: Original image description lesson_slug: Slug of the lesson run_id: Run identifier index: Index of this image within the lesson Returns: Path to the saved image """ short_desc = description[:50] safe_desc = slugify(short_desc) output_dir = f"output/lessons/{run_id}/images/{lesson_slug}" ensure_dir(output_dir) filename = f"image_{index:02d}_{safe_desc}.png" filepath = os.path.join(output_dir, filename) return write_bytes(filepath, image_data, content_type="image/png") def process_lesson( lesson_path: str, run_id: str ) -> Dict[str, Any]: """ Process a single lesson to generate images for its screenshot descriptions. Args: lesson_path: Path to the lesson markdown file run_id: Run identifier Returns: Dictionary with lesson info and image paths """ lesson_filename = os.path.basename(lesson_path) lesson_slug = os.path.splitext(lesson_filename)[0] with open(lesson_path, "r", encoding="utf-8") as f: markdown_content = f.read() descriptions = extract_image_descriptions(markdown_content) if not descriptions: print(f"No image descriptions found in {lesson_path}") return {"lesson_slug": lesson_slug, "images": []} lesson_context = extract_lesson_context(markdown_content) image_results = [] for i, description in enumerate(descriptions): print(f"Generating image {i+1}/{len(descriptions)} for {lesson_slug}: {description[:50]}...") try: result = generate_image_for_description( description=description, lesson_context=lesson_context, run_id=run_id ) # Placeholder: assume result contains 'image_data' as bytes image_path = save_generated_image( image_data=result["image_data"], description=description, lesson_slug=lesson_slug, run_id=run_id, index=i+1 ) image_results.append({ "description": description, "path": image_path, "success": True }) except Exception as e: print(f"Error generating image for {description[:50]}: {e}") image_results.append({ "description": description, "error": str(e), "success": False }) return { "lesson_slug": lesson_slug, "images": image_results } def create_images(run_id: str) -> Dict[str, Any]: """ Generate images for all lessons in a run. Args: run_id: Run identifier Returns: Dictionary with results of image generation """ lesson_pattern = f"output/lessons/{run_id}/lessons-markdown/*.md" lesson_files = glob.glob(lesson_pattern) if not lesson_files: print(f"No lesson markdown files found for run ID: {run_id}") return {"status": "error", "message": "No lessons found"} print(f"Found {len(lesson_files)} lesson files for run ID: {run_id}") all_results = [] for lesson_path in lesson_files: result = process_lesson( lesson_path=lesson_path, run_id=run_id ) all_results.append(result) results = {"lessons": all_results} save_output( data=results, pipeline="lessons", step="image_generation", run_id=run_id ) print(f"Image generation complete for run ID: {run_id}") return results