Update app.py

app.py

@@ -2,84 +2,95 @@ import os
 import gradio as gr
 import spaces
 import zipfile
+import json
+import re
+import uuid
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
-from llama_cpp_agent import LlamaCppAgent, Tool, ToolCall, MessagesFormatterType
+from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
 from llama_cpp_agent.providers import LlamaCppPythonProvider
 
 MODEL_NAME = "neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf"
 MODEL_PATH = f"models/{MODEL_NAME}"
 
+# Download model ahead of time
 hf_hub_download(
     repo_id="tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF",
     filename=MODEL_NAME,
-    local_dir="./models"
+    local_dir="./models",
 )
 
-# Define tools for the agent
-def create_file(filename: str, content: str) -> str:
-    os.makedirs("generated", exist_ok=True)
-    filepath = os.path.join("generated", filename)
-    with open(filepath, "w", encoding="utf-8") as f:
-        f.write(content)
-    return f"Created {filename}"
-
-def package_zip(output_name="code_package.zip"):
-    zip_path = os.path.join("generated", output_name)
-    with zipfile.ZipFile(zip_path, "w") as zipf:
-        for root, _, files in os.walk("generated"):
-            for file in files:
-                if file != output_name:
-                    zipf.write(os.path.join(root, file), arcname=file)
-    return zip_path
-
-# Wrap tools as LLM-callable
-tools = [
-    Tool(
-        name="create_file",
-        description="Create a file from given filename and content.",
-        parameters={"filename": str, "content": str},
-        function=create_file
-    ),
-    Tool(
-        name="package_zip",
-        description="Package all created files into a zip and return path.",
-        parameters={},
-        function=package_zip
-    )
-]
-
 @spaces.GPU(duration=120)
-def respond(prompt):
+def process_and_package(ai_response: str):
+    # Initialize the LLM & agent
     llm = Llama(
         model_path=MODEL_PATH,
-        n_gpu_layers=81,
-        n_ctx=8192,
-        n_batch=1024,
+        # Use CPU / fallback settings (or GPU if supported)
+        n_gpu_layers=0,
+        n_batch=512,
+        n_ctx=4096,
         verbose=False
     )
     provider = LlamaCppPythonProvider(llm)
     agent = LlamaCppAgent(
         provider,
-        tools=tools,
-        system_prompt="You are an agent that receives instructions and uses tools like create_file and package_zip to generate and package code files into a downloadable zip.",
-        predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2
+        system_prompt=(
+            "You are a code parsing assistant. The user will paste an AI response "
+            "that mentions multiple code files. Output a JSON object where keys are filenames "
+            "and values are the full source code text. Make sure the JSON structure is valid."
+        ),
+        predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
+        debug_output=True
     )
 
-    output = ""
-    for chunk in agent.run(prompt):
-        output += chunk
-    return output, package_zip()
+    # Ask the agent to output the JSON mapping
+    full = ""
+    for chunk in agent.run(ai_response):
+        full += chunk
 
-with gr.Blocks() as demo:
-    gr.Markdown("""### Code Generator with GGUF Model
-    Paste your AI response below. The model will generate code files and return a downloadable zip.
-    """)
-    input_text = gr.Textbox(lines=20, label="AI Response or Prompt")
-    output_text = gr.Textbox(lines=10, label="Agent Response")
-    download_file = gr.File(label="Download Code Package (.zip)")
-    submit_btn = gr.Button("Generate and Package")
+    # Try to parse JSON
+    file_map = {}
+    try:
+        file_map = json.loads(full)
+    except Exception as e:
+        # Fallback: simple regex heuristics
+        # Match patterns like "Filename: x/y.py ``` ... ```"
+        pattern = r"Filename:\s*(?P<fname>[\w\.\-/]+)\s*```(?:[a-zA-Z0-9]*\n)?(?P<code>[\s\S]*?)```"
+        for m in re.finditer(pattern, full, re.DOTALL):
+            fname = m.group("fname").strip()
+            code = m.group("code").rstrip()
+            file_map[fname] = code
+
+        if not file_map:
+            # as ultimate fallback, put entire output in a single file
+            file_map["output.txt"] = full
 
-    submit_btn.click(respond, inputs=[input_text], outputs=[output_text, download_file])
+    # Write files and zip
+    temp_dir = f"/tmp/codeproj_{uuid.uuid4().hex}"
+    os.makedirs(temp_dir, exist_ok=True)
+
+    for fname, code in file_map.items():
+        dest = os.path.join(temp_dir, fname)
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+        with open(dest, "w", encoding="utf-8") as f:
+            f.write(code)
+
+    zip_name = f"code_bundle_{uuid.uuid4().hex}.zip"
+    zip_path = os.path.join("/tmp", zip_name)
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
+        for root, _, files in os.walk(temp_dir):
+            for fn in files:
+                fullp = os.path.join(root, fn)
+                rel = os.path.relpath(fullp, temp_dir)
+                zf.write(fullp, arcname=rel)
+
+    return zip_path
+
+with gr.Blocks() as demo:
+    gr.Markdown("### Paste AI response below; get downloadable zip of parsed files")
+    inp = gr.Textbox(lines=15, label="AI response text")
+    btn = gr.Button("Generate Zip")
+    out = gr.File(label="Download .zip")
+    btn.click(fn=process_and_package, inputs=inp, outputs=out)
 
-demo.launch()
+demo.launch()