Test/ask with image

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "llm_server"  # The pip install <name>.
-version = "0.2.0"
+version = "0.3.0"
 description = "Template for repos that are intended to be packaged libraries."
 readme = "README.md"
 authors = [{ name = "Eric Apgar" }]

tests/test_server.py

@@ -19,8 +19,10 @@
 
 import pytest
 import requests
+from PIL import Image as PillowImage
 
 import llm_server
+from llm_server.helper.helper import encode_image
 
 
 MODEL_CACHE = os.environ.get('LLM_MODEL_CACHE')
@@ -36,7 +38,7 @@ def server():
     s.stop()
 
 
-@pytest.fixture(scope='module')
+@pytest.fixture(scope='function')
 def server_with_model():
     if not MODEL_CACHE:
         pytest.skip('LLM_MODEL_CACHE environment variable not set.')
@@ -47,6 +49,7 @@ def server_with_model():
     s.start()
     time.sleep(1)  # Allow uvicorn to finish starting.
     yield s
+    s.del_model(tag='gpt')  # Free GPU memory before teardown.
     s.stop()
 
 
@@ -63,3 +66,36 @@ def test_ask(server_with_model):
     data = response.json()
     assert isinstance(data['text'], str)
     assert len(data['text']) > 0
+
+
+@pytest.fixture(scope='module')
+def server_with_multimodal_model():
+    if not MODEL_CACHE:
+        pytest.skip('LLM_MODEL_CACHE environment variable not set.')
+    s = llm_server.Server()
+    s.set_host(ip_address='127.0.0.1', port=8003)
+    s.add_model(tag='phi4', name='microsoft/Phi-4-multimodal-instruct')
+    s.load_model(tag='phi4', location=MODEL_CACHE)
+    s.start()
+    time.sleep(1)  # Allow uvicorn to finish starting.
+    yield s
+    s.stop()
+
+
+@pytest.fixture(scope='module')
+def red_square_image():
+    return PillowImage.new('RGB', (64, 64), color=(255, 0, 0))
+
+
+def test_ask_with_image(server_with_multimodal_model, red_square_image):
+    response = requests.post(
+        'http://127.0.0.1:8003/ask',
+        json={
+            'tag': 'phi4',
+            'prompt': 'Describe the image.',
+            'images': [encode_image(red_square_image)],
+        },
+        timeout=60)
+    data = response.json()
+    assert isinstance(data['text'], str)
+    assert len(data['text']) > 0