Update app.py

app.py

@@ -1,90 +1,85 @@
-import spaces
 import os
 import gradio as gr
+import spaces
 import zipfile
-import tempfile
+from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
-from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
+from llama_cpp_agent import LlamaCppAgent, Tool, ToolCall, MessagesFormatterType
 from llama_cpp_agent.providers import LlamaCppPythonProvider
-from llama_cpp_agent.chat_history import BasicChatHistory
-from llama_cpp_agent.chat_history.messages import Roles
-from huggingface_hub import hf_hub_download
 
-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}"
+MODEL_NAME = "neuraldaredevil-8b-abliterated-q4_k_m-imat.gguf"
+MODEL_PATH = f"models/{MODEL_NAME}"
 
-hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME, local_dir="./models")
+hf_hub_download(
+    repo_id="tHottie/NeuralDaredevil-8B-abliterated-Q4_K_M-GGUF",
+    filename=MODEL_NAME,
+    local_dir="./models"
+)
 
-@spaces.GPU(duration=120)
-def generate_code_zip(prompt):
-    chat_template = MessagesFormatterType.GEMMA_2
+# 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):
     llm = Llama(
         model_path=MODEL_PATH,
-        flash_attn=True,
         n_gpu_layers=81,
-        n_batch=1024,
         n_ctx=8192,
+        n_batch=1024,
+        verbose=False
     )
     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
-    )
-
-    messages = BasicChatHistory()
-    messages.add_message({"role": Roles.user, "content": prompt})
-
-    stream = agent.get_chat_response(
-        prompt,
-        llm_sampling_settings=provider.get_provider_default_settings(),
-        chat_history=messages,
-        returns_streaming_generator=False,
-        print_output=True
+        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
     )
 
-    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")
-
-    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)
-
-    return zip_path
-
-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
-
-def interface_fn(prompt):
-    try:
-        zip_path = generate_code_zip(prompt)
-        return zip_path
-    except Exception as e:
-        return f"Error: {e}"
+    output = ""
+    for chunk in agent.run(prompt):
+        output += chunk
+    return output, package_zip()
 
 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)
+    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")
+
+    submit_btn.click(respond, inputs=[input_text], outputs=[output_text, download_file])
 
-demo.launch()
+demo.launch()