Implement image.generate

tests/client/test_images.py

@@ -0,0 +1,101 @@
+"""Tests for the Images class."""
+
+import pytest
+import requests
+from unittest.mock import patch
+from pathlib import Path
+
+from vlmrun.client import Client
+
+
+def test_image_generate_with_path():
+    """Test image.generate with a file path."""
+    client = Client(api_key="test-key")
+
+    with patch("vlmrun.client.images.requests.post") as mock_post:
+        # Setup mock response
+        mock_post.return_value.status_code = 200
+        mock_post.return_value.json.return_value = {
+            "id": "123",
+            "status": "success",
+            "result": {"text": "sample prediction"}
+        }
+
+        # Test with minimal parameters
+        response = client.image.generate(
+            image="tests/fixtures/sample.jpg",
+            model="vlm-1"
+        )
+
+        # Verify response
+        assert response["id"] == "123"
+        assert response["status"] == "success"
+        assert response["result"]["text"] == "sample prediction"
+
+        # Verify request
+        mock_post.assert_called_once()
+        args, kwargs = mock_post.call_args
+        assert args[0] == "https://api.vlm.run/v1/image/generate"
+        assert kwargs["headers"]["Authorization"] == "Bearer test-key"
+        assert kwargs["headers"]["Content-Type"] == "application/json"
+        assert "image" in kwargs["json"]
+        assert kwargs["json"]["model"] == "vlm-1"
+
+
+def test_image_generate_with_all_params():
+    """Test image.generate with all optional parameters."""
+    client = Client(api_key="test-key")
+
+    with patch("vlmrun.client.images.requests.post") as mock_post:
+        mock_post.return_value.status_code = 200
+        mock_post.return_value.json.return_value = {"id": "456", "status": "success"}
+
+        # Test with all parameters
+        response = client.image.generate(
+            image="tests/fixtures/sample.jpg",
+            model="vlm-1",
+            domain="document.invoice",
+            detail="hi",
+            json_schema={"type": "object"},
+            callback_url="https://example.com/callback",
+            metadata={"session_id": "test-session"}
+        )
+
+        # Verify response
+        assert response["id"] == "456"
+        assert response["status"] == "success"
+
+        # Verify all parameters were sent
+        args, kwargs = mock_post.call_args
+        payload = kwargs["json"]
+        assert payload["model"] == "vlm-1"
+        assert payload["domain"] == "document.invoice"
+        assert payload["detail"] == "hi"
+        assert payload["json_schema"] == {"type": "object"}
+        assert payload["callback_url"] == "https://example.com/callback"
+        assert payload["metadata"] == {"session_id": "test-session"}
+
+
+def test_image_generate_error_handling():
+    """Test error handling in image.generate."""
+    client = Client(api_key="test-key")
+
+    with patch("vlmrun.client.images.requests.post") as mock_post, \
+         pytest.raises(FileNotFoundError):
+        # Test with non-existent file
+        client.image.generate(
+            image="non_existent.jpg",
+            model="vlm-1"
+        )
+
+    with patch("vlmrun.client.images.requests.post") as mock_post:
+        # Test API error
+        mock_post.return_value.status_code = 400
+        mock_post.return_value.raise_for_status.side_effect = \
+            requests.exceptions.HTTPError("Bad Request")
+
+        with pytest.raises(requests.exceptions.HTTPError):
+            client.image.generate(
+                image="tests/fixtures/sample.jpg",
+                model="vlm-1"
+            )

tests/create_test_image.py

@@ -0,0 +1,13 @@
+"""Script to create a sample test image for unit tests."""
+
+from PIL import Image
+import numpy as np
+from pathlib import Path
+
+# Ensure fixtures directory exists
+fixtures_dir = Path(__file__).parent / "fixtures"
+fixtures_dir.mkdir(exist_ok=True)
+
+# Create a small test image (100x100 black square)
+img = Image.fromarray(np.zeros((100, 100, 3), dtype=np.uint8))
+img.save(fixtures_dir / "sample.jpg")

tests/test_client.py

@@ -1,5 +1,8 @@
 def test_client():
+    """Test client initialization."""
     from vlmrun.client import Client
 
-    client = Client()
+    client = Client(api_key="test-key")
+    assert client.api_key == "test-key"
+    assert hasattr(client, "image")
     assert client is not None

vlmrun/client/client.py

@@ -1,5 +1,24 @@
 from dataclasses import dataclass
+from typing import Any
+
+from .images import Images
 
 
 @dataclass
-class Client: ...
+class Client:
+    """VLMRun API client."""
+
+    api_key: str
+
+    def __post_init__(self) -> None:
+        """Initialize client resources after instantiation."""
+        self._image = Images(self)
+
+    @property
+    def image(self) -> Images:
+        """Get the image resource interface.
+
+        Returns:
+            Images: The image resource interface
+        """
+        return self._image

vlmrun/client/images.py

@@ -0,0 +1,93 @@
+"""Image-related functionality for VLMRun."""
+
+from typing import Any, Dict, Literal, Optional, Union
+from pathlib import Path
+
+from PIL.Image import Image
+import requests
+
+from ..common.image import encode_image
+
+
+class Images:
+    """Class for handling image-related operations."""
+
+    def __init__(self, client: Any) -> None:
+        """Initialize Images instance.
+
+        Args:
+            client: The VLMRun client instance
+        """
+        self._client = client
+
+    def generate(
+        self,
+        image: Union[str, Path, Image],
+        *,
+        domain: Optional[Literal[
+            "document.generative",
+            "document.presentation",
+            "document.invoice",
+            "document.receipt",
+            "document.markdown",
+            "video.tv-news",
+            "video.tv-intelligence"
+        ]] = None,
+        model: str = "vlm-1",
+        detail: Optional[Literal["auto", "hi", "lo"]] = None,
+        json_schema: Optional[Dict[str, Any]] = None,
+        callback_url: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Generate structured prediction for the given image.
+
+        Args:
+            image: Image to analyze (PIL Image, path to image, or Path object)
+            domain: Domain identifier for the prediction
+            model: Model to use for generating the response (default: vlm-1)
+            detail: Detail level to use for the model
+            json_schema: JSON schema to use for the model
+            callback_url: URL to call when the request is completed
+            metadata: Optional metadata to pass to the model
+
+        Returns:
+            Dict containing the prediction response
+
+        Raises:
+            FileNotFoundError: If image path doesn't exist
+            ValueError: If image type is invalid
+            requests.RequestException: If the API request fails
+        """
+        # Convert image to base64
+        image_data = encode_image(image)
+
+        # Build request payload
+        payload = {
+            "image": image_data,
+            "model": model,
+        }
+
+        if domain is not None:
+            payload["domain"] = domain
+        if detail is not None:
+            payload["detail"] = detail
+        if json_schema is not None:
+            payload["json_schema"] = json_schema
+        if callback_url is not None:
+            payload["callback_url"] = callback_url
+        if metadata is not None:
+            payload["metadata"] = metadata
+
+        # Make API request
+        headers = {
+            "Authorization": f"Bearer {self._client.api_key}",
+            "Content-Type": "application/json",
+        }
+        response = requests.post(
+            "https://api.vlm.run/v1/image/generate",
+            json=payload,
+            headers=headers,
+        )
+        response.raise_for_status()
+
+        return response.json()