# 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. """ import gradio as gr 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 ( 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 # --- Type aliases --- History = List[Tuple[str, str]] Model = Dict[str, Any] # --- 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" ] 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: Optional[History], ) -> Tuple[str, History, str, List[Dict[str, str]]]: query = query or "" history = history or [] try: # 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) # Process model 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", "")) 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 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=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("# 🚀 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( choices=[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()