feat: LLM - Exposed the TextGenerationResponse.raw_prediction_response

When a model returns some extra attributes that are not explicitly exposed by the SDK, these attributes will now be accessible via the `repsonse.raw_prediction_response` attribute. PiperOrigin-RevId: 557276971
googleapis · Aug 15, 2023 · f8f2b9c · f8f2b9c
1 parent e99f366
commit f8f2b9c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/tests/unit/aiplatform/test_language_models.py b/tests/unit/aiplatform/test_language_models.py
@@ -577,6 +577,10 @@ def test_text_generation(self):
             )
 
         assert response.text == _TEST_TEXT_GENERATION_PREDICTION["content"]
+        assert (
+            response.raw_prediction_response.predictions[0]
+            == _TEST_TEXT_GENERATION_PREDICTION
+        )
         assert (
             response.safety_attributes["Violent"]
             == _TEST_TEXT_GENERATION_PREDICTION["safetyAttributes"]["scores"][0]

diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
@@ -222,6 +222,11 @@ class TextGenerationResponse:
     def __repr__(self):
         return self.text
 
+    @property
+    def raw_prediction_response(self) -> aiplatform.models.Prediction:
+        """Raw prediction response."""
+        return self._prediction_response
+
 
 class _TextGenerationModel(_LanguageModel):
     """TextGenerationModel represents a general language model.