fix(bedrock): use dict literal instead of set in AI21 response parse default (#5527)
This commit is contained in:
@@ -409,7 +409,7 @@ class AWSBedrockLLM(LLMBase):
|
|||||||
elif self.provider == "cohere":
|
elif self.provider == "cohere":
|
||||||
return response_json.get("generations", [{"text": ""}])[0].get("text", "")
|
return response_json.get("generations", [{"text": ""}])[0].get("text", "")
|
||||||
elif self.provider == "ai21":
|
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:
|
else:
|
||||||
# Generic parsing - try common response fields
|
# Generic parsing - try common response fields
|
||||||
for field in ["content", "text", "completion", "generation"]:
|
for field in ["content", "text", "completion", "generation"]:
|
||||||
|
|||||||
@@ -424,3 +424,31 @@ class TestMiniMaxProvider:
|
|||||||
assert msg["role"] != "system"
|
assert msg["role"] != "system"
|
||||||
# user message must be present
|
# user message must be present
|
||||||
assert kwargs["messages"][0]["role"] == "user"
|
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"
|
||||||
|
|||||||
Reference in New Issue
Block a user