import argparse import json import logging import os import threading from concurrent.futures import ThreadPoolExecutor, as_completed, TimeoutError as FuturesTimeoutError from typing import Optional from dotenv import load_dotenv from openai import OpenAI from tqdm import tqdm from tqdm.contrib.logging import logging_redirect_tqdm load_dotenv() def load_posts_from_json(file_path: str, n: int | None = None, where: dict | None = None) -> list[dict]: logging.info(f"Loading posts from {file_path}") with open(file_path, "r") as file: posts = json.load(file) total_posts = len(posts) # Filter posts if where dict is provided if where: filtered_posts = [] for post in posts: match = True for key, value in where.items(): if post.get(key) != value: match = False break if match: filtered_posts.append(post) posts = filtered_posts logging.info(f"Filtered to {len(posts)} posts (out of {total_posts} total) matching where criteria: {where}") # Limit posts if n is provided if n is not None and n > 0: original_count = len(posts) posts = posts[:n] logging.info(f"Limited to first {n} posts (out of {original_count} after filtering)") logging.info(f"Loaded {len(posts)} post(s) from {file_path}") return posts def format_code_blocks_for_posts( posts: list[dict], openai_client: OpenAI, sql_file_path: Optional[str] = None, sql_write_lock: Optional[threading.Lock] = None, max_workers: Optional[int] = None, ) -> list[dict]: """ Format code blocks for multiple posts using multi-threaded concurrent processing. Args: posts: List of post dictionaries to format openai_client: OpenAI client instance sql_file_path: Optional path to SQL file for incremental writing sql_write_lock: Optional lock for SQL file writing max_workers: Maximum number of worker threads (default: CPU count) Returns: List of formatted post dictionaries """ logging.info(f"Starting to format code blocks for {len(posts)} post(s)") if max_workers is None: max_workers = os.cpu_count() or 1 logging.info(f"Using {max_workers} worker threads for concurrent processing") def process_single_post(idx: int, post: dict) -> tuple[int, dict]: """Process a single post and return its index and formatted result.""" post_id = post.get("id", "unknown") post_title = post.get("title", "untitled")[:50] logging.info(f"Starting to process post {idx + 1}/{len(posts)} (ID: {post_id}, Title: {post_title}...)") try: formatted_post = format_code_blocks_for_post(post, openai_client) logging.info(f"Completed post {idx + 1}/{len(posts)} (ID: {post_id}, Title: {post_title}...)") # Write SQL immediately if file path and lock are provided if sql_file_path and sql_write_lock: sql_statement = generate_sql_update_statement(formatted_post) if sql_statement: append_sql_statement(sql_file_path, sql_statement, sql_write_lock) return idx, formatted_post except Exception as e: logging.error(f"Error processing post {idx + 1} (ID: {post_id}): {e}", exc_info=True) return idx, post # Return original post on error # Process all posts concurrently with progress bar with logging_redirect_tqdm(): # Create progress bar pbar = tqdm(total=len(posts), desc="Formatting posts", unit="post") formatted_posts = [None] * len(posts) # Use ThreadPoolExecutor for concurrent post processing with ThreadPoolExecutor(max_workers=max_workers) as executor: # Submit all tasks future_to_idx = { executor.submit(process_single_post, idx, post): idx for idx, post in enumerate(posts) } # Process completed tasks with progress bar completed_count = 0 for future in as_completed(future_to_idx, timeout=None): idx = future_to_idx[future] try: # Add timeout to prevent hanging (5 minutes per post should be plenty) idx_result, formatted_post = future.result(timeout=300) formatted_posts[idx_result] = formatted_post completed_count += 1 pbar.update(1) logging.info(f"Progress: {completed_count}/{len(posts)} posts completed") except FuturesTimeoutError: completed_count += 1 logging.error(f"Post {idx + 1} timed out after 5 minutes, using original post") formatted_posts[idx] = posts[idx] # Keep original post on timeout pbar.update(1) logging.info(f"Progress: {completed_count}/{len(posts)} posts completed (after timeout)") except Exception as e: completed_count += 1 logging.error(f"Post {idx + 1} failed: {e}", exc_info=True) formatted_posts[idx] = posts[idx] # Keep original post on error pbar.update(1) logging.info(f"Progress: {completed_count}/{len(posts)} posts completed (after error)") pbar.close() # Fill any None values with original posts (shouldn't happen, but safety check) for idx, post in enumerate(formatted_posts): if post is None: logging.warning(f"Post {idx + 1} is None, using original") formatted_posts[idx] = posts[idx] logging.info(f"Completed formatting {len(formatted_posts)} post(s)") return formatted_posts def format_code_blocks_for_post(post: dict, openai_client: OpenAI) -> dict: """Format code blocks in a single post with concurrent processing of code blocks.""" if "content" not in post: logging.debug("Post has no content field, skipping") return post # Parse the content JSON string try: content = json.loads(post["content"]) except (json.JSONDecodeError, TypeError) as e: logging.warning(f"Failed to parse content JSON: {e}") return post # Collect all code blocks with their paths code_blocks_info: list[tuple[tuple, dict]] = [] # (path, block) def collect_code_blocks(node: dict | list, path: tuple = ()): """Collect all code blocks with their paths for concurrent processing.""" if isinstance(node, list): for idx, item in enumerate(node): collect_code_blocks(item, path + (idx,)) elif isinstance(node, dict): if node.get("type") == "code": code_blocks_info.append((path, node)) # Recursively process children for key, value in node.items(): if isinstance(value, (dict, list)): collect_code_blocks(value, path + (key,)) collect_code_blocks(content) if not code_blocks_info: # No code blocks found, return original post return post # Format all code blocks concurrently using ThreadPoolExecutor def format_code_block_sync(path_and_block: tuple[tuple, dict]) -> tuple[tuple, dict]: """Format a single code block.""" path, code_block = path_and_block formatted_block = format_code_block(code_block, openai_client) return path, formatted_block logging.info(f"Found {len(code_blocks_info)} code block(s) in this post, formatting concurrently...") # Use ThreadPoolExecutor for concurrent code block formatting formatted_blocks = [] with ThreadPoolExecutor(max_workers=min(len(code_blocks_info), 20)) as executor: future_to_info = { executor.submit(format_code_block_sync, info): info for info in code_blocks_info } for future in as_completed(future_to_info): path, original_block = future_to_info[future] try: # Add timeout to prevent hanging (2 minutes per code block) result = future.result(timeout=120) formatted_blocks.append(result) except FuturesTimeoutError: logging.error("Code block formatting timed out after 2 minutes, using original block") # Keep original block on timeout formatted_blocks.append((path, original_block)) except Exception as e: logging.error(f"Error formatting code block: {e}", exc_info=True) # Keep original block on error formatted_blocks.append((path, original_block)) # Create mapping from path to formatted block formatted_map = {path: formatted for path, formatted in formatted_blocks} # Update the content with formatted code blocks using path-based lookup def update_node_with_formatted_blocks(node: dict | list, path: tuple = ()) -> dict | list: """Update the node tree with formatted code blocks using path mapping.""" if path in formatted_map: # This node should be replaced with formatted block return formatted_map[path] if isinstance(node, list): return [ update_node_with_formatted_blocks(item, path + (idx,)) for idx, item in enumerate(node) ] elif isinstance(node, dict): # Recursively update children processed = {} for key, value in node.items(): if isinstance(value, (dict, list)): processed[key] = update_node_with_formatted_blocks(value, path + (key,)) else: processed[key] = value return processed return node formatted_content = update_node_with_formatted_blocks(content) post = post.copy() post["content"] = json.dumps(formatted_content, ensure_ascii=False) logging.info(f"Formatted {len(code_blocks_info)} code block(s) in this post") return post def escape_sql_string(value: str) -> str: r""" Escape special characters in SQL string for longtext field. For MySQL/TiDB string literals: - Single quote (') -> '' (double single quote) - Backslash (\) -> \\ (double backslash) - for compatibility across SQL_MODE """ # First escape backslashes, then escape single quotes # Order matters: backslashes first, then single quotes return value.replace("\\", "\\\\").replace("'", "''") def generate_sql_update_statement(post: dict) -> str | None: """Generate SQL UPDATE statement for a single post.""" slug = post.get("slug") content = post.get("content") if not slug: logging.warning(f"Skipping post with missing slug: {post.get('id', 'unknown')}") return None if "content" not in post: logging.warning(f"Skipping post {slug} with missing content") return None # Escape single quotes in the content and slug escaped_content = escape_sql_string(content) escaped_slug = escape_sql_string(slug) sql = f"UPDATE posts SET content = '{escaped_content}' WHERE slug = '{escaped_slug}';" return sql def init_sql_file(file_path: str) -> None: """Initialize SQL file with header comments.""" # Create directory if it doesn't exist directory = os.path.dirname(file_path) if directory: os.makedirs(directory, exist_ok=True) header = "-- SQL script to update posts with formatted code blocks\n" header += "-- Generated automatically by format_code script\n" header += "\n" with open(file_path, "w", encoding="utf-8") as f: f.write(header) logging.info(f"Initialized SQL file: {file_path}") def append_sql_statement(file_path: str, sql_statement: str, lock: threading.Lock) -> None: """Append SQL statement to file with lock protection.""" with lock: with open(file_path, "a", encoding="utf-8") as f: f.write(sql_statement + "\n") logging.debug(f"Appended SQL statement to {file_path}") def extract_text_from_node(node: dict | list | str) -> str: """Recursively extract text from a node structure.""" if isinstance(node, str): return node elif isinstance(node, dict): # If it has a "text" field, use it (but handle if it's a list) if "text" in node: text_value = node["text"] if isinstance(text_value, str): return text_value elif isinstance(text_value, list): # Recursively extract from list return "".join(extract_text_from_node(item) for item in text_value) else: return str(text_value) # If it has "children", extract from children elif "children" in node: return extract_text_from_node(node["children"]) else: # No text or children, return empty string return "" elif isinstance(node, list): # Recursively extract from each item in the list return "".join(extract_text_from_node(item) for item in node) else: # Fallback: convert to string return str(node) def format_code_block(code_block: dict, openai_client: OpenAI) -> dict: """Format a code block using OpenAI to detect language and format code.""" # Extract code text from children code_text = "" if "children" in code_block and isinstance(code_block["children"], list): code_text = "".join( extract_text_from_node(child) for child in code_block["children"] ) if not code_text.strip(): logging.debug("Code block is empty, skipping") return code_block original_lang = code_block.get("lang") code_preview = code_text[:50].replace("\n", " ") + ("..." if len(code_text) > 50 else "") logging.debug(f"Formatting code block (current lang: {original_lang}, preview: {code_preview})") try: # Call OpenAI to detect language and format code with timeout response = openai_client.chat.completions.create( model="gpt-4o-mini", messages=[ { "role": "system", "content": "You are a code analyzer. Analyze the provided code and determine its programming language. Then format the code according to the language's conventions. Return your response in JSON format with two fields: 'language' (the detected language name like 'python', 'bash', 'shell', 'rust', 'go', 'javascript', 'typescript', 'java', 'sql', 'yaml', 'json', etc.) and 'formatted_code' (the formatted code)." }, { "role": "user", "content": f"Analyze and format this code:\n\n```\n{code_text}\n```" } ], response_format={"type": "json_object"}, temperature=0.1, timeout=60.0, # 60 seconds timeout per API call ) result = json.loads(response.choices[0].message.content) detected_lang = result.get("language", "").lower() formatted_code = result.get("formatted_code", code_text) # Update the code block formatted_block = code_block.copy() formatted_block["lang"] = detected_lang formatted_block["children"] = [{"text": formatted_code}] lang_change = f" ({original_lang} -> {detected_lang})" if original_lang != detected_lang else f" (kept as {detected_lang})" logging.debug(f"Code block formatted, detected language: {detected_lang}{lang_change}") return formatted_block except Exception as e: # If OpenAI call fails, return the original block logging.error(f"Error formatting code block: {e}") return code_block def main(): parser = argparse.ArgumentParser(description="Format code blocks in blog posts") parser.add_argument( "-n", "--number", type=int, default=None, help="Number of posts to process (useful for testing). If not specified, all posts will be processed.", ) parser.add_argument( "-f", "--file", type=str, default="data/posts.json", help="Path to the JSON file containing posts. Default: data/posts.json", ) parser.add_argument( "-v", "--verbose", action="store_true", help="Enable verbose logging (debug level)", ) parser.add_argument( "-o", "--output-sql-file", type=str, default="data/format_code_update.sql", help="Path to output SQL file with UPDATE statements. Default: data/format_code_update.sql", ) parser.add_argument( "-w", "--where", type=str, default=None, help="JSON string to filter posts (e.g., '{\"has_code_block\": 1}'). Filters posts where field values match exactly.", ) args = parser.parse_args() # Parse where filter if provided where_filter = None if args.where: try: where_filter = json.loads(args.where) if not isinstance(where_filter, dict): raise ValueError("Where filter must be a JSON object") except json.JSONDecodeError as e: logging.error(f"Invalid JSON in --where argument: {e}") return except ValueError as e: logging.error(f"Invalid where filter: {e}") return # Configure logging log_level = logging.DEBUG if args.verbose else logging.INFO logging.basicConfig( level=log_level, format="%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) # Run main function run_main(args, where_filter) def run_main(args: argparse.Namespace, where_filter: dict | None): """Main function.""" logging.info("Starting code block formatting process") posts = load_posts_from_json(args.file, n=args.number, where=where_filter) # Create OpenAI client openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) # Initialize SQL file and create lock for concurrent writing sql_write_lock = threading.Lock() init_sql_file(args.output_sql_file) try: format_code_blocks_for_posts( posts, openai_client, sql_file_path=args.output_sql_file, sql_write_lock=sql_write_lock, ) logging.info("Formatting process completed successfully") logging.info(f"SQL update script written incrementally to {args.output_sql_file}") # logging.info(f"Formatted posts:\n{json.dumps(formatted_posts, indent=2, ensure_ascii=False)}") finally: # OpenAI client doesn't need explicit closing in sync version pass if __name__ == "__main__": main()