Update app.py

app.py

@@ -1,188 +1,201 @@
+# app.py
+
 """
-app.py – AnyCoder AI (Gradio)
+Main application file for SHASHA AI, a Gradio-based AI code generation tool.
 
-* Logo: assets/logo.png
-* Models: full list from constants.AVAILABLE_MODELS
-* No height= arg on gr.Code (Gradio ≥5)
+Provides a UI for generating code in many languages using various AI models.
+Supports text prompts, file uploads, website scraping, optional web search,
+and live previews of HTML output.
 """
 
-from __future__ import annotations
 import gradio as gr
-from typing import List, Tuple, Dict, Optional
+from typing import Optional, Dict, List, Tuple, Any
 
-# ── local modules ──────────────────────────────────────────────────────────
+# --- Local module imports ---
 from constants import (
-    HTML_SYSTEM_PROMPT, HTML_SYSTEM_PROMPT_WITH_SEARCH,
-    TRANSFORMERS_JS_SYSTEM_PROMPT, TRANSFORMERS_JS_SYSTEM_PROMPT_WITH_SEARCH,
-    GENERIC_SYSTEM_PROMPT, GENERIC_SYSTEM_PROMPT_WITH_SEARCH,
-    TransformersJSFollowUpSystemPrompt, FollowUpSystemPrompt,
-    AVAILABLE_MODELS, DEMO_LIST, get_gradio_language,
+    HTML_SYSTEM_PROMPT,
+    TRANSFORMERS_JS_SYSTEM_PROMPT,
+    AVAILABLE_MODELS,
+    DEMO_LIST,
 )
-from hf_client     import get_inference_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,
-)
-from search_replace import (         # <-- moved here
+from utils import (
+    extract_text_from_file,
+    extract_website_content,
     apply_search_replace_changes,
-    apply_transformers_js_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
+from deploy import send_to_sandbox
 
-# ── aliases ────────────────────────────────────────────────────────────────
+# --- Type aliases ---
 History = List[Tuple[str, str]]
-Model   = Dict[str, str]
-
-# ── code generation core ───────────────────────────────────────────────────
-def generate_code(
-    prompt: str,
-    file_path: Optional[str],
+Model = Dict[str, Any]
+
+# --- Supported languages for dropdown ---
+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"
+]
+
+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],
-    model: Model,
-    language: 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 []
-    prompt  = prompt or ""
-
-    # choose system prompt
-    if history:
-        system_prompt = (
-            TransformersJSFollowUpSystemPrompt if language == "transformers.js"
-            else FollowUpSystemPrompt
-        )
-    else:
+    try:
+        # Choose system prompt based on language
         if language == "html":
-            system_prompt = HTML_SYSTEM_PROMPT_WITH_SEARCH if enable_search else HTML_SYSTEM_PROMPT
+            system_prompt = HTML_SYSTEM_PROMPT
         elif language == "transformers.js":
-            system_prompt = TRANSFORMERS_JS_SYSTEM_PROMPT_WITH_SEARCH if enable_search else TRANSFORMERS_JS_SYSTEM_PROMPT
+            system_prompt = TRANSFORMERS_JS_SYSTEM_PROMPT
         else:
+            # Generic fallback prompt
             system_prompt = (
-                GENERIC_SYSTEM_PROMPT_WITH_SEARCH.format(language=language)
-                if enable_search else GENERIC_SYSTEM_PROMPT.format(language=language)
+                f"You are an expert {language} developer. "
+                f"Write clean, idiomatic {language} code based on the user's request."
             )
 
-    messages = history_to_messages(history, system_prompt)
-
-    # attach context
-    if file_path:
-        prompt += f"\n\n[File]\n{extract_text_from_file(file_path)[:5000]}"
-    if website_url:
-        prompt += f"\n\n[Website]\n{extract_website_content(website_url)[:8000]}"
-
-    messages.append({"role": "user", "content": enhance_query_with_search(prompt, enable_search)})
-
-    # call model
-    client = get_inference_client(model["id"])
-    try:
+        model_id = current_model["id"]
+        # Determine provider
+        if model_id.startswith("openai/") or model_id in {"gpt-4", "gpt-3.5-turbo"}:
+            provider = "openai"
+        elif model_id.startswith("gemini/") or model_id.startswith("google/"):
+            provider = "gemini"
+        elif model_id.startswith("fireworks-ai/"):
+            provider = "fireworks-ai"
+        else:
+            provider = "auto"
+
+        # Build message history
+        msgs = history_to_messages(history, system_prompt)
+        context = query
+        if file:
+            ftext = extract_text_from_file(file)
+            context += f"\n\n[Attached file]\n{ftext[:5000]}"
+        if website_url:
+            wtext = extract_website_content(website_url)
+            if not wtext.startswith("Error"):
+                context += f"\n\n[Website content]\n{wtext[:8000]}"
+        final_q = enhance_query_with_search(context, enable_search)
+        msgs.append({"role": "user", "content": final_q})
+
+        # Call the model
+        client = get_inference_client(model_id, provider)
         resp = client.chat.completions.create(
-            model=model["id"],
-            messages=messages,
+            model=model_id,
+            messages=msgs,
             max_tokens=16000,
-            temperature=0.1,
+            temperature=0.1
         )
-        answer = resp.choices[0].message.content
+        content = resp.choices[0].message.content
+
     except Exception as e:
         err = f"❌ **Error:**\n```\n{e}\n```"
-        history.append((prompt, err))
-        return err, history, "", history_to_chatbot_messages(history)
+        history.append((query, err))
+        return "", history, "", history_to_chatbot_messages(history)
 
-    # post‑process
+    # Process model output
     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 ""
+        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(answer)
-        if history and not history[-1][1].startswith("❌"):
-            clean = apply_search_replace_changes(history[-1][1], clean)
-        code = clean
+        cleaned = remove_code_block(content)
+        if history and history[-1][1] and not history[-1][1].startswith("❌"):
+            code = apply_search_replace_changes(history[-1][1], cleaned)
+        else:
+            code = cleaned
         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")
+    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-top: 1.5rem; }
+#subtitle { text-align: center; color: #4a5568; margin-bottom: 2.5rem; }
+.gradio-container { background-color: #f7fafc; }
+#gen_btn { box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
+"""
 
-with gr.Blocks(theme=theme, title="AnyCoder AI") as demo:
-    st_hist  = gr.State([])
-    st_model = gr.State(AVAILABLE_MODELS[0])
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=CUSTOM_CSS, title="Shasha AI") as demo:
+    history_state = gr.State([])
+    initial_model = AVAILABLE_MODELS[0]
+    model_state = gr.State(initial_model)
 
-    # header with logo
-    gr.HTML(
-        '<div style="text-align:center;margin:1rem 0;">'
-        '<img src="assets/logo.png" alt="logo" style="width:120px;"><br>'
-        '<h1 style="margin:0.4rem 0 0">AnyCoder AI</h1>'
-        '<p style="color:#555">Your AI partner for generating, modifying &amp; understanding code.</p>'
-        '</div>'
-    )
+    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():
         with gr.Column(scale=1):
-            gr.Markdown("### 1 · Model")
-            dd_model = gr.Dropdown([m["name"] for m in AVAILABLE_MODELS],
-                                   value=AVAILABLE_MODELS[0]["name"],
-                                   label="AI Model")
+            gr.Markdown("### 1. Select Model")
+            model_dd = gr.Dropdown(
+                choices=[m["name"] for m in AVAILABLE_MODELS],
+                value=initial_model["name"],
+                label="AI Model"
+            )
 
-            gr.Markdown("### 2 · Context")
+            gr.Markdown("### 2. Provide Context")
             with gr.Tabs():
-                with gr.Tab("Prompt"):
-                    tb_prompt = gr.Textbox(lines=6, placeholder="Describe what you want…")
-                with gr.Tab("File"):
-                    fi_file = gr.File()
-                with gr.Tab("Website"):
-                    tb_url = gr.Textbox(placeholder="https://example.com")
-
-            gr.Markdown("### 3 · Output")
-            dd_lang = gr.Dropdown(
-                [l for l in get_gradio_language.__defaults__[0] if l],  # supported list
-                value="html",
-                label="Target language"
-            )
-            cb_search = gr.Checkbox(label="Enable Tavily Web Search")
+                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():
-                btn_clear = gr.Button("Clear", variant="secondary")
-                btn_gen   = gr.Button("Generate Code", variant="primary")
+                clr_btn = gr.Button("Clear Session", variant="secondary")
+                gen_btn = 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="html", lines=25)
-                with gr.Tab("Preview"):
-                    html_prev = gr.HTML()
-                with gr.Tab("History"):
-                    chat_out = gr.Chatbot(type="messages", height=400)
-
-    # quick demos
-    gr.Markdown("#### Quick Start")
-    with gr.Row():
-        for d in DEMO_LIST[:6]:
-            gr.Button(d["title"], size="sm").click(
-                lambda desc=d["description"]: desc, outputs=tb_prompt
-            )
-
-    # callbacks
-    dd_model.change(lambda n: next(m for m in AVAILABLE_MODELS if m["name"] == n),
-                    dd_model, st_model)
-
-    btn_gen.click(
-        generate_code,
-        inputs=[tb_prompt, fi_file, tb_url, st_model, dd_lang, cb_search, st_hist],
-        outputs=[code_out, st_hist, html_prev, chat_out],
+                with gr.Tab("💻 Code"):
+                    code_out = gr.Code(language="html", 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_model, inputs=[model_dd], outputs=[model_state])
+
+    gen_btn.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],
     )
 
-    btn_clear.click(
-        lambda: ("", None, "", [], [], "", ""),
-        outputs=[tb_prompt, fi_file, tb_url, st_hist, chat_out, code_out, html_prev],
+    clr_btn.click(
+        lambda: ("", None, "", [], "", "", []),
+        outputs=[prompt_in, file_in, url_in, history_state, code_out, preview_out, chat_out],
         queue=False,
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch()