Update app.py

app.py

@@ -3,15 +3,22 @@
 """
 Main application file for SHASHA AI, a Gradio-based AI code generation tool.
 
-This application provides a user interface for generating code in various languages
-using different AI models. It supports inputs from text prompts, files, images,
-and websites, and includes features like web search enhancement and live code previews.
+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.
 """
 
 import gradio as gr
 from typing import Optional, Dict, List, Tuple, Any
 
-from constants import SYSTEM_PROMPTS, AVAILABLE_MODELS, DEMO_LIST
+# --- Local module imports ---
+from constants import (
+    HTML_SYSTEM_PROMPT,
+    TRANSFORMERS_JS_SYSTEM_PROMPT,
+    GENERIC_SYSTEM_PROMPT,
+    AVAILABLE_MODELS,
+    DEMO_LIST,
+)
 from hf_client import get_inference_client
 from tavily_search import enhance_query_with_search
 from utils import (
@@ -22,14 +29,15 @@ from utils import (
     history_to_chatbot_messages,
     remove_code_block,
     parse_transformers_js_output,
-    format_transformers_js_output
+    format_transformers_js_output,
 )
 from deploy import send_to_sandbox
 
+# --- Type aliases ---
 History = List[Tuple[str, str]]
 Model = Dict[str, Any]
 
-# Full list of supported languages for syntax highlighting & generation
+# --- Supported languages dropdown & syntax highlighting ---
 SUPPORTED_LANGUAGES = [
     "python", "c", "cpp", "markdown", "latex", "json", "html", "css",
     "javascript", "jinja2", "typescript", "yaml", "dockerfile", "shell",
@@ -38,13 +46,6 @@ SUPPORTED_LANGUAGES = [
     "sql-gpSQL", "sql-sparkSQL", "sql-esper"
 ]
 
-DEFAULT_SYSTEM_PROMPT = """
-You are a helpful AI coding assistant. Generate clean, correct, and efficient code based on the user's request.
-- Follow requirements precisely.
-- Enclose final code in a single ```code``` block of the target language.
-- Do not include any explanations outside the code block.
-"""
-
 def get_model_details(name: str) -> Optional[Model]:
     for m in AVAILABLE_MODELS:
         if m["name"] == name:
@@ -63,32 +64,39 @@ def generation_code(
     query = query or ""
     history = history or []
     try:
-        system_prompt = SYSTEM_PROMPTS.get(language, DEFAULT_SYSTEM_PROMPT)
-        model_id = current_model["id"]
+        # choose system prompt
+        if language == "html":
+            system_prompt = HTML_SYSTEM_PROMPT
+        elif language == "transformers.js":
+            system_prompt = TRANSFORMERS_JS_SYSTEM_PROMPT
+        else:
+            system_prompt = GENERIC_SYSTEM_PROMPT.format(language=language)
 
+        model_id = current_model["id"]
         # pick provider
-        if model_id.startswith("openai/") or model_id in ("gpt-4", "gpt-3.5-turbo"):
+        if model_id.startswith("openai/") or model_id in {"gpt-4","gpt-3.5-turbo"}:
             provider = "openai"
-        elif model_id.startswith("gemini/"):
+        elif model_id.startswith("gemini/") or model_id.startswith("google/"):
             provider = "gemini"
         elif model_id.startswith("fireworks-ai/"):
             provider = "fireworks-ai"
         else:
-            provider = "huggingface"
+            provider = "auto"
 
         # assemble messages
         msgs = history_to_messages(history, system_prompt)
         ctx = query
         if file:
-            txt = extract_text_from_file(file)
-            ctx += f"\n\n[File]\n{txt[:5000]}"
+            ftext = extract_text_from_file(file)
+            ctx += f"\n\n[Attached file]\n{ftext[:5000]}"
         if website_url:
-            txt = extract_website_content(website_url)
-            if not txt.startswith("Error"):
-                ctx += f"\n\n[Website]\n{txt[:8000]}"
+            wtext = extract_website_content(website_url)
+            if not wtext.startswith("Error"):
+                ctx += f"\n\n[Website content]\n{wtext[:8000]}"
         final_q = enhance_query_with_search(ctx, enable_search)
         msgs.append({"role": "user", "content": final_q})
 
+        # call model
         client = get_inference_client(model_id, provider)
         resp = client.chat.completions.create(
             model=model_id,
@@ -107,41 +115,44 @@ def generation_code(
     if language == "transformers.js":
         files = parse_transformers_js_output(content)
         code = format_transformers_js_output(files)
-        preview = send_to_sandbox(files.get("index.html",""))
+        preview = send_to_sandbox(files.get("index.html", ""))
     else:
-        clean = remove_code_block(content)
+        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], clean)
+            code = apply_search_replace_changes(history[-1][1], cleaned)
         else:
-            code = clean
+            code = cleaned
         preview = send_to_sandbox(code) if language == "html" else ""
 
     new_hist = history + [(query, code)]
     chat = history_to_chatbot_messages(new_hist)
     return code, new_hist, preview, chat
 
-# custom CSS
+# --- 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: 1.5rem 0 0.5rem; }
+#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: #f7fafc; }
+.gradio-container { background-color: #f7fafc; }
 #gen_btn { box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
 """
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=CUSTOM_CSS, title="Shasha AI") as demo:
     history_state = gr.State([])
-    initial = AVAILABLE_MODELS[0]
-    model_state = gr.State(initial)
+    initial_model = AVAILABLE_MODELS[0]
+    model_state = gr.State(initial_model)
 
     gr.Markdown("# 🚀 Shasha AI", elem_id="main_title")
-    gr.Markdown("Your personal AI partner for generating, modifying, and understanding code.", elem_id="subtitle")
+    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. Select Model")
-            names = [m["name"] for m in AVAILABLE_MODELS]
-            model_dd = gr.Dropdown(names, value=initial["name"], label="AI Model")
+            model_dd = gr.Dropdown(
+                choices=[m["name"] for m in AVAILABLE_MODELS],
+                value=initial_model["name"],
+                label="AI Model"
+            )
 
             gr.Markdown("### 2. Provide Context")
             with gr.Tabs():
@@ -157,28 +168,28 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=CUSTOM_CSS, title="
             search_chk = gr.Checkbox(label="Enable Web Search")
 
             with gr.Row():
-                clr = gr.Button("Clear Session", variant="secondary")
-                gen = gr.Button("Generate Code", variant="primary", elem_id="gen_btn")
+                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=lambda: lang_dd.value, interactive=True)
+                    code_out = gr.Code(language=None, 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, inputs=[model_dd], outputs=[model_state])
+    model_dd.change(lambda n: get_model_details(n) or initial_model, inputs=[model_dd], outputs=[model_state])
 
-    gen.click(
+    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],
     )
 
-    clr.click(
-        fn=lambda: ("", None, "", [], "", "", []),
+    clr_btn.click(
+        lambda: ("", None, "", [], "", "", []),
         outputs=[prompt_in, file_in, url_in, history_state, code_out, preview_out, chat_out],
         queue=False,
     )