Update app.py

app.py

@@ -1,31 +1,44 @@
+# app.py
 import gradio as gr
 import re
 import json
-import requests
+# Remove direct requests import, will use model_logic
+# import requests
 import os
 import tempfile
 
 # --- build_logic.py is now a hard requirement ---
 from build_logic import (
     create_space as build_logic_create_space,
-    _get_api_token as build_logic_get_api_token,
-    whoami as build_logic_whoami,
+    _get_api_token as build_logic_get_api_token, # Keep this for HF Hub token logic
+    whoami as build_logic_whoami, # Keep this for HF user info
     list_space_files_for_browsing,
     get_space_repository_info,
     get_space_file_content,
     update_space_file,
     parse_markdown as build_logic_parse_markdown,
     delete_space_file as build_logic_delete_space_file,
-    get_space_runtime_status # Import the new function
+    get_space_runtime_status
 )
 print("build_logic.py loaded successfully.")
-# --- End build_logic import ---
+
+# --- model_logic.py is now a hard requirement ---
+from model_logic import (
+    get_available_providers,
+    get_models_for_provider,
+    get_default_model_for_provider,
+    get_model_id_from_display_name, # Might not be strictly needed in app.py, but good practice
+    generate_stream # This is the core function we'll use
+)
+print("model_logic.py loaded successfully.")
+# --- End imports ---
 
 
 bbb = chr(96) * 3
 parsed_code_blocks_state_cache = []
 BOT_ROLE_NAME = "assistant"
-GROQ_API_ENDPOINT = "https://api.groq.com/openai/v1/chat/completions"
+# Removed GROQ_API_ENDPOINT as it's now in model_logic
+
 
 DEFAULT_SYSTEM_PROMPT = f"""You are an expert AI programmer. Your role is to generate code and file structures based on user requests, or to modify existing code provided by the user.
 When you provide NEW code for a file, or MODIFIED code for an existing file, use the following format exactly:
@@ -55,7 +68,7 @@ If the user asks to delete a file, simply omit it from your next full ### File:
 If no code is provided, assist the user with their tasks.
 """
 
-# --- Core Utility, Parsing, Export, API Call functions ---
+# --- Core Utility, Parsing, Export functions (mostly unchanged) ---
 def escape_html_for_markdown(text):
     if not isinstance(text, str): return ""
     # Minimal escaping, expand if needed
@@ -115,90 +128,79 @@ def _clean_filename(filename_line_content):
     return filename_candidate if filename_candidate else text.strip()
 
 
-def _parse_chat_stream_logic(chat_json_string, existing_files_state=None):
-    global parsed_code_blocks_state_cache
+def _parse_chat_stream_logic(latest_bot_message_content, existing_files_state=None):
+    """
+    Parses a single bot message content string to find file blocks and updates the state.
+    Assumes existing_files_state is the current state *before* this message.
+    """
     latest_blocks_dict = {}
     if existing_files_state:
-        for block in existing_files_state: latest_blocks_dict[block["filename"]] = block.copy()
-
-    results = {"parsed_code_blocks": [], "preview_md": "", "default_selected_filenames": [], "error_message": None}
-    try:
-        ai_chat_history = json.loads(chat_json_string)
-        if not isinstance(ai_chat_history, list): raise ValueError("JSON input must be a list of chat messages.")
-    except json.JSONDecodeError as e: results["error_message"] = f"JSON Parsing Error: {e}."; return results
-    except ValueError as e: results["error_message"] = str(e); return results
-
-    # Ensure the bot's message is the last one and is the only one being parsed for new/updated files
-    # This prevents reprocessing the entire history on every new AI turn
-    message_obj = None
-    if ai_chat_history and isinstance(ai_chat_history[-1], dict) and ai_chat_history[-1].get("role", "").lower() == BOT_ROLE_NAME:
-         message_obj = ai_chat_history[-1]
+        # Copy existing blocks, except for potential structure blocks that might be overwritten
+        for block in existing_files_state:
+             if not block.get("is_structure_block"):
+                 latest_blocks_dict[block["filename"]] = block.copy()
+             # Keep existing structure block for now, it might be replaced below
 
-    if not message_obj:
-         # If the last message isn't the bot's or is malformed, don't update state based on it
-         results["parsed_code_blocks"] = list(latest_blocks_dict.values()) # Return existing state
-         results["default_selected_filenames"] = [b["filename"] for b in results["parsed_code_blocks"] if not b.get("is_structure_block")]
-         # results["error_message"] = "No bot message found in the last entry for parsing." # Optional: add a debug message
-         return results
 
-
-    role, content = message_obj.get("role", "").lower(), message_obj.get("content", "")
+    results = {"parsed_code_blocks": [], "preview_md": "", "default_selected_filenames": [], "error_message": None}
+    content = latest_bot_message_content or ""
 
     file_pattern = re.compile(r"### File:\s*(?P<filename_line>[^\n]+)\n(?:```(?P<lang>[\w\.\-\+]*)\n(?P<code>[\s\S]*?)\n```|(?P<binary_msg>\[Binary file(?: - [^\]]+)?\]))")
     structure_pattern = re.compile(r"## File Structure\n```(?:(?P<struct_lang>[\w.-]*)\n)?(?P<structure_code>[\s\S]*?)\n```")
 
-    # Process only the *last* bot message for updates to file blocks
-    if role == BOT_ROLE_NAME:
-        structure_match = structure_pattern.search(content)
-        if structure_match:
-            # Overwrite or add the structure block from the latest response
-            latest_blocks_dict["File Structure (original)"] = {"filename": "File Structure (original)", "language": structure_match.group("struct_lang") or "plaintext", "code": structure_match.group("structure_code").strip(), "is_binary": False, "is_structure_block": True}
-
-        # Find all file blocks in the latest message
-        current_message_file_blocks = {}
-        for match in file_pattern.finditer(content):
-            filename = _clean_filename(match.group("filename_line"))
-            if not filename: continue
-            lang, code_block, binary_msg = match.group("lang"), match.group("code"), match.group("binary_msg")
-            item_data = {"filename": filename, "is_binary": False, "is_structure_block": False}
-            if code_block is not None:
-                item_data["code"], item_data["language"] = code_block.strip(), (lang.strip().lower() if lang else _infer_lang_from_filename(filename))
-            elif binary_msg is not None:
-                item_data["code"], item_data["language"], item_data["is_binary"] = binary_msg.strip(), "binary", True
-            else: continue
-            current_message_file_blocks[filename] = item_data
-
-        # Update latest_blocks_dict with blocks from the current message
-        # Any file mentioned in the latest message replaces its old version
-        # Any file *not* mentioned in the latest message but present previously is kept
-        # If the AI response implies deletion (e.g., by omitting a file it previously listed),
-        # the user needs to signal this intention in a way the prompt covers (e.g., "omit it from your next full ### File: list").
-        # Our current logic keeps files unless explicitly removed from the *entire* chat history's final list of blocks.
-        # A more sophisticated approach might look at AI language indicating deletion, but simple omission
-        # in a single turn's output isn't a reliable delete signal with complex interaction.
-        # For this implementation, we'll just update based on the latest message's defined blocks.
-        latest_blocks_dict.update(current_message_file_blocks)
-
-
+    # Process the latest bot message for updates to file blocks
+    structure_match = structure_pattern.search(content)
+    if structure_match:
+        # Add/Overwrite the structure block from the latest response
+        latest_blocks_dict["File Structure (original)"] = {"filename": "File Structure (original)", "language": structure_match.group("struct_lang") or "plaintext", "code": structure_match.group("structure_code").strip(), "is_binary": False, "is_structure_block": True}
+    else:
+         # If the latest message *doesn't* have a structure block, keep the previous one if it existed
+         existing_structure_block = next((b for b in (existing_files_state or []) if b.get("is_structure_block")), None)
+         if existing_structure_block:
+              latest_blocks_dict["File Structure (original)"] = existing_structure_block.copy()
+
+
+    # Find all file blocks in the latest message
+    current_message_file_blocks = {}
+    for match in file_pattern.finditer(content):
+        filename = _clean_filename(match.group("filename_line"))
+        if not filename: continue
+        lang, code_block, binary_msg = match.group("lang"), match.group("code"), match.group("binary_msg")
+        item_data = {"filename": filename, "is_binary": False, "is_structure_block": False}
+        if code_block is not None:
+            item_data["code"], item_data["language"] = code_block.strip(), (lang.strip().lower() if lang else _infer_lang_from_filename(filename))
+        elif binary_msg is not None:
+            item_data["code"], item_data["language"], item_data["is_binary"] = binary_msg.strip(), "binary", True
+        else: continue # Should not happen with the regex
+        current_message_file_blocks[filename] = item_data
+
+    # Update latest_blocks_dict with blocks from the current message
+    # Any file mentioned in the latest message replaces its old version
+    latest_blocks_dict.update(current_message_file_blocks)
+
+
+    # Convert dictionary values back to a list
     current_parsed_blocks = list(latest_blocks_dict.values())
+    # Sort: structure block first, then files alphabetically
     current_parsed_blocks.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
-    parsed_code_blocks_state_cache = current_parsed_blocks
+
+    # Update the global cache outside this function if needed, or pass it back
+    # For now, let's return the new state and let the caller update the cache.
     results["parsed_code_blocks"] = current_parsed_blocks
     results["default_selected_filenames"] = [b["filename"] for b in current_parsed_blocks if not b.get("is_structure_block")]
     return results
 
 def _export_selected_logic(selected_filenames, space_line_name_for_md, parsed_blocks_for_export):
+    # This function remains largely the same, using the provided parsed_blocks_for_export
     results = {"output_str": "", "error_message": None, "download_filepath": None}
-    global parsed_code_blocks_state_cache
-    # Use parsed_blocks_for_export if provided, otherwise use cache
-    if parsed_blocks_for_export is None:
-        parsed_blocks_for_export = parsed_code_blocks_state_cache
+    # Filter out structure blocks for file listing/export content
+    exportable_blocks_content = [b for b in parsed_blocks_for_export if not b.get("is_structure_block") and not b.get("is_binary") and not (b.get("code", "").startswith("[Error loading content:") or b.get("code", "").startswith("[Binary or Skipped file]"))]
+    binary_blocks_content = [b for b in parsed_blocks_for_export if b.get("is_binary") or b.get("code", "").startswith("[Binary or Skipped file]")]
 
-    # Filter out structure blocks and binary/error files for file listing/export
-    exportable_blocks = [b for b in parsed_blocks_for_export if not b.get("is_structure_block") and not b.get("is_binary") and not (b.get("code", "").startswith("[Error loading content:") or b.get("code", "").startswith("[Binary or Skipped file]"))]
-    binary_blocks = [b for b in parsed_blocks_for_export if b.get("is_binary") or b.get("code", "").startswith("[Binary or Skipped file]")]
+    # Collect all filenames (including binary/error ones) for the structure list
+    all_filenames_in_state = sorted(list(set(b["filename"] for b in parsed_blocks_for_export if not b.get("is_structure_block"))))
 
-    if not exportable_blocks and not binary_blocks and not any(b.get("is_structure_block") for b in parsed_blocks_for_export):
+    if not all_filenames_in_state and not any(b.get("is_structure_block") for b in parsed_blocks_for_export):
         results["output_str"] = f"# Space: {space_line_name_for_md}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files to list in structure or export.*"
         try:
             with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md", encoding='utf-8') as tmpfile:
@@ -215,34 +217,46 @@ def _export_selected_logic(selected_filenames, space_line_name_for_md, parsed_bl
     else:
         # If no structure block from AI, generate a simple one from detected files
         output_lines.extend(["## File Structure", bbb, "📁 Root"])
-        filenames_for_structure_list = sorted(list(set(b["filename"] for b in exportable_blocks + binary_blocks))) # Include binary files in structure list
-        if filenames_for_structure_list:
-            for fname in filenames_for_structure_list: output_lines.append(f"  📄 {fname}")
+        if all_filenames_in_state:
+            for fname in all_filenames_in_state: output_lines.append(f"  📄 {fname}")
         output_lines.extend([bbb, ""])
 
     output_lines.append("Below are the contents of all files in the space:\n")
     exported_content = False
 
     # Determine which files to export content for based on selection or default
+    # Exportable content blocks only
     files_to_export_content = []
     if selected_filenames:
-        files_to_export_content = [b for b in parsed_blocks_for_export if b["filename"] in selected_filenames and not b.get("is_structure_block")]
+        files_to_export_content = [b for b in exportable_blocks_content if b["filename"] in selected_filenames]
     else:
-        files_to_export_content = [b for b in parsed_blocks_for_export if not b.get("is_structure_block")]
+        files_to_export_content = exportable_blocks_content # Export all content blocks by default
+
+    # Add binary/error blocks if they were selected or if exporting all (and they exist)
+    # Binary/error blocks are listed in the structure, but their *content* is just the marker string
+    binary_error_blocks_to_export = []
+    if selected_filenames:
+         binary_error_blocks_to_export = [b for b in binary_blocks_content if b["filename"] in selected_filenames]
+    elif binary_blocks_content:
+         binary_error_blocks_to_export = binary_blocks_content # Include all binary/error if exporting all
+
+    # Combine and sort all blocks whose content/marker should be included
+    all_blocks_to_export_content = sorted(files_to_export_content + binary_error_blocks_to_export, key=lambda b: b["filename"])
 
-    # Sort blocks for consistent output order
-    files_to_export_content.sort(key=lambda b: (0, b["filename"]) if b.get("is_binary") else (1, b["filename"]))
 
-    for block in files_to_export_content:
+    for block in all_blocks_to_export_content:
         output_lines.append(f"### File: {block['filename']}")
-        if block.get('is_binary') or block.get("code", "").startswith("["): # Treat errors/skipped as binary for output format
-            output_lines.append(block.get('code',''))
+        if block.get('is_binary') or block.get("code", "").startswith("[Binary file") or block.get("code", "").startswith("[Error loading content:") or block.get("code", "").startswith("[Binary or Skipped file]"):
+            # For binary/error placeholders, just add the marker line
+            output_lines.append(block.get('code','[Binary or Skipped file]'))
         else:
+            # For actual code/text content
             output_lines.extend([f"{bbb}{block.get('language', 'plaintext') or 'plaintext'}", block.get('code',''), bbb])
         output_lines.append(""); exported_content = True
 
-    if not exported_content and not filenames_for_structure_list: output_lines.append("*No file content selected or available for export.*")
-    elif not exported_content and filenames_for_structure_list:  output_lines.append("*Selected files have no content blocks defined by AI, are binary, or encountered loading errors.*")
+    if not exported_content and not all_filenames_in_state: output_lines.append("*No files in state.*")
+    elif not exported_content: output_lines.append("*No files with editable content are in the state or selected.*") # Message updated
+
     final_output_str = "\n".join(output_lines)
     results["output_str"] = final_output_str
     try:
@@ -251,7 +265,9 @@ def _export_selected_logic(selected_filenames, space_line_name_for_md, parsed_bl
     except Exception as e: print(f"Error creating temp file: {e}"); results["error_message"] = "Could not prepare file for download."
     return results
 
+
 def _convert_gr_history_to_api_messages(system_prompt, gr_history, current_user_message=None):
+    # This function is fine as is, it produces standard OpenAI format
     messages = [{"role": "system", "content": system_prompt}] if system_prompt else []
     for user_msg, bot_msg in gr_history:
         if user_msg: messages.append({"role": "user", "content": user_msg})
@@ -260,22 +276,13 @@ def _convert_gr_history_to_api_messages(system_prompt, gr_history, current_user_
     if current_user_message: messages.append({"role": "user", "content": current_user_message})
     return messages
 
-def get_latest_bot_message_as_json(gr_history):
-    # Find the last non-empty bot message
-    latest_bot_msg_content = None
-    for user_msg, bot_msg in reversed(gr_history):
-        if bot_msg and isinstance(bot_msg, str):
-            latest_bot_msg_content = bot_msg
-            break # Found the latest, stop searching
-
-    if latest_bot_msg_content is None:
-        return json.dumps([]) # Return empty list if no bot message found
-
-    # Return JSON containing only the latest bot message for parsing
-    return json.dumps([{"role": BOT_ROLE_NAME, "content": latest_bot_msg_content}], indent=2)
+# This function is no longer needed as we process the *latest* message content directly from the stream handler
+# def get_latest_bot_message_as_json(gr_history):
+#    ...
 
 
 def _generate_ui_outputs_from_cache(owner, space_name):
+    # This function remains largely the same, generating UI previews and the export MD
     global parsed_code_blocks_state_cache
     preview_md_val = "*No files in cache to display.*"
     formatted_md_val = f"# Space: {owner}/{space_name}\n## File Structure\n{bbb}\n📁 Root\n{bbb}\n\n*No files in cache.*" if owner or space_name else "*Load or define a Space to see its Markdown structure.*"
@@ -292,7 +299,7 @@ def _generate_ui_outputs_from_cache(owner, space_name):
 
             # Handle content display
             content = block.get('code', '')
-            if block.get('is_binary') or content.startswith("[") : # Treat errors/skipped as binary for preview display
+            if block.get('is_binary') or content.startswith("["): # Treat errors/skipped as binary for preview display
                  preview_md_lines.append(f"\n`{escape_html_for_markdown(content)}`\n")
             elif block.get('is_structure_block'):
                  preview_md_lines.append(f"\n{bbb}{block.get('language', 'plaintext') or 'plaintext'}\n{content}\n{bbb}\n")
@@ -304,15 +311,20 @@ def _generate_ui_outputs_from_cache(owner, space_name):
         space_line_name = f"{owner}/{space_name}" if owner and space_name else (owner or space_name or "your-space")
 
         # _export_selected_logic handles selecting which files to include in the export MD
-        export_result = _export_selected_logic(None, space_line_name, parsed_code_blocks_state_cache) # Passing None means export all non-structure/non-binary/non-error content
+        # Passing None means export all non-structure/non-binary/non-error content + list all files in structure
+        export_result = _export_selected_logic(None, space_line_name, parsed_code_blocks_state_cache)
         formatted_md_val = export_result["output_str"]
         download_file = export_result["download_filepath"]
 
     return formatted_md_val, preview_md_val, gr.update(value=download_file, interactive=download_file is not None)
 
-def handle_groq_chat_submit(user_message, chat_history, api_key_input, model_select, system_prompt, hf_owner_name, hf_repo_name, _current_formatted_markdown):
+
+# --- Refactored Chat Submit Handler ---
+def handle_chat_submit(user_message, chat_history, api_key_input, provider_select, model_select, system_prompt, hf_owner_name, hf_repo_name, _current_formatted_markdown):
     global parsed_code_blocks_state_cache
-    _chat_msg_in, _chat_hist, _status = "", list(chat_history), "Initializing..."
+    _chat_msg_in = ""
+    _chat_hist = list(chat_history)
+    _status = "Initializing..."
     _detected_files_update, _formatted_output_update, _download_btn_update = gr.update(), gr.update(), gr.update(interactive=False, value=None)
 
     # --- Before sending to AI: Parse existing files from the current formatted markdown ---
@@ -322,27 +334,47 @@ def handle_groq_chat_submit(user_message, chat_history, api_key_input, model_sel
     if user_message and _current_formatted_markdown:
         try:
             parsed_from_md = build_logic_parse_markdown(_current_formatted_markdown)
-            # Update cache only with files parsed from the markdown, preserving structure block if any
-            structure_block = next((b for b in parsed_code_blocks_state_cache if b.get("is_structure_block")), None)
-            parsed_code_blocks_state_cache = []
-            if structure_block:
-                parsed_code_blocks_state_cache.append(structure_block)
+            # Update cache with files parsed from the markdown.
+            # Structure block from AI is volatile, always prefer structure from latest AI message.
+            # Files from markdown overwrite any previous file blocks.
+            new_cache_state = []
+            # Add structure block from *current cache* if it exists, it will be replaced if the AI provides a new one
+            existing_structure_block = next((b for b in parsed_code_blocks_state_cache if b.get("is_structure_block")), None)
+            if existing_structure_block:
+                 new_cache_state.append(existing_structure_block.copy()) # Add copy
+
 
             for f_info in parsed_from_md.get("files", []):
-                 # Only add if it has content and isn't the structure block representation
-                 if f_info.get("path") and f_info.get("path") != "File Structure (original)":
+                 # Only add if it has a path and isn't the structure block representation placeholder
+                 if f_info.get("path") and f_info["path"] != "File Structure (original)":
                       # Check if it's a binary representation string
                       is_binary_repr = isinstance(f_info.get("content"), str) and (f_info["content"].startswith("[Binary file") or f_info["content"].startswith("[Error loading content:") or f_info["content"].startswith("[Binary or Skipped file]"))
-                      parsed_code_blocks_state_cache.append({
-                           "filename": f_info["path"],
-                           "code": f_info.get("content", ""),
-                           "language": "binary" if is_binary_repr else _infer_lang_from_filename(f_info["path"]),
-                           "is_binary": is_binary_repr,
-                           "is_structure_block": False
-                      })
-            parsed_code_blocks_state_cache.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
-            # print(f"Parsed {len(parsed_code_blocks_state_cache)} blocks from current Markdown.")
-            # print(parsed_code_blocks_state_cache)
+                      # Check if a block with this filename already exists in new_cache_state and replace it
+                      found_existing = False
+                      for i, block in enumerate(new_cache_state):
+                           if block["filename"] == f_info["path"] and not block.get("is_structure_block"): # Only replace non-structure blocks
+                                new_cache_state[i] = {
+                                     "filename": f_info["path"],
+                                     "code": f_info.get("content", ""),
+                                     "language": "binary" if is_binary_repr else _infer_lang_from_filename(f_info["path"]),
+                                     "is_binary": is_binary_repr,
+                                     "is_structure_block": False
+                                }
+                                found_existing = True
+                                break
+                      if not found_existing:
+                            new_cache_state.append({
+                                 "filename": f_info["path"],
+                                 "code": f_info.get("content", ""),
+                                 "language": "binary" if is_binary_repr else _infer_lang_from_filename(f_info["path"]),
+                                 "is_binary": is_binary_repr,
+                                 "is_structure_block": False
+                            })
+
+            # Sort the updated cache state
+            new_cache_state.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
+            parsed_code_blocks_state_cache = new_cache_state # Update global cache
+
 
         except Exception as e:
             # Log error but don't block chat submission
@@ -354,18 +386,19 @@ def handle_groq_chat_submit(user_message, chat_history, api_key_input, model_sel
 
     if not user_message.strip():
         _status = "Cannot send an empty message."
-        yield (user_message, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update); return
-    _chat_hist.append((user_message, None)); _status = "Sending to AI..."
+        yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update); return
+    _chat_hist.append((user_message, None)); _status = f"Sending to {model_select} via {provider_select}..."
     yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
 
-    effective_api_key = api_key_input or os.environ.get("GROQ_API_KEY")
-    if not effective_api_key:
-        _chat_hist[-1] = (user_message, "Error: Groq API Key not set."); _status = "API Key missing."
-        yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update); return
+    # Pass the API key from the UI directly to model_logic
+    api_key_override = api_key_input
+    # model_id = get_model_id_from_display_name(provider_select, model_select) # model_logic handles display name to ID
+
 
     current_sys_prompt = system_prompt.strip() or DEFAULT_SYSTEM_PROMPT
 
-    # Include current file contents in the prompt as context
+    # Include current file contents in the prompt as context for the AI
+    # This context is built from the *current cache state* (which was just updated from the formatted markdown)
     current_files_context = ""
     if parsed_code_blocks_state_cache:
         current_files_context = "\n\n## Current Files in Space\n"
@@ -374,90 +407,114 @@ def handle_groq_chat_submit(user_message, chat_history, api_key_input, model_sel
                  current_files_context += f"### File: {block['filename']}\n{bbb}\n{block['code']}\n{bbb}\n"
             else:
                 current_files_context += f"### File: {block['filename']}\n"
-                if block.get("is_binary"):
+                if block.get("is_binary") or block.get("code", "").startswith("["): # Include binary/error markers
                     current_files_context += f"{block['code']}\n" # e.g. [Binary file...]
                 else:
                     current_files_context += f"{bbb}{block.get('language', 'plaintext') or 'plaintext'}\n{block.get('code','')}\n{bbb}\n"
         current_files_context += "\n"
 
     # Append current file context to the user message
+    # This combined message structure helps the model understand the context and the expected output format
     user_message_with_context = user_message.strip()
     if current_files_context.strip():
-         user_message_with_context = user_message_with_context + current_files_context + "Based on the current files above and our chat history, please provide updated file contents using the `### File: ...\n```...\n```\n` format for any files you are creating, modifying, or want to include in the final output. Omit files you want to delete from your response."
+         user_message_with_context = user_message_with_context + current_files_context + "\nBased on the current files above and our chat history, please provide updated file contents using the `### File: ...\n```...\n```\n` format for any files you are creating, modifying, or want to include in the final output. If you are providing a file structure, use the `## File Structure\n```\n...\n```\n` format. Omit files you want to delete from your response."
 
 
+    # Convert history to API messages, including the user message with context
     api_msgs = _convert_gr_history_to_api_messages(current_sys_prompt, _chat_hist[:-1], user_message_with_context)
 
-    headers = {"Authorization": f"Bearer {effective_api_key}", "Content-Type": "application/json"}
-    payload = {"model": model_select, "messages": api_msgs}
-
+    # --- Call the new model_logic streaming function ---
     try:
-        _status = f"Waiting for {model_select}..."; yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
-        response = requests.post(GROQ_API_ENDPOINT, headers=headers, json=payload, timeout=180)
-        response.raise_for_status()
-        api_resp_json = response.json()
-
-        bot_response_actual = None
-        # --- Corrected logic to safely access the message content ---
-        if api_resp_json and isinstance(api_resp_json, dict) and api_resp_json.get("choices") and isinstance(api_resp_json["choices"], list) and len(api_resp_json["choices"]) > 0:
-            first_choice = api_resp_json["choices"][0] # Access the first item in the 'choices' list
-            if isinstance(first_choice, dict) and first_choice.get("message"): # Ensure it's a dict and has 'message' key
-                message_obj = first_choice["message"] # Get the 'message' dictionary from the choice
-                if isinstance(message_obj, dict): # Ensure message_obj is a dictionary
-                    bot_response_actual = message_obj.get("content") # Get content from the message dictionary
-        # --- End corrected logic ---
-
-        if bot_response_actual:
-            # If bot_response_actual was successfully extracted
-            _chat_hist[-1] = (user_message, bot_response_actual); _status = "AI response received. Processing files..."
-            yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
-
-            # Pass the *existing* state (which was updated from markdown at the start of the function)
-            # to the parsing logic, so it knows about files already present.
-            latest_bot_message_json = json.dumps([{"role": BOT_ROLE_NAME, "content": bot_response_actual}], indent=2)
-            parsing_res = _parse_chat_stream_logic(latest_bot_message_json, existing_files_state=parsed_code_blocks_state_cache)
+        _status = f"Waiting for {model_select} via {provider_select}..."; yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
+
+        # Accumulate the full response content for parsing *after* streaming
+        full_bot_response_content = ""
+        error_during_stream = None
+
+        # Generate stream from model_logic
+        for chunk in generate_stream(provider_select, model_select, api_key_override, api_msgs):
+             if chunk is None: continue # Skip None chunks if any
+             if chunk.startswith("Error: ") or chunk.startswith("API HTTP Error"):
+                  # If an error chunk is received, treat it as the final output and stop
+                  full_bot_response_content = chunk
+                  error_during_stream = chunk
+                  break # Stop processing stream
+             else:
+                  # Accumulate response and update the last message in chat_hist
+                  full_bot_response_content += chunk
+                  _chat_hist[-1] = (user_message, full_bot_response_content)
+                  _status = f"Streaming from {model_select}..."
+                  # Yield update immediately after receiving chunk
+                  yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
+
+        # After the stream finishes or breaks
+        if error_during_stream:
+             _status = error_during_stream # Set status to the error message
+        elif full_bot_response_content:
+            _status = f"Streaming complete. Processing files from {model_select} response..."
+
+            # Pass the *current state* (updated from markdown at the start)
+            # and the *latest bot message content* to the parsing logic.
+            # _parse_chat_stream_logic will merge and update based on the bot's response.
+            parsing_res = _parse_chat_stream_logic(full_bot_response_content, existing_files_state=parsed_code_blocks_state_cache)
 
             if parsing_res["error_message"]:
                 _status = f"Parsing Error: {parsing_res['error_message']}"
-                # Append parsing error to the bot's response in chat? Or show separately?
-                # For now, update status and detected files area with error
+                # Append parsing error to the bot's response in chat for visibility? Or just status?
+                # For now, update status and detected files area with error message
                 _detected_files_update = gr.Markdown(f"## Parsing Error\n`{escape_html_for_markdown(parsing_res['error_message'])}`")
             else:
-                # parsed_code_blocks_state_cache is updated inside _parse_chat_stream_logic
+                # Update the global cache with the new state returned by the parser
+                global parsed_code_blocks_state_cache
+                parsed_code_blocks_state_cache = parsing_res["parsed_code_blocks"]
+
+                # Regenerate UI outputs from the *updated* cache
                 _formatted_output_update, _detected_files_update, _download_btn_update = _generate_ui_outputs_from_cache(hf_owner_name, hf_repo_name)
                 _status = "Processing complete. Previews updated."
-
-            yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update); return
         else:
-            # Handle cases where bot_response_actual is None (e.g., malformed API response)
-            error_msg = f"API Error: Unexpected response structure or empty message content. Response: {json.dumps(api_resp_json)}"
-            # This error_msg will be handled by the final error block below
-
-    except requests.exceptions.HTTPError as e: error_msg = f"API HTTP Error: {e} - {e.response.text if e.response else 'No details'}"
-    except requests.exceptions.RequestException as e: error_msg = f"API Request Error: {e}"
-    except Exception as e: error_msg = f"Unexpected error in chat submit: {e}"
-
-    # If an error occurred (and bot_response_actual was not set or an exception happened),
-    # update chat history with the error message.
-    # Assumes the last entry is the one for the current turn, added as (user_message, None).
-    # It's safer to append a new message if the last one isn't the one we expected to update.
-    if _chat_hist and len(_chat_hist) > 0 and _chat_hist[-1][1] is None:
-         _chat_hist[-1] = (_chat_hist[-1][0], error_msg) # Keep user message, add error as bot message
-    else:
-         _chat_hist.append((user_message, error_msg)) # Append as a new user/bot turn if structure unexpected
+            # Handle cases where the stream finished but yielded no content (e.g., filter)
+            _status = "AI response complete, but returned no content."
+            # Keep existing previews/markdown if no content was generated to parse
+
+    except Exception as e:
+        # Catch any errors that occurred *before* or *during* the stream setup/iteration
+        error_msg = f"An unexpected error occurred during AI generation: {e}"
+        print(f"Unexpected error in chat submit stream: {e}")
+        # Update the last chat message with the error
+        if _chat_hist and len(_chat_hist) > 0 and _chat_hist[-1][1] is None:
+             _chat_hist[-1] = (_chat_hist[-1][0], error_msg) # Keep user message, add error as bot message
+        else:
+             _chat_hist.append((user_message, error_msg)) # Append as a new user/bot turn if structure unexpected
+        _status = error_msg
+        # Previews and markdown might not be affected by a generation error, keep existing state
 
-    _status = error_msg
+    # Final yield to update UI after all processing
     yield (_chat_msg_in, _chat_hist, _status, _detected_files_update, _formatted_output_update, _download_btn_update)
 
 
+# --- Handler to update model dropdown based on provider selection ---
+def update_models_dropdown(provider_select):
+    """Updates the model dropdown choices and selects the default model."""
+    if not provider_select:
+        return gr.update(choices=[], value=None)
+    models = get_models_for_provider(provider_select)
+    default_model = get_default_model_for_provider(provider_select)
+    return gr.update(choices=models, value=default_model if default_model in models else (models[0] if models else None))
+
+
+# --- Existing handlers for Load, Build, Edit, Delete, Status (Mostly unchanged) ---
+# Ensure they correctly use hf_api_key_ui and potentially owner_name_input/space_name_input
+
 def handle_load_existing_space(hf_api_key_ui, ui_owner_name, ui_space_name):
+    # ... (rest of this function is largely the same, but needs to update global cache)
     global parsed_code_blocks_state_cache
     _formatted_md_val, _detected_preview_val, _status_val = "*Loading files...*", "*Loading files...*", f"Loading Space: {ui_owner_name}/{ui_space_name}..."
     _file_browser_update, _iframe_html_update, _download_btn_update = gr.update(visible=False, choices=[], value=None), gr.update(value=None, visible=False), gr.update(interactive=False, value=None)
     _build_status_clear, _edit_status_clear, _runtime_status_clear = "*Build status will appear here.*", "*Select a file to load or delete.*", "*Space runtime status will appear here after refresh.*"
+    _chat_history_clear = [] # Clear chat history on loading a new space
 
     # Yield initial state to update UI
-    yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, gr.update(value=ui_owner_name), gr.update(value=ui_space_name), _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+    yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, gr.update(value=ui_owner_name), gr.update(value=ui_space_name), _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear, _chat_history_clear)
 
     owner_to_use, updated_owner_name_val = ui_owner_name, ui_owner_name
     error_occurred = False
@@ -481,7 +538,8 @@ def handle_load_existing_space(hf_api_key_ui, ui_owner_name, ui_space_name):
 
     if error_occurred:
          # Yield error state
-         yield (f"*Error: {_status_val}*", f"*Error: {_status_val}*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+         yield (f"*Error: {_status_val}*", f"*Error: {_status_val}*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear, _chat_history_clear)
+         parsed_code_blocks_state_cache = [] # Clear cache on error
          return # Stop execution
 
     sdk_for_iframe, file_list, err_list_files = get_space_repository_info(hf_api_key_ui, ui_space_name, owner_to_use)
@@ -498,7 +556,7 @@ def handle_load_existing_space(hf_api_key_ui, ui_owner_name, ui_space_name):
         parsed_code_blocks_state_cache = [] # Clear cache on error
         _formatted_md_val, _detected_preview_val, _download_btn_update = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name)
         _file_browser_update = gr.update(visible=True, choices=[], value="Error loading files") # Update file browser with error state
-        yield (f"*Error: {err_list_files}*", "*Error loading files*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+        yield (f"*Error: {err_list_files}*", "*Error loading files*", _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear, _chat_history_clear)
         return # Stop execution
 
     if not file_list:
@@ -506,14 +564,14 @@ def handle_load_existing_space(hf_api_key_ui, ui_owner_name, ui_space_name):
         parsed_code_blocks_state_cache = []
         _formatted_md_val, _detected_preview_val, _download_btn_update = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name)
         _file_browser_update = gr.update(visible=True, choices=[], value="No files found")
-        yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+        yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear, _chat_history_clear)
         return # Stop execution
 
 
-    loaded_files_for_parse = []
+    loaded_files_for_cache = [] # Build a list to become the new cache state
     _status_val = f"Loading {len(file_list)} files from {owner_to_use}/{ui_space_name} (SDK: {sdk_for_iframe or 'unknown'})...";
     # Yield intermediate status while loading files
-    yield (_formatted_md_val, _detected_preview_val, _status_val, gr.update(visible=True, choices=sorted(file_list or []), value=None), updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+    yield (_formatted_md_val, _detected_preview_val, _status_val, gr.update(visible=True, choices=sorted(file_list or []), value=None), updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear, _chat_history_clear)
 
 
     for file_path in file_list:
@@ -527,35 +585,43 @@ def handle_load_existing_space(hf_api_key_ui, ui_owner_name, ui_space_name):
            file_path.startswith("__pycache__/") or "/__pycache__/" in file_path or \
            file_path.startswith("node_modules/") or "/node_modules/" in file_path or \
            file_path.startswith("venv/") or "/venv/" in file_path or \
-           file_path.startswith(".venv/") or "/.venv/" in file_path:
-            loaded_files_for_parse.append({"filename": file_path, "code": "[Binary or Skipped file]", "language": "binary", "is_binary": True, "is_structure_block": False}); continue
+           file_path.startswith(".venv/") or "/.venv/" in file_path or \
+           file_path == "README.md" or file_path == "LICENSE": # Optionally skip common non-code files like README/LICENSE
+            loaded_files_for_cache.append({"filename": file_path, "code": "[Binary or Skipped file]", "language": "binary", "is_binary": True, "is_structure_block": False}); continue
 
         # Handle potential issues with reading large files or non-utf8 files
         try:
             content, err_get = get_space_file_content(hf_api_key_ui, ui_space_name, owner_to_use, file_path)
             if err_get:
                  # If there's an error getting content, record it but don't fail the whole load
-                 loaded_files_for_parse.append({"filename": file_path, "code": f"[Error loading content: {err_get}]", "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
+                 loaded_files_for_cache.append({"filename": file_path, "code": f"[Error loading content: {err_get}]", "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
                  print(f"Error loading {file_path}: {err_get}");
                  continue
             # If content is successfully loaded
-            loaded_files_for_parse.append({"filename": file_path, "code": content, "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
+            loaded_files_for_cache.append({"filename": file_path, "code": content, "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
         except Exception as content_ex:
             # Catch any other unexpected exceptions during file content fetching
-            loaded_files_for_parse.append({"filename": file_path, "code": f"[Unexpected error loading content: {content_ex}]", "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
+            loaded_files_for_cache.append({"filename": file_path, "code": f"[Unexpected error loading content: {content_ex}]", "language": _infer_lang_from_filename(file_path), "is_binary": False, "is_structure_block": False})
             print(f"Unexpected error loading {file_path}: {content_ex}")
             continue
 
+    # Add a placeholder structure block if none was loaded (AI will generate one later if needed)
+    # This ensures the cache isn't empty except for files
+    # structure_block = next((b for b in loaded_files_for_cache if b.get("is_structure_block")), None)
+    # if not structure_block:
+    #      loaded_files_for_cache.insert(0, {"filename": "File Structure (original)", "code": "📁 Root\n  ...\n", "language": "plaintext", "is_binary": False, "is_structure_block": True})
 
-    parsed_code_blocks_state_cache = loaded_files_for_parse
+    parsed_code_blocks_state_cache = loaded_files_for_cache
     _formatted_md_val, _detected_preview_val, _download_btn_update = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name)
-    _status_val = f"Successfully loaded Space: {owner_to_use}/{ui_space_name}. Markdown ready."
+    _status_val = f"Successfully loaded Space: {owner_to_use}/{ui_space_name}. Markdown ready. {len(file_list)} files listed."
     _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value=None) # Use the full file list for the dropdown
 
-    yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear)
+    # Final yield with updated state
+    yield (_formatted_md_val, _detected_preview_val, _status_val, _file_browser_update, updated_owner_name_val, ui_space_name, _iframe_html_update, _download_btn_update, _build_status_clear, _edit_status_clear, _runtime_status_clear, _chat_history_clear)
 
 
 def handle_build_space_button(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, space_sdk_ui, formatted_markdown_content):
+    # ... (this function calls build_logic_create_space and refreshes file list)
     _build_status, _iframe_html, _file_browser_update = "Starting space build process...", gr.update(value=None, visible=False), gr.update(visible=False, choices=[], value=None)
     yield _build_status, _iframe_html, _file_browser_update # Yield initial status
     if not ui_space_name_part or "/" in ui_space_name_part: _build_status = f"Build Error: HF Space Name '{ui_space_name_part}' must be repo name only (no '/')."; yield _build_status, _iframe_html, _file_browser_update; return
@@ -573,25 +639,73 @@ def handle_build_space_button(hf_api_key_ui, ui_space_name_part, ui_owner_name_p
 
     if not final_owner_for_build: _build_status = "Build Error: HF Owner Name could not be determined. Please specify it."; yield _build_status, _iframe_html, _file_browser_update; return
 
+    # Before building, parse the markdown to ensure the cache reflects exactly what's being built
+    # This prevents inconsistencies if the user manually edited the markdown output
+    try:
+        parsed_from_md_for_build = build_logic_parse_markdown(formatted_markdown_content)
+        # Replace the global cache state with the state derived from the markdown being built
+        global parsed_code_blocks_state_cache
+        parsed_code_blocks_state_cache = []
+        if parsed_from_md_for_build.get("owner_md"): # Update UI owner/space name if present in MD
+             ui_owner_name_part = parsed_from_md_for_build["owner_md"]
+        if parsed_from_md_for_build.get("repo_name_md"):
+             ui_space_name_part = parsed_from_md_for_build["repo_name_md"]
+
+        # Rebuild cache from parsed markdown files + structure block
+        structure_block_md = next((f for f in parsed_from_md_for_build.get("files", []) if f.get("path") == "File Structure (original)"), None)
+        if structure_block_md:
+             parsed_code_blocks_state_cache.append({
+                 "filename": structure_block_md["path"],
+                 "code": structure_block_md["content"],
+                 "language": "plaintext", # Markdown parser doesn't detect lang for structure block ```
+                 "is_binary": False,
+                 "is_structure_block": True
+             })
+
+        for f_info in parsed_from_md_for_build.get("files", []):
+             if f_info.get("path") and f_info["path"] != "File Structure (original)":
+                  is_binary_repr = isinstance(f_info.get("content"), str) and (f_info["content"].startswith("[Binary file") or f_info["content"].startswith("[Error loading content:") or f_info["content"].startswith("[Binary or Skipped file]"))
+                  parsed_code_blocks_state_cache.append({
+                       "filename": f_info["path"],
+                       "code": f_info.get("content", ""),
+                       "language": "binary" if is_binary_repr else _infer_lang_from_filename(f_info["path"]),
+                       "is_binary": is_binary_repr,
+                       "is_structure_block": False
+                  })
+        parsed_code_blocks_state_cache.sort(key=lambda b: (0, b["filename"]) if b.get("is_structure_block") else (1, b["filename"]))
+
+    except Exception as e:
+         _build_status = f"Build Error: Failed to parse Markdown structure before building: {e}";
+         yield _build_status, _iframe_html, _file_browser_update; return # Stop build on parse error
+
+
     result_message = build_logic_create_space(hf_api_key_ui, ui_space_name_part, final_owner_for_build, space_sdk_ui, formatted_markdown_content)
     _build_status = f"Build Process: {result_message}"
 
     if "Successfully" in result_message:
-        sub_owner = re.sub(r'[^a-z0-9\-]+', '-', final_owner_for_build.lower()).strip('-') or 'owner'
+        # Use potentially updated owner/space name from markdown parsing
+        sub_owner = re.sub(r'[^a-z0-9\-]+', '-', ui_owner_name_part.lower()).strip('-') or 'owner'
         sub_repo = re.sub(r'[^a-z0-9\-]+', '-', ui_space_name_part.lower()).strip('-') or 'space'
         iframe_url = f"https://{sub_owner}-{sub_repo}{'.static.hf.space' if space_sdk_ui == 'static' else '.hf.space'}"
         _iframe_html = gr.update(value=f'<iframe src="{iframe_url}?__theme=light&embed=true" width="100%" height="700px" style="border:1px solid #eee; border-radius:8px;"></iframe>', visible=True)
-        _build_status += f"\nSpace live at: [Link]({iframe_url}) (Repo: https://huggingface.co/spaces/{final_owner_for_build}/{ui_space_name_part})"
+        _build_status += f"\nSpace live at: [Link]({iframe_url}) (Repo: https://huggingface.co/spaces/{ui_owner_name_part}/{ui_space_name_part})"
 
         # Refresh file list after successful build
-        file_list, err_list = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, final_owner_for_build)
-        if err_list: _build_status += f"\nFile list refresh error after build: {err_list}"; _file_browser_update = gr.update(visible=True, choices=[], value="Error refreshing files")
+        file_list, err_list = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, ui_owner_name_part)
+        if err_list: _build_status += f"\nFile list refresh error after build: {err_list}"; _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value="Error refreshing files")
         else: _file_browser_update = gr.update(visible=True, choices=sorted(file_list or []), value=None if file_list else "No files found")
 
-    yield _build_status, _iframe_html, _file_browser_update
+    # Update UI with owner/space names extracted from markdown if present
+    owner_name_output = gr.update(value=ui_owner_name_part)
+    space_name_output = gr.update(value=ui_space_name_part)
+
 
+    yield _build_status, _iframe_html, _file_browser_update, owner_name_output, space_name_output
 
+
+# File editing handlers are okay, just need to ensure they update the cache properly after commit/delete
 def handle_load_file_for_editing(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, selected_file_path):
+    # ... (rest of this function is the same)
     _file_content_val, _edit_status_val, _commit_msg_val, _lang_update = "", "Error: No file selected.", gr.update(value=""), gr.update(language="python") # Reset values
     if not selected_file_path or selected_file_path in ["No files found", "Error loading files", "Error refreshing files"]:
          yield _file_content_val, "Select a file from the dropdown.", _commit_msg_val, _lang_update # Clear editor and status
@@ -619,7 +733,7 @@ def handle_load_file_for_editing(hf_api_key_ui, ui_space_name_part, ui_owner_nam
         _edit_status_val = f"Error loading '{selected_file_path}': {err}"
         _commit_msg_val = f"Error loading {selected_file_path}"
         _file_content_val = f"Error loading {selected_file_path}:\n{err}"
-        _lang_update = gr.update(language="python")
+        _lang_update = gr.update(language="python") # Default language for error display
         yield _file_content_val, _edit_status_val, _commit_msg_val, _lang_update
         return
 
@@ -633,11 +747,15 @@ def handle_load_file_for_editing(hf_api_key_ui, ui_space_name_part, ui_owner_nam
 def handle_commit_file_changes(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, file_to_edit_path, edited_content, commit_message):
     global parsed_code_blocks_state_cache
     _edit_status_val = "Processing commit..."
-    _file_browser_update_val = gr.update()
-    _formatted_md_out, _detected_preview_out, _download_btn_out = gr.update(), gr.update(), gr.update()
+    # Initialize updates for components that might change
+    _file_browser_update_val = gr.update() # Will update choices or value
+    _formatted_md_out = gr.update() # Will update markdown
+    _detected_preview_out = gr.update() # Will update markdown preview
+    _download_btn_out = gr.update() # Will update download button
 
     yield _edit_status_val, _file_browser_update_val, _formatted_md_out, _detected_preview_out, _download_btn_out # Yield initial status
 
+
     if not file_to_edit_path or file_to_edit_path in ["No files found", "Error loading files", "Error refreshing files"]:
         _edit_status_val = "Error: No valid file selected for commit.";
         yield _edit_status_val, gr.update(), gr.update(), gr.update(), gr.update(); return
@@ -670,7 +788,10 @@ def handle_commit_file_changes(hf_api_key_ui, ui_space_name_part, ui_owner_name_
                 found_in_cache = True
                 break
         if not found_in_cache:
-             # If file was added via editor or wasn't in initial load cache (e.g. binary/error), add it
+             # If file was added/edited via editor and wasn't in initial load cache (e.g. binary/error placeholder), add/replace it
+             # First remove any existing placeholder for this file
+             parsed_code_blocks_state_cache = [b for b in parsed_code_blocks_state_cache if b["filename"] != file_to_edit_path]
+             # Then add the new text content block
              parsed_code_blocks_state_cache.append({
                  "filename": file_to_edit_path,
                  "code": edited_content,
@@ -698,12 +819,16 @@ def handle_commit_file_changes(hf_api_key_ui, ui_space_name_part, ui_owner_name_
 def handle_delete_file(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, file_to_delete_path):
     global parsed_code_blocks_state_cache
     _edit_status_val = "Processing deletion..."
-    _file_browser_choices_update = gr.update()
-    _file_browser_value_update = None # Clear selected file
+    # Initialize updates for components that might change/clear
+    _file_browser_choices_update = gr.update() # Update choices
+    _file_browser_value_update = None # Clear selected file value
     _file_content_editor_update = gr.update(value="") # Clear editor content
     _commit_msg_update = gr.update(value="") # Clear commit message
     _lang_update = gr.update(language="plaintext") # Reset editor language
-    _formatted_md_out, _detected_preview_out, _download_btn_out = gr.update(), gr.update(), gr.update()
+    _formatted_md_out = gr.update() # Update markdown
+    _detected_preview_out = gr.update() # Update markdown preview
+    _download_btn_out = gr.update() # Update download button
+
 
     yield (_edit_status_val, _file_browser_choices_update, _file_browser_value_update, _file_content_editor_update, _commit_msg_update, _lang_update, _formatted_md_out, _detected_preview_out, _download_btn_out) # Yield initial status
 
@@ -729,6 +854,9 @@ def handle_delete_file(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, fi
     deletion_status_msg = build_logic_delete_space_file(hf_api_key_ui, ui_space_name_part, owner_to_use, file_to_delete_path)
     _edit_status_val = deletion_status_msg
 
+    # Always refresh the file list dropdown choices after a delete attempt, successful or not
+    new_file_list, err_list = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, owner_to_use)
+
     if "Successfully deleted" in deletion_status_msg:
         # Remove the file from the cache
         parsed_code_blocks_state_cache = [b for b in parsed_code_blocks_state_cache if b["filename"] != file_to_delete_path]
@@ -736,26 +864,35 @@ def handle_delete_file(hf_api_key_ui, ui_space_name_part, ui_owner_name_part, fi
         # Regenerate markdown and preview from the updated cache
         _formatted_md_out, _detected_preview_out, _download_btn_out = _generate_ui_outputs_from_cache(owner_to_use, ui_space_name_part)
 
-        # Refresh file list choices and clear selected file
-        new_file_list, err_list = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, owner_to_use)
+
         if err_list:
              _edit_status_val += f"\nFile list refresh error: {err_list}"
-             _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value="Error refreshing files")
+             _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value="Error refreshing files") # Set value to error state
         else:
-             _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=None) # Clear selection
+             _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=None) # Clear selection visually and internally
 
         _file_browser_value_update = None # Explicitly set value to None to clear selection visual
 
-    else:
-        # If deletion failed, refresh the file list but keep the file selected
-        new_file_list, _ = list_space_files_for_browsing(hf_api_key_ui, ui_space_name_part, owner_to_use)
-        _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=file_to_delete_path)
-        _file_browser_value_update = file_to_delete_path # Keep the file selected
+
+    else: # If deletion failed
+        if err_list:
+            _edit_status_val += f"\nFile list refresh error: {err_list}"
+            _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value="Error refreshing files")
+            _file_browser_value_update = "Error refreshing files" # Keep error state in value if list failed
+        else:
+             # If list refresh succeeded but delete failed, refresh choices and keep the *failed-to-delete* file selected
+             _file_browser_choices_update = gr.update(choices=sorted(new_file_list or []), value=file_to_delete_path)
+             _file_browser_value_update = file_to_delete_path # Keep the file selected visually
+
+        # Markdown and preview are not changed if deletion failed, keep current updates as gr.update()
+
 
     yield (_edit_status_val, _file_browser_choices_update, _file_browser_value_update, _file_content_editor_update, _commit_msg_update, _lang_update, _formatted_md_out, _detected_preview_out, _download_btn_out)
 
 
+# Space status handler is okay
 def handle_refresh_space_status(hf_api_key_ui, ui_owner_name, ui_space_name):
+    # ... (rest of this function is the same)
     yield "*Fetching space status...*" # Initial feedback
     owner_to_use = ui_owner_name
     if not owner_to_use:
@@ -832,7 +969,7 @@ body {
 .gr-markdown {
     background-color: rgba(44, 62, 80, 0.7) !important; /* Transparent dark background */
     padding: 10px; /* Add some padding */
-    border-radius: 5px; /* Rounded corners */
+    border-radius: 5 al;
 }
 /* Style markdown headers for better contrast */
 .gr-markdown h1, .gr-markdown h2, .gr-markdown h3, .gr-markdown h4, .gr-markdown h5, .gr-markdown h6 {
@@ -844,7 +981,6 @@ body {
     background-color: rgba(52, 73, 94, 0.95) !important; /* Darker code background */
     border-color: rgba(189, 195, 199, 0.3) !important;
 }
-
 /* Chatbot specific styling */
 .gr-chatbot {
     background-color: rgba(44, 62, 80, 0.7) !important;
@@ -859,9 +995,17 @@ body {
     background-color: rgba(46, 204, 113, 0.9) !important; /* Greenish background for user messages */
     color: black !important; /* Dark text for green background */
 }
+"""
 
 
-"""
+# Get initial providers and models for UI setup
+available_providers = get_available_providers()
+default_provider = available_providers[0] if available_providers else None
+initial_models = get_models_for_provider(default_provider) if default_provider else []
+initial_default_model = get_default_model_for_provider(default_provider) if default_provider else None
+# Ensure initial_default_model is in the initial_models list, fallback if not
+if initial_default_model not in initial_models and initial_models:
+    initial_default_model = initial_models[0]
 
 
 with gr.Blocks(theme=custom_theme, css=custom_css) as demo:
@@ -870,58 +1014,88 @@ with gr.Blocks(theme=custom_theme, css=custom_css) as demo:
     with gr.Row():
         with gr.Sidebar():
             gr.Markdown("## ⚙️ Configuration")
-            with gr.Group(): gr.Markdown("### API Keys & Tokens"); groq_api_key_input = gr.Textbox(label="Groq API Key", type="password", placeholder="gsk_..."); hf_api_key_input = gr.Textbox(label="Hugging Face Token", type="password", placeholder="hf_...")
+            with gr.Group(): gr.Markdown("### API Keys & Tokens");
+            # Single API key input, model_logic decides which env var to check or uses this override
+            api_key_input = gr.Textbox(label="AI Provider API Key (Optional Override)", type="password", placeholder="Paste key here or set env var (e.g., GROQ_API_KEY)");
+            hf_api_key_input = gr.Textbox(label="Hugging Face Token (for building/loading)", type="password", placeholder="hf_...")
             with gr.Group(): gr.Markdown("### Hugging Face Space"); owner_name_input = gr.Textbox(label="HF Owner Name", placeholder="e.g., your-username"); space_name_input = gr.Textbox(label="HF Space Name", value="my-ai-space", placeholder="e.g., my-cool-app"); space_sdk_select = gr.Dropdown(label="Space SDK", choices=["gradio", "streamlit", "docker", "static"], value="gradio", info="Used for new/build."); load_space_button = gr.Button("🔄 Load Existing Space", variant="secondary", size="sm")
-            with gr.Group(): gr.Markdown("### AI Model Settings"); groq_model_select = gr.Dropdown(label="Groq Model", choices=["mixtral-8x7b-32768", "llama3-8b-8192", "llama3-70b-8192", "gemma-7b-it"], value="llama3-8b-8192"); groq_system_prompt_input = gr.Textbox(label="System Prompt", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True)
+            with gr.Group(): gr.Markdown("### AI Model Settings");
+            provider_select = gr.Dropdown(label="AI Provider", choices=available_providers, value=default_provider, info="Select an AI model provider.");
+            model_select = gr.Dropdown(label="AI Model", choices=initial_models, value=initial_default_model, info="Select a model.");
+            system_prompt_input = gr.Textbox(label="System Prompt", lines=8, value=DEFAULT_SYSTEM_PROMPT, interactive=True)
         with gr.Column(scale=3):
             gr.Markdown("## 💬 AI Chat & Code Generation")
-            groq_chatbot_display = gr.Chatbot(label="Groq AI Chat", height=400, bubble_full_width=False, avatar_images=(None, "https://raw.githubusercontent.com/groq/groq-api-cookbook/main/groq.png"))
-            with gr.Row(): groq_chat_message_input = gr.Textbox(show_label=False, placeholder="Your Message...", scale=7); groq_send_chat_button = gr.Button("Send", variant="primary", scale=1, size="lg")
-            groq_status_output = gr.Textbox(label="Chat/Process Status", interactive=False, lines=1, value="Ready.")
+            # Updated chatbot avatar
+            chatbot_display = gr.Chatbot(label="AI Chat", height=400, bubble_full_width=False, avatar_images=(None, "https://huggingface.co/datasets/huggingface/badges/resolve/main/huggingface-bot-avatar.svg"))
+            with gr.Row(): chat_message_input = gr.Textbox(show_label=False, placeholder="Your Message...", scale=7); send_chat_button = gr.Button("Send", variant="primary", scale=1, size="lg")
+            status_output = gr.Textbox(label="Chat/Process Status", interactive=False, lines=1, value="Ready.")
             gr.Markdown("---")
             with gr.Tabs():
                 with gr.TabItem("📝 Formatted Space Markdown"): gr.Markdown("Complete Markdown definition for your Space."); formatted_space_output_display = gr.Textbox(label="Current Space Definition", lines=15, interactive=True, show_copy_button=True, value="*Space definition...*"); download_button = gr.DownloadButton(label="Download .md", interactive=False, size="sm")
                 with gr.TabItem("🔍 Detected Files Preview"):
-                    # --- CORRECTED LINE ---
                     detected_files_preview = gr.Markdown(value="*Files preview...*")
-                    # --- END CORRECTED LINE ---
+
             gr.Markdown("---")
             with gr.Tabs():
                 with gr.TabItem("🚀 Build & Preview Space"):
                     with gr.Row(): build_space_button = gr.Button("Build / Update Space on HF", variant="primary", scale=2); refresh_status_button = gr.Button("🔄 Refresh Space Status", scale=1)
+                    # Build status outputs also include updating owner/space names in the textboxes
                     build_status_display = gr.Textbox(label="Build Operation Status", interactive=False, lines=2, value="*Build status will appear here.*"); gr.Markdown("---"); space_runtime_status_display = gr.Markdown("*Space runtime status will appear here after refresh.*"); gr.Markdown("---"); space_iframe_display = gr.HTML(value="<!-- Space Iframe -->", visible=False)
                 with gr.TabItem("✏️ Edit Space Files"):
                     gr.Markdown("Select a file to view, edit, or delete. Changes are committed to HF Hub.")
                     file_browser_dropdown = gr.Dropdown(label="Select File in Space", choices=[], interactive=True, visible=False, info="Load/build Space first.")
-                    file_content_editor = gr.Code(label="File Content Editor", language="python", lines=15, interactive=True) # Changed language to python from plaintext
+                    file_content_editor = gr.Code(label="File Content Editor", language="python", lines=15, interactive=True)
                     commit_message_input = gr.Textbox(label="Commit Message", placeholder="e.g., Updated app.py", value="Update via AI Space Editor")
                     with gr.Row(): update_file_button = gr.Button("Commit Changes", variant="primary", scale=2); delete_file_button = gr.Button("🗑️ Delete Selected File", variant="stop", scale=1)
                     edit_status_display = gr.Textbox(label="File Edit/Delete Status", interactive=False, lines=2, value="*Select file...*")
 
-    chat_outputs = [groq_chat_message_input, groq_chatbot_display, groq_status_output, detected_files_preview, formatted_space_output_display, download_button]
-    chat_inputs = [groq_chat_message_input, groq_chatbot_display, groq_api_key_input, groq_model_select, groq_system_prompt_input, owner_name_input, space_name_input, formatted_space_output_display] # Added formatted_space_output_display to inputs
-    groq_send_chat_button.click(fn=handle_groq_chat_submit, inputs=chat_inputs, outputs=chat_outputs)
-    groq_chat_message_input.submit(fn=handle_groq_chat_submit, inputs=chat_inputs, outputs=chat_outputs)
-
-    # Added outputs for clearing/updating edit status when loading a space
-    load_space_outputs = [formatted_space_output_display, detected_files_preview, groq_status_output, file_browser_dropdown, owner_name_input, space_name_input, space_iframe_display, download_button, build_status_display, edit_status_display, space_runtime_status_display]
+    # --- Event Handlers ---
+
+    # Provider dropdown change event to update model dropdown
+    provider_select.change(
+        fn=update_models_dropdown,
+        inputs=provider_select,
+        outputs=model_select
+    )
+
+    # Chat submit handler outputs
+    chat_outputs = [chat_message_input, chatbot_display, status_output, detected_files_preview, formatted_space_output_display, download_button]
+    # Chat submit handler inputs
+    chat_inputs = [chat_message_input, chatbot_display, api_key_input, provider_select, model_select, system_prompt_input, owner_name_input, space_name_input, formatted_space_output_display] # Pass current formatted markdown as context
+
+    # Wire chat buttons
+    send_chat_button.click(
+        fn=handle_chat_submit,
+        inputs=chat_inputs,
+        outputs=chat_outputs
+    )
+    chat_message_input.submit( # Allow submitting with Enter key
+        fn=handle_chat_submit,
+        inputs=chat_inputs,
+        outputs=chat_outputs
+    )
+
+    # Load space outputs include clearing chat history
+    load_space_outputs = [formatted_space_output_display, detected_files_preview, status_output, file_browser_dropdown, owner_name_input, space_name_input, space_iframe_display, download_button, build_status_display, edit_status_display, space_runtime_status_display, chatbot_display] # Added chatbot_display
     load_space_button.click(fn=handle_load_existing_space, inputs=[hf_api_key_input, owner_name_input, space_name_input], outputs=load_space_outputs)
 
-    build_outputs = [build_status_display, space_iframe_display, file_browser_dropdown]
+    # Build outputs now include updating owner/space name textboxes
+    build_outputs = [build_status_display, space_iframe_display, file_browser_dropdown, owner_name_input, space_name_input] # Added owner_name_input, space_name_input
     build_space_button.click(fn=handle_build_space_button, inputs=[hf_api_key_input, space_name_input, owner_name_input, space_sdk_select, formatted_space_output_display], outputs=build_outputs)
 
-    # Added commit message and language update to file edit load outputs
+    # File edit load outputs include clearing/setting commit message and language
     file_edit_load_outputs = [file_content_editor, edit_status_display, commit_message_input, file_content_editor]
     file_browser_dropdown.change(fn=handle_load_file_for_editing, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown], outputs=file_edit_load_outputs)
 
-    # Added formatted_space_output_display, detected_files_preview, download_button to commit outputs to refresh previews
+    # Commit file outputs include refreshing previews and file browser state
     commit_file_outputs = [edit_status_display, file_browser_dropdown, formatted_space_output_display, detected_files_preview, download_button]
     update_file_button.click(fn=handle_commit_file_changes, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown, file_content_editor, commit_message_input], outputs=commit_file_outputs)
 
-    # Added file_content_editor, commit_message_input, language update, formatted_space_output_display, detected_files_preview, download_button to delete outputs
-    delete_file_outputs = [edit_status_display, file_browser_dropdown, file_browser_dropdown, file_content_editor, commit_message_input, file_content_editor, formatted_space_output_display, detected_files_preview, download_button]
+    # Delete file outputs include refreshing previews, file browser state, and clearing editor
+    delete_file_outputs = [edit_status_display, file_browser_dropdown, file_browser_dropdown, file_content_editor, commit_message_input, file_content_editor, formatted_space_output_display, detected_files_preview, download_button] # Two file_browser_dropdown outputs: choices and value
     delete_file_button.click(fn=handle_delete_file, inputs=[hf_api_key_input, space_name_input, owner_name_input, file_browser_dropdown], outputs=delete_file_outputs)
 
+    # Refresh status handler is okay
     refresh_status_button.click(fn=handle_refresh_space_status, inputs=[hf_api_key_input, owner_name_input, space_name_input], outputs=[space_runtime_status_display])
 
 if __name__ == "__main__":