builder / chat_processing.py
mgbam's picture
Upload 4 files
20491fe verified
import re
from typing import Dict, List, Optional, Tuple
import base64
import numpy as np
from PIL import Image
import gradio as gr
from config import GRADIO_SUPPORTED_LANGUAGES, SEARCH_START, DIVIDER, REPLACE_END
History = List[Tuple[str, str]]
Messages = List[Dict[str, str]]
def get_gradio_language(language):
return language if language in GRADIO_SUPPORTED_LANGUAGES else None
def history_to_messages(history: History, system: str) -> Messages:
messages = [{'role': 'system', 'content': system}]
for h in history:
# Handle multimodal content in history
user_content = h[0]
if isinstance(user_content, list):
# Extract text from multimodal content
text_content = ""
for item in user_content:
if isinstance(item, dict) and item.get("type") == "text":
text_content += item.get("text", "")
user_content = text_content if text_content else str(user_content)
messages.append({'role': 'user', 'content': user_content})
messages.append({'role': 'assistant', 'content': h[1]})
return messages
def messages_to_history(messages: Messages) -> History:
assert messages[0]['role'] == 'system'
history = []
for q, r in zip(messages[1::2], messages[2::2]):
# Extract text content from multimodal messages for history
user_content = q['content']
if isinstance(user_content, list):
text_content = ""
for item in user_content:
if isinstance(item, dict) and item.get("type") == "text":
text_content += item.get("text", "")
user_content = text_content if text_content else str(user_content)
history.append((user_content, r['content']))
return history
def history_to_chatbot_messages(history: History) -> List[Dict[str, str]]:
"""Convert history tuples to chatbot message format"""
messages = []
for user_msg, assistant_msg in history:
# Handle multimodal content
if isinstance(user_msg, list):
text_content = ""
for item in user_msg:
if isinstance(item, dict) and item.get("type") == "text":
text_content += item.get("text", "")
user_msg = text_content if text_content else str(user_msg)
messages.append({"role": "user", "content": user_msg})
messages.append({"role": "assistant", "content": assistant_msg})
return messages
def remove_code_block(text):
# Try to match code blocks with language markers
patterns = [
r'```(?:html|HTML)\n([\s\S]+?)\n```', # Match ```html or ```HTML
r'```\n([\s\S]+?)\n```', # Match code blocks without language markers
r'```([\s\S]+?)```' # Match code blocks without line breaks
]
for pattern in patterns:
match = re.search(pattern, text, re.DOTALL)
if match:
extracted = match.group(1).strip()
return extracted
# If no code block is found, check if the entire text is HTML
if text.strip().startswith('<!DOCTYPE html>') or text.strip().startswith('<html') or text.strip().startswith('<'):
return text.strip()
return text.strip()
def clear_history():
return [], [], None, "" # Empty lists for both tuple format and chatbot messages, None for file, empty string for website URL
def update_image_input_visibility(model):
"""Update image input visibility based on selected model"""
is_ernie_vl = model.get("id") == "baidu/ERNIE-4.5-VL-424B-A47B-Base-PT"
is_glm_vl = model.get("id") == "THUDM/GLM-4.1V-9B-Thinking"
return gr.update(visible=is_ernie_vl or is_glm_vl)
def update_submit_button(query):
"""Enable submit button if query is not empty"""
return gr.update(interactive=bool(query))
def create_multimodal_message(text, image=None):
"""Create a multimodal message with text and optional image"""
if image is None:
return {"role": "user", "content": text}
from file_processing import process_image_for_model
content = [
{
"type": "text",
"text": text
},
{
"type": "image_url",
"image_url": {
"url": process_image_for_model(image)
}
}
]
return {"role": "user", "content": content}
def apply_search_replace_changes(original_html: str, changes_text: str) -> str:
"""Apply search/replace changes to HTML content"""
if not changes_text.strip():
return original_html
# Split the changes text into individual search/replace blocks
blocks = []
current_block = ""
lines = changes_text.split('\n')
for line in lines:
if line.strip() == SEARCH_START:
if current_block.strip():
blocks.append(current_block.strip())
current_block = line + '\n'
elif line.strip() == REPLACE_END:
current_block += line + '\n'
blocks.append(current_block.strip())
current_block = ""
else:
current_block += line + '\n'
if current_block.strip():
blocks.append(current_block.strip())
modified_html = original_html
for block in blocks:
if not block.strip():
continue
# Parse the search/replace block
lines = block.split('\n')
search_lines = []
replace_lines = []
in_search = False
in_replace = False
for line in lines:
if line.strip() == SEARCH_START:
in_search = True
in_replace = False
elif line.strip() == DIVIDER:
in_search = False
in_replace = True
elif line.strip() == REPLACE_END:
in_replace = False
elif in_search:
search_lines.append(line)
elif in_replace:
replace_lines.append(line)
# Apply the search/replace
if search_lines:
search_text = '\n'.join(search_lines).strip()
replace_text = '\n'.join(replace_lines).strip()
if search_text in modified_html:
modified_html = modified_html.replace(search_text, replace_text)
else:
print(f"Warning: Search text not found in HTML: {search_text[:100]}...")
return modified_html
def send_to_sandbox(code):
# Add a wrapper to inject necessary permissions and ensure full HTML
wrapped_code = f"""
<!DOCTYPE html>
<html>
<head>
<meta charset=\"UTF-8\">
<meta name=\"viewport\" content=\"width=device-width, initial-scale=1.0\">
<script>
// Safe localStorage polyfill
const safeStorage = {{
_data: {{}},
getItem: function(key) {{ return this._data[key] || null; }},
setItem: function(key, value) {{ this._data[key] = value; }},
removeItem: function(key) {{ delete this._data[key]; }},
clear: function() {{ this._data = {{}}; }}
}};
Object.defineProperty(window, 'localStorage', {{
value: safeStorage,
writable: false
}});
window.onerror = function(message, source, lineno, colno, error) {{
console.error('Error:', message);
}};
</script>
</head>
<body>
{code}
</body>
</html>
"""
encoded_html = base64.b64encode(wrapped_code.encode('utf-8')).decode('utf-8')
data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
iframe = f'<iframe src="{data_uri}" width="100%" height="920px" sandbox="allow-scripts allow-same-origin allow-forms allow-popups allow-modals allow-presentation" allow="display-capture"></iframe>'
return iframe