Update app.py

app.py

@@ -1,182 +1,216 @@
 # app.py
 """
-ShashaCode Builder – AI code‑generation playground.
+Main application file for SHASHA AI, a Gradio‑based AI code‑generation tool.
 
-• Hugging Face Spaces + Gradio front‑end
-• Supports prompts, file upload, web‑site scraping, optional web search
-• Streams code back, shows live HTML preview, can deploy to a user Space
+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.
 """
 
-# ─────────────────────────────────────────  Imports
 import gradio as gr
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Any
+from typing import Optional, Dict, List, Tuple, Any
 
-from constants import (                               # ← all constants live here
+# ─── Local module imports ────────────────────────────────────────────────
+from constants import (
     HTML_SYSTEM_PROMPT,
     TRANSFORMERS_JS_SYSTEM_PROMPT,
-    SYSTEM_PROMPTS,
     AVAILABLE_MODELS,
     DEMO_LIST,
-    GRADIO_SUPPORTED_LANGUAGES,                      # ← new import
-    SEARCH_START, DIVIDER, REPLACE_END,
 )
-
-from hf_client      import get_inference_client
-from tavily_search  import enhance_query_with_search
-from utils          import (                          # helpers split into utils.py
+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,
-    parse_svelte_output,
-    format_svelte_output,
-    apply_search_replace_changes,
-    apply_transformers_js_search_replace_changes,
-    extract_text_from_file,
-    extract_website_content,
-    get_gradio_language,
 )
-from deploy         import send_to_sandbox
+from deploy import send_to_sandbox
 
-# ─────────────────────────────────────────  Type Aliases
+# ─── Type aliases ───────────────────────────────────────────────────────
 History = List[Tuple[str, str]]
-ModelInfo = Dict[str, Any]
-
-# ─────────────────────────────────────────  Core Function
-def generate_code(
-    query: str,
-    file_path: Optional[str],
+Model   = Dict[str, Any]
+
+# ─── Supported languages (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
+
+# ─── Core generation function ───────────────────────────────────────────
+def generation_code(
+    query: Optional[str],
+    file: Optional[str],
     website_url: Optional[str],
-    model: ModelInfo,
+    current_model: Model,
     enable_search: bool,
     language: str,
     history: Optional[History],
 ) -> Tuple[str, History, str, List[Dict[str, str]]]:
-    """Main inference pipeline: build prompt → call model → post‑process."""
-    query = query or ""
-    history = history or []
-
-    # 1. pick system prompt
-    if   language == "html":             system = HTML_SYSTEM_PROMPT
-    elif language == "transformers.js":  system = TRANSFORMERS_JS_SYSTEM_PROMPT
-    else:                                system = SYSTEM_PROMPTS.get(language, HTML_SYSTEM_PROMPT)
-
-    # 2. build message list
-    messages = history_to_messages(history, system)
-
-    ctx_parts = [query.strip()]
-
-    if file_path:    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_parts += ["[Website]", html[:8000]]
-
-    user_query = "\n\n".join(ctx_parts)
-    user_query = enhance_query_with_search(user_query, enable_search)
-    messages.append({"role": "user", "content": user_query})
-
-    # 3. call model
-    client = get_inference_client(model["id"])
-    resp   = client.chat.completions.create(
-        model=model["id"],
-        messages=messages,
-        max_tokens=16000,
-        temperature=0.15,
-    )
-    answer = resp.choices[0].message.content
+    query    = query or ""
+    history  = history or []
+    try:
+        # Choose system prompt based on language
+        if language == "html":
+            system_prompt = HTML_SYSTEM_PROMPT
+        elif language == "transformers.js":
+            system_prompt = TRANSFORMERS_JS_SYSTEM_PROMPT
+        else:
+            system_prompt = (
+                f"You are an expert {language} developer. "
+                f"Write clean, idiomatic {language} code based on the user's request."
+            )
 
-    # 4. post‑process
+        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=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)
+
+    # Process model output
     if language == "transformers.js":
-        files = parse_transformers_js_output(answer)
-        code  = format_transformers_js_output(files)
+        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((query, code))
-    chat_msgs = history_to_chatbot_messages(history)
-    return code, history, preview, chat_msgs
-
-
-# ─────────────────────────────────────────  UI
-LOGO_PATH = "assets/logo.png"                       # ensure this file exists
+    new_hist = history + [(query, code)]
+    chat     = history_to_chatbot_messages(new_hist)
+    return code, new_hist, preview, chat
 
+# ─── Custom CSS (added #logo rule) ───────────────────────────────────────
 CUSTOM_CSS = """
-body {font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;}
-#logo {max-height:64px;margin:auto;}
+body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; }
+#logo      { display:block; margin:20px auto; max-height:80px; }
+#main_title{ text-align:center; font-size:2.5rem; margin-top:0.5rem; }
+#subtitle  { text-align:center; color:#4a5568; margin-bottom:2.0rem; }
+.gradio-container { background-color:#f7fafc; }
+#gen_btn   { box-shadow:0 4px 6px rgba(0,0,0,0.1); }
 """
 
-with gr.Blocks(css=CUSTOM_CSS, title="ShashaCode Builder") as demo:
-    state_history = gr.State([])
-    state_model   = gr.State(AVAILABLE_MODELS[0])
+# ─── Gradio UI ───────────────────────────────────────────────────────────
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"),
+               css=CUSTOM_CSS,
+               title="Shasha AI") as demo:
 
-    # Header
-    with gr.Row():
-        gr.Image(LOGO_PATH, elem_id="logo", show_label=False, height=64)
-        gr.Markdown("## **AnyCoder AI**\nYour AI partner for generating, modifying & understanding code.")
+    history_state = gr.State([])
+    initial_model = AVAILABLE_MODELS[0]
+    model_state   = gr.State(initial_model)
+
+    # Logo • Title • Subtitle
+    gr.Image("assets/logo.png", elem_id="logo", show_label=False)
+    gr.Markdown("# 🚀 Shasha AI", elem_id="main_title")
+    gr.Markdown("Your AI partner for generating, modifying, and understanding code.",
+                elem_id="subtitle")
 
-    # Sidebar (inputs)
     with gr.Row():
-        with gr.Column(scale=1, min_width=300):
-            # Model
-            dd_model = gr.Dropdown(
-                label="AI Model",
+        # ── Left column (inputs)
+        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"],
+                value=initial_model["name"],
+                label="AI Model",
             )
 
-            # Prompt / File / Website tabs
+            gr.Markdown("### 2. Provide Context")
             with gr.Tabs():
-                with gr.Tab("Prompt"):
-                    tb_prompt = gr.Textbox(label="Describe what you'd like to build…", lines=6)
-                with gr.Tab("File"):
-                    inp_file  = gr.File(label="Reference file", type="filepath")
-                with gr.Tab("Website"):
-                    tb_url    = gr.Textbox(label="URL to redesign")
-
-            # Output config
-            dd_lang = gr.Dropdown(
-                label="Target language",
-                choices=[l for l in GRADIO_SUPPORTED_LANGUAGES if l],  # ← fixed list
-                value="html",
-            )
-            chk_search = gr.Checkbox(label="Enable Tavily Web Search")
-
-            # Buttons
-            btn_generate = gr.Button("Generate Code", variant="primary")
-            btn_clear    = gr.Button("Clear Session",  variant="secondary")
-
-        # Main panel (outputs)
+                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_btn = gr.Button("Clear Session", variant="secondary")
+                gen_btn = gr.Button("Generate Code", variant="primary", elem_id="gen_btn")
+
+        # ── Right column (outputs)
         with gr.Column(scale=2):
             with gr.Tabs():
-                with gr.Tab("Code"):
-                    out_code = gr.Code(language="html", show_label=False)
-                with gr.Tab("Preview"):
-                    out_prev = gr.HTML()
-                with gr.Tab("History"):
-                    out_hist = gr.Chatbot(type="messages")
-
-    # ─── Callbacks ─────────────────────────────────────────────
-    def _model_from_name(name):
-        return next((m for m in AVAILABLE_MODELS if m["name"] == name), AVAILABLE_MODELS[0])
-
-    dd_model.change(lambda n: _model_from_name(n), inputs=dd_model, outputs=state_model)
-
-    btn_generate.click(
-        fn=generate_code,
-        inputs=[tb_prompt, inp_file, tb_url, state_model, chk_search, dd_lang, state_history],
-        outputs=[out_code, state_history, out_prev, out_hist],
+                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")
+
+    # ── Callbacks
+    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, inp_file, tb_url, state_history, out_code, out_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.queue().launch()