app.py

# app.py

"""
Main application file for SHASHA AI, 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.
"""

import gradio as gr
from typing import Optional, Dict, List, Tuple, Any

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

History = List[Tuple[str, str]]
Model = Dict[str, Any]

# Full list of supported languages for syntax highlighting & generation
SUPPORTED_LANGUAGES = [
    "python", "c", "cpp", "markdown", "latex", "json", "html", "css",
    "javascript", "jinja2", "typescript", "yaml", "dockerfile", "shell",
    "r", "sql", "sql-msSQL", "sql-mySQL", "sql-mariaDB", "sql-sqlite",
    "sql-cassandra", "sql-plSQL", "sql-hive", "sql-pgSQL", "sql-gql",
    "sql-gpSQL", "sql-sparkSQL", "sql-esper"
]

DEFAULT_SYSTEM_PROMPT = """
You are a helpful AI coding assistant. Generate clean, correct, and efficient code based on the user's request.
- Follow requirements precisely.
- Enclose final code in a single ```code``` block of the target language.
- Do not include any explanations outside the code block.
"""

def get_model_details(name: str) -> Optional[Model]:
    for m in AVAILABLE_MODELS:
        if m["name"] == name:
            return m
    return None

def generation_code(
    query: Optional[str],
    file: Optional[str],
    website_url: Optional[str],
    current_model: Model,
    enable_search: bool,
    language: str,
    history: Optional[History],
) -> Tuple[str, History, str, List[Dict[str, str]]]:
    query = query or ""
    history = history or []
    try:
        system_prompt = SYSTEM_PROMPTS.get(language, DEFAULT_SYSTEM_PROMPT)
        model_id = current_model["id"]

        # pick provider
        if model_id.startswith("openai/") or model_id in ("gpt-4", "gpt-3.5-turbo"):
            provider = "openai"
        elif model_id.startswith("gemini/"):
            provider = "gemini"
        elif model_id.startswith("fireworks-ai/"):
            provider = "fireworks-ai"
        else:
            provider = "huggingface"

        # assemble messages
        msgs = history_to_messages(history, system_prompt)
        ctx = query
        if file:
            txt = extract_text_from_file(file)
            ctx += f"\n\n[File]\n{txt[:5000]}"
        if website_url:
            txt = extract_website_content(website_url)
            if not txt.startswith("Error"):
                ctx += f"\n\n[Website]\n{txt[:8000]}"
        final_q = enhance_query_with_search(ctx, enable_search)
        msgs.append({"role": "user", "content": final_q})

        client = get_inference_client(model_id, provider)
        resp = client.chat.completions.create(
            model=model_id,
            messages=msgs,
            max_tokens=16000,
            temperature=0.1
        )
        content = resp.choices[0].message.content

    except Exception as e:
        err = f"❌ **Error:**\n```\n{e}\n```"
        history.append((query, err))
        return "", history, "", history_to_chatbot_messages(history)

    # post-process
    if language == "transformers.js":
        files = parse_transformers_js_output(content)
        code = format_transformers_js_output(files)
        preview = send_to_sandbox(files.get("index.html",""))
    else:
        clean = remove_code_block(content)
        if history and history[-1][1] and not history[-1][1].startswith("❌"):
            code = apply_search_replace_changes(history[-1][1], clean)
        else:
            code = clean
        preview = send_to_sandbox(code) if language == "html" else ""

    new_hist = history + [(query, code)]
    chat = history_to_chatbot_messages(new_hist)
    return code, new_hist, preview, chat

# custom CSS
CUSTOM_CSS = """
body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; }
#main_title { text-align: center; font-size: 2.5rem; margin: 1.5rem 0 0.5rem; }
#subtitle { text-align: center; color: #4a5568; margin-bottom: 2.5rem; }
.gradio-container { background: #f7fafc; }
#gen_btn { box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=CUSTOM_CSS, title="Shasha AI") as demo:
    history_state = gr.State([])
    initial = AVAILABLE_MODELS[0]
    model_state = gr.State(initial)

    gr.Markdown("# 🚀 Shasha AI", elem_id="main_title")
    gr.Markdown("Your personal AI partner for generating, modifying, and understanding code.", elem_id="subtitle")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Select Model")
            names = [m["name"] for m in AVAILABLE_MODELS]
            model_dd = gr.Dropdown(names, value=initial["name"], label="AI Model")

            gr.Markdown("### 2. Provide Context")
            with gr.Tabs():
                with gr.Tab("📝 Prompt"):
                    prompt_in = gr.Textbox(lines=7, placeholder="Describe your request...", show_label=False)
                with gr.Tab("📄 File"):
                    file_in = gr.File(type="filepath")
                with gr.Tab("🌐 Website"):
                    url_in = gr.Textbox(placeholder="https://example.com")

            gr.Markdown("### 3. Configure Output")
            lang_dd = gr.Dropdown(SUPPORTED_LANGUAGES, value="html", label="Target Language")
            search_chk = gr.Checkbox(label="Enable Web Search")

            with gr.Row():
                clr = gr.Button("Clear Session", variant="secondary")
                gen = gr.Button("Generate Code", variant="primary", elem_id="gen_btn")

        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("💻 Code"):
                    code_out = gr.Code(language=lambda: lang_dd.value, interactive=True)
                with gr.Tab("👁️ Live Preview"):
                    preview_out = gr.HTML()
                with gr.Tab("📜 History"):
                    chat_out = gr.Chatbot(type="messages")

    model_dd.change(lambda n: get_model_details(n) or initial, inputs=[model_dd], outputs=[model_state])

    gen.click(
        fn=generation_code,
        inputs=[prompt_in, file_in, url_in, model_state, search_chk, lang_dd, history_state],
        outputs=[code_out, history_state, preview_out, chat_out],
    )

    clr.click(
        fn=lambda: ("", None, "", [], "", "", []),
        outputs=[prompt_in, file_in, url_in, history_state, code_out, preview_out, chat_out],
        queue=False,
    )

if __name__ == "__main__":
    demo.queue().launch()