diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py index 783cb25..b95d40a 100644 --- a/tests/test_artifacts.py +++ b/tests/test_artifacts.py @@ -1,6 +1,10 @@ """Tests for artifacts operations.""" +import io +import warnings + import pytest +from PIL import Image from vlmrun.client.types import ArtifactListResponse @@ -83,3 +87,75 @@ def test_list_artifacts_rejects_both_ids(mock_client): session_id="sess-123", execution_id="exec-456", ) + + +@pytest.mark.parametrize( + "raw,expected_type,expected_id", + [ + ("img_a1b2c3", "img", "img_a1b2c3"), + ("img_4c129a.jpg", "img", "img_4c129a"), + ("vid_4d0e56.mp4", "vid", "vid_4d0e56"), + ("doc_abcdef.PDF", "doc", "doc_abcdef"), + ("IMG_ABCDEF.PNG", "img", "img_abcdef"), + ], +) +def test_normalize_object_id(raw, expected_type, expected_id): + from vlmrun.client.artifacts import normalize_object_id + + obj_type, normalized = normalize_object_id(raw) + assert obj_type == expected_type + assert normalized == expected_id + + +def test_normalize_object_id_invalid(): + from vlmrun.client.artifacts import normalize_object_id + + with pytest.raises(ValueError): + normalize_object_id("not-an-id") + + +def test_artifacts_get_strips_extension_and_accepts_octet_stream(monkeypatch, tmp_path): + """artifacts.get should strip .jpg and not crash on application/octet-stream.""" + from vlmrun.client import artifacts as artifacts_mod + from vlmrun.client.artifacts import Artifacts + + class FakeClient: + api_key = "test" + base_url = "https://agent.vlm.run/v1" + max_retries = 1 + + buf = io.BytesIO() + Image.new("RGB", (2, 2), color=(0, 0, 255)).save(buf, format="PNG") + png_bytes = buf.getvalue() + + def fake_request(self, method, url, params=None, raw_response=False, **kwargs): + assert params["object_id"] == "img_4c129a" + return png_bytes, 200, {"Content-Type": "application/octet-stream"} + + monkeypatch.setattr(artifacts_mod.APIRequestor, "request", fake_request) + monkeypatch.setattr(artifacts_mod, "VLMRUN_ARTIFACTS_CACHE_DIR", tmp_path) + + art = Artifacts(FakeClient()) + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + image = art.get( + object_id="img_4c129a.jpg", + session_id="550e8400-e29b-41d4-a716-446655440000", + ) + assert isinstance(image, Image.Image) + assert image.size == (2, 2) + + +def test_long_request_pending_and_result_helpers(): + from vlmrun.client.long_request import is_pending_status, is_result_status + + assert is_pending_status(204) + assert is_pending_status(303) + assert is_pending_status(500) + assert is_pending_status(0) + assert not is_pending_status(400) + + assert is_result_status(201, b'{"id":"x"}') + assert is_result_status(200, b'{"id":"x"}') + assert not is_result_status(201, b"") + assert not is_result_status(204, b"") diff --git a/vlmrun/client/agent.py b/vlmrun/client/agent.py index a1ad669..c2166d7 100644 --- a/vlmrun/client/agent.py +++ b/vlmrun/client/agent.py @@ -26,33 +26,124 @@ # via `extra_body`. _VLM_EXTRA_KEYS: frozenset[str] = frozenset({"skills", "toolsets", "models"}) - -def _patch_create(create_fn: Any) -> Any: - """Wrap an OpenAI ``create`` callable to accept VLM Run-specific kwargs. +# Default wall-clock budget for following a Modal 303 long-request poll URL. +_DEFAULT_LONG_REQUEST_TIMEOUT = 900.0 + + +def _pop_vlm_extra_kwargs(kwargs: dict[str, Any]) -> None: + """Move VLM-specific kwargs into ``extra_body`` in-place.""" + vlm_kwargs = {k: kwargs.pop(k) for k in _VLM_EXTRA_KEYS if k in kwargs} + if vlm_kwargs: + kwargs["extra_body"] = {**(kwargs.get("extra_body") or {}), **vlm_kwargs} + + +def _status_code_from_exc(exc: BaseException) -> int | None: + """Best-effort extract of an HTTP status code from an OpenAI/httpx error.""" + status = getattr(exc, "status_code", None) + if isinstance(status, int): + return status + response = getattr(exc, "response", None) + if response is not None: + status = getattr(response, "status_code", None) + if isinstance(status, int): + return status + return None + + +def _headers_from_exc(exc: BaseException) -> dict[str, Any]: + """Best-effort extract of response headers from an OpenAI/httpx error.""" + response = getattr(exc, "response", None) + if response is None: + return {} + headers = getattr(response, "headers", None) or {} + try: + return dict(headers) + except Exception: + return {} + + +def _parse_chat_completion(body: bytes) -> Any: + """Parse poll-response bytes into an OpenAI ChatCompletion object.""" + import json + + from openai.types.chat import ChatCompletion + + payload = json.loads(body) + return ChatCompletion.model_validate(payload) + + +def _follow_long_request_303( + exc: BaseException, + *, + api_key: str | None, + timeout: float, +) -> Any: + """Poll the Modal ``Location`` URL from a 303 and return the ChatCompletion.""" + from vlmrun.client.long_request import ( + extract_location, + poll_location, + ) + + location = extract_location(_headers_from_exc(exc)) + if not location: + raise exc + + headers: dict[str, str] = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + body, _status, _resp_headers = poll_location( + location, + headers=headers, + timeout=timeout, + ) + return _parse_chat_completion(body) + + +def _patch_create(create_fn: Any, *, api_key: str | None, timeout: float) -> Any: + """Wrap an OpenAI ``create`` callable for VLM kwargs + Modal 303 polling. ``skills``, ``toolsets``, and ``models`` are popped from ``kwargs`` and - merged into ``extra_body`` before the underlying call is made. + merged into ``extra_body`` before the underlying call is made. If the + gateway returns ``303 See Other`` (long request), the Location URL is + polled until the chat completion is ready. """ @functools.wraps(create_fn) def _create(*args: Any, **kwargs: Any) -> Any: - vlm_kwargs = {k: kwargs.pop(k) for k in _VLM_EXTRA_KEYS if k in kwargs} - if vlm_kwargs: - kwargs["extra_body"] = {**(kwargs.get("extra_body") or {}), **vlm_kwargs} - return create_fn(*args, **kwargs) + _pop_vlm_extra_kwargs(kwargs) + try: + return create_fn(*args, **kwargs) + except Exception as exc: + if _status_code_from_exc(exc) != 303: + raise + return _follow_long_request_303( + exc, + api_key=api_key, + timeout=max(float(timeout or 0), _DEFAULT_LONG_REQUEST_TIMEOUT), + ) return _create -def _patch_async_create(create_fn: Any) -> Any: +def _patch_async_create(create_fn: Any, *, api_key: str | None, timeout: float) -> Any: """Async variant of :func:`_patch_create`.""" + import asyncio @functools.wraps(create_fn) async def _create(*args: Any, **kwargs: Any) -> Any: - vlm_kwargs = {k: kwargs.pop(k) for k in _VLM_EXTRA_KEYS if k in kwargs} - if vlm_kwargs: - kwargs["extra_body"] = {**(kwargs.get("extra_body") or {}), **vlm_kwargs} - return await create_fn(*args, **kwargs) + _pop_vlm_extra_kwargs(kwargs) + try: + return await create_fn(*args, **kwargs) + except Exception as exc: + if _status_code_from_exc(exc) != 303: + raise + return await asyncio.to_thread( + _follow_long_request_303, + exc, + api_key=api_key, + timeout=max(float(timeout or 0), _DEFAULT_LONG_REQUEST_TIMEOUT), + ) return _create @@ -325,7 +416,11 @@ def completions(self): ) completions = openai_client.chat.completions - completions.create = _patch_create(completions.create) + completions.create = _patch_create( + completions.create, + api_key=self._client.api_key, + timeout=float(self._client.timeout or _DEFAULT_LONG_REQUEST_TIMEOUT), + ) return completions @cached_property @@ -379,5 +474,9 @@ async def main(): ) completions = async_openai_client.chat.completions - completions.create = _patch_async_create(completions.create) + completions.create = _patch_async_create( + completions.create, + api_key=self._client.api_key, + timeout=float(self._client.timeout or _DEFAULT_LONG_REQUEST_TIMEOUT), + ) return completions diff --git a/vlmrun/client/artifacts.py b/vlmrun/client/artifacts.py index 70c4772..d3c58f9 100644 --- a/vlmrun/client/artifacts.py +++ b/vlmrun/client/artifacts.py @@ -3,10 +3,12 @@ from __future__ import annotations import io -import requests +import re +import warnings from pathlib import Path from typing import TYPE_CHECKING, Union +import requests from PIL import Image from pydantic import AnyHttpUrl @@ -15,10 +17,70 @@ from vlmrun.common.utils import _HEADERS from vlmrun.constants import VLMRUN_ARTIFACTS_CACHE_DIR - if TYPE_CHECKING: from vlmrun.types.abstract import VLMRunProtocol +# Object refs are `_<6hex>` and may optionally include a file extension +# (e.g. `img_4c129a.jpg`, `vid_4d0e56.mp4`) in preview tags / markdown. +_OBJECT_ID_RE = re.compile( + r"^(?P[a-z]+)_(?P[0-9a-f]{6})(?:\.[a-z0-9]+)?$", + re.IGNORECASE, +) + +_IMAGE_CONTENT_TYPES = frozenset( + { + "image/jpeg", + "image/jpg", + "image/png", + "image/webp", + "image/gif", + "application/octet-stream", + } +) + +_CONTENT_TYPE_MAPPING = { + "vid": frozenset({"video/mp4", "application/octet-stream", "binary/octet-stream"}), + "aud": frozenset( + {"audio/mpeg", "audio/mp3", "application/octet-stream", "binary/octet-stream"} + ), + "doc": frozenset( + {"application/pdf", "application/octet-stream", "binary/octet-stream"} + ), + "recon": frozenset({"application/octet-stream", "binary/octet-stream"}), +} + +_EXT_MAPPING = {"vid": "mp4", "aud": "mp3", "doc": "pdf", "recon": "spz"} + + +def normalize_object_id(object_id: str) -> tuple[str, str]: + """Normalize an artifact object ID, stripping any file extension. + + Args: + object_id: Raw object ID, optionally with an extension + (e.g. ``img_a1b2c3`` or ``img_a1b2c3.jpg``). + + Returns: + Tuple of ``(obj_type, normalized_object_id)`` where the ID has no extension. + + Raises: + ValueError: If the object ID does not match ``_<6hex>[.ext]``. + """ + match = _OBJECT_ID_RE.match(object_id.strip()) + if not match: + raise ValueError( + f"Invalid object ID: {object_id}, expected format: " + "_<6-digit-hex-string> with optional file extension" + ) + obj_type = match.group("prefix").lower() + hex_id = match.group("hex").lower() + return obj_type, f"{obj_type}_{hex_id}" + + +def _content_type_base(headers: dict[str, str]) -> str: + """Return the Content-Type header without parameters (e.g. charset).""" + raw = headers.get("Content-Type") or headers.get("content-type") or "" + return raw.split(";", 1)[0].strip().lower() + class Artifacts: """Artifacts resource for VLM Run API.""" @@ -43,7 +105,7 @@ def get( """Get an artifact by session ID or execution ID and object ID or filename. Supported artifact types: - - img: Returns PIL.Image.Image (JPEG) + - img: Returns PIL.Image.Image - url: Returns AnyHttpUrl - vid: Returns Path to MP4 file - aud: Returns Path to MP3 file @@ -51,7 +113,8 @@ def get( - recon: Returns Path to SPZ file Args: - object_id: Object ID for the artifact (format: _<6-hex-chars>). + object_id: Object ID for the artifact (format: ``_<6-hex-chars>``, + optionally with a file extension such as ``.jpg`` / ``.mp4``). Mutually exclusive with filename. session_id: Session ID for the artifact (mutually exclusive with execution_id) execution_id: Execution ID for the artifact (mutually exclusive with session_id) @@ -79,10 +142,12 @@ def get( "Only one of `object_id` or `filename` is allowed, not both" ) - # Build query parameters, filtering out None values query_params: dict[str, str] = {} + normalized_id: str | None = None + obj_type: str | None = None if object_id is not None: - query_params["object_id"] = object_id + obj_type, normalized_id = normalize_object_id(object_id) + query_params["object_id"] = normalized_id if filename is not None: query_params["filename"] = filename if session_id is not None: @@ -100,7 +165,6 @@ def get( if not isinstance(response, bytes): raise TypeError("Expected bytes response") - # If raw response is requested, return the raw response as bytes if raw_response: return response @@ -108,71 +172,57 @@ def get( if object_id is None: return response - # Otherwise, return the appropriate type based on the content type - obj_type, _obj_id = object_id.split("_") - if len(_obj_id) != 6: - raise ValueError( - f"Invalid object ID: {object_id}, expected format: _<6-digit-hex-string>" - ) + assert obj_type is not None and normalized_id is not None - # Create artifacts directory with session_id subdirectory - sess_id: str = session_id or execution_id + sess_id: str = session_id or execution_id # type: ignore[assignment] artifacts_dir: Path = VLMRUN_ARTIFACTS_CACHE_DIR / sess_id artifacts_dir.mkdir(parents=True, exist_ok=True) - # Extension and content-type mappings for file-based artifacts - ext_mapping = {"vid": "mp4", "aud": "mp3", "doc": "pdf", "recon": "spz"} - content_type_mapping = { - "vid": "video/mp4", - "aud": "audio/mpeg", - "doc": "application/pdf", - "recon": "application/octet-stream", - } + content_type = _content_type_base(headers) if obj_type == "img": - assert headers["Content-Type"] in ( - "image/jpeg", - "image/png", - ), f"Expected image/jpeg or image/png, got {headers['Content-Type']}" + if content_type and content_type not in _IMAGE_CONTENT_TYPES: + warnings.warn( + f"Unexpected Content-Type for image artifact {normalized_id}: " + f"{content_type!r}; attempting to decode as an image anyway", + UserWarning, + stacklevel=2, + ) return Image.open(io.BytesIO(response)).convert("RGB") elif obj_type == "url": - # Get the filename including extension frm the URL by stripping any query parameters url: AnyHttpUrl = AnyHttpUrl(response.decode("utf-8")) path: Path = Path(str(url)) - filename: str = path.name.split("?")[0] - ext: str = filename.split(".")[-1].lower() - tmp_path: Path = artifacts_dir / f"{filename}.{ext}" + url_filename: str = path.name.split("?")[0] + ext: str = url_filename.split(".")[-1].lower() + tmp_path: Path = artifacts_dir / f"{url_filename}.{ext}" if tmp_path.exists(): return tmp_path - # Download the file, and move it to the appropriate path - with requests.get(url, headers=_HEADERS, stream=True) as r: + with requests.get(str(url), headers=_HEADERS, stream=True) as r: r.raise_for_status() with tmp_path.open("wb") as f: for chunk in r.iter_content(chunk_size=8192): f.write(chunk) return tmp_path elif obj_type in ("vid", "aud", "doc", "recon"): - # Validate content type - expected_content_type = content_type_mapping[obj_type] - actual_content_type = headers.get("Content-Type") - assert ( - actual_content_type == expected_content_type - ), f"Expected {expected_content_type}, got {actual_content_type}" - - # Build file path with appropriate extension - ext = ext_mapping.get(obj_type, None) - if ext is None: - raise IOError( - f"Unsupported file type [file_type={filename}, object_id={object_id}]" + expected = _CONTENT_TYPE_MAPPING[obj_type] + if content_type and content_type not in expected: + warnings.warn( + f"Unexpected Content-Type for {obj_type} artifact {normalized_id}: " + f"expected one of {sorted(expected)}, got {content_type!r}; " + "saving bytes anyway", + UserWarning, + stacklevel=2, ) - tmp_path: Path = artifacts_dir / f"{object_id}.{ext}" - # Return cached version if it exists + ext = _EXT_MAPPING.get(obj_type) + if ext is None: + raise IOError(f"Unsupported file type [object_id={normalized_id}]") + tmp_path = artifacts_dir / f"{normalized_id}.{ext}" + if tmp_path.exists(): return tmp_path - # Write the binary response to file with tmp_path.open("wb") as f: f.write(response) return tmp_path diff --git a/vlmrun/client/long_request.py b/vlmrun/client/long_request.py new file mode 100644 index 0000000..d464409 --- /dev/null +++ b/vlmrun/client/long_request.py @@ -0,0 +1,138 @@ +"""Helpers for Modal's long-request 303 See Other poll flow. + +When an Orion chat-completions request exceeds the Modal gateway idle timeout +(~150s), the gateway returns ``303 See Other`` with a ``Location`` URL. The +job keeps running server-side; clients must poll ``Location`` until the result +is ready. + +Poll semantics observed on the live API: + - connection hold-then-timeout, ``202``, ``204``, ``303``, ``5xx`` + → keep waiting (job still running) + - ``200`` / ``201`` with a non-empty body → result +""" + +from __future__ import annotations + +import time +from typing import Any, Mapping, Optional +from urllib.parse import urljoin + +import requests + +# Status codes that mean "job still running — keep polling". +_PENDING_STATUS_CODES = frozenset({0, 202, 204, 303}) + + +class LongRequestTimeoutError(TimeoutError): + """Raised when polling a long-request Location URL times out.""" + + +def is_pending_status(status_code: int) -> bool: + """Return True if ``status_code`` means the long job is still running.""" + if status_code in _PENDING_STATUS_CODES: + return True + # Transient 5xx mid-job does not mean failure on Modal poll URLs. + if 500 <= status_code < 600: + return True + return False + + +def is_result_status(status_code: int, body: bytes | str | None) -> bool: + """Return True if the poll response is the finished result.""" + if status_code not in (200, 201): + return False + if body is None: + return False + if isinstance(body, bytes): + return len(body) > 0 + return len(body) > 0 + + +def poll_location( + location: str, + *, + headers: Optional[Mapping[str, str]] = None, + timeout: float = 900.0, + poll_interval: float = 2.0, + request_timeout: float = 130.0, + session: Optional[requests.Session] = None, +) -> tuple[bytes, int, dict[str, str]]: + """Poll a Modal long-request ``Location`` URL until the result is ready. + + Args: + location: Absolute or relative URL from the ``303`` ``Location`` header. + headers: Optional request headers (Authorization, etc.). + timeout: Maximum wall-clock seconds to wait for the result. + poll_interval: Seconds to sleep between polls after a pending response. + request_timeout: Per-request timeout. Modal may hold the connection open + for up to ~2 minutes before returning; keep this above that. + session: Optional ``requests.Session`` to reuse. + + Returns: + Tuple of ``(body_bytes, status_code, response_headers)``. + + Raises: + LongRequestTimeoutError: If ``timeout`` elapses before a result arrives. + requests.HTTPError: For non-pending client errors. + """ + http = session or requests.Session() + own_session = session is None + started = time.monotonic() + _headers = dict(headers or {}) + + try: + while True: + elapsed = time.monotonic() - started + if elapsed >= timeout: + raise LongRequestTimeoutError( + f"Long request did not complete within {timeout:.0f}s " + f"(last Location={location!r})" + ) + + try: + response = http.get( + location, + headers=_headers, + timeout=min(request_timeout, max(1.0, timeout - elapsed)), + allow_redirects=False, + ) + except (requests.exceptions.Timeout, requests.exceptions.ConnectionError): + # Connection hold / network blip while the job is still running. + time.sleep(poll_interval) + continue + + status = response.status_code + body = response.content + resp_headers = dict(response.headers) + + if is_result_status(status, body): + return body, status, resp_headers + + if status in (301, 302, 303, 307, 308): + next_location = resp_headers.get("Location") or resp_headers.get( + "location" + ) + if next_location: + location = urljoin(location, next_location) + time.sleep(poll_interval) + continue + + if is_pending_status(status): + time.sleep(poll_interval) + continue + + response.raise_for_status() + time.sleep(poll_interval) + finally: + if own_session: + http.close() + + +def extract_location(headers: Mapping[str, Any] | None) -> Optional[str]: + """Extract a ``Location`` header value (case-insensitive).""" + if not headers: + return None + for key, value in headers.items(): + if str(key).lower() == "location" and value: + return str(value) + return None