Added a new test mechanism for shuffle tools

Author: frikky

email/1.3.0/src/app.py

@@ -427,6 +427,9 @@ def remove_similar_items(self, items):
         return result
 
     def parse_eml(self, filedata, extract_attachments=False):
+        if filedata.startswith("file_"):
+            return self.parse_email_file(filedata, extract_attachments)
+
         parsedfile = {
             "success": True,
             "filename": "email.eml",

http/1.4.0/api.yaml

@@ -88,7 +88,20 @@ actions:
     returns:
       schema:
         type: string
-      example: "404 NOT FOUND"
+      example: |
+        {
+          "status": 200,
+          "body": {
+            "example": "json"
+            "data": "json"
+          },
+          "url": "https://example.com",
+          "headers": {
+            "Content-Type": "application/json"
+          },
+          "cookies": {},
+          "success": true
+        }
   - name: POST 
     description: Runs a POST request towards the specified endpoint
     parameters:

shuffle-ai/1.0.0/Dockerfile

@@ -15,16 +15,34 @@ RUN apt install -y clang g++ make automake autoconf libtool cmake
 RUN apt install -y jq git curl
 RUN apt install -y file openssl bash tini libpng-dev aspell-en
 RUN apt install -y git clang g++ make automake autoconf libtool cmake 
-RUN apt install -y autoconf-archive wget
-RUN mkdir -p /models
+RUN apt install -y autoconf-archive wget 
+
+# Install cuda toolkit
+#RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-archive-keyring.gpg
+#RUN dpkg -i cuda-archive-keyring.gpg
+#RUN rm cuda-archive-keyring.gpg
+#RUN apt update
+#RUN apt install -y cuda
+#RUN echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
+#RUN echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
+#RUN source ~/.bashrc
 
 # Larger model 
-RUN wget https://huggingface.co/unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF/resolve/main/DeepSeek-R1-Distill-Llama-8B-Q2_K.gguf
-ENV MODEL_PATH="/models/DeepSeek-R1-Distill-Llama-8B-Q2_K.gguf"
+RUN mkdir -p /models
+
+# Fails. 6 bit, 8B model. 
+#RUN wget https://huggingface.co/RichardErkhov/meta-llama_-_Meta-Llama-3-8B-gguf/blob/main/Meta-Llama-3-8B.Q6_K.gguf?download=true -O /models/Meta-Llama-3-8B.Q6_K.gguf
+#ENV MODEL_PATH="/models/Meta-Llama-3-8B.Q6_K.gguf"
+
+# Simple small Llama wrapper
+RUN wget https://huggingface.co/unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF/resolve/main/DeepSeek-R1-Distill-Llama-8B-Q2_K.gguf?download=true -O /models/DeepSeek-R1-Distill-Llama.gguf
+# Larger one
+#RUN wget https://huggingface.co/unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF/resolve/main/DeepSeek-R1-Distill-Llama-8B-Q8_0.gguf?download=true -O /models/DeepSeek-R1-Distill-Llama.gguf
+ENV MODEL_PATH="/models/DeepSeek-R1-Distill-Llama.gguf"
+
+# Failing? Bad magic bytes.
+#RUN wget https://huggingface.co/QuantFactory/Llama-3.2-3B-GGUF/resolve/main/Llama-3.2-3B.Q2_K.gguf?download=true -O /models/Llama-3.2-3B.Q2_K.gguf
 
-# https://huggingface.co/unsloth/DeepSeek-R1-Distill-Llama-8B-GGUF/resolve/main/DeepSeek-R1-Distill-Llama-8B-Q8_0.gguf
-#RUN wget https://huggingface.co/QuantFactory/Llama-3.2-3B-GGUF/resolve/main/Llama-3.2-3B.Q2_K.gguf?download=true -O /models/Llama-3.2-3B.Q8_0.gguf
-#RUN wget https://huggingface.co/QuantFactory/Llama-3.2-3B-GGUF/resolve/main/Llama-3.2-3B.Q2_K.gguf?download=true -O /models/Llama-3.2-3B.Q8_0.gguf
 
 # Install all of our pip packages in a single directory that we can copy to our base image later
 RUN mkdir /install
@@ -33,6 +51,8 @@ WORKDIR /install
 # Switch back to our base image and copy in all of our built packages and source code
 COPY requirements.txt /requirements.txt
 RUN python3 -m pip install -r /requirements.txt
+RUN CMAKE_ARGS="-DLLAMA_CUBLAS=on" python3 -m pip install llama-cpp-python --upgrade --force-reinstall --no-cache-dir
+
 
 # Install any binary dependencies needed in our final image
 

shuffle-ai/1.0.0/docker-compose.yml

@@ -18,23 +18,29 @@
 except Exception as e:
     print("Skipping pdf2image import: %s" % e)
 
+
 try:
     import llama_cpp 
 except Exception as e:
     print("Skipping llama_cpp import: %s" % e)
 
+print("LD Library: '%s'" % os.environ.get("LD_LIBRARY_PATH", ""))
+
 from shuffle_sdk import AppBase
 
 #model = "/models/Llama-3.2-3B.Q8_0.gguf" # Larger 
 #model = "/models/Llama-3.2-3B.Q2_K.gguf" # Smol
 
-#model = "/models/DeepSeek-R1-Distill-Llama-8B-Q8_0.gguf" # Larger 8-bit
-model = "/models/DeepSeek-R1-Distill-Llama-8B-Q2_K.gguf" # Smaller
+#model = "/models/DeepSeek-R1-Distill-Llama-8B-Q2_K.gguf" # Smaller
+#model = "/models/Meta-Llama-3-8B.Q6_K.gguf"
+model = "/models/DeepSeek-R1-Distill-Llama.gguf"
 if os.getenv("MODEL_PATH"):
     model = os.getenv("MODEL_PATH")
 
 def load_llm_model(model):
+    print("Using model path '%s'" % model)
     if not os.path.exists(model):
+        print("Could not find model at path %s" % model)
         model_name = model.split("/")[-1]
         # Check $HOME/downloads/{model}
 
@@ -54,20 +60,34 @@ def load_llm_model(model):
     innerllm = None
     gpu_layers = os.getenv("GPU_LAYERS")
     if gpu_layers:
+        print("GPU Layers: %s" % gpu_layers)
+
         gpu_layers = int(gpu_layers)
         if gpu_layers > 0:
-            print("GPU Layers: %s" % gpu_layers)
             innerllm = llama_cpp.Llama(model_path=model, n_gpu_layers=gpu_layers)
         else:
-            innerllm = llama_cpp.Llama(model_path=model)
+            innerllm = llama_cpp.Llama(model_path=model, n_gpu_layers=8)
     else:
         # Check if GPU available
-        #print("No GPU layers set.")
-        innerllm = llama_cpp.Llama(model_path=model)
+        print("No GPU layers set.")
+        #innerllm = llama_cpp.Llama(model_path=model)
+
+        return {
+            "success": False,
+            "reason": "GPU layers not set",
+            "details": "Set GPU_LAYERS environment variable to the number of GPU layers to use (e.g. 8)."
+        }
 
     return innerllm
 
-llm = load_llm_model(model)
+try:
+    llm = load_llm_model(model)
+except Exception as e:
+    print("[ERROR] Failed to load LLM model: %s" % e)
+    llm = {
+        "success": False,
+        "reason": "Failed to load LLM model %s" % model,
+    }
 
 class Tools(AppBase):
     __version__ = "1.0.0"
@@ -80,13 +100,35 @@ def run_llm(self, input, system_message=""):
         global llm
         global model
 
+        self.logger.info("[DEBUG] LD LIbrary: '%s'. If this is empty, GPU's may not work." % os.environ.get("LD_LIBRARY_PATH", ""))
+
         if not system_message:
-            system_message = "Be a friendly assistant",
+            system_message = "Answer their question directly. Don't use HTML or Markdown",
 
         self.logger.info("[DEBUG] Running LLM with model '%s'. To overwrite path, use environment variable MODEL_PATH=<path>" % model)
 
+        # Check if llm is a dict or not and look for success and reason in it
+        if not llm:
+            return {
+                "success": False,
+                "reason": "LLM model not loaded",
+                "details": "Ensure the LLM model is loaded",
+                "gpu_layers": os.getenv("GPU_LAYERS"),
+            }
+
+        if isinstance(llm, dict):
+            if "success" in llm and not llm["success"]:
+                # List files in /model folder
+                llm["folder"] = os.listdir("/models")
+                llm["gpu_layers"] = os.getenv("GPU_LAYERS")
+                return llm
+
+        self.logger.info("[DEBUG] Running LLM with input '%s' and system message '%s'. GPU Layers: %s" % (input, system_message, os.getenv("GPU_LAYERS")))
+
         # https://github.com/abetlen/llama-cpp-python 
         try:
+            print("LLM: ", llm)
+
             self.logger.info("[DEBUG] LLM: %s" % llm)
             output = llm.create_chat_completion(
                 max_tokens=100,
@@ -117,14 +159,19 @@ def run_llm(self, input, system_message=""):
         parsed_output = {
             "success": True,
             "model": output["model"],
-            "tokens": output["tokens"],
             "output": new_message,
         }
 
+        if "tokens" in output:
+            parsed_output["tokens"] = output["tokens"]
+
+        if "usage" in output:
+            parsed_output["tokens"] = output["usage"]
+
         if not os.getenv("GPU_LAYERS"):
             parsed_output["debug"] = "GPU_LAYERS not set. Running on CPU. Set GPU_LAYERS to the number of GPU layers to use (e.g. 8)."
 
-        return output
+        return parsed_output
 
     def security_assistant(self):
         # Currently testing outside the Shuffle environment

shuffle-ai/1.0.0/src/app.py

@@ -1,7 +1,16 @@
+gcloud config set project shuffler 
 
-gcloud run deploy shuffle-ai-1-0-0 \
-	--region=europe-west2 \
-	--max-instances=5 \
-	--memory=2Gi \
-	--set-env-vars=SHUFFLE_APP_EXPOSED_PORT=8080,SHUFFLE_SWARM_CONFIG=run,SHUFFLE_LOGS_DISABLED=true,SHUFFLE_APP_SDK_TIMEOUT=120 --source=./  \
+gcloud beta run deploy shuffle-ai-1-0-0 \
+	--project=shuffler \
+	--region=europe-west4 \
+	--source=./ \
+	--max-instances=1 \
+	--concurrency=64 \
+	--gpu 1 --gpu-type=nvidia-l4 \
+	--cpu 4 \
+	--memory=16Gi \
+	--no-cpu-throttling \
+	--set-env-vars=MODEL_PATH=/models/DeepSeek-R1-Distill-Llama.gguf,GPU_LAYERS=64,SHUFFLE_APP_EXPOSED_PORT=8080,SHUFFLE_SWARM_CONFIG=run,SHUFFLE_LOGS_DISABLED=true,SHUFFLE_APP_SDK_TIMEOUT=300,LD_LIBRARY_PATH=/usr/local/lib:/usr/local/nvidia/lib64:$LD_LIBRARY_PATH \
+	--source=./  \
+	[email protected] \
 	--timeout=120s

shuffle-ai/1.0.0/upload.sh

@@ -2931,6 +2931,9 @@ def merge_incoming_branches(self, input_type="list"):
 
         return results
 
+    def bodyparse_test(self, body):
+        return body
+
     def list_cidr_ips(self, cidr):
         defaultreturn = {
             "success": False,

shuffle-tools/1.2.0/src/app.py

@@ -26,6 +26,7 @@ gcloud run deploy yara \
 	--region=europe-west2 \
 	--max-instances=1 \
 	--set-env-vars=SHUFFLE_APP_EXPOSED_PORT=8080,SHUFFLE_SWARM_CONFIG=run,SHUFFLE_LOGS_DISABLED=true --source=./  \
+	[email protected] \
 	--timeout=1800s
 
 # With image