fix: LLM - Fixed the handling of the TextEmbeddingInput.task_type p…

…arameter. PiperOrigin-RevId: 560854460
googleapis · Aug 28, 2023 · 2e3090b · 2e3090b
1 parent 4f66472
commit 2e3090b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/tests/unit/aiplatform/test_language_models.py b/tests/unit/aiplatform/test_language_models.py
@@ -2359,12 +2359,12 @@ def test_text_embedding(self):
                 {"content": "What is life?"},
                 {
                     "content": "Foo",
-                    "taskType": "RETRIEVAL_DOCUMENT",
+                    "task_type": "RETRIEVAL_DOCUMENT",
                     "title": "Bar",
                 },
                 {
                     "content": "Baz",
-                    "taskType": "CLASSIFICATION",
+                    "task_type": "CLASSIFICATION",
                 },
             ]
             prediction_parameters = mock_predict.call_args[1]["parameters"]

diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
@@ -1084,7 +1084,7 @@ def get_embeddings(self,
             if isinstance(text, TextEmbeddingInput):
                 instance = {"content": text.text}
                 if text.task_type:
-                    instance["taskType"] = text.task_type
+                    instance["task_type"] = text.task_type
                 if text.title:
                     instance["title"] = text.title
             elif isinstance(text, str):