Update app.py

app.py

@@ -1,52 +1,40 @@
 # app.py
-# ---------------------------------------------------------------------
-# AnyCoder / Shasha AI  –  Gradio front‑end (no external static files)
-# ---------------------------------------------------------------------
 """
-Interactive UI for generating / modifying / previewing code with multiple
-LLM back‑ends.  Relies on:
-
- • models.py       – AVAILABLE_MODELS registry + find_model()
- • inference.py    – chat_completion()  (provider routing handled there)
- • utils.py        – file / website extractors, history helpers, etc.
- • tavily_search.py (optional) – enhance_query_with_search()
- • deploy.py       – send_to_sandbox() for live HTML preview
+Shasha / AnyCoder AI  –  Gradio front‑end + back‑end
+
+Generate or refactor code with multiple foundation models.
+Uploads, website scrape, optional Tavily search enrichment, live HTML preview.
 """
 
 from __future__ import annotations
-
-from pathlib import Path
 from typing import List, Tuple, Dict, Any, Optional
 
 import gradio as gr
 
-# ---------- local helpers --------------------------------------------------
-from models          import AVAILABLE_MODELS, find_model, ModelInfo
-from inference       import chat_completion
-from tavily_search   import enhance_query_with_search
-from utils           import (
+# ──────────────────────────────────────────────────────────────────────────────
+# Local helpers (unchanged)
+# ──────────────────────────────────────────────────────────────────────────────
+from constants import (
+    HTML_SYSTEM_PROMPT,
+    TRANSFORMERS_JS_SYSTEM_PROMPT,
+    AVAILABLE_MODELS,
+)
+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,
-    apply_search_replace_changes,
     remove_code_block,
     parse_transformers_js_output,
     format_transformers_js_output,
 )
-from deploy          import send_to_sandbox
-
-# ---------- constants ------------------------------------------------------
-SYSTEM_PROMPTS = {
-    "html": (
-        "ONLY USE HTML, CSS AND JAVASCRIPT.  Produce ONE complete html file "
-        "wrapped in ```html … ```."
-    ),
-    "transformers.js": (
-        "Generate THREE fenced blocks: index.html, index.js, style.css "
-        "for a transformers.js demo."
-    ),
-}
+from deploy         import send_to_sandbox
+
+History = List[Tuple[str, str]]
+
 SUPPORTED_LANGUAGES = [
     "python", "c", "cpp", "markdown", "latex", "json", "html", "css",
     "javascript", "jinja2", "typescript", "yaml", "dockerfile", "shell",
@@ -54,149 +42,168 @@ SUPPORTED_LANGUAGES = [
     "sql-cassandra", "sql-plSQL", "sql-hive", "sql-pgSQL", "sql-gql",
     "sql-gpSQL", "sql-sparkSQL", "sql-esper",
 ]
-History = List[Tuple[str, str]]  # [(prompt, code/result)]
-
 
-# ---------- core generation callback --------------------------------------
-def generate(
-    prompt: str,
-    file_path: str | None,
-    website_url: str | None,
-    model_name: str,
+# ──────────────────────────────────────────────────────────────────────────────
+# Generation callback
+# ──────────────────────────────────────────────────────────────────────────────
+def generation_code(
+    query: Optional[str],
+    file_path: Optional[str],
+    website_url: Optional[str],
+    current_model: Dict[str, Any],
+    enable_search: bool,
     language: str,
-    use_search: bool,
-    hist: History | None,
+    history: History | None,
 ) -> Tuple[str, History, str, List[Dict[str, str]]]:
-    """Main callback wired to the “Generate Code” button."""
-    hist = hist or []
-    user_prompt = (prompt or "").strip()
-
-    # 1 · system + previous messages
-    sys = SYSTEM_PROMPTS.get(language, f"You are an expert {language} developer.")
-    messages = history_to_messages(hist, sys)
+    """Core back‑end that the “Generate Code” button calls."""
+    query   = (query or "").strip()
+    history = history or []
+
+    # Pick proper system prompt
+    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} based on the user request."
+        )
+
+    # Route to provider
+    model_id = current_model["id"]
+    if model_id.startswith(("openai/", "gpt-")):
+        provider = "openai"
+    elif model_id.startswith(("gemini/", "google/")):
+        provider = "gemini"
+    elif model_id.startswith("fireworks-ai/"):
+        provider = "fireworks-ai"
+    else:
+        provider = "auto"
 
-    # 2 · gather extra context
-    ctx_parts: list[str] = [user_prompt]
+    # Build message list
+    messages = history_to_messages(history, system_prompt)
+    context_parts: List[str] = [query]
 
     if file_path:
-        ctx_parts.append("[Reference file]")
-        ctx_parts.append(extract_text_from_file(file_path)[:5000])
+        context_parts.append("[Attached file]")
+        context_parts.append(extract_text_from_file(file_path)[:5_000])
 
     if website_url:
         html = extract_website_content(website_url)
-        if not html.lower().startswith("error"):
-            ctx_parts.append("[Website content]")
-            ctx_parts.append(html[:8000])
+        if not html.startswith("Error"):
+            context_parts.append("[Website content]")
+            context_parts.append(html[:8_000])
 
-    user_query = "\n\n".join(filter(None, ctx_parts))
-    user_query = enhance_query_with_search(user_query, use_search)
-    messages.append({"role": "user", "content": user_query})
+    user_msg = "\n\n".join(context_parts)
+    user_msg = enhance_query_with_search(user_msg, enable_search)
+    messages.append({"role": "user", "content": user_msg})
 
-    # 3 · call the model
-    model: ModelInfo = find_model(model_name) or AVAILABLE_MODELS[0]
+    # Call model
     try:
-        assistant = chat_completion(model.id, messages)
+        client  = get_inference_client(model_id, provider)
+        resp    = client.chat.completions.create(
+            model      = model_id,
+            messages   = messages,
+            max_tokens = 16_000,
+            temperature= 0.1,
+        )
+        answer = resp.choices[0].message.content
     except Exception as exc:
-        err = f"❌ **Error**\n```{exc}```"
-        hist.append((user_prompt, err))
-        return "", hist, "", history_to_chatbot_messages(hist)
+        err = f"❌ **Error**\n```\n{exc}\n```"
+        history.append((query, err))
+        return "", history, "", history_to_chatbot_messages(history)
 
-    # 4 · post‑process
+    # Post‑process
     if language == "transformers.js":
-        files = parse_transformers_js_output(assistant)
-        code = format_transformers_js_output(files)
+        files   = parse_transformers_js_output(answer)
+        code    = format_transformers_js_output(files)
         preview = send_to_sandbox(files.get("index.html", ""))
     else:
-        cleaned = remove_code_block(assistant)
-        if hist and not hist[-1][1].startswith("❌"):
-            cleaned = apply_search_replace_changes(hist[-1][1], cleaned)
-        code = cleaned
+        cleaned = remove_code_block(answer)
+        if history and not history[-1][1].startswith("❌"):
+            cleaned = apply_search_replace_changes(history[-1][1], cleaned)
+        code    = cleaned
         preview = send_to_sandbox(code) if language == "html" else ""
 
-    # 5 · update history & chat
-    hist.append((user_prompt, code))
-    chat_messages = history_to_chatbot_messages(hist)
-    return code, hist, preview, chat_messages
-
-
-# ---------- UI (Gradio 5) --------------------------------------------------
-THEME = gr.themes.Soft(primary_hue="indigo")
-CSS = """
-.gradio-container {max-width: 1450px !important;}
-@media (min-width: 960px){
-  .layout {display:flex; gap:32px}
-  .inputs  {flex:1 0 360px}
-  .outputs {flex:2 0 0}
-}
+    history.append((query, code))
+    chat_history = history_to_chatbot_messages(history)
+    return code, history, preview, chat_history
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Gradio UI
+# ──────────────────────────────────────────────────────────────────────────────
+CUSTOM_CSS = """
+body{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif;}
+.gradio-container{background:#f9fafb;}
+#gen-btn{box-shadow:0 2px 4px rgba(0,0,0,.08);}
 """
 
-with gr.Blocks(title="AnyCoder AI", theme=THEME, css=CSS) as demo:
-    state_hist: gr.State = gr.State([])
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo"),
+               title="AnyCoder AI",
+               css=CUSTOM_CSS) as demo:
+
+    state_history: gr.State = gr.State([])
+    state_model  : gr.State = gr.State(AVAILABLE_MODELS[0])
 
-    # ---- header ----------------------------------------------------------
-    gr.Markdown(
-        "## 🚀 **AnyCoder AI**  \n"
-        "Your AI partner for generating, modifying & understanding code."
+    gr.Markdown("## 🏢 **AnyCoder AI**\nYour AI partner for generating, modifying  and  understanding code.",
+                elem_id="header")
+
+    # 1 · MODEL
+    gr.Markdown("### 1 · Model")
+    model_dd = gr.Dropdown(
+        choices=[m["name"] for m in AVAILABLE_MODELS],
+        value=AVAILABLE_MODELS[0]["name"],
+        label="AI Model",
     )
 
-    with gr.Row(elem_classes="layout"):
-        # -------- left column (inputs) ------------------------------------
-        with gr.Column(elem_classes="inputs"):
-            gr.Markdown("### 1 · Model")
-            model_dd = gr.Dropdown(
-                choices=[m.name for m in AVAILABLE_MODELS],
-                value=AVAILABLE_MODELS[0].name,
-                label="AI Model",
-            )
-
-            gr.Markdown("### 2 · Context")
-            with gr.Tabs():
-                with gr.Tab("Prompt"):
-                    prompt_box = gr.Textbox(lines=6, placeholder="Describe what you need…")
-                with gr.Tab("File"):
-                    file_box = gr.File(type="filepath")
-                with gr.Tab("Website"):
-                    url_box = gr.Textbox(placeholder="https://example.com")
-
-            gr.Markdown("### 3 · Output")
-            lang_dd = gr.Dropdown(SUPPORTED_LANGUAGES, value="html", label="Target Language")
-            search_chk = gr.Checkbox(label="Enable Tavily Web Search")
-
-            with gr.Row():
-                clear_btn = gr.Button("Clear Session", variant="secondary")
-                gen_btn   = gr.Button("Generate Code", variant="primary")
-
-        # -------- right column (outputs) ----------------------------------
-        with gr.Column(elem_classes="outputs"):
-            with gr.Tabs():
-                with gr.Tab("Code"):
-                    code_out = gr.Code(height=500, language="html")
-                with gr.Tab("Live Preview"):
-                    preview_out = gr.HTML()
-                with gr.Tab("History"):
-                    chat_out = gr.Chatbot(type="messages")
-
-    # ---- callbacks -------------------------------------------------------
+    # 2 · CONTEXT
+    gr.Markdown("### 2 · Context")
+    with gr.Tabs():
+        with gr.Tab("Prompt"):
+            prompt_tb = gr.Textbox(lines=6, placeholder="Describe what you need…")
+        with gr.Tab("File"):
+            file_upl  = gr.File(label="Reference file")
+            import_btn= gr.Button("Import Project", variant="secondary")
+            import_btn.click(lambda: print("import clicked"), outputs=[])
+        with gr.Tab("Website"):
+            url_tb    = gr.Textbox(placeholder="https://example.com")
+
+    # 3 · OUTPUT
+    gr.Markdown("### 3 · Output")
+    lang_dd   = gr.Dropdown(SUPPORTED_LANGUAGES, value="html", label="Target Language")
+    search_ck = gr.Checkbox(label="Enable Tavily Web Search")
+
+    with gr.Row():
+        clear_btn = gr.Button("Clear Session", variant="secondary")
+        gen_btn   = gr.Button("Generate Code", elem_id="gen-btn")
+
+    # OUTPUT PANELS
+    with gr.Tabs():
+        with gr.Tab("Code"):
+            code_out = gr.Code(language="html", interactive=True, lines=25)
+        with gr.Tab("Live Preview"):
+            preview  = gr.HTML()
+        with gr.Tab("History"):
+            chatbox  = gr.Chatbot(type="messages")
+
+    # ── CALLBACKS ──────────────────────────────────────────────
+    model_dd.change(lambda n: next(m for m in AVAILABLE_MODELS if m["name"] == n),
+                    inputs=model_dd, outputs=state_model)
+
     gen_btn.click(
-        generate,
-        inputs=[
-            prompt_box, file_box, url_box,
-            model_dd, lang_dd, search_chk, state_hist
-        ],
-        outputs=[code_out, state_hist, preview_out, chat_out],
+        fn=generation_code,
+        inputs=[prompt_tb, file_upl, url_tb,
+                state_model, search_ck, lang_dd, state_history],
+        outputs=[code_out, state_history, preview, chatbox],
     )
 
     clear_btn.click(
-        lambda: ("", None, "", "html", False, [], [], "", []),
-        outputs=[
-            prompt_box, file_box, url_box,
-            lang_dd, search_chk, state_hist,
-            code_out, preview_out, chat_out,
-        ],
+        lambda: ("", None, "", [], "", "", []),
+        outputs=[prompt_tb, file_upl, url_tb, state_history,
+                 code_out, preview, chatbox],
         queue=False,
     )
 
-
-# ---------- launch ---------------------------------------------------------
 if __name__ == "__main__":
     demo.queue().launch()