diff --git a/src/celeste/http.py b/src/celeste/http.py index 856b016..204c99a 100644 --- a/src/celeste/http.py +++ b/src/celeste/http.py @@ -3,7 +3,7 @@ import asyncio import json import logging -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any import httpx @@ -16,6 +16,25 @@ MAX_CONNECTIONS = 20 MAX_KEEPALIVE_CONNECTIONS = 10 DEFAULT_TIMEOUT = 180.0 +MAX_RETRIES = 2 +RETRY_BASE_DELAY = 0.5 +RETRYABLE_STATUS = frozenset({408, 429, 500, 502, 503, 504}) + + +async def _retry_request( + send: Callable[[], Awaitable[httpx.Response]], +) -> httpx.Response: + """Retry `send` on transient failures (network errors + retryable status) with backoff, then fail hard.""" + for attempt in range(MAX_RETRIES): + try: + response = await send() + except (httpx.TimeoutException, httpx.NetworkError): + pass # transient — retry after backoff + else: + if response.status_code not in RETRYABLE_STATUS: + return response + await asyncio.sleep(RETRY_BASE_DELAY * 2**attempt) + return await send() class HTTPClient: @@ -81,11 +100,13 @@ async def post( raise ValueError("URL cannot be empty") client = await self._get_client() - return await client.post( - url, - headers=headers, - json=json_body, - timeout=timeout, + return await _retry_request( + lambda: client.post( + url, + headers=headers, + json=json_body, + timeout=timeout, + ) ) async def post_multipart( @@ -116,12 +137,14 @@ async def post_multipart( raise ValueError("URL cannot be empty") client = await self._get_client() - return await client.post( - url, - headers=headers, - files=files, - data=data, - timeout=timeout, + return await _retry_request( + lambda: client.post( + url, + headers=headers, + files=files, + data=data, + timeout=timeout, + ) ) async def get( @@ -150,11 +173,13 @@ async def get( raise ValueError("URL cannot be empty") client = await self._get_client() - return await client.get( - url, - headers=headers or {}, - timeout=timeout, - follow_redirects=follow_redirects, + return await _retry_request( + lambda: client.get( + url, + headers=headers or {}, + timeout=timeout, + follow_redirects=follow_redirects, + ) ) async def stream_post( @@ -285,6 +310,7 @@ def clear_http_clients() -> None: "DEFAULT_TIMEOUT", "MAX_CONNECTIONS", "MAX_KEEPALIVE_CONNECTIONS", + "MAX_RETRIES", "HTTPClient", "clear_http_clients", "close_all_http_clients", diff --git a/src/celeste/structured_outputs.py b/src/celeste/structured_outputs.py index e1f9a02..ab6f584 100644 --- a/src/celeste/structured_outputs.py +++ b/src/celeste/structured_outputs.py @@ -7,7 +7,7 @@ class StrictJsonSchemaGenerator(GenerateJsonSchema): - """Adds additionalProperties: false to all objects (OpenAI, Anthropic, xAI).""" + """Forces additionalProperties: false and all properties required on objects (OpenAI, Anthropic, xAI).""" def generate( self, schema: CoreSchema, mode: JsonSchemaMode = "validation" @@ -21,6 +21,8 @@ def _add_strict_props(self, schema: dict) -> None: return if schema.get("type") == "object" and "additionalProperties" not in schema: schema["additionalProperties"] = False + if "properties" in schema: + schema["required"] = list(schema["properties"]) for key in ("properties", "$defs", "items", "allOf", "anyOf", "oneOf"): if key in schema: if isinstance(schema[key], dict): diff --git a/tests/unit_tests/test_http.py b/tests/unit_tests/test_http.py index 264ea95..2c06bc4 100644 --- a/tests/unit_tests/test_http.py +++ b/tests/unit_tests/test_http.py @@ -8,6 +8,7 @@ from celeste.core import Modality, Provider from celeste.http import ( + MAX_RETRIES, HTTPClient, clear_http_clients, close_all_http_clients, @@ -243,7 +244,7 @@ async def test_get_request_with_all_parameters( async def test_post_propagates_httpx_errors( self, mock_httpx_client: AsyncMock ) -> None: - """POST method must propagate httpx.HTTPError to caller.""" + """POST must retry a transient error, then propagate it after MAX_RETRIES.""" # Arrange http_client = HTTPClient() mock_httpx_client.post.side_effect = httpx.TimeoutException("Request timeout") @@ -251,6 +252,7 @@ async def test_post_propagates_httpx_errors( # Act & Assert with ( patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + patch("celeste.http.asyncio.sleep", new=AsyncMock()), pytest.raises(httpx.TimeoutException, match="Request timeout"), ): await http_client.post( @@ -258,11 +260,12 @@ async def test_post_propagates_httpx_errors( headers={"Authorization": "Bearer test"}, json_body={"key": "value"}, ) + assert mock_httpx_client.post.call_count == MAX_RETRIES + 1 async def test_get_propagates_httpx_errors( self, mock_httpx_client: AsyncMock ) -> None: - """GET method must propagate httpx.HTTPError to caller.""" + """GET must retry a transient error, then propagate it after MAX_RETRIES.""" # Arrange http_client = HTTPClient() mock_httpx_client.get.side_effect = httpx.ConnectError("Connection failed") @@ -270,12 +273,14 @@ async def test_get_propagates_httpx_errors( # Act & Assert with ( patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + patch("celeste.http.asyncio.sleep", new=AsyncMock()), pytest.raises(httpx.ConnectError, match="Connection failed"), ): await http_client.get( url="https://api.example.com/test", headers={"Authorization": "Bearer test"}, ) + assert mock_httpx_client.get.call_count == MAX_RETRIES + 1 async def test_post_uses_default_timeout_when_not_specified( self, mock_httpx_client: AsyncMock @@ -339,6 +344,51 @@ async def test_custom_timeout_passed_to_httpx( assert call_kwargs["timeout"] == custom_timeout +class TestHTTPClientRetry: + """Transient-failure retry behavior (MAX_RETRIES attempts with backoff).""" + + @pytest.mark.parametrize( + ("outcomes", "want_status", "want_calls"), + [ + ([httpx.ReadTimeout("x"), httpx.Response(200)], 200, 2), + ([httpx.Response(503), httpx.Response(200)], 200, 2), + ([httpx.Response(401)], 401, 1), + ([httpx.Response(503)] * (MAX_RETRIES + 1), 503, MAX_RETRIES + 1), + ], + ) + async def test_retry_then_return( + self, + mock_httpx_client: AsyncMock, + outcomes: list[object], + want_status: int, + want_calls: int, + ) -> None: + """Transient errors and retryable statuses are retried; others return immediately.""" + mock_httpx_client.post.side_effect = outcomes + with ( + patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + patch("celeste.http.asyncio.sleep", new=AsyncMock()), + ): + response = await HTTPClient().post( + url="https://x", headers={}, json_body={} + ) + assert response.status_code == want_status + assert mock_httpx_client.post.call_count == want_calls + + async def test_reraises_after_exhausting_retries( + self, mock_httpx_client: AsyncMock + ) -> None: + """A persistent transient error is re-raised after MAX_RETRIES + 1 attempts.""" + mock_httpx_client.post.side_effect = httpx.ConnectError("down") + with ( + patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + patch("celeste.http.asyncio.sleep", new=AsyncMock()), + pytest.raises(httpx.ConnectError), + ): + await HTTPClient().post(url="https://x", headers={}, json_body={}) + assert mock_httpx_client.post.call_count == MAX_RETRIES + 1 + + class TestHTTPClientRegistry: """Test get_http_client registry and singleton behavior.""" diff --git a/tests/unit_tests/test_structured_outputs.py b/tests/unit_tests/test_structured_outputs.py index e2fbe0e..5de0268 100644 --- a/tests/unit_tests/test_structured_outputs.py +++ b/tests/unit_tests/test_structured_outputs.py @@ -24,6 +24,19 @@ class SimpleModel(BaseModel): assert schema.get("additionalProperties") is False + def test_forces_all_properties_required(self) -> None: + """Strict mode requires every property in `required`, even fields with defaults.""" + + class Model(BaseModel): + required_field: str + defaulted_field: str = "x" + + schema = TypeAdapter(Model).json_schema( + schema_generator=StrictJsonSchemaGenerator + ) + + assert set(schema["required"]) == {"required_field", "defaulted_field"} + def test_preserves_existing_additional_properties(self) -> None: """Doesn't override existing additionalProperties when already set.""" generator = StrictJsonSchemaGenerator(ref_template="{model}")