app.py

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

# --- 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 (
    extract_text_from_file,
    extract_website_content,
    apply_search_replace_changes,
    history_to_messages,
    history_to_chatbot_messages,
    remove_code_block,
    parse_transformers_js_output,
    format_transformers_js_output
)
from deploy import send_to_sandbox, load_project_from_url

# --- 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
# ==============================================================================

def generation_code(
    query: Optional[str],
    file: Optional[str],
    website_url: Optional[str],
    current_model: Model,
    enable_search: bool,
    language: str,
    history: Optional[History],
    hf_token: str,
) -> Tuple[str, History, str, List[Dict[str, str]]]:
    """
    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.get('index.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)

    return code_str, updated_history, preview_html, chat_messages


# ==============================================================================
# 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([])
    # 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")
        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", 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", 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__':
    # Launch the Gradio app with queuing enabled for handling multiple users.
    demo.queue().launch()