mgbam commited on
Commit
31f94f1
Β·
verified Β·
1 Parent(s): eb96582

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +87 -164
app.py CHANGED
@@ -1,27 +1,16 @@
1
- """
2
- app.py
3
 
4
- Main application file for AnyCoder, 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
- 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, load_project_from_url
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. Your primary goal is to generate clean, correct, and efficient code based on the user's request.
45
- - Follow the user's requirements precisely.
46
- - If the user asks for a specific language, provide the code in that language.
47
- - Enclose the final code in a single markdown code block (e.g., ```html ... ```).
48
- - Do not include any conversational text, apologies, or explanations outside of the code block in your final response.
49
  """
50
 
51
- # ==============================================================================
52
- # HELPER FUNCTIONS
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
- # Robustly determine the provider based on ID, falling back to a default
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
- # 3. --- Assemble Full Context for the AI ---
100
- messages = history_to_messages(history, system_prompt)
101
- context_query = query
102
-
103
  if file:
104
- text = extract_text_from_file(file)
105
- context_query += f"\n\n[Attached File Content]\n{text[:5000]}"
106
-
107
  if website_url:
108
- text = extract_website_content(website_url)
109
- if not text.startswith('Error'):
110
- context_query += f"\n\n[Scraped Website Content]\n{text[:8000]}"
111
-
112
- final_query = enhance_query_with_search(context_query, enable_search)
113
- messages.append({'role': 'user', 'content': final_query})
114
 
115
- # 4. --- AI Model Inference with Robust Error Handling ---
116
- client = get_inference_client(model_id, provider, user_token=hf_token)
117
  resp = client.chat.completions.create(
118
  model=model_id,
119
- messages=messages,
120
- max_tokens=16384,
121
  temperature=0.1
122
  )
123
  content = resp.choices[0].message.content
124
 
125
  except Exception as e:
126
- error_message = f"❌ **An error occurred:**\n\n```\n{str(e)}\n```\n\nPlease check your API keys, model selection, or try again."
127
- history.append((query, error_message))
128
  return "", history, "", history_to_chatbot_messages(history)
129
 
130
- # 5. --- Post-process the AI's Output ---
131
- if language == 'transformers.js':
132
  files = parse_transformers_js_output(content)
133
- code_str = format_transformers_js_output(files)
134
- preview_html = send_to_sandbox(files.get('index.html', ''))
135
  else:
136
- clean_code = remove_code_block(content)
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
- code_str = apply_search_replace_changes(history[-1][1], clean_code)
140
  else:
141
- code_str = clean_code
142
- preview_html = send_to_sandbox(code_str) if language == 'html' else ''
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
- # UI LAYOUT & EVENT WIRING
153
- # ==============================================================================
154
 
155
- # Custom CSS for a more professional and modern look
156
  CUSTOM_CSS = """
157
- body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif; }
158
- #main_title { text-align: center; font-size: 2.5rem; font-weight: 700; color: #1a202c; margin: 1.5rem 0 0.5rem 0; }
159
- #subtitle { text-align: center; color: #4a5568; margin-bottom: 2.5rem; font-size: 1.1rem; }
160
- .gradio-container { background-color: #f7fafc; }
161
- /* Custom styling for the generate button to make it stand out */
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", secondary_hue="sky"), title="AnyCoder - AI Code Generator", css=CUSTOM_CSS) as demo:
166
- # --- State Management ---
167
  history_state = gr.State([])
168
- initial_model = AVAILABLE_MODELS[0]
169
- model_state = gr.State(initial_model)
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(equal_height=False):
176
- # Left column for controls and inputs
177
  with gr.Column(scale=1):
178
  gr.Markdown("### 1. Select Model")
179
- model_choices = [model["name"] for model in AVAILABLE_MODELS]
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(label="Attach File (Optional)", type="filepath")
198
  with gr.Tab("🌐 Website"):
199
- url_site = gr.Textbox(label="Scrape Website (Optional)", placeholder="https://example.com")
200
 
201
  gr.Markdown("### 3. Configure Output")
202
- language_dd = gr.Dropdown(
203
- choices=["html", "python", "transformers.js", "sql", "javascript", "css"],
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
- clr_btn = gr.Button("Clear Session", variant="secondary")
211
- gen_btn = gr.Button("Generate Code", variant="primary", elem_id="gen_btn")
212
 
213
- # Right column for outputs
214
  with gr.Column(scale=2):
215
- with gr.Tabs() as main_tabs:
216
- with gr.Tab("πŸ’» Code", id="code_tab"):
217
- code_out = gr.Code(label="Generated Code", language="html", interactive=True)
218
- with gr.Tab("πŸ‘οΈ Live Preview", id="preview_tab"):
219
- preview_out = gr.HTML(label="Live Preview")
220
- with gr.Tab("πŸ“œ History", id="history_tab"):
221
- chat_out = gr.Chatbot(label="Conversation History", type="messages")
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(fn=on_model_change, inputs=[model_dd], outputs=[model_state])
230
- language_dd.change(fn=lambda lang: gr.update(language=lang), inputs=[language_dd], outputs=[code_out])
231
 
232
- gen_btn.click(
233
  fn=generation_code,
234
- inputs=[
235
- prompt_in, file_in, url_site,
236
- model_state, search_chk, language_dd, history_state
237
- ],
238
- outputs=[code_out, history_state, preview_out, chat_out]
239
  )
240
 
241
- def clear_session():
242
- """Resets all UI components and state to their initial values."""
243
- return (
244
- "", # prompt_in
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
- # APPLICATION ENTRY POINT
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()