# app.py """ Main application file for SHASHA AI (Gradio). Only change: enlarge logo to 120 × 120 px. """ import gradio as gr from typing import Optional, Dict, List, Tuple, Any import os # --- Local 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 ( 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, ) from deploy import send_to_sandbox History = List[Tuple[str, str]] Model = Dict[str, Any] 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]: return next((m for m in AVAILABLE_MODELS if m["name"]==name), None) def generation_code( query:Optional[str], file:Optional[str], website_url:Optional[str], current_model:Model, enable_search:bool, language:str, history:Optional[History], )->Tuple[str,History,str,List[Dict[str,str]]]: query = query or "" history = history or [] try: system_prompt = ( HTML_SYSTEM_PROMPT if language=="html" else TRANSFORMERS_JS_SYSTEM_PROMPT if language=="transformers.js" else f"You are an expert {language} developer. Write clean, idiomatic {language} code." ) model_id = current_model["id"] 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" ) msgs = history_to_messages(history, system_prompt) context = query if file: 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]}" msgs.append({"role":"user","content":enhance_query_with_search(context, enable_search)}) 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) if language=="transformers.js": 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(content) 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)] return code, new_hist, preview, history_to_chatbot_messages(new_hist) # --- 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:.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);} """ LOGO_PATH = "assets/logo.png" 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) # enlarged logo (120×120 px) if os.path.exists(LOGO_PATH): gr.Image(value=LOGO_PATH, height=120, width=120, show_label=False, container=False) 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(): with gr.Column(scale=1): gr.Markdown("### 1. Select Model") model_dd = gr.Dropdown([m["name"] for m in AVAILABLE_MODELS], value=initial_model["name"], label="AI Model") 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) 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") with gr.Column(scale=2): with gr.Tabs(): 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") 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], ) 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()