Skip to content

Commit 5294972

Browse files
Ark-kuncopybara-github
authored andcommitted
chore: LLM - TextEmbeddingModel - Added argument type checking
PiperOrigin-RevId: 580022424
1 parent d4667f2 commit 5294972

File tree

2 files changed

+8
-0
lines changed

2 files changed

+8
-0
lines changed

tests/unit/aiplatform/test_language_models.py

+4
Original file line numberDiff line numberDiff line change
@@ -3567,6 +3567,10 @@ def test_text_embedding_ga(self):
35673567
assert len(vector) == _TEXT_EMBEDDING_VECTOR_LENGTH
35683568
assert vector == _TEST_TEXT_EMBEDDING_PREDICTION["embeddings"]["values"]
35693569

3570+
# Validating that a single string is not accepted.
3571+
with pytest.raises(TypeError):
3572+
model.get_embeddings("What is life?")
3573+
35703574
def test_batch_prediction(
35713575
self,
35723576
get_endpoint_mock,

vertexai/language_models/_language_models.py

+4
Original file line numberDiff line numberDiff line change
@@ -1520,6 +1520,10 @@ def _prepare_text_embedding_request(
15201520
Returns:
15211521
A `_MultiInstancePredictionRequest` object.
15221522
"""
1523+
if isinstance(texts, str) or not isinstance(texts, Sequence):
1524+
raise TypeError(
1525+
"The `texts` argument must be a list, not a single string."
1526+
)
15231527
instances = []
15241528
for text in texts:
15251529
if isinstance(text, TextEmbeddingInput):

0 commit comments

Comments
 (0)