[Bug] Fix regex for urls in image function

guidance/library/_image.py

@@ -1,22 +1,24 @@
-from .._guidance import guidance
-import urllib
-import typing
 import http
+import pathlib
 import re
+import typing
+import urllib
+
+from .._guidance import guidance
 
 
 @guidance
-def image(lm, src, allow_local=True):
+def image(lm, src: typing.Union[str, pathlib.Path, bytes], allow_local: bool = True):
 
     # load the image bytes
     # ...from a url
-    if isinstance(src, str) and re.match(r"$[^:/]+://", src):
+    if isinstance(src, str) and re.match(r"[^:/]+://", src):
         with urllib.request.urlopen(src) as response:
             response = typing.cast(http.client.HTTPResponse, response)
             bytes_data = response.read()
 
     # ...from a local path
-    elif allow_local and isinstance(src, str):
+    elif allow_local and (isinstance(src, str) or isinstance(src, pathlib.Path)):
         with open(src, "rb") as f:
             bytes_data = f.read()
 

tests/conftest.py

@@ -1,8 +1,12 @@
 import os
+import pathlib
 import random
+import tempfile
 import time
+import uuid
 
 import pytest
+import requests
 
 from guidance import models
 
@@ -11,14 +15,11 @@
 
 from .utils import get_model
 
-
 SELECTED_MODEL_ENV_VARIABLE = "GUIDANCE_SELECTED_MODEL"
 
 AVAILABLE_MODELS = {
     "gpt2cpu": dict(name="transformers:gpt2", kwargs=dict()),
-    "phi2cpu": dict(
-        name="transformers:microsoft/phi-2", kwargs={"trust_remote_code": True}
-    ),
+    "phi2cpu": dict(name="transformers:microsoft/phi-2", kwargs={"trust_remote_code": True}),
     "azure_guidance": dict(
         name="azure_guidance:",
         kwargs={},
@@ -41,9 +42,7 @@
         name="huggingface_hubllama:TheBloke/Llama-2-7B-GGUF:llama-2-7b.Q5_K_M.gguf",
         kwargs={"verbose": True, "n_ctx": 4096},
     ),
-    "transformers_mistral_7b": dict(
-        name="transformers:mistralai/Mistral-7B-v0.1", kwargs=dict()
-    ),
+    "transformers_mistral_7b": dict(name="transformers:mistralai/Mistral-7B-v0.1", kwargs=dict()),
     "hfllama_mistral_7b": dict(
         name="huggingface_hubllama:TheBloke/Mistral-7B-Instruct-v0.2-GGUF:mistral-7b-instruct-v0.2.Q8_0.gguf",
         kwargs={"verbose": True},
@@ -113,3 +112,26 @@ def rate_limiter() -> int:
     delay_secs = random.randint(10, 30)
     time.sleep(delay_secs)
     return delay_secs
+
+
+@pytest.fixture(scope="session")
+def remote_image_url():
+    return "https://picsum.photos/300/200"
+
+
+@pytest.fixture(scope="session")
+def local_image_path(remote_image_url):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        td = pathlib.Path(temp_dir)
+        filename = f"{str(uuid.uuid4())}.jpg"
+        with open(td / filename, "wb") as file:
+            response = requests.get(remote_image_url)
+            file.write(response.content)
+        assert (td / filename).exists()
+        yield td / filename
+
+
+@pytest.fixture(scope="session")
+def local_image_bytes(local_image_path):
+    with open(local_image_path, "rb") as f:
+        return f.read()

tests/unit/library/test_image.py

@@ -0,0 +1,36 @@
+import pytest
+
+from urllib.error import HTTPError
+from guidance import models, image
+
+
+def test_local_image(local_image_path):
+    model = models.Mock()
+    model += image(local_image_path)
+
+    assert str(model).startswith("<|_image:")
+
+
+def test_local_image_not_found():
+    model = models.Mock()
+    with pytest.raises(FileNotFoundError):
+        model += image("not_found.jpg")
+
+
+def test_remote_image(remote_image_url):
+    model = models.Mock()
+    model += image(remote_image_url)
+
+    assert str(model).startswith("<|_image:")
+
+
+def test_remote_image_not_found():
+    model = models.Mock()
+    with pytest.raises(HTTPError):
+        model += image("https://example.com/not_found.jpg")
+
+
+def test_image_from_bytes(local_image_bytes):
+    model = models.Mock()
+    model += image(local_image_bytes)
+    assert str(model).startswith("<|_image:")