client: add support for passing in Image type to generate (#408)

tests/test_client.py

@@ -11,6 +11,7 @@ from pytest_httpserver import HTTPServer, URIPattern
 from werkzeug.wrappers import Request, Response
 
 from ollama._client import CONNECTION_ERROR_MESSAGE, AsyncClient, Client, _copy_tools
+from ollama._types import Image
 
 PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC'
 PNG_BYTES = base64.b64decode(PNG_BASE64)
@@ -286,6 +287,46 @@ def test_client_generate(httpserver: HTTPServer):
   assert response['response'] == 'Because it is.'
 
 
+def test_client_generate_with_image_type(httpserver: HTTPServer):
+  httpserver.expect_ordered_request(
+    '/api/generate',
+    method='POST',
+    json={
+      'model': 'dummy',
+      'prompt': 'What is in this image?',
+      'stream': False,
+      'images': [PNG_BASE64],
+    },
+  ).respond_with_json(
+    {
+      'model': 'dummy',
+      'response': 'A blue sky.',
+    }
+  )
+
+  client = Client(httpserver.url_for('/'))
+  response = client.generate('dummy', 'What is in this image?', images=[Image(value=PNG_BASE64)])
+  assert response['model'] == 'dummy'
+  assert response['response'] == 'A blue sky.'
+
+
+def test_client_generate_with_invalid_image(httpserver: HTTPServer):
+  httpserver.expect_ordered_request(
+    '/api/generate',
+    method='POST',
+    json={
+      'model': 'dummy',
+      'prompt': 'What is in this image?',
+      'stream': False,
+      'images': ['invalid_base64'],
+    },
+  ).respond_with_json({'error': 'Invalid image data'}, status=400)
+
+  client = Client(httpserver.url_for('/'))
+  with pytest.raises(ValueError):
+    client.generate('dummy', 'What is in this image?', images=[Image(value='invalid_base64')])
+
+
 def test_client_generate_stream(httpserver: HTTPServer):
   def stream_handler(_: Request):
     def generate():