fix(python): preserve valid embeddings when an OpenAI batch is rejected (#4005)

Retries individual inputs one at a time when the OpenAI API rejects a
batch with BadRequestError, instead of discarding the whole batch.
Valid inputs now still get embeddings; only the actually-bad input
resolves to None, and the log names the exact failing input.

Fixes #1677

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

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()