Update app.py

app.py

@@ -1,174 +1,210 @@
-# app.py
-# ────────────────────────────────────────────────────────────────
-"""
-AnyCoder / Shasha AI – Gradio interface
-• Works with the model registry in models.py
-• Uses inference.chat_completion for provider‑aware calls
-• Supports prompt / file / website context
-• 25+ target languages
-• Optional Tavily web‑search enrichment
-• Live HTML preview + chat history
 """
+Main application file for SHASHA AI, a Gradio-based AI code generation tool.
 
-from __future__ import annotations
-from typing import List, Tuple, Dict, Any
+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
-
-# local helpers
-from models        import AVAILABLE_MODELS, find_model, ModelInfo
-from inference     import chat_completion
+from typing import Optional, Dict, List, Tuple, Any
+
+# --- Local module imports ---
+from constants import (
+    HTML_SYSTEM_PROMPT,
+    TRANSFORMERS_JS_SYSTEM_PROMPT,
+    AVAILABLE_MODELS,
+    DEMO_LIST,
+)
+from hf_client import get_inference_client
 from tavily_search import enhance_query_with_search
-from utils         import (
+from utils import (
     extract_text_from_file,
     extract_website_content,
+    apply_search_replace_changes,
     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
+from deploy import send_to_sandbox, load_project_from_url
+
+# --- Type aliases ---
+History = List[Tuple[str, str]]
+Model = Dict[str, Any]
 
-# ── constants ───────────────────────────────────────────────────
-SUPPORTED_LANGS = [
+# --- 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",
+    "sql-gpSQL", "sql-sparkSQL", "sql-esper"
 ]
 
-SYSTEM_PROMPTS = {
-    "html": (
-        "ONLY USE HTML, CSS AND JAVASCRIPT. "
-        "Return exactly one complete page wrapped in ```html ...```."
-    ),
-    "transformers.js": (
-        "Generate THREE fenced blocks in order: "
-        "index.html, index.js, style.css for a transformers.js web‑app."
-    ),
-}
-
-History = List[Tuple[str, str]]
-
-# ── backend callback ────────────────────────────────────────────
-def generate(
-    prompt: str,
-    file_path: str | None,
-    website_url: str | None,
-    model_name: str,
+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],
+    current_model: Model,
     enable_search: bool,
     language: str,
-    history: History | None,
+    history: Optional[History],
 ) -> Tuple[str, History, str, List[Dict[str, str]]]:
-    """Main generation routine – called when the user clicks **Generate Code**."""
+    query = query or ""
     history = history or []
-    prompt  = (prompt or "").strip()
-
-    # 1 system + conversation
-    sys_p   = SYSTEM_PROMPTS.get(language, f"You are an expert {language} developer.")
-    msgs    = history_to_messages(history, sys_p)
-
-    ctx: List[str] = [prompt] if prompt else []
-    if file_path:
-        ctx += ["[File]", extract_text_from_file(file_path)[:5_000]]
-    if website_url:
-        html = extract_website_content(website_url)
-        if not html.startswith("Error"):
-            ctx += ["[Website]", html[:8_000]]
-
-    user_msg = enhance_query_with_search("\n\n".join(ctx), enable_search)
-    msgs.append({"role": "user", "content": user_msg})
-
-    # 2 model call
-    model: ModelInfo = find_model(model_name) or AVAILABLE_MODELS[0]
     try:
-        raw = chat_completion(model.id, msgs)
-    except Exception as exc:
-        err = f"❌ **Error**\n```{exc}```"
-        history.append((prompt, err))
+        # 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:
+            # Generic fallback prompt
+            system_prompt = (
+                f"You are an expert {language} developer. "
+                f"Write clean, idiomatic {language} code based on the user's request."
+            )
+
+        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)
 
-    # 3 post‑process
+    # Process model output
     if language == "transformers.js":
-        files   = parse_transformers_js_output(raw)
-        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:
-        cleaned = remove_code_block(raw)
-        if history and not history[-1][1].startswith("❌"):
-            cleaned = apply_search_replace_changes(history[-1][1], cleaned)
-        code    = cleaned
-        preview = send_to_sandbox(cleaned) if language == "html" else ""
-
-    history.append((prompt, code))
-    return code, history, preview, history_to_chatbot_messages(history)
-
-
-# ── Gradio UI ───────────────────────────────────────────────────
-THEME = gr.themes.Soft(primary_hue="indigo")
-CSS = """
-body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,Arial,sans-serif}
-#title{font-size:2.4rem;text-align:center;margin-top:1rem}
-#subtitle{color:#56606d;text-align:center;margin-bottom:1.8rem}
+        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 ""
+
+    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, css=CSS, title="ShashaCode_Builder") as demo:
-    state_hist: gr.State[History] = gr.State([])
+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)
 
-    gr.Markdown("## 🚀 AnyCoder AI", elem_id="title")
-    gr.Markdown("Your AI partner for generating, modifying & understanding code.", elem_id="subtitle")
+    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():
-        # ── inputs sidebar
         with gr.Column(scale=1):
-            gr.Markdown("### 1 · Model")
+            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",
+                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", id="tab-prompt"):
-                    prompt_box = gr.Textbox(lines=6, placeholder="Describe what you need…")
-                with gr.Tab("File", id="tab-file"):
-                    file_box   = gr.File(type="filepath")
-                with gr.Tab("Website", id="tab-site"):
-                    url_box    = gr.Textbox(placeholder="https://example.com")
+                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 · Output")
-            lang_dd   = gr.Dropdown(SUPPORTED_LANGS, value="html", label="Target Language")
-            search_ck = gr.Checkbox(label="Enable Tavily 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():
-                clear_btn = gr.Button("Clear Session", variant="secondary")
-                gen_btn   = 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")
+
+            # --- Hugging Face Integration ---
+            gr.Markdown("---")
+            gr.Markdown("#### Hugging Face Integration")
+            hf_url = gr.Textbox(label="Load HF Space URL")
+            hf_load_btn = gr.Button("Import Project")
+            hf_status = gr.Markdown(visible=False)
+
+            hf_load_btn.click(
+                fn=lambda u: load_project_from_url(u),
+                inputs=hf_url,
+                outputs=[hf_status, code_out, preview_out, hf_url, history_state, chat_out]
+            )
 
-        # ── outputs
         with gr.Column(scale=2):
             with gr.Tabs():
-                with gr.Tab("Code"):
-                    code_out    = gr.Code(interactive=True)
-                with gr.Tab("Live Preview"):
+                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")
+                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])
 
-    # ── wiring
     gen_btn.click(
-        generate,
-        inputs=[prompt_box, file_box, url_box, model_dd, search_ck, lang_dd, state_hist],
-        outputs=[code_out, state_hist, preview_out, chat_out],
+        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],
     )
 
-    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],
+    clr_btn.click(
+        lambda: ("", None, "", [], "", "", []),
+        outputs=[prompt_in, file_in, url_in, history_state, code_out, preview_out, chat_out],
         queue=False,
     )