Update app.py

app.py

@@ -1,140 +1,180 @@
 # app.py
-# ──────────────────────────────────────────────────────────────────────────────
+# ────────────────────────────────────────────────────────────────
 """
-AnyCoder / Shasha AI – Gradio back‑end
+AnyCoder / Shasha AI – lightweight Gradio front‑end
 
-• Serves the custom front‑end located in static/index.html (+ static/style.css,
-  static/index.js).
-• Exposes ONE http POST endpoint  →  /run/predict   (Gradio api_name="predict")
-  that the browser JS calls to run the model and get generated code.
-
-All heavy lifting (model registry, provider routing, web‑search, etc.) lives in
-  • models.py, inference.py, utils.py, deploy.py …
+• Pick an AI model from models.py ➜ AVAILABLE_MODELS
+• Provide context (prompt / file / website)
+• Choose a target language from 25+ options
+• Optional Tavily web‑search enrichment
+• Generate code & live‑preview HTML
 """
 
+from __future__ import annotations
+
 from pathlib import Path
-from typing import List, Tuple, Dict
+from typing import List, Tuple, Dict, Any, Optional
 
 import gradio as gr
 
-# ── Local helpers ────────────────────────────────────────────────────────────
-from inference   import chat_completion
+# ── local helpers ───────────────────────────────────────────────
+from models       import AVAILABLE_MODELS, find_model, ModelInfo
+from inference    import chat_completion
 from tavily_search import enhance_query_with_search
-from deploy      import send_to_sandbox
-from models      import AVAILABLE_MODELS, find_model, ModelInfo
-from utils       import (
+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_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",
+]
 
-# ── System prompts keyed by language ────────────────────────────────────────
 SYSTEM_PROMPTS = {
     "html": (
-        "ONLY USE HTML, CSS AND JAVASCRIPT. Return exactly ONE file wrapped in "
-        "```html ...```."
+        "ONLY USE HTML, CSS AND JAVASCRIPT. Produce ONE complete HTML file "
+        "wrapped in ```html ...```."
     ),
     "transformers.js": (
-        "Generate THREE separate files (index.html • index.js • style.css) each "
-        "inside its own fenced block."
+        "Generate THREE fenced blocks (index.html / index.js / style.css) "
+        "for a transformers.js web‑app."
     ),
 }
 
-# ── Output‑history data structure ───────────────────────────────────────────
-History = List[Tuple[str, str]]  # [(user_query, generated_code), …]
+History = List[Tuple[str, str]]
 
-# ════════════════════════════════════════════════════════════════════════════
-# 1.  Backend callback hit by the JS front‑end
-# ════════════════════════════════════════════════════════════════════════════
-def generate(
+
+# ── core callback ───────────────────────────────────────────────
+def generate_code(
     prompt: str,
     file_path: str | None,
     website_url: str | None,
-    model_id: str,
+    model_name: str,
+    enable_search: bool,
     language: str,
-    web_search: bool,
-    history: History | None,
-) -> Dict[str, str]:
-    """
-    The only public API.  Returns: { "code": <string> }
-    """
-    history = history or []
-
-    # ---- Build system + user messages --------------------------------------
-    sys_prompt = SYSTEM_PROMPTS.get(language, f"You are an expert {language} developer.")
-    messages  = history_to_messages(history, sys_prompt)
+    hist: History | None,
+) -> Tuple[str, History, str, List[Dict[str, str]]]:
+    """Back‑end for the ‘Generate Code’ button."""
+    hist = hist or []
+    prompt = (prompt or "").strip()
 
-    ctx: list[str] = [prompt.strip()]
+    # 1 build messages
+    sys_prompt = SYSTEM_PROMPTS.get(language, f"You are an expert {language} developer.")
+    messages   = history_to_messages(hist, sys_prompt)
 
+    ctx_parts: list[str] = [prompt]
     if file_path:
-        ctx.append("[File]")
-        ctx.append(extract_text_from_file(file_path)[:5_000])
-
+        ctx_parts += ["[File]", extract_text_from_file(file_path)[:5000]]
     if website_url:
         html = extract_website_content(website_url)
         if not html.startswith("Error"):
-            ctx.append("[Website]")
-            ctx.append(html[:8_000])
+            ctx_parts += ["[Website]", html[:8000]]
 
-    user_query = "\n\n".join(filter(None, ctx))
-    user_query = enhance_query_with_search(user_query, web_search)
-    messages.append({"role": "user", "content": user_query})
+    user_msg = enhance_query_with_search("\n\n".join(ctx_parts), enable_search)
+    messages.append({"role": "user", "content": user_msg})
 
-    # ---- Call model --------------------------------------------------------
-    model: ModelInfo = find_model(model_id) or AVAILABLE_MODELS[0]
-    assistant = chat_completion(model.id, messages)
+    # 2 run model (provider selection handled in inference.chat_completion)
+    model: ModelInfo = find_model(model_name) or AVAILABLE_MODELS[0]
+    try:
+        raw_out = chat_completion(model.id, messages)
+    except Exception as exc:  # pragma: no cover
+        err = f"❌ **Error**\n```{exc}```"
+        hist.append((prompt, err))
+        return "", hist, "", history_to_chatbot_messages(hist)
 
-    # ---- Post‑process output ----------------------------------------------
+    # 3 post‑process
     if language == "transformers.js":
-        files = parse_transformers_js_output(assistant)
+        files = parse_transformers_js_output(raw_out)
         code  = format_transformers_js_output(files)
+        preview = send_to_sandbox(files.get("index.html", ""))
     else:
-        clean = remove_code_block(assistant)
-        if history and not history[-1][1].startswith("❌"):
-            clean = apply_search_replace_changes(history[-1][1], clean)
-        code = clean
-
-    # (preview iframe is rendered entirely client‑side from code)
-    return {"code": code}
-
-
-# ════════════════════════════════════════════════════════════════════════════
-# 2.  Gradio wrapper
-# ════════════════════════════════════════════════════════════════════════════
-# Read static front‑end
-INDEX_HTML = Path("static/index.html").read_text(encoding="utf-8")
-
-with gr.Blocks(css="body{margin:0}", title="AnyCoder AI") as demo:
-    # 2‑a  Serve the custom UI
-    gr.HTML(INDEX_HTML)
-
-    # 2‑b  Invisible components acting as our JSON REST API
-    with gr.Group(visible=False) as backend:
-        prompt_in   = gr.Textbox()
-        file_in     = gr.File()
-        url_in      = gr.Textbox()
-        model_in    = gr.Textbox()
-        lang_in     = gr.Textbox()
-        search_in   = gr.Checkbox()
-        hist_state  = gr.State([])      # persists conversation on server
-
-        code_out = gr.JSON(label="code")  # returns {"code": "..."} to the JS
-
-    trigger = gr.Button(visible=False)
-    trigger.click(
-        fn=generate,
-        inputs=[
-            prompt_in, file_in, url_in,
-            model_in, lang_in, search_in, hist_state
-        ],
-        outputs=[code_out],
-        api_name="predict",   #  <─ POST /run/predict
+        cleaned = remove_code_block(raw_out)
+        if hist and not hist[-1][1].startswith("❌"):
+            cleaned = apply_search_replace_changes(hist[-1][1], cleaned)
+        code    = cleaned
+        preview = send_to_sandbox(cleaned) if language == "html" else ""
+
+    hist.append((prompt, code))
+    chat_view = history_to_chatbot_messages(hist)
+    return code, hist, preview, chat_view
+
+
+# ── UI ──────────────────────────────────────────────────────────
+THEME = gr.themes.Soft(primary_hue="indigo")
+custom_css = """
+body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; }
+#main_title { text-align:center;font-size:2.4rem;margin-top:1rem }
+#subtitle   { text-align:center;color:#5a6475;margin-bottom:2rem }
+"""
+
+with gr.Blocks(title="AnyCoder AI", theme=THEME, css=custom_css) as demo:
+    state_hist: gr.State[History] = gr.State([])
+
+    gr.Markdown("## 🚀 AnyCoder AI", elem_id="main_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 · Select 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"):
+                    prompt_box = gr.Textbox(lines=6, placeholder="Describe what you want…")
+                with gr.Tab("File"):
+                    file_box   = gr.File(type="filepath")
+                with gr.Tab("Website"):
+                    url_box    = gr.Textbox(placeholder="https://example.com")
+
+            gr.Markdown("#### 3 · Configure Output")
+            lang_dd   = gr.Dropdown(SUPPORTED_LANGUAGES, 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 (main panel) ─────
+        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_code,
+        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()