From 513b56159fb05d4733a0859b8a994c375aeb809a Mon Sep 17 00:00:00 2001
From: Lucas Kim <41357160+kimnamu@users.noreply.github.com>
Date: Mon, 22 Jun 2026 15:19:23 +0900
Subject: [PATCH] fix(embeddings): forward embedding_dims to Titan V2 in AWS
 Bedrock embedder (#5671)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 mem0/embeddings/aws_bedrock.py                |  5 ++
 .../embeddings/test_aws_bedrock_embeddings.py | 47 +++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/mem0/embeddings/aws_bedrock.py b/mem0/embeddings/aws_bedrock.py
index a6fe3c224..9672625c0 100644
--- a/mem0/embeddings/aws_bedrock.py
+++ b/mem0/embeddings/aws_bedrock.py
@@ -70,6 +70,11 @@ class AWSBedrockEmbedding(EmbeddingBase):
         else:
             # Amazon and other providers
             input_body["inputText"] = text
+            # Titan Text Embeddings V2 accepts an optional output dimension
+            # (256/512/1024). Only forward embedding_dims when the user set it,
+            # mirroring the OpenAI embedder's guarded `dimensions` pass-through.
+            if self.config.embedding_dims is not None and "v2" in self.config.model:
+                input_body["dimensions"] = self.config.embedding_dims
 
         body = json.dumps(input_body)
 
diff --git a/tests/embeddings/test_aws_bedrock_embeddings.py b/tests/embeddings/test_aws_bedrock_embeddings.py
index b1924d43f..6bcdbfcf6 100644
--- a/tests/embeddings/test_aws_bedrock_embeddings.py
+++ b/tests/embeddings/test_aws_bedrock_embeddings.py
@@ -1,3 +1,4 @@
+import json
 from unittest.mock import Mock, patch
 
 import pytest
@@ -58,3 +59,49 @@ def test_no_session_token_passes_none(mock_boto3_client):
 
     _, kwargs = mock_boto3_client.call_args
     assert kwargs["aws_session_token"] is None
+
+
+def _captured_request_body(mock_boto3_client, config):
+    """Run a single embed call and return the JSON body sent to invoke_model."""
+    runtime = mock_boto3_client.return_value
+    response_stream = Mock()
+    response_stream.read.return_value = json.dumps({"embedding": [0.0, 0.1, 0.2]}).encode()
+    runtime.invoke_model.return_value = {"body": response_stream}
+
+    embedder = AWSBedrockEmbedding(config)
+    embedder.embed("hello world")
+
+    _, call_kwargs = runtime.invoke_model.call_args
+    return json.loads(call_kwargs["body"])
+
+
+def test_titan_v2_forwards_embedding_dims_as_dimensions(mock_boto3_client):
+    """Titan Text Embeddings V2 must receive embedding_dims as `dimensions`.
+
+    Titan V2 supports an optional output size (256/512/1024); without forwarding
+    it the model returns its default 1024-d vector and ignores the user's request.
+    """
+    with patch("mem0.embeddings.aws_bedrock.os.environ", {}):
+        config = BaseEmbedderConfig(model="amazon.titan-embed-text-v2:0", embedding_dims=512)
+        body = _captured_request_body(mock_boto3_client, config)
+
+    assert body["dimensions"] == 512
+    assert body["inputText"] == "hello world"
+
+
+def test_titan_v2_without_embedding_dims_omits_dimensions(mock_boto3_client):
+    """When embedding_dims is unset, no `dimensions` key is sent (model default)."""
+    with patch("mem0.embeddings.aws_bedrock.os.environ", {}):
+        config = BaseEmbedderConfig(model="amazon.titan-embed-text-v2:0")
+        body = _captured_request_body(mock_boto3_client, config)
+
+    assert "dimensions" not in body
+
+
+def test_titan_v1_ignores_embedding_dims(mock_boto3_client):
+    """Titan V1 has no configurable output size, so `dimensions` must not be sent."""
+    with patch("mem0.embeddings.aws_bedrock.os.environ", {}):
+        config = BaseEmbedderConfig(model="amazon.titan-embed-text-v1", embedding_dims=512)
+        body = _captured_request_body(mock_boto3_client, config)
+
+    assert "dimensions" not in body