|
|
|
|
|
""" |
|
A collection of utility functions for data manipulation and formatting. |
|
|
|
This module provides helpers for tasks like converting chat history formats, |
|
processing images for multimodal models, cleaning model outputs, and |
|
applying code modifications. |
|
""" |
|
import base64 |
|
import io |
|
import re |
|
from typing import Dict, List, Optional, Tuple |
|
|
|
import numpy as np |
|
from PIL import Image |
|
|
|
from config import SEARCH_START, DIVIDER, REPLACE_END, GRADIO_SUPPORTED_LANGUAGES |
|
|
|
|
|
History = List[Tuple[Optional[str], Optional[str]]] |
|
Messages = List[Dict[str, any]] |
|
|
|
|
|
|
|
def history_to_messages(history: History, system_prompt: str) -> Messages: |
|
"""Converts Gradio's history format to the list of messages format for an API call.""" |
|
messages = [{'role': 'system', 'content': system_prompt}] |
|
for user_msg, assistant_msg in history: |
|
|
|
if isinstance(user_msg, list): |
|
|
|
text_content = next((item.get("text", "") for item in user_msg if isinstance(item, dict) and item.get("type") == "text"), "") |
|
messages.append({'role': 'user', 'content': text_content}) |
|
elif user_msg: |
|
messages.append({'role': 'user', 'content': user_msg}) |
|
|
|
if assistant_msg: |
|
messages.append({'role': 'assistant', 'content': assistant_msg}) |
|
return messages |
|
|
|
def messages_to_history(messages: Messages) -> History: |
|
"""Converts a list of messages back to Gradio's history format.""" |
|
history = [] |
|
|
|
for i in range(1, len(messages), 2): |
|
user_msg = messages[i]['content'] |
|
assistant_msg = messages[i+1]['content'] if (i+1) < len(messages) else "" |
|
history.append((user_msg, assistant_msg)) |
|
return history |
|
|
|
|
|
|
|
def process_image_for_model(image_data: np.ndarray) -> str: |
|
"""Converts a NumPy image array to a base64-encoded string.""" |
|
pil_img = Image.fromarray(image_data) |
|
buffer = io.BytesIO() |
|
pil_img.save(buffer, format="PNG") |
|
img_str = base64.b64encode(buffer.getvalue()).decode("utf-8") |
|
return f"data:image/png;base64,{img_str}" |
|
|
|
|
|
|
|
def remove_code_block(text: str) -> str: |
|
"""Extracts code from a markdown-style code block.""" |
|
pattern = r'```(?:[a-zA-Z]+)?\n(.*?)\n```' |
|
match = re.search(pattern, text, re.DOTALL) |
|
if match: |
|
return match.group(1).strip() |
|
return text.strip() |
|
|
|
def apply_search_replace(original_code: str, change_block: str) -> str: |
|
"""Applies a single search-and-replace block to the code.""" |
|
try: |
|
parts = re.split(f"^{DIVIDER}$", change_block, flags=re.MULTILINE) |
|
if len(parts) != 2: return original_code |
|
|
|
search_part, replace_part = parts |
|
search_content = search_part.replace(SEARCH_START, "").strip() |
|
replace_content = replace_part.replace(REPLACE_END, "").strip() |
|
|
|
|
|
if not search_content: |
|
|
|
return replace_content + "\n" + original_code |
|
|
|
if search_content in original_code: |
|
return original_code.replace(search_content, replace_content) |
|
else: |
|
|
|
|
|
|
|
|
|
print(f"Warning: Search block not found:\n---\n{search_content}\n---") |
|
return original_code |
|
|
|
except Exception as e: |
|
print(f"Error applying changes: {e}") |
|
return original_code |
|
|
|
def get_gradio_language(language: str) -> Optional[str]: |
|
"""Returns the language name if supported by Gradio, otherwise None.""" |
|
return language if language in GRADIO_SUPPORTED_LANGUAGES else None |