Skip to content

Commit 222f222

Browse files
Ark-kuncopybara-github
authored andcommitted
feat: LLM - Improved representation for blocked responses
PiperOrigin-RevId: 569046322
1 parent c22220e commit 222f222

File tree

2 files changed

+22
-1
lines changed

2 files changed

+22
-1
lines changed

tests/unit/aiplatform/test_language_models.py

+11
Original file line numberDiff line numberDiff line change
@@ -1354,6 +1354,17 @@ def test_text_generation_model_predict_streaming(self):
13541354
):
13551355
assert len(response.text) > 10
13561356

1357+
def test_text_generation_response_repr(self):
1358+
response = language_models.TextGenerationResponse(
1359+
text="",
1360+
is_blocked=True,
1361+
safety_attributes={"Violent": 0.1},
1362+
_prediction_response=None,
1363+
)
1364+
response_repr = repr(response)
1365+
assert "blocked" in response_repr
1366+
assert "Violent" in response_repr
1367+
13571368
@pytest.mark.parametrize(
13581369
"job_spec",
13591370
[_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_JOB],

vertexai/language_models/_language_models.py

+11-1
Original file line numberDiff line numberDiff line change
@@ -613,7 +613,17 @@ class TextGenerationResponse:
613613
safety_attributes: Dict[str, float] = dataclasses.field(default_factory=dict)
614614

615615
def __repr__(self):
616-
return self.text
616+
if self.text:
617+
return self.text
618+
else:
619+
# Falling back to the full representation
620+
return (
621+
"TextGenerationResponse("
622+
f"text={self.text!r}"
623+
f", is_blocked={self.is_blocked!r}"
624+
f", safety_attributes={self.safety_attributes!r}"
625+
")"
626+
)
617627

618628
@property
619629
def raw_prediction_response(self) -> aiplatform.models.Prediction:

0 commit comments

Comments
 (0)