@@ -267,6 +267,79 @@ def test_embed_cohere_multiple_embedding_types(self, mock_boto3_session):
267267 with pytest .raises (ValueError ):
268268 AmazonBedrockDocumentImageEmbedder (model = "cohere.embed-english-v3" , embedding_types = ["float" , "int8" ])
269269
270+ def test_run_cohere_end_to_end (self , mock_boto3_session , image_paths ):
271+ embedder = AmazonBedrockDocumentImageEmbedder (model = "cohere.embed-english-v3" )
272+
273+ def mock_invoke_model (* args , ** kwargs ):
274+ return {"body" : io .StringIO ('{"embeddings": [[0.1, 0.2, 0.3]]}' )}
275+
276+ with patch .object (embedder , "_client" ) as mock_client :
277+ mock_client .invoke_model .side_effect = mock_invoke_model
278+ docs = [Document (content = "cat" , meta = {"file_path" : str (image_paths [0 ])})]
279+ result = embedder .run (documents = docs )
280+
281+ assert len (result ["documents" ]) == 1
282+ assert result ["documents" ][0 ].embedding == [0.1 , 0.2 , 0.3 ]
283+ assert result ["documents" ][0 ].meta ["embedding_source" ] == {
284+ "type" : "image" ,
285+ "file_path_meta_field" : "file_path" ,
286+ }
287+
288+ def test_run_titan_end_to_end (self , mock_boto3_session , image_paths ):
289+ embedder = AmazonBedrockDocumentImageEmbedder (model = "amazon.titan-embed-image-v1" )
290+
291+ def mock_invoke_model (* args , ** kwargs ):
292+ return {"body" : io .StringIO ('{"embedding": [0.4, 0.5]}' )}
293+
294+ with patch .object (embedder , "_client" ) as mock_client :
295+ mock_client .invoke_model .side_effect = mock_invoke_model
296+ docs = [Document (content = "apple" , meta = {"file_path" : str (image_paths [0 ])})]
297+ result = embedder .run (documents = docs )
298+
299+ assert result ["documents" ][0 ].embedding == [0.4 , 0.5 ]
300+
301+ def test_run_with_pdf_cohere (self , mock_boto3_session , image_paths ):
302+ embedder = AmazonBedrockDocumentImageEmbedder (model = "cohere.embed-english-v3" )
303+
304+ def mock_invoke_model (* args , ** kwargs ):
305+ return {"body" : io .StringIO ('{"embeddings": [[0.9]]}' )}
306+
307+ with patch .object (embedder , "_client" ) as mock_client :
308+ mock_client .invoke_model .side_effect = mock_invoke_model
309+ pdf_doc = Document (content = "pdf" , meta = {"file_path" : str (image_paths [2 ]), "page_number" : 1 })
310+ result = embedder .run (documents = [pdf_doc ])
311+
312+ assert result ["documents" ][0 ].embedding == [0.9 ]
313+ body_sent = mock_client .invoke_model .call_args .kwargs ["body" ]
314+ assert "data:application/pdf;base64," in body_sent
315+
316+ def test_embed_titan_with_embedding_config (self , mock_boto3_session , image_paths ):
317+ embedder = AmazonBedrockDocumentImageEmbedder (
318+ model = "amazon.titan-embed-image-v1" ,
319+ embeddingConfig = {"outputEmbeddingLength" : 256 },
320+ )
321+
322+ def mock_invoke_model (* args , ** kwargs ):
323+ return {"body" : io .StringIO ('{"embedding": [0.1]}' )}
324+
325+ with patch .object (embedder , "_client" ) as mock_client :
326+ mock_client .invoke_model .side_effect = mock_invoke_model
327+ embedder ._embed_titan (images = ["fake_base64" ])
328+
329+ body_sent = mock_client .invoke_model .call_args .kwargs ["body" ]
330+ assert '"embeddingConfig": {"outputEmbeddingLength": 256}' in body_sent
331+
332+ def test_embed_titan_invocation_error (self , mock_boto3_session ):
333+ embedder = AmazonBedrockDocumentImageEmbedder (model = "amazon.titan-embed-image-v1" )
334+
335+ with patch .object (embedder , "_client" ) as mock_client :
336+ mock_client .invoke_model .side_effect = ClientError (
337+ error_response = {"Error" : {"Code" : "x" , "Message" : "y" }},
338+ operation_name = "invoke_model" ,
339+ )
340+ with pytest .raises (AmazonBedrockInferenceError ):
341+ embedder ._embed_titan (images = ["fake_base64" ])
342+
270343 @pytest .mark .integration
271344 @pytest .mark .skipif (
272345 not os .getenv ("AWS_ACCESS_KEY_ID" )
0 commit comments