File size: 7,513 Bytes
3a43332 f22daae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import os
from typing import Dict, List, Optional, Tuple
import gradio as gr
from huggingface_hub import InferenceClient
from tavily import TavilyClient
from config import (
HTML_SYSTEM_PROMPT, GENERIC_SYSTEM_PROMPT, HTML_SYSTEM_PROMPT_WITH_SEARCH,
GENERIC_SYSTEM_PROMPT_WITH_SEARCH, FollowUpSystemPrompt
)
from chat_processing import (
history_to_messages, messages_to_history,
remove_code_block, apply_search_replace_changes, send_to_sandbox,
history_to_chatbot_messages, get_gradio_language
)
from file_processing import ( # file_processing.py
extract_text_from_file, create_multimodal_message,
)
from web_extraction import extract_website_content, enhance_query_with_search
# HF Inference Client
HF_TOKEN = os.getenv('HF_TOKEN')
def get_inference_client(model_id):
"""Return an InferenceClient with provider based on model_id."""
provider = "groq" if model_id == "moonshotai/Kimi-K2-Instruct" else "auto"
return InferenceClient(
provider=provider,
api_key=HF_TOKEN,
bill_to="huggingface"
)
# Tavily Search Client
TAVILY_API_KEY = os.getenv('TAVILY_API_KEY')
tavily_client = None
if TAVILY_API_KEY:
try:
tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
except Exception as e:
print(f"Failed to initialize Tavily client: {e}")
tavily_client = None
def generation_code(query: Optional[str], image: Optional[gr.Image], file: Optional[str], website_url: Optional[str], _setting: Dict[str, str], _history: Optional[List[Tuple[str, str]]], _current_model: Dict, enable_search: bool = False, language: str = "html"):
if query is None:
query = ''
if _history is None:
_history = []
# Check if there's existing HTML content in history to determine if this is a modification request
has_existing_html = False
if _history:
# Check the last assistant message for HTML content
last_assistant_msg = _history[-1][1] if len(_history) > 0 else ""
if '<!DOCTYPE html>' in last_assistant_msg or '<html' in last_assistant_msg:
has_existing_html = True
# Choose system prompt based on context
if has_existing_html:
# Use follow-up prompt for modifying existing HTML
system_prompt = FollowUpSystemPrompt
else:
# Use language-specific prompt
if language == "html":
system_prompt = HTML_SYSTEM_PROMPT_WITH_SEARCH if enable_search else HTML_SYSTEM_PROMPT
else:
system_prompt = GENERIC_SYSTEM_PROMPT_WITH_SEARCH.format(language=language) if enable_search else GENERIC_SYSTEM_PROMPT.format(language=language)
messages = history_to_messages(_history, system_prompt)
# Extract file text and append to query if file is present
file_text = ""
if file:
file_text = extract_text_from_file(file)
if file_text:
file_text = file_text[:5000] # Limit to 5000 chars for prompt size
query = f"{query}\n\n[Reference file content below]\n{file_text}"
# Extract website content and append to query if website URL is present
website_text = ""
if website_url and website_url.strip():
website_text = extract_website_content(website_url.strip())
if website_text and not website_text.startswith("Error"):
website_text = website_text[:8000] # Limit to 8000 chars for prompt size
query = f"{query}\n\n[Website content to redesign below]\n{website_text}"
elif website_text.startswith("Error"):
# Provide helpful guidance when website extraction fails
fallback_guidance = """
Since I couldn't extract the website content, please provide additional details about what you'd like to build:
1. What type of website is this? (e.g., e-commerce, blog, portfolio, dashboard)
2. What are the main features you want?
3. What's the target audience?
4. Any specific design preferences? (colors, style, layout)
This will help me create a better design for you."""
query = f"{query}\n\n[Error extracting website: {website_text}]{fallback_guidance}"
# Enhance query with search if enabled
enhanced_query = enhance_query_with_search(query, enable_search)
# Use dynamic client based on selected model
client = get_inference_client(_current_model["id"])
if image is not None:
messages.append(create_multimodal_message(enhanced_query, image))
else:
messages.append({'role': 'user', 'content': enhanced_query})
try:
completion = client.chat.completions.create(
model=_current_model["id"],
messages=messages,
stream=True,
max_tokens=5000
)
content = ""
for chunk in completion:
if chunk.choices[0].delta.content:
content += chunk.choices[0].delta.content
clean_code = remove_code_block(content)
if has_existing_html:
# Fallback: If the model returns a full HTML file, use it directly
if clean_code.strip().startswith("<!DOCTYPE html>") or clean_code.strip().startswith("<html"):
yield {
"code_output": gr.update(value=clean_code, language=get_gradio_language(language)),
"history_output": history_to_chatbot_messages(_history),
"sandbox": send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
}
else:
last_html = _history[-1][1] if _history else ""
modified_html = apply_search_replace_changes(last_html, clean_code)
clean_html = remove_code_block(modified_html)
yield {
"code_output": gr.update(value=clean_html, language=get_gradio_language(language)),
"history_output": history_to_chatbot_messages(_history),
"sandbox": send_to_sandbox(clean_html) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
}
else:
yield {
"code_output": gr.update(value=clean_code, language=get_gradio_language(language)),
"history_output": history_to_chatbot_messages(_history),
"sandbox": send_to_sandbox(clean_code) if language == "html" else "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML. Please download your code using the download button above.</div>",
}
# Final update
_history = messages_to_history(messages + [{'role': 'assistant', 'content': content}])
yield {
"code_output": remove_code_block(content),
"history": _history,
"sandbox": send_to_sandbox(remove_code_block(content)),
"history_output": history_to_chatbot_messages(_history),
}
except Exception as e:
error_message = f"Error: {str(e)}"
yield {
"code_output": error_message,
"history_output": history_to_chatbot_messages(_history),
} |