added server tests for allowed local media path and size args

Author: cchadowitz

tools/server/tests/unit/test_vision_api.py

@@ -2,12 +2,15 @@
 from utils import *
 import base64
 import requests
+from pathlib import Path
 
 server: ServerProcess
 
-def get_img_url(id: str) -> str:
+
+def get_img_url(id: str, tmp_path: str | None = None) -> str:
     IMG_URL_0 = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/11_truck.png"
     IMG_URL_1 = "https://huggingface.co/ggml-org/tinygemma3-GGUF/resolve/main/test/91_cat.png"
+    IMG_FILE_2 = "https://picsum.photos/id/237/5000"
     if id == "IMG_URL_0":
         return IMG_URL_0
     elif id == "IMG_URL_1":
@@ -28,6 +31,46 @@ def get_img_url(id: str) -> str:
         response = requests.get(IMG_URL_1)
         response.raise_for_status() # Raise an exception for bad status codes
         return base64.b64encode(response.content).decode("utf-8")
+    elif id == "IMG_FILE_0":
+        if tmp_path is None:
+            raise RuntimeError("get_img_url must be called with a tmp_path if using local files")
+        image_name = IMG_URL_0.split('/')[-1]
+        file_name: Path = Path(tmp_path) / image_name
+        if file_name.exists():
+            return f"file://{file_name}"
+        else:
+            response = requests.get(IMG_URL_0)
+            response.raise_for_status() # Raise an exception for bad status codes
+            with open(file_name, 'wb') as f:
+                f.write(response.content)
+            return f"file://{file_name}"
+    elif id == "IMG_FILE_1":
+        if tmp_path is None:
+            raise RuntimeError("get_img_url must be called with a tmp_path if using local files")
+        image_name = IMG_URL_1.split('/')[-1]
+        file_name: Path = Path(tmp_path) / image_name
+        if file_name.exists():
+            return f"file://{file_name}"
+        else:
+            response = requests.get(IMG_URL_1)
+            response.raise_for_status() # Raise an exception for bad status codes
+            with open(file_name, 'wb') as f:
+                f.write(response.content)
+            return f"file://{file_name}"
+    elif id == "IMG_FILE_2":
+        if tmp_path is None:
+            raise RuntimeError("get_img_url must be called with a tmp_path if using local files")
+        image_name = "dog.jpg"
+        file_name: Path = Path(tmp_path) / image_name
+        if file_name.exists():
+            return f"file://{file_name}"
+        else:
+            response = requests.get(IMG_FILE_2)
+            response.raise_for_status() # Raise an exception for bad status codes
+            with open(file_name, 'wb') as f:
+                f.write(response.content)
+            return f"file://{file_name}"
+
     else:
         return id
 
@@ -70,6 +113,9 @@ def test_v1_models_supports_multimodal_capability():
         ("What is this:\n", "malformed",              False, None),
         ("What is this:\n", "https://google.com/404", False, None), # non-existent image
         ("What is this:\n", "https://ggml.ai",        False, None), # non-image data
+        ("What is this:\n", "IMG_FILE_0",             False, None),
+        ("What is this:\n", "IMG_FILE_1",             False, None),
+        ("What is this:\n", "IMG_FILE_2",             False, None),
         # TODO @ngxson : test with multiple images, no images and with audio
     ]
 )
@@ -83,7 +129,7 @@ def test_vision_chat_completion(prompt, image_url, success, re_content):
             {"role": "user", "content": [
                 {"type": "text", "text": prompt},
                 {"type": "image_url", "image_url": {
-                    "url": get_img_url(image_url),
+                    "url": get_img_url(image_url, "./tmp"),
                 }},
             ]},
         ],
@@ -97,6 +143,45 @@ def test_vision_chat_completion(prompt, image_url, success, re_content):
         assert res.status_code != 200
 
 
+@pytest.mark.parametrize(
+    "allowed_mb_size, allowed_path, img_dir_path, prompt, image_url, success, re_content",
+    [
+        # test model is trained on CIFAR-10, but it's quite dumb due to small size
+        (0, "./tmp", "./tmp", "What is this:\n", "IMG_FILE_0",                          True, "(cat)+"),
+        (0, "./tmp", "./tmp", "What is this:\n", "IMG_FILE_1",                          True, "(frog)+"),
+        (1, "./tmp", "./tmp", "What is this:\n", "IMG_FILE_2",                          False, None),
+        (0, "./tmp/allowed", "./tmp", "What is this:\n", "IMG_FILE_0",                  False, None),
+        (0, "./tm", "./tmp", "What is this:\n", "IMG_FILE_0",                           False, None),
+        (0, "./tmp/allowed", "./tmp/allowed/..", "What is this:\n", "IMG_FILE_0",       False, None),
+        (0, "./tmp/allowed", "./tmp/allowed/../.", "What is this:\n", "IMG_FILE_0",     False, None),
+    ]
+)
+def test_vision_chat_completion_local_files(allowed_mb_size, allowed_path, img_dir_path, prompt, image_url, success, re_content):
+    global server
+    server.local_media_max_size_mb = allowed_mb_size
+    server.allowed_local_media_path = allowed_path
+    Path(allowed_path).mkdir(exist_ok=True)
+    server.start()
+    res = server.make_request("POST", "/chat/completions", data={
+        "temperature": 0.0,
+        "top_k": 1,
+        "messages": [
+            {"role": "user", "content": [
+                {"type": "text", "text": prompt},
+                {"type": "image_url", "image_url": {
+                    "url": get_img_url(image_url, img_dir_path),
+                }},
+            ]},
+        ],
+    })
+    if success:
+        assert res.status_code == 200
+        choice = res.body["choices"][0]
+        assert "assistant" == choice["message"]["role"]
+        assert match_regex(re_content, choice["message"]["content"])
+    else:
+        assert res.status_code != 200
+
 @pytest.mark.parametrize(
     "prompt, image_data, success, re_content",
     [

tools/server/tests/utils.py

@@ -95,6 +95,8 @@ class ServerProcess:
     chat_template_file: str | None = None
     server_path: str | None = None
     mmproj_url: str | None = None
+    local_media_max_size_mb: int | None = None
+    allowed_local_media_path: str | None = None
 
     # session variables
     process: subprocess.Popen | None = None
@@ -215,6 +217,10 @@ def start(self, timeout_seconds: int | None = DEFAULT_HTTP_TIMEOUT) -> None:
             server_args.extend(["--chat-template-file", self.chat_template_file])
         if self.mmproj_url:
             server_args.extend(["--mmproj-url", self.mmproj_url])
+        if self.local_media_max_size_mb:
+            server_args.extend(["--local-media-max-size-mb", self.local_media_max_size_mb])
+        if self.allowed_local_media_path:
+            server_args.extend(["--allowed-local-media-path", self.allowed_local_media_path])
 
         args = [str(arg) for arg in [server_path, *server_args]]
         print(f"tests: starting server with: {' '.join(args)}")