Update app.py

app.py

@@ -1,169 +1,207 @@
-# app.py
 """
-Main application file for Shasha AI, an AI‑assisted code‑generation tool built
-with Gradio.  
-• Generates code in dozens of languages with multiple OSS / proprietary models  
-• Accepts plain prompts, reference files, or a web‑site URL to redesign  
-• Optional Tavily web‑search augmentation  
-• Live HTML preview, import‑existing‑project, and one‑click Space deploy
+app.py  –  Gradio front‑end for “AnyCoder AI” (a.k.a. Shasha AI)
+
+UI   : single‑page, 3‑column layout
+Logo : assets/logo.png   (120 px wide, centred)
+SDK  : Gradio 5.38.2     (no `height=` arg on gr.Code)
 """
 
 from __future__ import annotations
 
-import os, time, urllib.parse, tempfile, webbrowser
-from typing import Optional, Dict, List, Tuple, Any
-
 import gradio as gr
-from huggingface_hub import HfApi
-from tavily import TavilyClient
+from typing import List, Tuple, Dict, Optional, Any
 
-# ──────────────────────────── local modules ──────────────────────────────
-from constants      import (
+# ── local helpers ----------------------------------------------------------
+from constants import (                           # all kept in one place
+    SEARCH_START, DIVIDER, REPLACE_END,
     HTML_SYSTEM_PROMPT, HTML_SYSTEM_PROMPT_WITH_SEARCH,
     TRANSFORMERS_JS_SYSTEM_PROMPT, TRANSFORMERS_JS_SYSTEM_PROMPT_WITH_SEARCH,
-    SVELTE_SYSTEM_PROMPT, SVELTE_SYSTEM_PROMPT_WITH_SEARCH,
     GENERIC_SYSTEM_PROMPT, GENERIC_SYSTEM_PROMPT_WITH_SEARCH,
-    TransformersJSFollowUpSystemPrompt, FollowUpSystemPrompt,
-    SEARCH_START, DIVIDER, REPLACE_END,
-    AVAILABLE_MODELS, DEMO_LIST, GRADIO_SUPPORTED_LANGUAGES,
+    SYSTEM_PROMPTS, FollowUpSystemPrompt,
+    TransformersJSFollowUpSystemPrompt,
+    AVAILABLE_MODELS, DEMO_LIST,
+    get_gradio_language,
 )
-from hf_client      import get_inference_client, HF_TOKEN
-from tavily_search  import enhance_query_with_search, tavily_client
+
+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,
     history_to_messages, history_to_chatbot_messages,
-    remove_code_block, parse_transformers_js_output,
-    format_transformers_js_output, parse_svelte_output,
-    format_svelte_output, apply_search_replace_changes,
-    apply_transformers_js_search_replace_changes, get_gradio_language,
+    remove_code_block, parse_transformers_js_output, format_transformers_js_output,
+    apply_search_replace_changes, apply_transformers_js_search_replace_changes,
 )
-from web_scraper    import extract_website_content
-from search_replace import apply_search_replace_changes  # alias kept for clarity
-from deploy         import send_to_sandbox, deploy_to_user_space
-from web_scraper    import extract_text_from_file, extract_text_from_image
-# ─────────────────────────────────────────────────────────────────────────
+from deploy         import send_to_sandbox
+from search_replace import SEARCH_START as SR_START  # just to avoid name clash
+                                                     # (optional import)
 
-
-# ==== type aliases ====
+# ── type aliases -----------------------------------------------------------
 History = List[Tuple[str, str]]
-Model   = Dict[str, Any]
-
-# ==== helpers ====
-def get_model_details(name: str) -> Model:
-    return next((m for m in AVAILABLE_MODELS if m["name"] == name), AVAILABLE_MODELS[0])
-
+ModelInfo = Dict[str, str]
+
+# ── generation core --------------------------------------------------------
+def generate_code(
+    prompt: str,
+    file_path: Optional[str],
+    website_url: Optional[str],
+    model: ModelInfo,
+    language: str,
+    enable_search: bool,
+    history: Optional[History],
+) -> Tuple[str, History, str, List[Dict[str, str]]]:
+
+    history = history or []
+    prompt = prompt or ""
+
+    # 1. choose system prompt  ------------------------------------------------
+    if history:
+        # modification request
+        if language == "transformers.js":
+            system_prompt = TransformersJSFollowUpSystemPrompt
+        else:
+            system_prompt = FollowUpSystemPrompt
+    else:
+        # fresh generation
+        if language == "html":
+            system_prompt = HTML_SYSTEM_PROMPT_WITH_SEARCH if enable_search else HTML_SYSTEM_PROMPT
+        elif language == "transformers.js":
+            system_prompt = TRANSFORMERS_JS_SYSTEM_PROMPT_WITH_SEARCH if enable_search else TRANSFORMERS_JS_SYSTEM_PROMPT
+        else:
+            system_prompt = (
+                GENERIC_SYSTEM_PROMPT_WITH_SEARCH.format(language=language)
+                if enable_search else GENERIC_SYSTEM_PROMPT.format(language=language)
+            )
 
-# --------------------  Gradio UI  ----------------------------------------
-CUSTOM_CSS = """
-#brand_logo{margin-right:.5rem;border-radius:8px}
-body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif}
-#main_title{font-size:2rem;margin:0}
-#subtitle{color:#4a5568;margin-bottom:2rem}
-.gradio-container{background:#f7fafc}
-#gen_btn{box-shadow:0 4px 6px rgba(0,0,0,.1)}
-"""
+    messages = history_to_messages(history, system_prompt)
+
+    # 2. augment prompt with file / website  ---------------------------------
+    if file_path:
+        file_txt = extract_text_from_file(file_path)[:5000]
+        prompt += f"\n\n[Reference file]\n{file_txt}"
+
+    if website_url:
+        site_ctx = extract_website_content(website_url.strip())
+        prompt += f"\n\n[Website]\n{site_ctx[:8000]}"
+
+    # 3. optional web‑search enrichment --------------------------------------
+    user_query = enhance_query_with_search(prompt, enable_search)
+    messages.append({"role": "user", "content": user_query})
+
+    # 4. call model -----------------------------------------------------------
+    client = get_inference_client(model["id"])
+    try:
+        resp = client.chat.completions.create(
+            model=model["id"],
+            messages=messages,
+            max_tokens=16_000,
+            temperature=0.1,
+        )
+        answer = resp.choices[0].message.content
+    except Exception as e:
+        err = f"❌ **Error:**\n```\n{e}\n```"
+        history.append((prompt, err))
+        return "", history, "", history_to_chatbot_messages(history)
+
+    # 5. post‑processing ------------------------------------------------------
+    if language == "transformers.js":
+        files = parse_transformers_js_output(answer)
+        code = format_transformers_js_output(files)
+        preview = send_to_sandbox(files["index.html"]) if files["index.html"] else ""
+    else:
+        clean = remove_code_block(answer)
+        if history and not history[-1][1].startswith("❌"):
+            clean = apply_search_replace_changes(history[-1][1], clean)
+        code = clean
+        preview = send_to_sandbox(code) if language == "html" else ""
+
+    history.append((prompt, code))
+    chat_msgs = history_to_chatbot_messages(history)
+
+    return code, history, preview, chat_msgs
+
+
+# ── UI ---------------------------------------------------------------------
+THEME = gr.themes.Base(primary_hue="indigo", font="Inter")
+
+with gr.Blocks(theme=THEME, title="AnyCoder AI") as demo:
+    state_hist  = gr.State([])             # History list
+    state_model = gr.State(AVAILABLE_MODELS[0])
+
+    # ––– Header with logo –––
+    with gr.Row():
+        gr.HTML(
+            '<div style="text-align:center; margin:1.2rem 0;">'
+            '<img src="assets/logo.png" alt="AnyCoder logo" style="width:120px;"><br>'
+            '<h1 style="margin:0.4rem 0 0; font-size:1.9rem;">AnyCoder AI</h1>'
+            '<p style="color:#555;">Your AI partner for generating, modifying &amp; understanding code.</p>'
+            '</div>'
+        )
 
-with gr.Blocks(
-    theme=gr.themes.Soft(primary_hue="blue"),
-    css=CUSTOM_CSS,
-    title="Shasha AI"
-) as demo:
-
-    # ────────── states ──────────
-    history_state : gr.State = gr.State([])
-    model_state   : gr.State = gr.State(AVAILABLE_MODELS[0])
-    provider_state: gr.State = gr.State("auto")
-
-    # ────────── header with logo ──────────
-    with gr.Row(elem_id="header"):
-        gr.Image(value="assets/logo.png",
-                 width=48, height=48,
-                 show_label=False, container=False,
-                 elem_id="brand_logo")
-        with gr.Column():
-            gr.Markdown("## 🚀 Shasha AI", elem_id="main_title")
-            gr.Markdown(
-                "Your AI partner for generating, modifying, and understanding code.",
-                elem_id="subtitle"
+    with gr.Row():
+        # ── Sidebar (column‑1) ───────────────────────────────────────────
+        with gr.Column(scale=1):
+            gr.Markdown("### 1 · Select Model")
+            dd_model = gr.Dropdown(
+                [m["name"] for m in AVAILABLE_MODELS],
+                value=AVAILABLE_MODELS[0]["name"],
+                label="AI Model",
             )
 
-    # ────────── sidebar (inputs) ──────────
-    with gr.Sidebar():
-        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 · Provide Context")
+            with gr.Tabs():
+                with gr.Tab("Prompt"):
+                    tb_prompt = gr.Textbox(lines=6, placeholder="Describe what you want to build…")
+                with gr.Tab("File"):
+                    fi_file = gr.File()
+                with gr.Tab("Website"):
+                    tb_url = gr.Textbox(placeholder="https://example.com")
+
+            gr.Markdown("### 3 · Configure Output")
+            dd_lang = gr.Dropdown(
+                GRADIO_SUPPORTED_LANGUAGES[:-1],  # drop trailing None
+                value="html",
+                label="Target Language",
+            )
+            cb_search = gr.Checkbox(label="Enable Tavily Web Search")
+
+            with gr.Row():
+                btn_clear = gr.Button("Clear Session", variant="secondary")
+                btn_gen   = gr.Button("Generate Code", variant="primary")
+
+        # ── Output / preview (column‑2) ──────────────────────────────────
+        with gr.Column(scale=2):
+            with gr.Tabs():
+                with gr.Tab("Code"):
+                    code_out = gr.Code(language="html", lines=25, label="Generated code")
+                with gr.Tab("Preview"):
+                    html_prev = gr.HTML()
+                with gr.Tab("History"):
+                    chat_out = gr.Chatbot(type="messages", height=400)
+
+    # ––– Quick‑start buttons –––
+    gr.Markdown("#### Quick Start Examples")
+    with gr.Row():
+        for demo in DEMO_LIST[:6]:
+            gr.Button(demo["title"], size="sm").click(
+                lambda d=demo: d["description"], outputs=tb_prompt
+            )
 
-        gr.Markdown("### 2 · Context")
-        with gr.Tabs():
-            with gr.Tab("📝 Prompt"):
-                prompt_in = gr.Textbox(lines=7, placeholder="Describe what you’d like…")
-            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 · Output")
-        lang_dd = gr.Dropdown(
-            choices=GRADIO_SUPPORTED_LANGUAGES,
-            value="html",
-            label="Target Language"
-        )
-        search_chk = gr.Checkbox(label="Enable Tavily Web Search")
-
-        with gr.Row():
-            clr_btn = gr.Button("Clear Session", variant="secondary")
-            gen_btn = gr.Button("Generate Code",  variant="primary", elem_id="gen_btn")
-
-    # ────────── main panel (outputs) ──────────
-    with gr.Tabs():
-        with gr.Tab("💻 Code"):
-            code_out = gr.Code(language="html", lines=25, interactive=True)
-        with gr.Tab("👁️ Live Preview"):
-            preview_out = gr.HTML()
-        with gr.Tab("📜 History"):
-            chat_out = gr.Chatbot(type="messages")
-
-    # ────────── callbacks ──────────
-    def generation_code(
-        query        : Optional[str],
-        file_path    : Optional[str],
-        website_url  : Optional[str],
-        current_model: Model,
-        enable_search: bool,
-        language     : str,
-        history      : Optional[History],
-    ):
-        # (implementation identical to previous working version…)
-        # For brevity, assume the body of generation_code remains unchanged.
-        ...
-
-    # dropdown change
-    def _on_model_change(name): return get_model_details(name)
-
-    model_dd.change(
-        _on_model_change,
-        inputs=model_dd,
-        outputs=model_state
-    )
+    # ── Callbacks -----------------------------------------------------------
+    def _select_model(name: str) -> ModelInfo:
+        return next((m for m in AVAILABLE_MODELS if m["name"] == name), AVAILABLE_MODELS[0])
 
-    # generate button
-    gen_btn.click(
-        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]
+    dd_model.change(_select_model, dd_model, state_model)
+    btn_gen.click(
+        generate_code,
+        inputs=[tb_prompt, fi_file, tb_url,
+                state_model, dd_lang, cb_search, state_hist],
+        outputs=[code_out, state_hist, html_prev, chat_out],
     )
-
-    # clear
-    def _reset(): return "", None, "", [], "", ""
-    clr_btn.click(
-        _reset,
-        outputs=[prompt_in, file_in, url_in,
-                 history_state, code_out, preview_out, chat_out],
-        queue=False
+    btn_clear.click(
+        lambda: ("", None, "", [], [], "", ""),
+        outputs=[tb_prompt, fi_file, tb_url, state_hist, chat_out, code_out, html_prev],
+        queue=False,
     )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch()