File tree 2 files changed +22
-1
lines changed
2 files changed +22
-1
lines changed Original file line number Diff line number Diff line change @@ -1354,6 +1354,17 @@ def test_text_generation_model_predict_streaming(self):
1354
1354
):
1355
1355
assert len (response .text ) > 10
1356
1356
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
+
1357
1368
@pytest .mark .parametrize (
1358
1369
"job_spec" ,
1359
1370
[_TEST_PIPELINE_SPEC_JSON , _TEST_PIPELINE_JOB ],
Original file line number Diff line number Diff line change @@ -613,7 +613,17 @@ class TextGenerationResponse:
613
613
safety_attributes : Dict [str , float ] = dataclasses .field (default_factory = dict )
614
614
615
615
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
+ )
617
627
618
628
@property
619
629
def raw_prediction_response (self ) -> aiplatform .models .Prediction :
You can’t perform that action at this time.
0 commit comments