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') if not HF_TOKEN: raise RuntimeError("HF_TOKEN environment variable is not set. Please set it to your Hugging Face API token.") def get_inference_client(model_id, provider="auto"): """Return an InferenceClient with provider based on model_id and user selection.""" 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 async 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", progress=gr.Progress(track_tqdm=True)): 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.", history_to_chatbot_messages(_history) ) # Final update _history = messages_to_history(messages + [{'role': 'assistant', 'content': content}]) final_code = remove_code_block(content) yield ( final_code, _history, send_to_sandbox(final_code), history_to_chatbot_messages(_history), ) except Exception as e: error_message = f"Error: {str(e)}" yield (error_message, _history, None, history_to_chatbot_messages(_history))