diff --git a/.sampo/changesets/prompts-refetch-cooldown.md b/.sampo/changesets/prompts-refetch-cooldown.md new file mode 100644 index 000000000..f8abad701 --- /dev/null +++ b/.sampo/changesets/prompts-refetch-cooldown.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +After a failed prompt refetch, the SDK now serves the stale cached prompt for a cooldown period (60 seconds by default) instead of retrying the network on every `prompts.get()` call. When the failure is a 429, the cooldown follows the `Retry-After` the server sends, capped at one hour. This keeps a rate-limited client from holding itself against the limit. Matches the behavior the JavaScript SDK already has. diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index d4bc2cbe4..2cbae9db4 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -6,6 +6,7 @@ import copy import logging +import math import re import time import urllib.parse @@ -13,6 +14,7 @@ from dataclasses import dataclass from typing import Any, Dict, List, Literal, Optional, Union, overload +from posthog.capture_v1 import _parse_retry_after from posthog.request import USER_AGENT, _get_session from posthog.utils import remove_trailing_slash @@ -24,6 +26,10 @@ # default page size covers 10,000 prompts; past that get_all raises rather than # returning a truncated result. _MAX_PROMPT_LIST_PAGES = 100 +# After a failed refetch the stale entry is served for this long before the next network attempt. +# The server's tightest prompt limit is per-minute, so a minute lets the bucket refill. +DEFAULT_REFETCH_COOLDOWN_SECONDS = 60 +MAX_REFETCH_COOLDOWN_SECONDS = 3600 PromptVariables = Dict[str, Union[str, int, float, bool]] PromptCacheKey = tuple[str, Optional[int], Optional[str]] @@ -70,6 +76,7 @@ def __init__( self.version = version self.label = label self.config = config + self.retry_not_before: Optional[float] = None def _cache_key( @@ -153,6 +160,22 @@ def _is_same_origin(url: str, host: str) -> bool: return (parsed.scheme, parsed.netloc) == (expected.scheme, expected.netloc) +def _parse_retry_after_seconds(value: Optional[str]) -> Optional[float]: + """Read a Retry-After header (delta-seconds or HTTP-date) as a bounded cooldown.""" + seconds = _parse_retry_after(value) + if seconds is None or not math.isfinite(seconds) or seconds <= 0: + return None + return min(seconds, MAX_REFETCH_COOLDOWN_SECONDS) + + +class PromptFetchError(Exception): + """Carries the server's own cooldown so a rate-limited client waits as long as it was told to.""" + + def __init__(self, message: str, retry_after_seconds: Optional[float] = None): + super().__init__(message) + self.retry_after_seconds = retry_after_seconds + + def _authentication_error(reference: str) -> Exception: return Exception( f"[PostHog Prompts] Authentication failed for {reference}. " @@ -525,6 +548,24 @@ def _get_internal( config=copy.deepcopy(cached.config), ) + # A failed refetch left this entry in cooldown. Serving it keeps one + # throttled client from turning every later get() into another + # request, which is what holds it against the limit. A zero TTL is + # an explicit request to refetch on every read, so it opts out. + if ( + ttl > 0 + and cached.retry_not_before is not None + and now < cached.retry_not_before + ): + return PromptResult( + source="stale_cache", + prompt=cached.prompt, + name=cached.name, + version=cached.version, + label=cached.label, + config=copy.deepcopy(cached.config), + ) + # Try to fetch from API try: data = self._fetch_prompt_from_api(name, version, label) @@ -568,6 +609,15 @@ def _get_internal( prompt_reference = _prompt_reference(name, version, label) # Return stale cache (with warning) if cached is not None: + cooldown_seconds: float = DEFAULT_REFETCH_COOLDOWN_SECONDS + if ( + isinstance(error, PromptFetchError) + and error.retry_after_seconds is not None + ): + cooldown_seconds = error.retry_after_seconds + cached.retry_not_before = max( + cached.retry_not_before or 0, time.time() + cooldown_seconds + ) log.warning( "[PostHog Prompts] Failed to fetch %s, using stale cache: %s", prompt_reference, @@ -723,8 +773,11 @@ def _fetch_prompt_list_from_api(self, label: Optional[str]) -> List[Dict[str, An raise _authentication_error(reference) if response.status_code == 403: raise _access_denied_error(reference) - raise Exception( - f"[PostHog Prompts] Failed to fetch {reference}: HTTP {response.status_code}" + raise PromptFetchError( + f"[PostHog Prompts] Failed to fetch {reference}: HTTP {response.status_code}", + _parse_retry_after_seconds(response.headers.get("Retry-After")) + if response.status_code == 429 + else None, ) try: @@ -811,8 +864,11 @@ def _fetch_prompt_from_api( if response.status_code == 403: raise _access_denied_error(prompt_reference) - raise Exception( - f"[PostHog Prompts] Failed to fetch {prompt_title}: HTTP {response.status_code}" + raise PromptFetchError( + f"[PostHog Prompts] Failed to fetch {prompt_title}: HTTP {response.status_code}", + _parse_retry_after_seconds(response.headers.get("Retry-After")) + if response.status_code == 429 + else None, ) try: diff --git a/posthog/test/ai/test_prompts.py b/posthog/test/ai/test_prompts.py index 22460d513..4e124f16d 100644 --- a/posthog/test/ai/test_prompts.py +++ b/posthog/test/ai/test_prompts.py @@ -10,10 +10,11 @@ class MockResponse: """Mock HTTP response for testing.""" - def __init__(self, json_data=None, status_code=200, ok=True): + def __init__(self, json_data=None, status_code=200, ok=True, headers=None): self._json_data = json_data self.status_code = status_code self.ok = ok + self.headers = headers or {} def json(self): if self._json_data is None: @@ -324,6 +325,63 @@ def test_use_stale_cache_on_fetch_failure_with_warning( warning_call = mock_log.warning.call_args self.assertIn("using stale cache", warning_call[0][0]) + @parameterized.expand( + [ + # The Retry-After cooldown (300s) outlives the 60s default, so the + # holds after second 1460 prove the server's value governs. + ("server_error", 500, None, 60.0), + ("rate_limited_retry_after", 429, {"Retry-After": "300"}, 300.0), + ] + ) + @patch("posthog.ai.prompts._get_session") + @patch("posthog.ai.prompts.time.time") + def test_hold_a_cooldown_after_each_failed_refetch( + self, _scenario, status, headers, cooldown, mock_time, mock_get_session + ): + # Without the cooldown, one throttled client turns every later get() + # into another network request until one succeeds. + mock_get = mock_get_session.return_value.get + mock_get.side_effect = [ + MockResponse(json_data=self.mock_prompt_response), + MockResponse(status_code=status, ok=False, headers=headers), + MockResponse(status_code=status, ok=False, headers=headers), + MockResponse(json_data=self.mock_prompt_response), + ] + mock_time.return_value = 1000.0 + + prompts = Prompts(self.create_mock_posthog()) + prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=False) + + # Past TTL: the refetch fails, stale cache is served, a cooldown starts. + mock_time.return_value = 1400.0 + result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True) + self.assertEqual(result.source, "stale_cache") + self.assertEqual(mock_get.call_count, 2) + + # Within the cooldown: stale cache again, no network attempt. + mock_time.return_value = 1400.0 + cooldown - 1 + result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True) + self.assertEqual(result.source, "stale_cache") + self.assertEqual(mock_get.call_count, 2) + + # Past the cooldown: the retry fails too and a new cooldown starts. + mock_time.return_value = 1400.0 + cooldown + 1 + result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True) + self.assertEqual(result.source, "stale_cache") + self.assertEqual(mock_get.call_count, 3) + + # Within the second cooldown: no network attempt. + mock_time.return_value = 1400.0 + 2 * cooldown + result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True) + self.assertEqual(result.source, "stale_cache") + self.assertEqual(mock_get.call_count, 3) + + # Past the second cooldown: the network is retried and the cache refreshed. + mock_time.return_value = 1400.0 + 2 * cooldown + 2 + result = prompts.get("test-prompt", cache_ttl_seconds=300, with_metadata=True) + self.assertEqual(result.source, "api") + self.assertEqual(mock_get.call_count, 4) + @patch("posthog.ai.prompts._get_session") @patch("posthog.ai.prompts.log") def test_use_fallback_when_no_cache_and_fetch_fails_with_warning( @@ -548,6 +606,18 @@ def test_default_cache_ttl_seconds_zero_disables_caching( prompts.get("test-prompt", with_metadata=False) self.assertEqual(mock_get.call_count, 2) + # A failed refetch must not start a cooldown here: a zero TTL is an + # explicit request to refetch on every read. + mock_get.side_effect = [ + MockResponse(status_code=500, ok=False), + MockResponse(json_data=self.mock_prompt_response), + ] + result = prompts.get("test-prompt", with_metadata=True) + self.assertEqual(result.source, "stale_cache") + result = prompts.get("test-prompt", with_metadata=True) + self.assertEqual(result.source, "api") + self.assertEqual(mock_get.call_count, 4) + @patch("posthog.ai.prompts._get_session") def test_url_encode_prompt_names_with_special_characters(self, mock_get_session): """Should URL-encode prompt names with special characters.""" diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 6599ea487..36142dac0 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -449,9 +449,13 @@ attribute posthog.ai.prompts.CachedPrompt.fetched_at = fetched_at attribute posthog.ai.prompts.CachedPrompt.label = label attribute posthog.ai.prompts.CachedPrompt.name = name attribute posthog.ai.prompts.CachedPrompt.prompt = prompt +attribute posthog.ai.prompts.CachedPrompt.retry_not_before: Optional[float] = None attribute posthog.ai.prompts.CachedPrompt.version = version attribute posthog.ai.prompts.DEFAULT_CACHE_TTL_SECONDS = 300 +attribute posthog.ai.prompts.DEFAULT_REFETCH_COOLDOWN_SECONDS = 60 +attribute posthog.ai.prompts.MAX_REFETCH_COOLDOWN_SECONDS = 3600 attribute posthog.ai.prompts.PromptCacheKey = tuple[str, Optional[int], Optional[str]] +attribute posthog.ai.prompts.PromptFetchError.retry_after_seconds = retry_after_seconds attribute posthog.ai.prompts.PromptResult.config: Optional[Dict[str, Any]] = None attribute posthog.ai.prompts.PromptResult.label: Optional[str] = None attribute posthog.ai.prompts.PromptResult.name: Optional[str] = None @@ -990,6 +994,7 @@ class posthog.ai.openai_agents.processor.PostHogTracingProcessor(client: Optiona class posthog.ai.otel.exporter.PostHogTraceExporter(api_key: str, host: str = DEFAULT_HOST) class posthog.ai.otel.processor.PostHogSpanProcessor(api_key: str, host: str = DEFAULT_HOST) class posthog.ai.prompts.CachedPrompt(prompt: str, fetched_at: float, name: str, version: int, label: Optional[str] = None, config: Optional[Dict[str, Any]] = None) +class posthog.ai.prompts.PromptFetchError(message: str, retry_after_seconds: Optional[float] = None) class posthog.ai.prompts.PromptResult(source: PromptSource, prompt: str, name: Optional[str] = None, version: Optional[int] = None, label: Optional[str] = None, config: Optional[Dict[str, Any]] = None) class posthog.ai.prompts.Prompts(posthog: Optional[Any] = None, *, personal_api_key: Optional[str] = None, project_api_key: Optional[str] = None, host: Optional[str] = None, default_cache_ttl_seconds: Optional[int] = None, capture_errors: bool = False) class posthog.ai.stream.AsyncStreamWrapper(generator: AsyncGenerator[T, None], stream: Optional[Any] = None)