diff --git a/README.md b/README.md index d74d912d..f18f03f9 100644 --- a/README.md +++ b/README.md @@ -527,6 +527,18 @@ result = await agent.run( print(result.output) ``` +## Uploading Files + +`files.upload` takes a file path, raw bytes, or a binary file object. Files over 100 MB go through a presigned storage URL. + +```python +file = metorial.files.upload("./report.pdf") + +logo = metorial.files.upload("./logo.png", purpose="organization_image") + +notes = metorial.files.upload(b"hello", file_name="notes.txt", title="Notes") +``` + ## Error Handling ```python diff --git a/src/metorial/_endpoint/endpoint_manager.py b/src/metorial/_endpoint/endpoint_manager.py index 0e3938b5..0eae0079 100644 --- a/src/metorial/_endpoint/endpoint_manager.py +++ b/src/metorial/_endpoint/endpoint_manager.py @@ -47,8 +47,12 @@ def _request(self, method: str, request: MetorialRequest, try_count: int = 0) -> has_body = method in {"POST", "PUT", "PATCH"} json_payload = None files_payload = None + data_payload = None - if has_body and request.body is not None: + if has_body and request.files is not None: + files_payload = request.files + data_payload = request.data + elif has_body and request.body is not None: # If body is a file-like object, send as multipart if hasattr(request.body, "read"): files_payload = request.body @@ -67,6 +71,7 @@ def _request(self, method: str, request: MetorialRequest, try_count: int = 0) -> headers=headers, json=json_payload, files=files_payload, + data=data_payload, allow_redirects=True, timeout=30, ) diff --git a/src/metorial/_endpoint/request.py b/src/metorial/_endpoint/request.py index 4dd141c6..b98de4ee 100644 --- a/src/metorial/_endpoint/request.py +++ b/src/metorial/_endpoint/request.py @@ -4,7 +4,10 @@ class MetorialRequest: - """Request configuration object for Metorial API calls.""" + """Request configuration object for Metorial API calls. + + ``files`` and ``data`` send a multipart body instead of JSON. + """ def __init__( self, @@ -12,8 +15,12 @@ def __init__( host: str | None = None, query: dict[str, Any] | None = None, body: dict[str, Any] | None = None, + files: dict[str, Any] | None = None, + data: dict[str, str] | None = None, ): self.path = path self.host = host self.query = query self.body = body + self.files = files + self.data = data diff --git a/src/metorial/_file_upload.py b/src/metorial/_file_upload.py new file mode 100644 index 00000000..850c673c --- /dev/null +++ b/src/metorial/_file_upload.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import io +import mimetypes +import os +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import IO, Any, Union + +import requests + +from metorial._endpoint import MetorialEndpointManager, MetorialRequest +from metorial._generated.magnetar.resources import ( + DashboardInstanceFilesGetOutput, + mapDashboardInstanceFilesGetOutput, +) +from metorial.exceptions import MetorialSDKError + +# Must stay in sync with maxDirectUploadSize in the files API. +MAX_DIRECT_UPLOAD_SIZE = 100 * 1024 * 1024 + +UPLOAD_TIMEOUT_SECONDS = 30 + +FileSource = Union[str, "os.PathLike[str]", bytes, IO[bytes]] + + +@dataclass +class _Source: + stream: IO[bytes] + name: str + content_type: str + size: int + + +def upload_file( + manager: MetorialEndpointManager, + file: FileSource, + *, + purpose: str, + title: str | None, + file_name: str | None, + content_type: str | None, + store_id: str | None, + store_path: str | None, + store_replace: bool, + instance_id: str | None, + direct_limit: int = MAX_DIRECT_UPLOAD_SIZE, +) -> DashboardInstanceFilesGetOutput: + if (store_id is None) != (store_path is None): + raise ValueError("store_id and store_path must be given together") + if store_replace and store_id is None: + raise ValueError("store_replace requires store_id and store_path") + + fields: dict[str, Any] = {"purpose": purpose} + if title is not None: + fields["title"] = title + if instance_id is not None: + fields["instance_id"] = instance_id + if store_id is not None: + fields["store_id"] = store_id + fields["path"] = store_path + fields["store_replace"] = store_replace + + with _open_source(file, file_name, content_type) as source: + if source.size <= direct_limit: + data = _upload_direct(manager, source, fields) + else: + data = _upload_with_url(manager, source, fields, instance_id) + + return mapDashboardInstanceFilesGetOutput.from_dict(data) + + +def _upload_direct( + manager: MetorialEndpointManager, source: _Source, fields: dict[str, Any] +) -> Any: + form = { + key: ("true" if value else "false") if isinstance(value, bool) else str(value) + for key, value in fields.items() + } + # Read into memory so the 429 retry in the endpoint manager can resend the body. + content = source.stream.read() + request = MetorialRequest( + path=["files"], + files={"file": (source.name, content, source.content_type)}, + data=form, + ) + return manager._request("POST", request) + + +def _upload_with_url( + manager: MetorialEndpointManager, + source: _Source, + fields: dict[str, Any], + instance_id: str | None, +) -> Any: + session = manager._request( + "POST", + MetorialRequest( + path=["files"], + body={ + "mode": "get_upload_url", + "file_name": source.name, + "file_size": source.size, + "file_type": source.content_type, + **fields, + }, + ), + ) + + upload = session.get("upload") if isinstance(session, dict) else None + if not isinstance(upload, dict) or not upload.get("url") or not session.get("id"): + raise MetorialSDKError( + { + "status": 0, + "code": "malformed_response", + "message": "The Metorial API did not return an upload URL.", + } + ) + + _put_to_storage(upload.get("method") or "PUT", upload["url"], source) + + complete: dict[str, Any] = {"mode": "complete", "file_upload_id": session["id"]} + if instance_id is not None: + complete["instance_id"] = instance_id + return manager._request("POST", MetorialRequest(path=["files"], body=complete)) + + +def _put_to_storage(method: str, url: str, source: _Source) -> None: + # The URL signature covers the content type, so it must match the declared + # file_type, and Metorial credentials must not be sent to the storage backend. + try: + response = requests.request( + method, + url, + data=source.stream, + headers={ + "Content-Type": source.content_type, + "Content-Length": str(source.size), + }, + timeout=UPLOAD_TIMEOUT_SECONDS, + ) + except Exception as error: + raise MetorialSDKError( + { + "status": 0, + "code": "network_error", + "message": f"Failed to upload {source.name} to storage", + "error": str(error), + } + ) from error + + if not response.ok: + raise MetorialSDKError( + { + "status": response.status_code, + "code": "upload_failed", + "message": f"Failed to upload {source.name}: storage responded with {response.status_code}", + "body_snippet": (response.text or "")[:1000], + } + ) + + +@contextmanager +def _open_source( + file: FileSource, file_name: str | None, content_type: str | None +) -> Iterator[_Source]: + if isinstance(file, (bytes, bytearray)): + if not file_name: + raise ValueError("file_name is required when uploading bytes") + yield _Source( + stream=io.BytesIO(file), + name=file_name, + content_type=content_type or _guess_content_type(file_name), + size=len(file), + ) + return + + if isinstance(file, (str, os.PathLike)): + path = os.path.abspath(os.fspath(file)) + if not os.path.isfile(path): + raise ValueError(f"{os.fspath(file)} is not a file") + + name = file_name or os.path.basename(path) + with open(path, "rb") as stream: + yield _Source( + stream=stream, + name=name, + content_type=content_type or _guess_content_type(name), + size=os.fstat(stream.fileno()).st_size, + ) + return + + name = file_name or os.path.basename(str(getattr(file, "name", "") or "")) + if not name: + raise ValueError("file_name is required when uploading a stream without a name") + + if file.seekable(): + start = file.tell() + size = file.seek(0, io.SEEK_END) - start + file.seek(start) + readable: IO[bytes] = file + else: + content = file.read() + readable = io.BytesIO(content) + size = len(content) + + yield _Source( + stream=readable, + name=name, + content_type=content_type or _guess_content_type(name), + size=size, + ) + + +def _guess_content_type(name: str) -> str: + return mimetypes.guess_type(name)[0] or "application/octet-stream" diff --git a/src/metorial/_magnetar_sdk.py b/src/metorial/_magnetar_sdk.py index 300712d6..003e13e3 100644 --- a/src/metorial/_magnetar_sdk.py +++ b/src/metorial/_magnetar_sdk.py @@ -18,6 +18,7 @@ import httpx from metorial._endpoint import MetorialEndpointManager +from metorial._file_upload import FileSource, upload_file from metorial._generated.magnetar.endpoints.callback_events import ( MetorialCallbackEventsEndpoint, ) @@ -336,6 +337,7 @@ from metorial._generated.magnetar.endpoints.tool_calls import ( MetorialToolCallsEndpoint, ) +from metorial._generated.magnetar.resources import DashboardInstanceFilesGetOutput from metorial._sdk_builder import MetorialSDKBuilder from metorial._sdk_shared import SDKConfig, get_api_host, get_headers from metorial.sandbox import ( @@ -717,6 +719,39 @@ class MagnetarFilesGroup(MetorialFilesEndpoint): def __init__(self, manager: MetorialEndpointManager) -> None: super().__init__(manager) self.links = MetorialFilesLinksEndpoint(manager) + self._manager = manager + + def upload( + self, + file: FileSource, + *, + purpose: str = "generic", + title: str | None = None, + file_name: str | None = None, + content_type: str | None = None, + store_id: str | None = None, + store_path: str | None = None, + store_replace: bool = False, + instance_id: str | None = None, + ) -> DashboardInstanceFilesGetOutput: + """Upload a file and return the created file. + + ``file`` is a file path, raw bytes, or a binary file object. Bytes and + unnamed streams need ``file_name``. Files over 100 MB are uploaded + through a presigned storage URL. + """ + return upload_file( + self._manager, + file, + purpose=purpose, + title=title, + file_name=file_name, + content_type=content_type, + store_id=store_id, + store_path=store_path, + store_replace=store_replace, + instance_id=instance_id, + ) class MagnetarSkillsGroup(MetorialSkillsEndpoint): diff --git a/tests/test_file_upload.py b/tests/test_file_upload.py new file mode 100644 index 00000000..626abba9 --- /dev/null +++ b/tests/test_file_upload.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +import io +import json +import threading +from collections.abc import Iterator +from email.parser import BytesParser +from email.policy import HTTP +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any + +import pytest + +from metorial._endpoint import MetorialEndpointManager +from metorial._file_upload import upload_file +from metorial.exceptions import MetorialSDKError + +FILE_RESPONSE = { + "object": "file", + "id": "file_1", + "status": "active", + "file_name": "notes.txt", + "file_size": 5, + "file_type": "text/plain", + "title": None, + "purpose": {"name": "Generic", "identifier": "generic"}, + "created_at": "2026-01-01T00:00:00.000Z", + "updated_at": "2026-01-01T00:00:00.000Z", +} + + +class FakeFilesApi: + def __init__(self) -> None: + self.requests: list[dict[str, Any]] = [] + self.storage_status = 200 + self.api_status = 200 + + def handle(self, handler: BaseHTTPRequestHandler) -> None: + length = int(handler.headers.get("Content-Length") or 0) + body = handler.rfile.read(length) + content_type = handler.headers.get("Content-Type") or "" + record: dict[str, Any] = { + "method": handler.command, + "path": handler.path, + "headers": dict(handler.headers), + "body": body, + } + self.requests.append(record) + + if handler.command == "PUT": + handler.send_response(self.storage_status) + handler.end_headers() + if self.storage_status != 200: + handler.wfile.write(b"SignatureDoesNotMatch") + return + + if self.api_status != 200: + self._json( + handler, self.api_status, {"status": "error", "message": "Unknown file purpose"} + ) + return + + if content_type.startswith("multipart/form-data"): + record["form"], record["files"] = parse_multipart(content_type, body) + self._json(handler, 200, FILE_RESPONSE) + return + + payload = json.loads(body) + record["json"] = payload + if payload["mode"] == "get_upload_url": + host = handler.headers["Host"] + self._json( + handler, + 200, + { + "object": "file_upload", + "id": "fup_1", + "upload": {"url": f"http://{host}/storage", "method": "PUT"}, + }, + ) + else: + self._json(handler, 200, {**FILE_RESPONSE, "id": "file_2"}) + + def _json(self, handler: BaseHTTPRequestHandler, status: int, value: Any) -> None: + data = json.dumps(value).encode() + handler.send_response(status) + handler.send_header("Content-Type", "application/json") + handler.send_header("Content-Length", str(len(data))) + handler.end_headers() + handler.wfile.write(data) + + +def parse_multipart( + content_type: str, body: bytes +) -> tuple[dict[str, str], dict[str, tuple[str, str, bytes]]]: + message = BytesParser(policy=HTTP).parsebytes( + f"Content-Type: {content_type}\r\n\r\n".encode() + body + ) + form: dict[str, str] = {} + files: dict[str, tuple[str, str, bytes]] = {} + for part in message.iter_parts(): + name = part.get_param("name", header="content-disposition") + payload = part.get_payload(decode=True) + filename = part.get_filename() + if filename: + files[str(name)] = (filename, part.get_content_type(), payload) + else: + form[str(name)] = payload.decode() + return form, files + + +@pytest.fixture +def api() -> Iterator[tuple[FakeFilesApi, MetorialEndpointManager]]: + fake = FakeFilesApi() + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + fake.handle(self) + + def do_PUT(self) -> None: + fake.handle(self) + + def log_message(self, *args: Any) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + + manager = MetorialEndpointManager( + {}, + f"http://127.0.0.1:{server.server_address[1]}", + lambda _: {"Authorization": "Bearer metorial_sk_test"}, + ) + try: + yield fake, manager + finally: + server.shutdown() + + +def upload(manager: MetorialEndpointManager, file: Any, **kwargs: Any) -> Any: + options: dict[str, Any] = { + "purpose": "generic", + "title": None, + "file_name": None, + "content_type": None, + "store_id": None, + "store_path": None, + "store_replace": False, + "instance_id": None, + } + options.update(kwargs) + return upload_file(manager, file, **options) + + +def test_small_files_are_sent_as_multipart( + api: tuple[FakeFilesApi, MetorialEndpointManager], tmp_path: Path +) -> None: + fake, manager = api + local = tmp_path / "notes.txt" + local.write_text("hello") + + result = upload( + manager, + local, + title="Notes", + instance_id="inst_1", + store_id="store_1", + store_path="docs/notes.txt", + store_replace=True, + ) + + assert result.id == "file_1" + [request] = fake.requests + assert request["path"] == "/files" + assert request["form"] == { + "purpose": "generic", + "title": "Notes", + "instance_id": "inst_1", + "store_id": "store_1", + "path": "docs/notes.txt", + "store_replace": "true", + } + assert request["files"]["file"] == ("notes.txt", "text/plain", b"hello") + + +def test_directories_are_rejected( + api: tuple[FakeFilesApi, MetorialEndpointManager], tmp_path: Path +) -> None: + fake, manager = api + (tmp_path / "notes.txt").write_text("hello") + + with pytest.raises(ValueError, match="is not a file"): + upload(manager, tmp_path) + assert fake.requests == [] + + +def test_bytes_and_streams_need_a_name( + api: tuple[FakeFilesApi, MetorialEndpointManager], +) -> None: + fake, manager = api + with pytest.raises(ValueError, match="file_name"): + upload(manager, b"hello") + + upload(manager, io.BytesIO(b"hello"), file_name="hello.json") + assert fake.requests[0]["files"]["file"] == ( + "hello.json", + "application/json", + b"hello", + ) + + +def test_large_files_use_an_upload_url( + api: tuple[FakeFilesApi, MetorialEndpointManager], tmp_path: Path +) -> None: + fake, manager = api + local = tmp_path / "archive.bin" + local.write_bytes(b"large contents") + + result = upload( + manager, local, purpose="skill_export", instance_id="inst_1", direct_limit=4 + ) + + assert result.id == "file_2" + get_url, put, complete = fake.requests + assert get_url["json"] == { + "mode": "get_upload_url", + "file_name": "archive.bin", + "file_size": len(b"large contents"), + "file_type": "application/octet-stream", + "purpose": "skill_export", + "instance_id": "inst_1", + } + assert put["method"] == "PUT" + assert put["body"] == b"large contents" + assert put["headers"]["Content-Type"] == "application/octet-stream" + assert "Authorization" not in put["headers"] + assert complete["json"] == { + "mode": "complete", + "file_upload_id": "fup_1", + "instance_id": "inst_1", + } + + +def test_storage_failures_are_reported( + api: tuple[FakeFilesApi, MetorialEndpointManager], tmp_path: Path +) -> None: + fake, manager = api + fake.storage_status = 403 + local = tmp_path / "archive.bin" + local.write_bytes(b"large contents") + + with pytest.raises(MetorialSDKError, match="storage responded with 403"): + upload(manager, local, direct_limit=4) + + +def test_api_errors_are_raised( + api: tuple[FakeFilesApi, MetorialEndpointManager], tmp_path: Path +) -> None: + fake, manager = api + fake.api_status = 400 + local = tmp_path / "notes.txt" + local.write_text("hello") + + with pytest.raises(MetorialSDKError, match="Unknown file purpose"): + upload(manager, local, purpose="nope") + + +def test_store_options_are_validated( + api: tuple[FakeFilesApi, MetorialEndpointManager], tmp_path: Path +) -> None: + _, manager = api + with pytest.raises(ValueError, match="together"): + upload(manager, b"x", file_name="x.txt", store_id="store_1") + with pytest.raises(ValueError, match="store_replace"): + upload(manager, b"x", file_name="x.txt", store_replace=True)