fix: LLM - Fixed the handling of the TextEmbeddingInput.task_type parameter.

vertex-sdk-bot · copybara-github · commit 2e3090b9f1f8 · 2023-08-28T16:49:16.000-07:00
PiperOrigin-RevId: 560854460
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):

Original file line number	Diff line number	Diff line change
`@@ -2359,12 +2359,12 @@ def test_text_embedding(self):`
`2359`	`2359`	`{"content": "What is life?"},`
`2360`	`2360`	`{`
`2361`	`2361`	`"content": "Foo",`
`2362`		`- "taskType": "RETRIEVAL_DOCUMENT",`
	`2362`	`+ "task_type": "RETRIEVAL_DOCUMENT",`
`2363`	`2363`	`"title": "Bar",`
`2364`	`2364`	`},`
`2365`	`2365`	`{`
`2366`	`2366`	`"content": "Baz",`
`2367`		`- "taskType": "CLASSIFICATION",`
	`2367`	`+ "task_type": "CLASSIFICATION",`
`2368`	`2368`	`},`
`2369`	`2369`	`]`
`2370`	`2370`	`prediction_parameters = mock_predict.call_args[1]["parameters"]`