utils/utils.py

import os
import re
import json
import logging

from PIL import Image
from typing import Any, Union


COLOR_GRAY = "\033[37m"
COLOR_BLUE = "\033[94m"
COLOR_RED = "\033[31m"
COLOR_GREEN = "\033[32m"
COLOR_RESET = "\033[0m"
JSON_PATTERN = r"```json\s*(.*?)```"

def extract_json_result(answer: str, logger: logging.Logger) -> Any:
    text_to_parse = answer
    match = re.search(JSON_PATTERN, answer, re.DOTALL)
    if match:
        text_to_parse = match.group(1).strip()

    try:
        parsed_json = json.loads(text_to_parse)
        return parsed_json
    except json.JSONDecodeError as e:
        logger.warning(f"{COLOR_RED}JSON parse error: {e}. Raw output: {answer[:500]}{COLOR_RESET}")
        return {"error": f"JSON parsing failed: {e}", "raw_output": answer.strip()}
    except Exception as e:
        logger.error(f"{COLOR_RED}Unexpected error during JSON extraction: {e}. Raw output: {answer[:500]}{COLOR_RESET}")
        return {"error": str(e), "raw_output": answer.strip()}


def load_image(image_or_path: Union[str, Image.Image]) -> Image.Image:
    if isinstance(image_or_path, str):
        if not os.path.exists(image_or_path):
            raise FileNotFoundError(f"Image file not found: {image_or_path}")
        return Image.open(image_or_path).convert("RGB")
    elif isinstance(image_or_path, Image.Image):
        return image_or_path.convert("RGB")
    else:
        raise TypeError(f"Invalid image input type: {type(image_or_path)}. Expected str or PIL.Image.Image.")


class PartialFormatter(dict):

    def __missing__(self, key):
        return "{" + key + "}"


def get_logger(name: str, log_file: str):
    # 1. Get a logger instance
    logger = logging.Logger(name)

    # 2. Create a console handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)  # Set the level for console output
    console_formatter = logging.Formatter(f"{COLOR_GREEN}[%(asctime)s][%(name)s][%(levelname)s]{COLOR_RESET} - %(message)s")
    console_handler.setFormatter(console_formatter)

    # 3. Create a file handler
    file_handler = logging.FileHandler(log_file)  # Specify the log file name
    file_handler.setLevel(logging.DEBUG)  # Set the level for file output (e.g., capture more details in the file)
    file_formatter = logging.Formatter("[%(asctime)s][%(name)s][%(levelname)s] - %(message)s")
    file_handler.setFormatter(file_formatter)

    # 4. Add handlers to the logger
    logger.addHandler(console_handler)
    logger.addHandler(file_handler)

    return logger