Update app.py

app.py

@@ -1,103 +1,90 @@
+import spaces
 import os
-import re
-import json
-import uuid
+import gradio as gr
 import zipfile
-from huggingface_hub import hf_hub_download
+import tempfile
 from llama_cpp import Llama
 from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
 from llama_cpp_agent.providers import LlamaCppPythonProvider
-import spaces
+from llama_cpp_agent.chat_history import BasicChatHistory
+from llama_cpp_agent.chat_history.messages import Roles
+from huggingface_hub import hf_hub_download
 
-LOG_PATH = "/tmp/codegen.log"
+MODEL_REPO = "tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF"
+MODEL_FILENAME = "neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf"
+MODEL_PATH = f"models/{MODEL_FILENAME}"
 
-def log(message: str):
-    with open(LOG_PATH, "a", encoding="utf-8") as log_file:
-        log_file.write(message + "\n")
-    print(message)
+hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME, local_dir="./models")
 
 @spaces.GPU(duration=120)
-def process_code_generation(ai_response: str):
-    try:
-        log("=== Started process_code_generation ===")
+def generate_code_zip(prompt):
+    chat_template = MessagesFormatterType.GEMMA_2
 
-        # Download model
-        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"
-        )
-        log(f"Downloaded model to: {model_path}")
+    llm = Llama(
+        model_path=MODEL_PATH,
+        flash_attn=True,
+        n_gpu_layers=81,
+        n_batch=1024,
+        n_ctx=8192,
+    )
+    provider = LlamaCppPythonProvider(llm)
+    agent = LlamaCppAgent(
+        provider,
+        system_prompt="You are an AI code generation agent. Respond with code blocks per file name.",
+        predefined_messages_formatter_type=chat_template,
+        debug_output=True
+    )
 
-        # Load LLM
-        llm = Llama(
-            model_path=model_path,
-            n_gpu_layers=0,
-            n_batch=512,
-            n_ctx=4096,
-            verbose=True
-        )
-        log("Loaded Llama model.")
+    messages = BasicChatHistory()
+    messages.add_message({"role": Roles.user, "content": prompt})
 
-        provider = LlamaCppPythonProvider(llm)
-        agent = LlamaCppAgent(
-            provider,
-            system_prompt="Convert an AI response with multiple code files into a JSON mapping file names to code.",
-            predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
-        )
+    stream = agent.get_chat_response(
+        prompt,
+        llm_sampling_settings=provider.get_provider_default_settings(),
+        chat_history=messages,
+        returns_streaming_generator=False,
+        print_output=True
+    )
 
-        prompt = f"""Given this AI response, output a JSON like: {{"main.py": "...", "utils/helper.js": "..."}} only:
-        
-{ai_response}
-"""
-        log("Generated prompt. Sending to model...")
+    response = stream if isinstance(stream, str) else "".join(stream)
+    file_map = parse_code_blocks(response)
+    
+    temp_dir = tempfile.mkdtemp()
+    zip_path = os.path.join(temp_dir, "generated_code.zip")
 
-        full_response = ""
-        for chunk in agent.get_chat_response(prompt, returns_streaming_generator=True):
-            full_response += chunk
-        log("Model response received:")
-        log(full_response)
+    with zipfile.ZipFile(zip_path, "w") as zipf:
+        for filename, code in file_map.items():
+            full_path = os.path.join(temp_dir, filename)
+            os.makedirs(os.path.dirname(full_path), exist_ok=True)
+            with open(full_path, "w", encoding="utf-8") as f:
+                f.write(code)
+            zipf.write(full_path, arcname=filename)
 
-        # Try parsing JSON
-        try:
-            parsed = json.loads(full_response)
-        except Exception as e:
-            log(f"JSON decode failed: {e}")
-            parsed = {}
-            matches = re.finditer(r"(?s)File(?:name)?:\s*(.+?)\n```(?:\w*\n)?(.*?)```", full_response)
-            for match in matches:
-                filename = match.group(1).strip()
-                code = match.group(2).strip()
-                parsed[filename] = code
-            if not parsed:
-                parsed = {"raw_output.txt": full_response}
+    return zip_path
 
-        log(f"Parsed files: {list(parsed.keys())}")
+def parse_code_blocks(text):
+    import re
+    file_blocks = {}
+    pattern = r"###\s*(.*?)\n```(?:[a-zA-Z0-9+]*\n)?(.*?)```"
+    matches = re.findall(pattern, text, re.DOTALL)
+    for filename, content in matches:
+        cleaned_filename = filename.strip()
+        cleaned_code = content.strip()
+        file_blocks[cleaned_filename] = cleaned_code
+    return file_blocks
 
-        temp_dir = f"/tmp/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)
-            log(f"Saved file: {file_path}")
-
-        zip_path = f"/tmp/code_bundle_{uuid.uuid4().hex}.zip"
-        with zipfile.ZipFile(zip_path, "w") 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, temp_dir)
-                    zipf.write(full_path, arcname)
-                    log(f"Added to ZIP: {arcname}")
-
-        log(f"Zip file created at: {zip_path}")
-        log("=== Finished process_code_generation ===\n\n")
+def interface_fn(prompt):
+    try:
+        zip_path = generate_code_zip(prompt)
         return zip_path
+    except Exception as e:
+        return f"Error: {e}"
+
+with gr.Blocks() as demo:
+    gr.Markdown("## AI Code Packager (GGUF-Powered)")
+    prompt = gr.Textbox(label="Enter AI response or code instructions")
+    output = gr.File(label="Download Zip")
+    btn = gr.Button("Generate Zip")
+    btn.click(fn=interface_fn, inputs=prompt, outputs=output)
 
-    except Exception as ex:
-        error_msg = f"Exception: {repr(ex)}"
-        log(error_msg)
-        return error_msg
+demo.launch()