Update app.py

app.py

@@ -1,7 +1,27 @@
-from typing import Optional, Dict, List, Tuple
+"""
+app.py
+
+Main application file for AnyCoder, a Gradio-based AI code generation tool.
+
+This application provides a user interface for generating code in various languages
+using different AI models. It supports inputs from text prompts, files, images,
+and websites, and includes features like web search enhancement and live code previews.
+
+Structure:
+- Imports & Configuration: Loads necessary libraries and constants.
+- Helper Functions: Small utility functions supporting the UI logic.
+- Core Application Logic: The main `generation_code` function that handles the AI interaction.
+- UI Layout: Defines the Gradio interface using `gr.Blocks`.
+- Event Wiring: Connects UI components to backend functions.
+- Application Entry Point: Launches the Gradio app.
+"""
+
 import gradio as gr
+from typing import Optional, Dict, List, Tuple, Any
 
-from constants import HTML_SYSTEM_PROMPT, AVAILABLE_MODELS, DEMO_LIST
+# --- Local Module Imports ---
+# These modules contain the application's configuration, clients, and utility functions.
+from constants import SYSTEM_PROMPTS, AVAILABLE_MODELS, DEMO_LIST
 from hf_client import get_inference_client
 from tavily_search import enhance_query_with_search
 from utils import (
@@ -16,141 +36,214 @@ from utils import (
 )
 from deploy import send_to_sandbox, load_project_from_url
 
-# Type aliases
+# --- Type Aliases for Readability ---
 History = List[Tuple[str, str]]
+Model = Dict[str, Any]
+
+
+# ==============================================================================
+# HELPER FUNCTIONS
+# ==============================================================================
+
+def get_model_details(model_name: str) -> Optional[Model]:
+    """Finds the full dictionary for a model given its name."""
+    for model in AVAILABLE_MODELS:
+        if model["name"] == model_name:
+            return model
+    return None
+
+
+# ==============================================================================
+# CORE APPLICATION LOGIC
+# ==============================================================================
 
-# Core generation function
-# This function signature correctly includes 'hf_token'. Gradio will automatically
-# provide this value if the user is logged in.
 def generation_code(
     query: Optional[str],
-    image: Optional[gr.Image],
     file: Optional[str],
     website_url: Optional[str],
-    _setting: Dict[str, str],
-    _history: Optional[History],
-    _current_model_name: str,
+    current_model: Model,
     enable_search: bool,
     language: str,
-    provider: str,
-    hf_token: str
+    history: Optional[History],
+    hf_token: str,
 ) -> Tuple[str, History, str, List[Dict[str, str]]]:
-    # Initialize inputs
-    if query is None:
-        query = ''
-    if _history is None:
-        _history = []
-
-    # System prompt and history
-    system_prompt = _setting.get('system', HTML_SYSTEM_PROMPT)
-    messages = history_to_messages(_history, system_prompt)
-
-    # File input
-    if file:
-        text = extract_text_from_file(file)
-        query += f"\n\n[File content]\n{text[:5000]}"
-
-    # Website input
-    if website_url:
-        text = extract_website_content(website_url)
-        if not text.startswith('Error'):
-            query += f"\n\n[Website content]\n{text[:8000]}"
-
-    # Web search enhancement
-    final_query = enhance_query_with_search(query, enable_search)
-    messages.append({'role': 'user', 'content': final_query})
-
-    # Model inference - using the provided user token for billing
-    client = get_inference_client(_current_model_name, provider, user_token=hf_token)
-    resp = client.chat.completions.create(
-        model=_current_model_name,
-        messages=messages,
-        max_tokens=10000
-    )
-    content = resp.choices[0].message.content
-
-    # Post-processing
-    has_existing = bool(_history)
+    """
+    The main function to handle a user's code generation request.
+
+    Args:
+        query: The user's text prompt.
+        file: An uploaded file for context.
+        website_url: A URL to scrape for context.
+        current_model: The dictionary of the currently selected AI model.
+        enable_search: Flag to enable web search for query enhancement.
+        language: The target programming language.
+        history: The existing conversation history.
+        hf_token: The logged-in user's Hugging Face token for billing.
+
+    Returns:
+        A tuple containing the generated code, updated history, preview HTML,
+        and formatted chatbot messages.
+    """
+    # 1. --- Initialization and Input Sanitization ---
+    query = query or ""
+    history = history or []
+
+    try:
+        # 2. --- System Prompt and Model Selection ---
+        system_prompt = SYSTEM_PROMPTS.get(language, SYSTEM_PROMPTS["default"])
+        model_id = current_model["id"]
+        provider = current_model["provider"]
+
+        # 3. --- Assemble Full Context for the AI ---
+        messages = history_to_messages(history, system_prompt)
+        context_query = query
+
+        if file:
+            text = extract_text_from_file(file)
+            context_query += f"\n\n[Attached File Content]\n{text[:5000]}"
+
+        if website_url:
+            text = extract_website_content(website_url)
+            if not text.startswith('Error'):
+                context_query += f"\n\n[Scraped Website Content]\n{text[:8000]}"
+
+        final_query = enhance_query_with_search(context_query, enable_search)
+        messages.append({'role': 'user', 'content': final_query})
+
+        # 4. --- AI Model Inference with Robust Error Handling ---
+        client = get_inference_client(model_id, provider, user_token=hf_token)
+        resp = client.chat.completions.create(
+            model=model_id,
+            messages=messages,
+            max_tokens=16384,  # Increased token limit for complex code
+            temperature=0.1   # Low temperature for more predictable, stable code
+        )
+        content = resp.choices[0].message.content
+
+    except Exception as e:
+        # If the API call fails, show a user-friendly error in the chat.
+        error_message = f"❌ **An error occurred:**\n\n```\n{str(e)}\n```\n\nPlease check your API keys, model selection, or try again."
+        history.append((query, error_message))
+        return "", history, "", history_to_chatbot_messages(history)
+
+    # 5. --- Post-process the AI's Output ---
     if language == 'transformers.js':
         files = parse_transformers_js_output(content)
         code_str = format_transformers_js_output(files)
-        preview_html = send_to_sandbox(files['index.html'])
+        preview_html = send_to_sandbox(files.get('index.html', ''))
     else:
-        clean = remove_code_block(content)
-        if has_existing:
-            clean = apply_search_replace_changes(_history[-1][1], clean)
-        code_str = clean
-        preview_html = send_to_sandbox(clean) if language == 'html' else ''
+        clean_code = remove_code_block(content)
+        if history and history[-1][1] not in (None, ""):
+            # Apply search/replace if a previous turn exists
+            code_str = apply_search_replace_changes(history[-1][1], clean_code)
+        else:
+            code_str = clean_code
+        preview_html = send_to_sandbox(code_str) if language == 'html' else ''
+
+    # 6. --- Update History and Final Outputs ---
+    updated_history = history + [(query, code_str)]
+    chat_messages = history_to_chatbot_messages(updated_history)
 
-    # Update history
-    new_history = _history + [(query, code_str)]
-    chat_msgs = history_to_chatbot_messages(new_history)
+    return code_str, updated_history, preview_html, chat_messages
 
-    return code_str, new_history, preview_html, chat_msgs
 
-# Build UI
-with gr.Blocks(theme=gr.themes.Base(), title="AnyCoder - AI Code Generator") as demo:
-    # State
+# ==============================================================================
+# UI LAYOUT
+# ==============================================================================
+
+with gr.Blocks(theme=gr.themes.Soft(), title="AnyCoder - AI Code Generator") as demo:
+    # --- State Management ---
+    # Using gr.State to hold non-visible data like conversation history
+    # and the full dictionary of the selected model.
     history_state = gr.State([])
-    setting_state = gr.State({'system': HTML_SYSTEM_PROMPT})
-    model_state = gr.State(AVAILABLE_MODELS[0]['id'])
+    # Initialize with the first model from our constants list
+    initial_model = AVAILABLE_MODELS[0]
+    model_state = gr.State(initial_model)
 
+    # --- UI Definition ---
     with gr.Sidebar():
-        gr.Markdown("## AnyCoder AI")
-        url_in = gr.Textbox(label="Load HF Space URL", placeholder="https://huggingface.co/spaces/user/project")
-        load_btn = gr.Button("Import Project")
-        load_status = gr.Markdown(visible=False)
-        gr.Markdown("---")
-        prompt_in = gr.Textbox(label="Prompt", lines=3)
-        file_in = gr.File(label="Reference file")
-        image_in = gr.Image(label="Design image")
-        url_site = gr.Textbox(label="Website URL")
-        search_chk = gr.Checkbox(label="Enable Web Search")
-        language_dd = gr.Dropdown(choices=["html", "python", "transformers.js"], value="html", label="Language")
-        model_dd = gr.Dropdown(choices=[m['name'] for m in AVAILABLE_MODELS], value=AVAILABLE_MODELS[0]['name'], label="Model")
-        gen_btn = gr.Button("Generate")
-        clr_btn = gr.Button("Clear")
+        gr.Markdown("## 🚀 AnyCoder AI")
+        gr.Markdown("Your personal AI partner for generating, modifying, and understanding code.")
+        
+        # Group models by category for a better user experience
+        model_choices = {}
+        for model in AVAILABLE_MODELS:
+            category = model.get("category", "Other")
+            if category not in model_choices:
+                model_choices[category] = []
+            model_choices[category].append(model["name"])
+
+        model_dd = gr.Dropdown(
+            choices=list(model_choices.values()),
+            value=initial_model["name"],
+            label="🤖 Select AI Model",
+            info="Different models have different strengths. Experiment!"
+        )
+
+        with gr.Accordion("🛠️ Inputs & Context", open=True):
+            prompt_in = gr.Textbox(label="Prompt", lines=3, placeholder="e.g., 'Create a dark-themed login form with a spinning loader.'")
+            file_in = gr.File(label="📎 Attach File (Optional)", type="filepath")
+            url_site = gr.Textbox(label="🌐 Scrape Website (Optional)", placeholder="https://example.com")
+        
+        with gr.Accordion("⚙️ Settings", open=False):
+            language_dd = gr.Dropdown(
+                choices=["html", "python", "transformers.js", "sql", "javascript", "css"],
+                value="html",
+                label="🎯 Target Language"
+            )
+            search_chk = gr.Checkbox(label="🧠 Enable Web Search", info="Enhances the AI's knowledge with real-time information.")
+        
+        with gr.Row():
+            gen_btn = gr.Button("Generate Code", variant="primary", scale=2)
+            clr_btn = gr.Button("Clear", variant="secondary", scale=1)
 
     with gr.Column():
         with gr.Tabs():
-            with gr.Tab("Code"):
-                code_out = gr.Code(label="Generated Code")
-            with gr.Tab("Preview"):
+            with gr.Tab("💻 Code", id="code_tab"):
+                code_out = gr.Code(label="Generated Code", language="html")
+            with gr.Tab("👁️ Live Preview", id="preview_tab"):
                 preview_out = gr.HTML(label="Live Preview")
-            with gr.Tab("History"):
-                # FIX: The chatbot now correctly uses the modern 'messages' format.
-                chat_out = gr.Chatbot(label="History", type="messages")
-
-    # Events
-    load_btn.click(
-        fn=lambda u: load_project_from_url(u),
-        inputs=[url_in],
-        outputs=[load_status, code_out, preview_out, url_in, history_state, chat_out]
-    )
-
-    def on_model_change(name: str) -> str:
-        for m in AVAILABLE_MODELS:
-            if m['name'] == name:
-                return m['id']
-        return AVAILABLE_MODELS[0]['id']
-
-    model_dd.change(fn=on_model_change, inputs=[model_dd], outputs=[model_state])
-
-    gen_btn.click(
-        fn=generation_code,
-        # FIX: The 'inputs' list now only contains Gradio components.
-        # The 'hf_token' is passed automatically by Gradio because it's
-        # in the function's signature.
-        inputs=[
-            prompt_in, image_in, file_in, url_site,
-            setting_state, history_state, model_state,
-            search_chk, language_dd, gr.State('auto')
-        ],
-        outputs=[code_out, history_state, preview_out, chat_out]
-    )
-
-    clr_btn.click(fn=lambda: ([], [], "", []), outputs=[history_state, chat_out, preview_out, code_out])
+            with gr.Tab("📜 History", id="history_tab"):
+                chat_out = gr.Chatbot(label="Conversation History", type="messages", bubble_full_width=False)
+
+
+# ==============================================================================
+# EVENT WIRING
+# ==============================================================================
+
+# Update the model_state when the user selects a new model from the dropdown.
+def on_model_change(model_name: str) -> Dict:
+    model_details = get_model_details(model_name)
+    return model_details or initial_model
+model_dd.change(fn=on_model_change, inputs=[model_dd], outputs=[model_state])
+
+# Update the syntax highlighting when the language changes.
+language_dd.change(fn=lambda lang: gr.Code(language=lang), inputs=[language_dd], outputs=[code_out])
+
+# The main event listener for the "Generate" button.
+gen_btn.click(
+    fn=generation_code,
+    # Note: `hf_token` is passed automatically by Gradio and is not listed here.
+    inputs=[
+        prompt_in, file_in, url_site,
+        model_state, search_chk, language_dd, history_state
+    ],
+    outputs=[code_out, history_state, preview_out, chat_out]
+)
+
+# Clear button functionality to reset the interface.
+def clear_session():
+    return "", [], "", [], None, ""
+clr_btn.click(
+    fn=clear_session,
+    outputs=[prompt_in, history_state, preview_out, chat_out, file_in, url_site]
+)
+
+
+# ==============================================================================
+# APPLICATION ENTRY POINT
+# ==============================================================================
 
 if __name__ == '__main__':
-    # FIX: The launch() method is now correct, without any extra arguments.
+    # Launch the Gradio app with queuing enabled for handling multiple users.
     demo.queue().launch()