fix(bedrock): reject s3 Marengo media without bucketOwner and skip items without an embedding

This commit is contained in:
mateo-berri 2026-09-08 12:53:01 -07:00
parent 6076e9f611
commit 384eee26a5
5 changed files with 55 additions and 18 deletions

View file

@ -91,19 +91,21 @@ class Marengo3Params(BaseModel):
return self.model_dump(include=MARENGO_2_7_ONLY_FIELDS, exclude_none=True)
def _s3_location(uri: str, bucket_owner: str | None) -> TwelveLabsS3Location:
def _require_bucket_owner(bucket_owner: str | None) -> str:
if bucket_owner is None:
unowned: Final[TwelveLabsS3Location] = {"uri": uri}
return unowned
owned: Final[TwelveLabsS3Location] = {"uri": uri, "bucketOwner": bucket_owner}
return owned
raise BedrockError(
status_code=400,
message="s3:// media requires the 'bucketOwner' parameter, the account id that owns the bucket",
)
return bucket_owner
def _media_source(media: str, bucket_owner: str | None) -> TwelveLabsMediaSource:
if not media.startswith(S3_URI_PREFIX):
inline: Final[TwelveLabsMediaSource] = {"base64String": get_base64_str(media)}
return inline
remote: Final[TwelveLabsMediaSource] = {"s3Location": _s3_location(media, bucket_owner)}
s3_location: Final[TwelveLabsS3Location] = {"uri": media, "bucketOwner": _require_bucket_owner(bucket_owner)}
remote: Final[TwelveLabsMediaSource] = {"s3Location": s3_location}
return remote

View file

@ -35,7 +35,7 @@ from litellm.types.utils import Embedding, EmbeddingResponse, PromptTokensDetail
class MarengoEmbeddingItem(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
embedding: tuple[float, ...]
embedding: tuple[float, ...] | None = None
class MarengoInvokeResponse(BaseModel):
@ -47,10 +47,10 @@ class MarengoInvokeResponse(BaseModel):
def vectors(self) -> tuple[tuple[float, ...], ...]:
if self.data:
return tuple(item.embedding for item in self.data)
return tuple(item.embedding for item in self.data if item.embedding is not None)
if self.embedding is not None:
return (self.embedding,)
return tuple(item.embedding for item in self.embeddings)
return tuple(item.embedding for item in self.embeddings if item.embedding is not None)
class MarengoBilledMultiInput(BaseModel):

View file

@ -204,6 +204,7 @@ class TestBedrockAsyncInvokeEmbedding:
input_type="video",
embeddingOption=["visual", "audio"],
segmentation={"method": "fixed", "fixed": {"durationSec": 6}},
bucketOwner="123456789012",
output_s3_uri="s3://test-bucket/async-invoke-output/",
)
@ -214,7 +215,7 @@ class TestBedrockAsyncInvokeEmbedding:
"modelInput": {
"inputType": "video",
"video": {
"mediaSource": {"s3Location": {"uri": "s3://test-bucket/clip.mp4"}},
"mediaSource": {"s3Location": {"uri": "s3://test-bucket/clip.mp4", "bucketOwner": "123456789012"}},
"segmentation": {"method": "fixed", "fixed": {"durationSec": 6}},
"embeddingOption": ["visual", "audio"],
},

View file

@ -1220,6 +1220,16 @@ def test_marengo_usage_without_request_data_bills_nothing():
assert response.usage.prompt_tokens_details is None
def test_marengo_response_items_without_an_embedding_are_skipped():
response = TwelveLabsMarengoEmbeddingConfig()._transform_response(
response_list=[{"data": [{"embeddingOption": "visual-text", "startSec": 0.0}, {"embedding": [0.1, 0.2, 0.3]}]}],
model="us.twelvelabs.marengo-embed-3-0-v1:0",
)
assert [item["embedding"] for item in response.data] == [[0.1, 0.2, 0.3]]
assert response.data[0]["index"] == 0
def test_marengo_3_text_image_without_media_source_is_a_bad_request():
with pytest.raises(litellm.BadRequestError, match=r"text_image.*media_source"):
litellm.embedding(

View file

@ -75,9 +75,22 @@ def test_image_request_from_s3_carries_bucket_owner():
}
def test_s3_media_without_bucket_owner_omits_the_key():
request = build_marengo_3_request("s3://media/duck.png", {"input_type": "image"})
assert request["image"]["mediaSource"] == {"s3Location": {"uri": "s3://media/duck.png"}}
@pytest.mark.parametrize(
"input_media,params",
[
("s3://media/duck.png", {"input_type": "image"}),
("s3://media/clip.mp4", {"input_type": "video"}),
("a duck", {"input_type": "text_image", "media_source": "s3://media/duck.png"}),
("a duck", {"input_type": "multi_input", "media_sources": {"img1": "s3://media/duck.png"}}),
],
)
def test_s3_media_without_bucket_owner_is_rejected_naming_it(input_media, params):
with pytest.raises(BedrockError) as excinfo:
build_marengo_3_request(input_media, params)
assert excinfo.value.status_code == 400
assert excinfo.value.message == (
"s3:// media requires the 'bucketOwner' parameter, the account id that owns the bucket"
)
def test_text_image_request_pairs_text_with_media_source():
@ -147,12 +160,13 @@ def test_timed_media_request_nests_every_option_under_the_media_key(input_type):
"embeddingType": ["fused_embedding"],
"embeddingScope": ["clip", "asset"],
"inferenceId": "req-42",
"bucketOwner": "123456789012",
},
)
assert wire(request) == {
"inputType": input_type,
input_type: {
"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4"}},
"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}},
"startSec": 2.0,
"endSec": 12.5,
"segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}},
@ -165,8 +179,10 @@ def test_timed_media_request_nests_every_option_under_the_media_key(input_type):
def test_timed_media_request_without_options_carries_only_the_media_source():
request = build_marengo_3_request("s3://media/clip.mp4", {"input_type": "video"})
assert request["video"] == {"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4"}}}
request = build_marengo_3_request("s3://media/clip.mp4", {"input_type": "video", "bucketOwner": "123456789012"})
assert request["video"] == {
"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}}
}
@pytest.mark.parametrize(
@ -211,7 +227,12 @@ def test_marengo_3_video_and_audio_still_require_the_async_route(input_type):
def test_marengo_3_async_invoke_wraps_the_nested_payload_with_the_base_model_id():
request = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request(
input="s3://media/clip.mp4",
inference_params={"input_type": "video", "embeddingOption": ["visual"], "output_s3_uri": OUTPUT_S3_URI},
inference_params={
"input_type": "video",
"embeddingOption": ["visual"],
"bucketOwner": "123456789012",
"output_s3_uri": OUTPUT_S3_URI,
},
async_invoke_route=True,
model_id="async_invoke%2Ftwelvelabs.marengo-embed-3-0-v1%3A0",
output_s3_uri=OUTPUT_S3_URI,
@ -220,7 +241,10 @@ def test_marengo_3_async_invoke_wraps_the_nested_payload_with_the_base_model_id(
"modelId": MARENGO_3_BASE,
"modelInput": {
"inputType": "video",
"video": {"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4"}}, "embeddingOption": ["visual"]},
"video": {
"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4", "bucketOwner": "123456789012"}},
"embeddingOption": ["visual"],
},
},
"outputDataConfig": {"s3OutputDataConfig": {"s3Uri": OUTPUT_S3_URI}},
}