Update app.py

app.py

@@ -5,90 +5,80 @@ import spaces
 
 @spaces.GPU(duration=120)
 def process_code_generation(ai_response: str):
-    import json
-    import re
-    from huggingface_hub import hf_hub_download
-    from llama_cpp import Llama
-    from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
-    from llama_cpp_agent.providers import LlamaCppPythonProvider
-
-    # Download the model (cached if already exists)
-    model_path = hf_hub_download(
-        repo_id="tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF",
-        filename="neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf",
-        local_dir="./models"
-    )
-
-    # Initialize LLaMA model (CPU-safe)
-    llm = Llama(
-        model_path=model_path,
-        n_gpu_layers=0,     # IMPORTANT: ZeroGPU uses CPU under the hood
-        n_batch=512,
-        n_ctx=4096,
-        verbose=True
-    )
-    provider = LlamaCppPythonProvider(llm)
-    agent = LlamaCppAgent(
-        provider,
-        system_prompt="You are an AI code parser. Given a pasted AI response with multiple code files, extract the individual files and output a JSON object like {\"file.py\": \"...\", \"utils.py\": \"...\"}.",
-        predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
-        debug_output=False
-    )
-
-    # Prompt to interpret the response
-    user_prompt = f"""
-Given the following AI output that includes multiple code files, extract and convert it into a JSON object that maps filenames to code contents.
-
-Example output: 
-{{
-  "app.py": "print('Hello World')",
-  "utils/helper.py": "def greet(): return 'hi'"
-}}
-
-AI Output:
+    try:
+        import json
+        import re
+        from huggingface_hub import hf_hub_download
+        from llama_cpp import Llama
+        from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
+        from llama_cpp_agent.providers import LlamaCppPythonProvider
+        import os
+        import zipfile
+        import uuid
+
+        model_path = hf_hub_download(
+            repo_id="tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF",
+            filename="neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf",
+            local_dir="./models"
+        )
+
+        llm = Llama(
+            model_path=model_path,
+            n_gpu_layers=0,
+            n_batch=512,
+            n_ctx=4096,
+            verbose=True
+        )
+
+        provider = LlamaCppPythonProvider(llm)
+        agent = LlamaCppAgent(
+            provider,
+            system_prompt="You are an AI code parser. Given a pasted AI response with multiple code files, output a JSON like {\"file.py\": \"code\"}.",
+            predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
+            debug_output=False
+        )
+
+        prompt = f"""Here is an AI output with multiple files. Convert it to a JSON object mapping filenames to code.
+
 {ai_response}
 """
-
-    full_response = ""
-    for chunk in agent.get_chat_response(user_prompt, returns_streaming_generator=True):
-        full_response += chunk
-
-    # Try to parse as JSON
-    try:
-        parsed = json.loads(full_response)
-    except Exception:
-        # Fallback regex parsing if model output wasn't pure JSON
-        pattern = r"Filename:\s*(?P<fname>[\w\.\-/]+)\s*```(?:[a-zA-Z0-9]*)\n(?P<code>[\s\S]*?)```"
-        parsed = {}
-        for m in re.finditer(pattern, full_response):
-            parsed[m.group("fname").strip()] = m.group("code")
-
-        if not parsed:
-            parsed = {"output.txt": full_response}
-
-    # Write files and zip them
-    import os
-    import zipfile
-    import uuid
-
-    temp_dir = f"/tmp/code_files_{uuid.uuid4().hex}"
-    os.makedirs(temp_dir, exist_ok=True)
-
-    for filename, content in parsed.items():
-        file_path = os.path.join(temp_dir, filename)
-        os.makedirs(os.path.dirname(file_path), exist_ok=True)
-        with open(file_path, "w", encoding="utf-8") as f:
-            f.write(content)
-
-    zip_path = f"/tmp/code_bundle_{uuid.uuid4().hex}.zip"
-    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
-        for root, _, files in os.walk(temp_dir):
-            for file in files:
-                full_path = os.path.join(root, file)
-                arcname = os.path.relpath(full_path, start=temp_dir)
-                zipf.write(full_path, arcname=arcname)
-
-    return zip_path
+        full_response = ""
+        for chunk in agent.get_chat_response(prompt, returns_streaming_generator=True):
+            full_response += chunk
+
+        try:
+            parsed = json.loads(full_response)
+        except Exception as e:
+            print("JSON parse failed:", e)
+            parsed = {}
+            pattern = r"Filename:\s*(?P<fname>[\w\.\-/]+)\s*```(?:[a-zA-Z0-9]*)\n(?P<code>[\s\S]*?)```"
+            for m in re.finditer(pattern, full_response):
+                parsed[m.group("fname").strip()] = m.group("code")
+            if not parsed:
+                parsed = {"output.txt": full_response}
+
+        temp_dir = f"/tmp/code_files_{uuid.uuid4().hex}"
+        os.makedirs(temp_dir, exist_ok=True)
+
+        for fname, content in parsed.items():
+            file_path = os.path.join(temp_dir, fname)
+            os.makedirs(os.path.dirname(file_path), exist_ok=True)
+            with open(file_path, "w", encoding="utf-8") as f:
+                f.write(content)
+
+        zip_path = f"/tmp/code_bundle_{uuid.uuid4().hex}.zip"
+        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+            for root, _, files in os.walk(temp_dir):
+                for file in files:
+                    full_path = os.path.join(root, file)
+                    arcname = os.path.relpath(full_path, temp_dir)
+                    zf.write(full_path, arcname)
+
+        return zip_path
+
+    except Exception as err:
+        print("🔥 Error in process_code_generation():", err)
+        raise err
 
 
 with gr.Blocks() as demo: