Update app.py
Browse files
app.py
CHANGED
@@ -1,27 +1,16 @@
|
|
1 |
-
|
2 |
-
app.py
|
3 |
|
4 |
-
|
|
|
5 |
|
6 |
This application provides a user interface for generating code in various languages
|
7 |
using different AI models. It supports inputs from text prompts, files, images,
|
8 |
and websites, and includes features like web search enhancement and live code previews.
|
9 |
-
|
10 |
-
Structure:
|
11 |
-
- Imports & Configuration: Loads necessary libraries and constants.
|
12 |
-
- Helper Functions: Small utility functions supporting the UI logic.
|
13 |
-
- Core Application Logic: The main `generation_code` function that handles the AI interaction.
|
14 |
-
- UI Layout: Defines the Gradio interface using `gr.Blocks`.
|
15 |
-
- Event Wiring: Connects UI components to backend functions.
|
16 |
-
- Application Entry Point: Launches the Gradio app.
|
17 |
"""
|
18 |
|
19 |
import gradio as gr
|
20 |
from typing import Optional, Dict, List, Tuple, Any
|
21 |
|
22 |
-
# --- Local Module Imports ---
|
23 |
-
# These modules contain the application's configuration, clients, and utility functions.
|
24 |
-
# Note: These files (hf_client.py, etc.) must exist in the same directory.
|
25 |
from constants import SYSTEM_PROMPTS, AVAILABLE_MODELS, DEMO_LIST
|
26 |
from hf_client import get_inference_client
|
27 |
from tavily_search import enhance_query_with_search
|
@@ -35,34 +24,33 @@ from utils import (
|
|
35 |
parse_transformers_js_output,
|
36 |
format_transformers_js_output
|
37 |
)
|
38 |
-
from deploy import send_to_sandbox
|
39 |
|
40 |
-
# --- Type Aliases and Constants ---
|
41 |
History = List[Tuple[str, str]]
|
42 |
Model = Dict[str, Any]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
43 |
DEFAULT_SYSTEM_PROMPT = """
|
44 |
-
You are a helpful AI coding assistant.
|
45 |
-
- Follow
|
46 |
-
-
|
47 |
-
-
|
48 |
-
- Do not include any conversational text, apologies, or explanations outside of the code block in your final response.
|
49 |
"""
|
50 |
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
def get_model_details(model_name: str) -> Optional[Model]:
|
56 |
-
"""Finds the full dictionary for a model given its name."""
|
57 |
-
for model in AVAILABLE_MODELS:
|
58 |
-
if model["name"] == model_name:
|
59 |
-
return model
|
60 |
return None
|
61 |
|
62 |
-
# ==============================================================================
|
63 |
-
# CORE APPLICATION LOGIC
|
64 |
-
# ==============================================================================
|
65 |
-
|
66 |
def generation_code(
|
67 |
query: Optional[str],
|
68 |
file: Optional[str],
|
@@ -71,194 +59,129 @@ def generation_code(
|
|
71 |
enable_search: bool,
|
72 |
language: str,
|
73 |
history: Optional[History],
|
74 |
-
hf_token: str,
|
75 |
) -> Tuple[str, History, str, List[Dict[str, str]]]:
|
76 |
-
"""
|
77 |
-
The main function to handle a user's code generation request.
|
78 |
-
"""
|
79 |
-
# 1. --- Initialization and Input Sanitization ---
|
80 |
query = query or ""
|
81 |
history = history or []
|
82 |
-
|
83 |
try:
|
84 |
-
# 2. --- System Prompt and Model Selection ---
|
85 |
system_prompt = SYSTEM_PROMPTS.get(language, DEFAULT_SYSTEM_PROMPT)
|
86 |
model_id = current_model["id"]
|
87 |
-
|
88 |
-
#
|
89 |
-
if model_id.startswith("openai/"):
|
90 |
provider = "openai"
|
91 |
elif model_id.startswith("gemini/"):
|
92 |
provider = "gemini"
|
93 |
elif model_id.startswith("fireworks-ai/"):
|
94 |
provider = "fireworks-ai"
|
95 |
else:
|
96 |
-
# Assume other models are served via standard Hugging Face TGI
|
97 |
provider = "huggingface"
|
98 |
|
99 |
-
#
|
100 |
-
|
101 |
-
|
102 |
-
|
103 |
if file:
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
if website_url:
|
108 |
-
|
109 |
-
if not
|
110 |
-
|
111 |
-
|
112 |
-
|
113 |
-
messages.append({'role': 'user', 'content': final_query})
|
114 |
|
115 |
-
|
116 |
-
client = get_inference_client(model_id, provider, user_token=hf_token)
|
117 |
resp = client.chat.completions.create(
|
118 |
model=model_id,
|
119 |
-
messages=
|
120 |
-
max_tokens=
|
121 |
temperature=0.1
|
122 |
)
|
123 |
content = resp.choices[0].message.content
|
124 |
|
125 |
except Exception as e:
|
126 |
-
|
127 |
-
history.append((query,
|
128 |
return "", history, "", history_to_chatbot_messages(history)
|
129 |
|
130 |
-
#
|
131 |
-
if language ==
|
132 |
files = parse_transformers_js_output(content)
|
133 |
-
|
134 |
-
|
135 |
else:
|
136 |
-
|
137 |
-
# Apply search/replace if a previous turn exists and contains valid code
|
138 |
if history and history[-1][1] and not history[-1][1].startswith("β"):
|
139 |
-
|
140 |
else:
|
141 |
-
|
142 |
-
|
143 |
-
|
144 |
-
# 6. --- Update History and Final Outputs ---
|
145 |
-
updated_history = history + [(query, code_str)]
|
146 |
-
chat_messages = history_to_chatbot_messages(updated_history)
|
147 |
-
|
148 |
-
return code_str, updated_history, preview_html, chat_messages
|
149 |
-
|
150 |
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
|
155 |
-
#
|
156 |
CUSTOM_CSS = """
|
157 |
-
body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
|
158 |
-
#main_title { text-align: center; font-size: 2.5rem;
|
159 |
-
#subtitle { text-align: center; color: #4a5568; margin-bottom: 2.5rem;
|
160 |
-
.gradio-container { background
|
161 |
-
|
162 |
-
#gen_btn { box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06); }
|
163 |
"""
|
164 |
|
165 |
-
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue",
|
166 |
-
# --- State Management ---
|
167 |
history_state = gr.State([])
|
168 |
-
|
169 |
-
model_state = gr.State(
|
170 |
|
171 |
-
# --- UI Definition ---
|
172 |
gr.Markdown("# π Shasha AI", elem_id="main_title")
|
173 |
gr.Markdown("Your personal AI partner for generating, modifying, and understanding code.", elem_id="subtitle")
|
174 |
|
175 |
-
with gr.Row(
|
176 |
-
# Left column for controls and inputs
|
177 |
with gr.Column(scale=1):
|
178 |
gr.Markdown("### 1. Select Model")
|
179 |
-
|
180 |
-
model_dd = gr.Dropdown(
|
181 |
-
choices=model_choices,
|
182 |
-
value=initial_model["name"],
|
183 |
-
label="AI Model",
|
184 |
-
info="Different models have different strengths."
|
185 |
-
)
|
186 |
|
187 |
gr.Markdown("### 2. Provide Context")
|
188 |
with gr.Tabs():
|
189 |
with gr.Tab("π Prompt"):
|
190 |
-
prompt_in = gr.Textbox(
|
191 |
-
label="Your Request",
|
192 |
-
lines=7,
|
193 |
-
placeholder="e.g., 'Create a modern, responsive landing page for a SaaS product.'",
|
194 |
-
show_label=False
|
195 |
-
)
|
196 |
with gr.Tab("π File"):
|
197 |
-
file_in = gr.File(
|
198 |
with gr.Tab("π Website"):
|
199 |
-
|
200 |
|
201 |
gr.Markdown("### 3. Configure Output")
|
202 |
-
|
203 |
-
|
204 |
-
value="html",
|
205 |
-
label="Target Language"
|
206 |
-
)
|
207 |
-
search_chk = gr.Checkbox(label="Enable Web Search", info="Enhances AI with real-time info.")
|
208 |
|
209 |
with gr.Row():
|
210 |
-
|
211 |
-
|
212 |
|
213 |
-
# Right column for outputs
|
214 |
with gr.Column(scale=2):
|
215 |
-
with gr.Tabs()
|
216 |
-
with gr.Tab("π» Code"
|
217 |
-
code_out = gr.Code(
|
218 |
-
with gr.Tab("ποΈ Live Preview"
|
219 |
-
preview_out = gr.HTML(
|
220 |
-
with gr.Tab("π History"
|
221 |
-
chat_out = gr.Chatbot(
|
222 |
-
|
223 |
-
# --- Event Wiring ---
|
224 |
-
def on_model_change(model_name: str) -> Dict:
|
225 |
-
"""Updates the model_state when the user selects a new model."""
|
226 |
-
model_details = get_model_details(model_name)
|
227 |
-
return model_details or initial_model
|
228 |
|
229 |
-
model_dd.change(
|
230 |
-
language_dd.change(fn=lambda lang: gr.update(language=lang), inputs=[language_dd], outputs=[code_out])
|
231 |
|
232 |
-
|
233 |
fn=generation_code,
|
234 |
-
inputs=[
|
235 |
-
|
236 |
-
model_state, search_chk, language_dd, history_state
|
237 |
-
],
|
238 |
-
outputs=[code_out, history_state, preview_out, chat_out]
|
239 |
)
|
240 |
|
241 |
-
|
242 |
-
""
|
243 |
-
|
244 |
-
|
245 |
-
None, # file_in
|
246 |
-
"", # url_site
|
247 |
-
[], # history_state
|
248 |
-
"", # code_out
|
249 |
-
"", # preview_out
|
250 |
-
[] # chat_out
|
251 |
-
)
|
252 |
-
|
253 |
-
clr_btn.click(
|
254 |
-
fn=clear_session,
|
255 |
-
outputs=[prompt_in, file_in, url_site, history_state, code_out, preview_out, chat_out],
|
256 |
-
queue=False
|
257 |
)
|
258 |
|
259 |
-
|
260 |
-
|
261 |
-
# ==============================================================================
|
262 |
-
|
263 |
-
if __name__ == '__main__':
|
264 |
-
demo.queue().launch()
|
|
|
1 |
+
# app.py
|
|
|
2 |
|
3 |
+
"""
|
4 |
+
Main application file for SHASHA AI, a Gradio-based AI code generation tool.
|
5 |
|
6 |
This application provides a user interface for generating code in various languages
|
7 |
using different AI models. It supports inputs from text prompts, files, images,
|
8 |
and websites, and includes features like web search enhancement and live code previews.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
9 |
"""
|
10 |
|
11 |
import gradio as gr
|
12 |
from typing import Optional, Dict, List, Tuple, Any
|
13 |
|
|
|
|
|
|
|
14 |
from constants import SYSTEM_PROMPTS, AVAILABLE_MODELS, DEMO_LIST
|
15 |
from hf_client import get_inference_client
|
16 |
from tavily_search import enhance_query_with_search
|
|
|
24 |
parse_transformers_js_output,
|
25 |
format_transformers_js_output
|
26 |
)
|
27 |
+
from deploy import send_to_sandbox
|
28 |
|
|
|
29 |
History = List[Tuple[str, str]]
|
30 |
Model = Dict[str, Any]
|
31 |
+
|
32 |
+
# Full list of supported languages for syntax highlighting & generation
|
33 |
+
SUPPORTED_LANGUAGES = [
|
34 |
+
"python", "c", "cpp", "markdown", "latex", "json", "html", "css",
|
35 |
+
"javascript", "jinja2", "typescript", "yaml", "dockerfile", "shell",
|
36 |
+
"r", "sql", "sql-msSQL", "sql-mySQL", "sql-mariaDB", "sql-sqlite",
|
37 |
+
"sql-cassandra", "sql-plSQL", "sql-hive", "sql-pgSQL", "sql-gql",
|
38 |
+
"sql-gpSQL", "sql-sparkSQL", "sql-esper"
|
39 |
+
]
|
40 |
+
|
41 |
DEFAULT_SYSTEM_PROMPT = """
|
42 |
+
You are a helpful AI coding assistant. Generate clean, correct, and efficient code based on the user's request.
|
43 |
+
- Follow requirements precisely.
|
44 |
+
- Enclose final code in a single ```code``` block of the target language.
|
45 |
+
- Do not include any explanations outside the code block.
|
|
|
46 |
"""
|
47 |
|
48 |
+
def get_model_details(name: str) -> Optional[Model]:
|
49 |
+
for m in AVAILABLE_MODELS:
|
50 |
+
if m["name"] == name:
|
51 |
+
return m
|
|
|
|
|
|
|
|
|
|
|
52 |
return None
|
53 |
|
|
|
|
|
|
|
|
|
54 |
def generation_code(
|
55 |
query: Optional[str],
|
56 |
file: Optional[str],
|
|
|
59 |
enable_search: bool,
|
60 |
language: str,
|
61 |
history: Optional[History],
|
|
|
62 |
) -> Tuple[str, History, str, List[Dict[str, str]]]:
|
|
|
|
|
|
|
|
|
63 |
query = query or ""
|
64 |
history = history or []
|
|
|
65 |
try:
|
|
|
66 |
system_prompt = SYSTEM_PROMPTS.get(language, DEFAULT_SYSTEM_PROMPT)
|
67 |
model_id = current_model["id"]
|
68 |
+
|
69 |
+
# pick provider
|
70 |
+
if model_id.startswith("openai/") or model_id in ("gpt-4", "gpt-3.5-turbo"):
|
71 |
provider = "openai"
|
72 |
elif model_id.startswith("gemini/"):
|
73 |
provider = "gemini"
|
74 |
elif model_id.startswith("fireworks-ai/"):
|
75 |
provider = "fireworks-ai"
|
76 |
else:
|
|
|
77 |
provider = "huggingface"
|
78 |
|
79 |
+
# assemble messages
|
80 |
+
msgs = history_to_messages(history, system_prompt)
|
81 |
+
ctx = query
|
|
|
82 |
if file:
|
83 |
+
txt = extract_text_from_file(file)
|
84 |
+
ctx += f"\n\n[File]\n{txt[:5000]}"
|
|
|
85 |
if website_url:
|
86 |
+
txt = extract_website_content(website_url)
|
87 |
+
if not txt.startswith("Error"):
|
88 |
+
ctx += f"\n\n[Website]\n{txt[:8000]}"
|
89 |
+
final_q = enhance_query_with_search(ctx, enable_search)
|
90 |
+
msgs.append({"role": "user", "content": final_q})
|
|
|
91 |
|
92 |
+
client = get_inference_client(model_id, provider)
|
|
|
93 |
resp = client.chat.completions.create(
|
94 |
model=model_id,
|
95 |
+
messages=msgs,
|
96 |
+
max_tokens=16000,
|
97 |
temperature=0.1
|
98 |
)
|
99 |
content = resp.choices[0].message.content
|
100 |
|
101 |
except Exception as e:
|
102 |
+
err = f"β **Error:**\n```\n{e}\n```"
|
103 |
+
history.append((query, err))
|
104 |
return "", history, "", history_to_chatbot_messages(history)
|
105 |
|
106 |
+
# post-process
|
107 |
+
if language == "transformers.js":
|
108 |
files = parse_transformers_js_output(content)
|
109 |
+
code = format_transformers_js_output(files)
|
110 |
+
preview = send_to_sandbox(files.get("index.html",""))
|
111 |
else:
|
112 |
+
clean = remove_code_block(content)
|
|
|
113 |
if history and history[-1][1] and not history[-1][1].startswith("β"):
|
114 |
+
code = apply_search_replace_changes(history[-1][1], clean)
|
115 |
else:
|
116 |
+
code = clean
|
117 |
+
preview = send_to_sandbox(code) if language == "html" else ""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
118 |
|
119 |
+
new_hist = history + [(query, code)]
|
120 |
+
chat = history_to_chatbot_messages(new_hist)
|
121 |
+
return code, new_hist, preview, chat
|
122 |
|
123 |
+
# custom CSS
|
124 |
CUSTOM_CSS = """
|
125 |
+
body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; }
|
126 |
+
#main_title { text-align: center; font-size: 2.5rem; margin: 1.5rem 0 0.5rem; }
|
127 |
+
#subtitle { text-align: center; color: #4a5568; margin-bottom: 2.5rem; }
|
128 |
+
.gradio-container { background: #f7fafc; }
|
129 |
+
#gen_btn { box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
|
|
|
130 |
"""
|
131 |
|
132 |
+
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=CUSTOM_CSS, title="Shasha AI") as demo:
|
|
|
133 |
history_state = gr.State([])
|
134 |
+
initial = AVAILABLE_MODELS[0]
|
135 |
+
model_state = gr.State(initial)
|
136 |
|
|
|
137 |
gr.Markdown("# π Shasha AI", elem_id="main_title")
|
138 |
gr.Markdown("Your personal AI partner for generating, modifying, and understanding code.", elem_id="subtitle")
|
139 |
|
140 |
+
with gr.Row():
|
|
|
141 |
with gr.Column(scale=1):
|
142 |
gr.Markdown("### 1. Select Model")
|
143 |
+
names = [m["name"] for m in AVAILABLE_MODELS]
|
144 |
+
model_dd = gr.Dropdown(names, value=initial["name"], label="AI Model")
|
|
|
|
|
|
|
|
|
|
|
145 |
|
146 |
gr.Markdown("### 2. Provide Context")
|
147 |
with gr.Tabs():
|
148 |
with gr.Tab("π Prompt"):
|
149 |
+
prompt_in = gr.Textbox(lines=7, placeholder="Describe your request...", show_label=False)
|
|
|
|
|
|
|
|
|
|
|
150 |
with gr.Tab("π File"):
|
151 |
+
file_in = gr.File(type="filepath")
|
152 |
with gr.Tab("π Website"):
|
153 |
+
url_in = gr.Textbox(placeholder="https://example.com")
|
154 |
|
155 |
gr.Markdown("### 3. Configure Output")
|
156 |
+
lang_dd = gr.Dropdown(SUPPORTED_LANGUAGES, value="html", label="Target Language")
|
157 |
+
search_chk = gr.Checkbox(label="Enable Web Search")
|
|
|
|
|
|
|
|
|
158 |
|
159 |
with gr.Row():
|
160 |
+
clr = gr.Button("Clear Session", variant="secondary")
|
161 |
+
gen = gr.Button("Generate Code", variant="primary", elem_id="gen_btn")
|
162 |
|
|
|
163 |
with gr.Column(scale=2):
|
164 |
+
with gr.Tabs():
|
165 |
+
with gr.Tab("π» Code"):
|
166 |
+
code_out = gr.Code(language=lambda: lang_dd.value, interactive=True)
|
167 |
+
with gr.Tab("ποΈ Live Preview"):
|
168 |
+
preview_out = gr.HTML()
|
169 |
+
with gr.Tab("π History"):
|
170 |
+
chat_out = gr.Chatbot(type="messages")
|
|
|
|
|
|
|
|
|
|
|
|
|
171 |
|
172 |
+
model_dd.change(lambda n: get_model_details(n) or initial, inputs=[model_dd], outputs=[model_state])
|
|
|
173 |
|
174 |
+
gen.click(
|
175 |
fn=generation_code,
|
176 |
+
inputs=[prompt_in, file_in, url_in, model_state, search_chk, lang_dd, history_state],
|
177 |
+
outputs=[code_out, history_state, preview_out, chat_out],
|
|
|
|
|
|
|
178 |
)
|
179 |
|
180 |
+
clr.click(
|
181 |
+
fn=lambda: ("", None, "", [], "", "", []),
|
182 |
+
outputs=[prompt_in, file_in, url_in, history_state, code_out, preview_out, chat_out],
|
183 |
+
queue=False,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
184 |
)
|
185 |
|
186 |
+
if __name__ == "__main__":
|
187 |
+
demo.queue().launch()
|
|
|
|
|
|
|
|