From a1eefc31bc8aae6d027a80558f12cdc4682f0013 Mon Sep 17 00:00:00 2001
From: Hrushikesh Yadav <136978914+HrushiYadav@users.noreply.github.com>
Date: Mon, 15 Jun 2026 12:09:34 +0530
Subject: [PATCH] fix(bedrock): use dict literal instead of set in AI21
 response parse default (#5527)

---
 mem0/llms/aws_bedrock.py       |  2 +-
 tests/llms/test_aws_bedrock.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/mem0/llms/aws_bedrock.py b/mem0/llms/aws_bedrock.py
index aa2c08101..d749ebb27 100644
--- a/mem0/llms/aws_bedrock.py
+++ b/mem0/llms/aws_bedrock.py
@@ -409,7 +409,7 @@ class AWSBedrockLLM(LLMBase):
             elif self.provider == "cohere":
                 return response_json.get("generations", [{"text": ""}])[0].get("text", "")
             elif self.provider == "ai21":
-                return response_json.get("completions", [{"data", {"text": ""}}])[0].get("data", {}).get("text", "")
+                return response_json.get("completions", [{"data": {"text": ""}}])[0].get("data", {}).get("text", "")
             else:
                 # Generic parsing - try common response fields
                 for field in ["content", "text", "completion", "generation"]:
diff --git a/tests/llms/test_aws_bedrock.py b/tests/llms/test_aws_bedrock.py
index 029c4a0b0..059cce650 100644
--- a/tests/llms/test_aws_bedrock.py
+++ b/tests/llms/test_aws_bedrock.py
@@ -424,3 +424,31 @@ class TestMiniMaxProvider:
             assert msg["role"] != "system"
         # user message must be present
         assert kwargs["messages"][0]["role"] == "user"
+
+
+# ---------------------------------------------------------------------------
+# _parse_response — legacy InvokeModel provider-specific parsing
+# ---------------------------------------------------------------------------
+
+class TestParseResponseLegacy:
+    def test_ai21_missing_completions_returns_empty(self, mock_boto3):
+        """When AI21 response lacks 'completions', the fallback default must
+        be a valid dict (not a set literal), returning empty string."""
+        llm = _make_llm("ai21.j2-mid-v1", mock_boto3)
+        import io
+        import json
+        body = io.BytesIO(json.dumps({"not_completions": True}).encode())
+        response = {"body": body}
+        result = llm._parse_response(response, tools=None)
+        assert result == ""
+
+    def test_ai21_normal_response(self, mock_boto3):
+        llm = _make_llm("ai21.j2-mid-v1", mock_boto3)
+        import io
+        import json
+        body = io.BytesIO(json.dumps({
+            "completions": [{"data": {"text": "hello from ai21"}}]
+        }).encode())
+        response = {"body": body}
+        result = llm._parse_response(response, tools=None)
+        assert result == "hello from ai21"