app.py

# app.py
# ────────────────────────────────────────────────────────────────
"""
AnyCoder / Shasha AI – Gradio interface
• Works with the model registry in models.py
• Uses inference.chat_completion for provider‑aware calls
• Supports prompt / file / website context
• 25+ target languages
• Optional Tavily web‑search enrichment
• Live HTML preview + chat history
"""

from __future__ import annotations
from typing import List, Tuple, Dict, Any

import gradio as gr

# local helpers
from models        import AVAILABLE_MODELS, find_model, ModelInfo
from inference     import chat_completion
from tavily_search import enhance_query_with_search
from utils         import (
    extract_text_from_file,
    extract_website_content,
    history_to_messages,
    history_to_chatbot_messages,
    apply_search_replace_changes,
    remove_code_block,
    parse_transformers_js_output,
    format_transformers_js_output,
)
from deploy        import send_to_sandbox

# ── constants ───────────────────────────────────────────────────
SUPPORTED_LANGS = [
    "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",
]

SYSTEM_PROMPTS = {
    "html": (
        "ONLY USE HTML, CSS AND JAVASCRIPT. "
        "Return exactly one complete page wrapped in ```html ...```."
    ),
    "transformers.js": (
        "Generate THREE fenced blocks in order: "
        "index.html, index.js, style.css for a transformers.js web‑app."
    ),
}

History = List[Tuple[str, str]]

# ── backend callback ────────────────────────────────────────────
def generate(
    prompt: str,
    file_path: str | None,
    website_url: str | None,
    model_name: str,
    enable_search: bool,
    language: str,
    history: History | None,
) -> Tuple[str, History, str, List[Dict[str, str]]]:
    """Main generation routine – called when the user clicks **Generate Code**."""
    history = history or []
    prompt  = (prompt or "").strip()

    # 1 system + conversation
    sys_p   = SYSTEM_PROMPTS.get(language, f"You are an expert {language} developer.")
    msgs    = history_to_messages(history, sys_p)

    ctx: List[str] = [prompt] if prompt else []
    if file_path:
        ctx += ["[File]", extract_text_from_file(file_path)[:5_000]]
    if website_url:
        html = extract_website_content(website_url)
        if not html.startswith("Error"):
            ctx += ["[Website]", html[:8_000]]

    user_msg = enhance_query_with_search("\n\n".join(ctx), enable_search)
    msgs.append({"role": "user", "content": user_msg})

    # 2 model call
    model: ModelInfo = find_model(model_name) or AVAILABLE_MODELS[0]
    try:
        raw = chat_completion(model.id, msgs)
    except Exception as exc:
        err = f"❌ **Error**\n```{exc}```"
        history.append((prompt, err))
        return "", history, "", history_to_chatbot_messages(history)

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

    history.append((prompt, code))
    return code, history, preview, history_to_chatbot_messages(history)


# ── Gradio UI ───────────────────────────────────────────────────
THEME = gr.themes.Soft(primary_hue="indigo")
CSS = """
body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,Arial,sans-serif}
#title{font-size:2.4rem;text-align:center;margin-top:1rem}
#subtitle{color:#56606d;text-align:center;margin-bottom:1.8rem}
"""

with gr.Blocks(theme=THEME, css=CSS, title="ShashaCode_Builder") as demo:
    state_hist: gr.State[History] = gr.State([])

    gr.Markdown("## 🚀 AnyCoder AI", elem_id="title")
    gr.Markdown("Your AI partner for generating, modifying & understanding code.", elem_id="subtitle")

    with gr.Row():
        # ── inputs sidebar
        with gr.Column(scale=1):
            gr.Markdown("### 1 · Model")
            model_dd = gr.Dropdown(
                choices=[m.name for m in AVAILABLE_MODELS],
                value=AVAILABLE_MODELS[0].name,
                label="AI Model",
            )

            gr.Markdown("### 2 · Context")
            with gr.Tabs():
                with gr.Tab("Prompt", id="tab-prompt"):
                    prompt_box = gr.Textbox(lines=6, placeholder="Describe what you need…")
                with gr.Tab("File", id="tab-file"):
                    file_box   = gr.File(type="filepath")
                with gr.Tab("Website", id="tab-site"):
                    url_box    = gr.Textbox(placeholder="https://example.com")

            gr.Markdown("### 3 · Output")
            lang_dd   = gr.Dropdown(SUPPORTED_LANGS, value="html", label="Target Language")
            search_ck = gr.Checkbox(label="Enable Tavily Web Search")

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

        # ── outputs
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("Code"):
                    code_out    = gr.Code(interactive=True)
                with gr.Tab("Live Preview"):
                    preview_out = gr.HTML()
                with gr.Tab("History"):
                    chat_out    = gr.Chatbot(type="messages")

    # ── wiring
    gen_btn.click(
        generate,
        inputs=[prompt_box, file_box, url_box, model_dd, search_ck, lang_dd, state_hist],
        outputs=[code_out, state_hist, preview_out, chat_out],
    )

    clear_btn.click(
        lambda: ("", None, "", "html", False, [], [], "", ""),
        outputs=[prompt_box, file_box, url_box, lang_dd, search_ck,
                 state_hist, code_out, preview_out, chat_out],
        queue=False,
    )

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