diff --git a/tests/recipes/wrangles/test_create.py b/tests/recipes/wrangles/test_create.py index 1f141f43..9c3e2ae8 100644 --- a/tests/recipes/wrangles/test_create.py +++ b/tests/recipes/wrangles/test_create.py @@ -1,8 +1,11 @@ import wrangles +import base64 +import json import pandas as pd import pytest import numpy as np import os +import requests import uuid import random from datetime import datetime @@ -1202,6 +1205,245 @@ class TestCreateEmbeddings: """ Test create.embeddings """ + @staticmethod + def _response(status_code, body, headers=None): + response = requests.Response() + response.status_code = status_code + response._content = json.dumps(body).encode("utf-8") + response.headers = headers or {} + return response + + @classmethod + def _openai_embedding_response(cls, values=None): + if values is None: + values = [0.1, 0.2, 0.3] + embedding = np.asarray(values, dtype=np.float32) + return cls._response( + 200, + { + "data": [{ + "embedding": base64.b64encode(embedding.tobytes()).decode("ascii"), + "index": 0, + }] + }, + ) + + def test_create_embeddings_retries_falsey_http_response_and_forwards_timeout(self): + """Retry a real 429 Response, which is false-valued in requests.""" + rate_limited = self._response( + 429, + {"error": {"message": "Rate limit exceeded"}}, + {"retry-after": "0.25"}, + ) + responses = [rate_limited, self._openai_embedding_response()] + + with ( + patch("wrangles.openai._requests.post", side_effect=responses) as mock_post, + patch("wrangles.openai._openai_responses._sleep_for_retry") as mock_sleep, + ): + result = wrangles.openai.embeddings( + ["test"], + api_key="fake-key", + retries=1, + timeout=7.5, + ) + + assert np.allclose(result[0], [0.1, 0.2, 0.3]) + assert mock_post.call_count == 2 + assert all(call.kwargs["timeout"] == 7.5 for call in mock_post.call_args_list) + assert mock_sleep.call_count == 1 + assert mock_sleep.call_args.args[0]["status_code"] == 429 + assert mock_sleep.call_args.args[0]["retry_after"] == 0.25 + + @pytest.mark.parametrize( + "transport_error", + [ + requests.exceptions.Timeout("timed out"), + requests.exceptions.ConnectionError("connection lost"), + requests.exceptions.SSLError("TLS failed"), + requests.exceptions.ProxyError("proxy failed"), + requests.exceptions.ChunkedEncodingError("response truncated"), + requests.exceptions.ContentDecodingError("response decoding failed"), + ], + ids=["timeout", "connection", "ssl", "proxy", "chunked", "content-decoding"], + ) + def test_create_embeddings_retries_transient_transport_errors(self, transport_error): + with ( + patch( + "wrangles.openai._requests.post", + side_effect=[transport_error, self._openai_embedding_response()], + ) as mock_post, + patch("wrangles.openai._openai_responses._sleep_for_retry") as mock_sleep, + ): + result = wrangles.openai.embeddings( + ["test"], + api_key="fake-key", + retries=1, + ) + + assert np.allclose(result[0], [0.1, 0.2, 0.3]) + assert mock_post.call_count == 2 + assert mock_sleep.call_count == 1 + + def test_create_embeddings_does_not_retry_permanent_request_error(self): + with ( + patch( + "wrangles.openai._requests.post", + side_effect=requests.exceptions.InvalidURL("invalid URL"), + ) as mock_post, + patch("wrangles.openai._openai_responses._sleep_for_retry") as mock_sleep, + ): + with pytest.raises(RuntimeError, match=r"after 1 attempt\(s\)") as exc_info: + wrangles.openai.embeddings( + ["test"], + api_key="fake-key", + retries=2, + ) + + assert isinstance(exc_info.value.__cause__, requests.exceptions.InvalidURL) + assert mock_post.call_count == 1 + mock_sleep.assert_not_called() + + def test_create_embeddings_does_not_retry_permanent_http_error(self): + bad_request = self._response( + 400, + {"error": {"message": "Invalid embedding request"}}, + ) + + with ( + patch("wrangles.openai._requests.post", return_value=bad_request) as mock_post, + patch("wrangles.openai._openai_responses._sleep_for_retry") as mock_sleep, + ): + with pytest.raises(RuntimeError, match=r"after 1 attempt\(s\)"): + wrangles.openai.embeddings( + ["test"], + api_key="fake-key", + retries=2, + ) + + assert mock_post.call_count == 1 + mock_sleep.assert_not_called() + + def test_create_embeddings_invalid_api_key_fails_fast(self): + unauthorized = self._response( + 401, + {"error": {"message": "Incorrect API key provided"}}, + ) + + with patch("wrangles.openai._requests.post", return_value=unauthorized) as mock_post: + with pytest.raises(ValueError, match="API Key provided is missing or invalid"): + wrangles.openai.embeddings( + ["test"], + api_key="invalid-key", + retries=2, + ) + + assert mock_post.call_count == 1 + + @pytest.mark.parametrize("retries, expected_attempts", [(0, 1), (2, 3)]) + def test_create_embeddings_transport_retry_exhaustion(self, retries, expected_attempts): + with ( + patch( + "wrangles.openai._requests.post", + side_effect=requests.exceptions.ConnectionError("connection lost"), + ) as mock_post, + patch("wrangles.openai._openai_responses._sleep_for_retry") as mock_sleep, + ): + with pytest.raises( + RuntimeError, + match=rf"after {expected_attempts} attempt\(s\)", + ): + wrangles.openai.embeddings( + ["test"], + api_key="fake-key", + retries=retries, + ) + + assert mock_post.call_count == expected_attempts + assert mock_sleep.call_count == expected_attempts - 1 + + def test_create_embeddings_model_not_found_fails_fast_for_openai(self): + model_not_found = self._response( + 404, + { + "error": { + "message": "The requested model does not exist", + "code": "model_not_found", + } + }, + ) + + with patch("wrangles.openai._requests.post", return_value=model_not_found) as mock_post: + with pytest.raises(ValueError, match="does not exist or is not accessible"): + wrangles.openai.embeddings( + ["test"], + api_key="fake-key", + model="missing-embedding-model", + retries=2, + ) + + assert mock_post.call_count == 1 + + def test_create_embeddings_jina_errors_do_not_use_openai_fatal_error(self): + model_not_found = self._response( + 404, + { + "error": { + "message": "The requested Jina model does not exist", + "code": "model_not_found", + } + }, + ) + + with ( + patch("wrangles.openai._requests.post", return_value=model_not_found), + patch("wrangles.openai._openai_responses._raise_for_fatal_error") as mock_raise, + ): + with pytest.raises(RuntimeError, match=r"after 1 attempt\(s\)"): + wrangles.openai.embeddings( + ["test"], + api_key="fake-key", + model="missing-jina-model", + provider="jina", + retries=2, + ) + + mock_raise.assert_not_called() + + def test_create_embeddings_recipe_forwards_timeout(self): + with patch( + "wrangles.openai._requests.post", + return_value=self._openai_embedding_response(), + ) as mock_post: + df = wrangles.recipe.run( + """ + wrangles: + - create.embeddings: + input: text + output: embedding + api_key: fake-key + timeout: 6.5 + """, + dataframe=pd.DataFrame({"text": ["test"]}), + ) + + assert len(df["embedding"][0]) == 3 + assert mock_post.call_args.kwargs["timeout"] == 6.5 + + @pytest.mark.parametrize("retries", [True, -1, 1.5, "1"]) + def test_create_embeddings_rejects_invalid_retries(self, retries): + with pytest.raises(ValueError, match="retries must be a non-negative integer"): + wrangles.openai.embeddings(["test"], api_key="fake-key", retries=retries) + + @pytest.mark.parametrize( + "timeout", + [True, 0, -1, float("inf"), float("nan"), "30"], + ids=["boolean", "zero", "negative", "infinite", "nan", "string"], + ) + def test_create_embeddings_rejects_invalid_timeout(self, timeout): + with pytest.raises(ValueError, match="timeout must be a positive finite number"): + wrangles.openai.embeddings(["test"], api_key="fake-key", timeout=timeout) + def test_create_embeddings(self): """ Test generating openai embeddings @@ -2017,4 +2259,4 @@ def test_create_hash_empty(self): """, dataframe=pd.DataFrame({'text': []}) ) - assert df.empty and list(df.columns) == ['text', 'hash'] \ No newline at end of file + assert df.empty and list(df.columns) == ['text', 'hash'] diff --git a/wrangles/openai.py b/wrangles/openai.py index d0c25736..489d1b75 100644 --- a/wrangles/openai.py +++ b/wrangles/openai.py @@ -21,6 +21,12 @@ "openai": "https://api.openai.com/v1/embeddings", "jina": "https://api.jina.ai/v1/embeddings", } +_RETRYABLE_EMBEDDING_TRANSPORT_ERRORS = ( + _requests.exceptions.Timeout, + _requests.exceptions.ConnectionError, + _requests.exceptions.ChunkedEncodingError, + _requests.exceptions.ContentDecodingError, +) def format_input_data(data: any) -> str: @@ -192,7 +198,8 @@ def _embedding_thread( retries: int = 0, request_params: dict = None, precision: str = "float32", - provider: str = "openai" + provider: str = "openai", + timeout: float = 30, ): """ Get embeddings @@ -205,6 +212,7 @@ def _embedding_thread( :param request_params: Additional request parameters to pass to the backend. :param precision: The precision of the embeddings. Default is float32. :param provider: The embedding provider to use. Default is openai. + :param timeout: Per-attempt request timeout in seconds. Default is 30. """ if request_params is None: request_params = {} @@ -228,53 +236,48 @@ def _embedding_thread( _logging.debug(f": Computing embeddings :: model :: {model}, record_count :: {len(input_list)}") response = None + transport_error = None backoff_time = 1 - while (retries + 1): + for attempt in range(retries + 1): + response = None + transport_error = None try: response = _requests.post( url=url, - headers={ - "Authorization": f"Bearer {api_key}" - }, + headers={"Authorization": f"Bearer {api_key}"}, json=request_body, - timeout=30 + timeout=timeout, ) - except Exception: - pass + except _requests.exceptions.RequestException as exc: + transport_error = exc - if response and response.ok: + if response is not None and response.ok: break - else: - if response is not None and response.status_code == 401: - raise ValueError("API Key provided is missing or invalid.") - try: - error_message = response.json().get('error').get('message') - except Exception: - error_message = "" - # Raise errors for fatal errors rather than continuing - if error_message: - if "Incorrect API key" in error_message: - raise ValueError("API Key provided is missing or invalid.") - context = _openai_responses._response_context( - response, - endpoint="embeddings", - model=model, - ) if response else {} - if retries == 0 or not _openai_responses._should_retry(context): - if response: - _openai_responses._log_api_error(context, final=True) - break - if response: - _openai_responses._log_api_error(context, final=False) - retries -= 1 - if response and not response.ok: - _openai_responses._sleep_for_retry(context, backoff_time) - else: - _time.sleep(backoff_time) + context = _openai_responses._response_context( + response, endpoint="embeddings", model=model, attempt=attempt + 1, + ) + if transport_error is not None: + context["message"] = f"{type(transport_error).__name__}: {transport_error}" + if response is not None and ( + response.status_code == 401 or "Incorrect API key" in context.get("message", "") + ): + raise ValueError("API Key provided is missing or invalid.") + if provider == "openai": + _openai_responses._raise_for_fatal_error(context) + + retryable = ( + isinstance(transport_error, _RETRYABLE_EMBEDDING_TRANSPORT_ERRORS) + or _openai_responses._should_retry(context) + ) + final = attempt == retries or not retryable + _openai_responses._log_api_error(context, final=final) + if final: + break + _openai_responses._sleep_for_retry(context, backoff_time) backoff_time *= 2 - if response and response.ok: + if response is not None and response.ok: if provider == "jina": try: return [ @@ -293,19 +296,11 @@ def _embedding_thread( for row in response.json()['data'] ] else: - try: - error_msg = _openai_responses._error_message( - _openai_responses._response_context( - response, - endpoint="embeddings", - model=model, - ) - ) - except Exception: - error_msg = 'Unknown error' + error_msg = _openai_responses._error_message(context) raise RuntimeError( - f"Failed to get embeddings: {error_msg}. Consider raising the number of retries." - ) + f"Failed to get embeddings after {attempt + 1} attempt(s): {error_msg}" + ) from transport_error + def embeddings( input_list, @@ -318,6 +313,7 @@ def embeddings( precision: str = "float32", provider: str = None, task: str = None, + timeout: float = 30, **kwargs ) -> list: """ @@ -334,8 +330,8 @@ def embeddings( :param batch_size: (Optional, default 100) The number of rows to submit per individual request. :param threads: (Optional, default 10) The number of requests to submit in parallel. \ Each request contains the number of rows set as batch_size. - :param retries: The number of times to retry. This will exponentially \ - backoff to assist with rate limiting + :param retries: Additional attempts after transient HTTP or transport failures. + Defaults to 0. Uses exponential backoff and Retry-After; permanent errors fail immediately. :param url: The endpoint to send requests to. Defaults to the standard endpoint for \ the resolved provider. Setting a Jina URL without an explicit provider will \ automatically use Jina's request/response format. @@ -346,8 +342,17 @@ def embeddings( Pass both only when using a custom endpoint with a non-default provider's API format. :param task: (Optional, Jina only) The task type for the embedding model. \ Valid values: retrieval.query, retrieval.passage, text-matching, classification, separation. + :param timeout: Per-attempt request timeout in seconds. Default is 30. :return: A list of embeddings corresponding to the input """ + if not isinstance(retries, int) or isinstance(retries, bool) or retries < 0: + raise ValueError("retries must be a non-negative integer.") + if ( + not isinstance(timeout, (int, float)) or isinstance(timeout, bool) + or not _np.isfinite(timeout) or timeout <= 0 + ): + raise ValueError("timeout must be a positive finite number of seconds.") + # Infer provider from URL when not explicitly set if provider is None: if "jina.ai" in url: @@ -397,7 +402,8 @@ def embeddings( [retries] * len(batches), [kwargs] * len(batches), [precision] * len(batches), - [provider] * len(batches) + [provider] * len(batches), + [timeout] * len(batches), )) results = list(_chain.from_iterable(results)) diff --git a/wrangles/recipe_wrangles/create.py b/wrangles/recipe_wrangles/create.py index c61c566e..ca7011b9 100644 --- a/wrangles/recipe_wrangles/create.py +++ b/wrangles/recipe_wrangles/create.py @@ -211,6 +211,7 @@ def embeddings( precision: str = "float32", provider: str = None, task: str = None, + timeout: float = 30, **kwargs ) -> _pd.DataFrame: """ @@ -255,9 +256,18 @@ def embeddings( - python list retries: type: integer + minimum: 0 description: >- - The number of times to retry if the request fails. - This will apply exponential backoff to help with rate limiting. + Additional attempts after transient transport or HTTP errors. + Defaults to 0. Retries use exponential backoff and respect Retry-After. + Permanent errors fail immediately. + timeout: + type: number + exclusiveMinimum: 0 + default: 30 + description: >- + Request timeout in seconds for each attempt. Defaults to 30. + Each retry receives the full timeout; this is not a total batch deadline. provider: type: string description: >- @@ -322,6 +332,7 @@ def embeddings( precision, provider=provider, task=task, + timeout=timeout, **kwargs ) @@ -564,4 +575,4 @@ def hash(df: _pd.DataFrame, input: _Union[str, int, list], output: _Union[str, l hash_fn = getattr(_hashlib, method) df[out_col] = [hash_fn(str(x).encode('utf-8')).hexdigest() for x in df[in_col]] - return df \ No newline at end of file + return df