diff --git a/python/python/lancedb/embeddings/openai.py b/python/python/lancedb/embeddings/openai.py index 1edfaa2ca..d4ecbdf3f 100644 --- a/python/python/lancedb/embeddings/openai.py +++ b/python/python/lancedb/embeddings/openai.py @@ -101,24 +101,35 @@ class OpenAIEmbeddings(TextEmbeddingFunction): valid_indices.append(idx) # TODO retry, rate limit, token limit - try: + def embed_batch(texts_batch, indices_batch): kwargs = { - "input": valid_texts, + "input": texts_batch, "model": self.name, } if self.name != "text-embedding-ada-002": kwargs["dimensions"] = self.dim rs = self._openai_client.embeddings.create(**kwargs) - valid_embeddings = { - idx: v.embedding for v, idx in zip(rs.data, valid_indices) - } + return {idx: v.embedding for v, idx in zip(rs.data, indices_batch)} + + def embed_one_or_log(idx, text): + # Retry a single rejected input + try: + return embed_batch([text], [idx]) + except openai.BadRequestError: + logging.exception("Bad request for input at index %d: %r", idx, text) + return {} + + try: + valid_embeddings = embed_batch(valid_texts, valid_indices) except openai.AuthenticationError: logging.error("Authentication failed: Invalid API key provided") raise except openai.BadRequestError: - logging.exception("Bad request: %s", texts) - return [None] * len(texts) + # Retry each input individually + valid_embeddings = {} + for idx, text in zip(valid_indices, valid_texts): + valid_embeddings.update(embed_one_or_log(idx, text)) except Exception: logging.exception("OpenAI embeddings error") raise diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 9434033e1..c42172f77 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -596,6 +596,51 @@ def test_openai_no_retry_on_401(mock_sleep): assert mock_sleep.call_count == 0 +def test_openai_bad_request_partial_batch_recovery(): + # Regression test for issue #1677 + + # Fake exception types, avoids requiring the optional openai/httpx deps + class FakeBadRequestError(Exception): + pass + + class FakeAuthenticationError(Exception): + pass + + fake_openai_module = MagicMock() + fake_openai_module.BadRequestError = FakeBadRequestError + fake_openai_module.AuthenticationError = FakeAuthenticationError + + registry = get_registry() + model = registry.get("openai").create(name="text-embedding-ada-002") + + def fake_create(input, model=None, **kwargs): + # Full batch always rejected + if len(input) > 1: + raise FakeBadRequestError("bad request") + if input[0] == "bad text": + raise FakeBadRequestError("bad request") + embedding = MagicMock() + embedding.embedding = [0.1, 0.2, 0.3] + result = MagicMock() + result.data = [embedding] + return result + + mock_client = MagicMock() + mock_client.embeddings.create.side_effect = fake_create + + with ( + patch( + "lancedb.embeddings.openai.attempt_import_or_raise", + return_value=fake_openai_module, + ), + patch.object(type(model), "_openai_client", new_callable=lambda: mock_client), + ): + embeddings = model.generate_embeddings(["good text", "bad text"]) + + assert embeddings[0] == [0.1, 0.2, 0.3] + assert embeddings[1] is None + + def test_ollama_embeddings_pickle(): """OllamaEmbeddings must pickle even after the cached client is created.""" registry = get_registry()