# app.py """ Main application file for SHASHA AI, a Gradio-based AI code generation tool. This application provides a user interface for generating code in various languages using different AI models. It supports inputs from text prompts, files, images, and websites, and includes features like web search enhancement and live code previews. """ import gradio as gr from typing import Optional, Dict, List, Tuple, Any from constants import SYSTEM_PROMPTS, 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] # Full list of supported languages for syntax highlighting & generation 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" ] DEFAULT_SYSTEM_PROMPT = """ You are a helpful AI coding assistant. Generate clean, correct, and efficient code based on the user's request. - Follow requirements precisely. - Enclose final code in a single ```code``` block of the target language. - Do not include any explanations outside the code block. """ 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: system_prompt = SYSTEM_PROMPTS.get(language, DEFAULT_SYSTEM_PROMPT) model_id = current_model["id"] # pick provider if model_id.startswith("openai/") or model_id in ("gpt-4", "gpt-3.5-turbo"): provider = "openai" elif model_id.startswith("gemini/"): provider = "gemini" elif model_id.startswith("fireworks-ai/"): provider = "fireworks-ai" else: provider = "huggingface" # assemble messages msgs = history_to_messages(history, system_prompt) ctx = query if file: txt = extract_text_from_file(file) ctx += f"\n\n[File]\n{txt[:5000]}" if website_url: txt = extract_website_content(website_url) if not txt.startswith("Error"): ctx += f"\n\n[Website]\n{txt[:8000]}" final_q = enhance_query_with_search(ctx, enable_search) msgs.append({"role": "user", "content": final_q}) 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) # post-process 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: clean = remove_code_block(content) if history and history[-1][1] and not history[-1][1].startswith("❌"): code = apply_search_replace_changes(history[-1][1], clean) else: code = clean 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: 1.5rem 0 0.5rem; } #subtitle { text-align: center; color: #4a5568; margin-bottom: 2.5rem; } .gradio-container { background: #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 = AVAILABLE_MODELS[0] model_state = gr.State(initial) gr.Markdown("# 🚀 Shasha AI", elem_id="main_title") gr.Markdown("Your personal AI partner for generating, modifying, and understanding code.", elem_id="subtitle") with gr.Row(): with gr.Column(scale=1): gr.Markdown("### 1. Select Model") names = [m["name"] for m in AVAILABLE_MODELS] model_dd = gr.Dropdown(names, value=initial["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 = gr.Button("Clear Session", variant="secondary") gen = 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=lambda: lang_dd.value, 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, inputs=[model_dd], outputs=[model_state]) gen.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.click( fn=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()