Update app.py

app.py

@@ -1,16 +1,14 @@
 # app.py
 """
-Main application file for SHASHA AI, a Gradio‑based AI code‑generation tool.
-
-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.
+Main application file for SHASHA AI, a Gradio‑based AI code‑generation tool.
+Adds a small banner logo (assets/logo.png) above the page title.
 """
 
 import gradio as gr
 from typing import Optional, Dict, List, Tuple, Any
+import os
 
-# ─── Local module imports ────────────────────────────────────────────────
+# --- Local module imports ---
 from constants import (
     HTML_SYSTEM_PROMPT,
     TRANSFORMERS_JS_SYSTEM_PROMPT,
@@ -31,17 +29,16 @@ from utils import (
 )
 from deploy import send_to_sandbox
 
-# ─── Type aliases ───────────────────────────────────────────────────────
+# --- Type aliases ---
 History = List[Tuple[str, str]]
 Model   = Dict[str, Any]
 
-# ─── Supported languages (dropdown) ─────────────────────────────────────
+# --- 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",
+    "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]:
@@ -50,7 +47,6 @@ def get_model_details(name: str) -> Optional[Model]:
             return m
     return None
 
-# ─── Core generation function ───────────────────────────────────────────
 def generation_code(
     query: Optional[str],
     file: Optional[str],
@@ -60,10 +56,10 @@ def generation_code(
     language: str,
     history: Optional[History],
 ) -> Tuple[str, History, str, List[Dict[str, str]]]:
-    query    = query or ""
-    history  = history or []
+    query   = query or ""
+    history = history or []
     try:
-        # Choose system prompt based on language
+        # Choose system prompt
         if language == "html":
             system_prompt = HTML_SYSTEM_PROMPT
         elif language == "transformers.js":
@@ -75,36 +71,28 @@ def generation_code(
             )
 
         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"
+        provider = (
+            "openai"       if model_id.startswith("openai/") or model_id in {"gpt-4","gpt-3.5-turbo"}
+            else "gemini"  if model_id.startswith(("gemini/","google/"))
+            else "fireworks-ai" if model_id.startswith("fireworks-ai/")
+            else "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]}"
+            context += f"\n\n[Attached file]\n{extract_text_from_file(file)[: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})
+        msgs.append({"role":"user","content":enhance_query_with_search(context, enable_search)})
 
         # 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,
+        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
 
@@ -113,80 +101,72 @@ def generation_code(
         history.append((query, err))
         return "", history, "", history_to_chatbot_messages(history)
 
-    # Process model output
+    # Process output
     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:
         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
+        code    = apply_search_replace_changes(history[-1][1], cleaned) if history and not history[-1][1].startswith("❌") else 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
+    return code, new_hist, preview, history_to_chatbot_messages(new_hist)
 
-# ─── Custom CSS (added #logo rule) ───────────────────────────────────────
+# --- Custom CSS ---
 CUSTOM_CSS = """
-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); }
+body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif;}
+#main_title{text-align:center;font-size:2.5rem;margin-top:0.5rem;}
+#subtitle{text-align:center;color:#4a5568;margin-bottom:2rem;}
+.gradio-container{background-color:#f7fafc;}
+#gen_btn{box-shadow:0 4px 6px rgba(0,0,0,0.1);}
 """
 
-# ─── Gradio UI ───────────────────────────────────────────────────────────
+LOGO_PATH = "assets/logo.png"
+logo_exists = os.path.exists(LOGO_PATH)
+
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"),
                css=CUSTOM_CSS,
-               title="Shasha AI") as demo:
-
+               title="Shasha AI") as demo:
     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")
+    if logo_exists:
+        # Small banner logo (height ~70px) centred
+        gr.Image(value=LOGO_PATH, height=70, show_label=False,
+                 container=False, elem_id="banner_logo")
+
+    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():
-        # ── Left column (inputs)
         with gr.Column(scale=1):
-            gr.Markdown("### 1. Select 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",
+                label="AI Model"
             )
 
-            gr.Markdown("### 2. Provide Context")
+            gr.Markdown("### 2. Provide Context")
             with gr.Tabs():
                 with gr.Tab("📝 Prompt"):
-                    prompt_in = gr.Textbox(lines=7,
-                                            placeholder="Describe your request…",
-                                            show_label=False)
+                    prompt_in = gr.Textbox(lines=7, placeholder="Describe your request...", show_label=False)
                 with gr.Tab("📄 File"):
-                    file_in   = gr.File(type="filepath")
+                    file_in = gr.File(type="filepath")
                 with gr.Tab("🌐 Website"):
-                    url_in    = gr.Textbox(placeholder="https://example.com")
+                    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")
+            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")
+                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"):
@@ -196,7 +176,6 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"),
                 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])