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 '' in last_assistant_msg or '") or clean_code.strip().startswith("Preview is only available for HTML. Please download your code using the download button above.", } 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 "

Preview is only available for HTML. Please download your code using the download button above.

", } 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 "

Preview is only available for HTML. Please download your code using the download button above.

", } # 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), }