diff --git a/.sampo/changesets/prompts-single-flight.md b/.sampo/changesets/prompts-single-flight.md new file mode 100644 index 000000000..6051ac9b1 --- /dev/null +++ b/.sampo/changesets/prompts-single-flight.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +`Prompts.get()` now makes one request when several threads ask for the same prompt at once, instead of one request per thread. The other threads wait for that result. With `capture_errors=True`, a shared failure is reported once instead of once per thread. diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index 2cbae9db4..3a785f0ce 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -7,10 +7,13 @@ import copy import logging import math +import os import re +import threading import time import urllib.parse import warnings +import weakref from dataclasses import dataclass from typing import Any, Dict, List, Literal, Optional, Union, overload @@ -37,6 +40,33 @@ PromptSource = Literal["api", "cache", "stale_cache", "code_fallback"] +class _InFlightFetch: + """One network request that several callers wait on.""" + + def __init__(self) -> None: + self.done = threading.Event() + self.result: Optional[CachedPrompt] = None + self.error: Optional[BaseException] = None + + +# Every live Prompts instance, so one fork hook can reset them all. Registered +# once here, because os.register_at_fork keeps each callback for the process +# lifetime and an app that builds many instances would accumulate them. +_instances: "weakref.WeakSet[Prompts]" = weakref.WeakSet() + + +def _reset_in_flight_after_fork() -> None: + # The child has no copy of the thread that was fetching, so an inherited + # in-flight entry would hold every later caller for that key. + for prompts in list(_instances): + prompts._in_flight = {} + prompts._in_flight_lock = threading.Lock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_in_flight_after_fork) + + @dataclass(frozen=True) class PromptResult: """Result of a prompt fetch with metadata about its source. @@ -270,6 +300,12 @@ def __init__( else DEFAULT_CACHE_TTL_SECONDS ) self._cache: Dict[PromptCacheKey, CachedPrompt] = {} + # One request per cache key at a time. Threads that arrive while it is in + # flight wait for that result, so a worker with many concurrent callers does + # not turn every cache miss or expiry into one request per caller. + self._in_flight: Dict[PromptCacheKey, _InFlightFetch] = {} + self._in_flight_lock = threading.Lock() + _instances.add(self) self._has_warned_deprecation = False self._client = posthog self._capture_errors = capture_errors @@ -568,44 +604,18 @@ def _get_internal( # Try to fetch from API try: - data = self._fetch_prompt_from_api(name, version, label) - - # An older PostHog server ignores the label param and returns the latest - # version with no label field — surface that instead of failing silently. - if label is not None and data.get("label") != label: - log.warning( - "[PostHog Prompts] Requested label %r for prompt %r but the server " - "resolved %r. It may not support prompt labels yet and returned the " - "latest version instead.", - label, - name, - data.get("label"), - ) - - config = _extract_config(data) - - # Update cache - self._cache[cache_key] = CachedPrompt( - prompt=data["prompt"], - fetched_at=time.time(), - name=data["name"], - version=data["version"], - label=data.get("label"), - config=config, - ) + fetched = self._fetch_prompt_once(name, version, label, ttl=ttl) return PromptResult( source="api", - prompt=data["prompt"], - name=data["name"], - version=data["version"], - label=data.get("label"), - config=copy.deepcopy(config), + prompt=fetched.prompt, + name=fetched.name, + version=fetched.version, + label=fetched.label, + config=copy.deepcopy(fetched.config), ) except Exception as error: - self._maybe_capture_error(error, name=name, version=version, label=label) - prompt_reference = _prompt_reference(name, version, label) # Return stale cache (with warning) if cached is not None: @@ -634,6 +644,79 @@ def _get_internal( raise + def _fetch_prompt_once( + self, + name: str, + version: Optional[int], + label: Optional[str], + *, + ttl: float, + ) -> CachedPrompt: + """Fetch a prompt, sharing the request with other threads that want the same key.""" + cache_key = _cache_key(name, version, label) + with self._in_flight_lock: + # Another thread may have finished a fetch between the caller's cache + # check and this lock. Its entry is fresh, so serve it instead of + # starting another request. A zero TTL is never fresh and still fetches. + cached = self._cache.get(cache_key) + if cached is not None and (time.time() - cached.fetched_at) < ttl: + return cached + + pending = self._in_flight.get(cache_key) + is_leader = pending is None + if pending is None: + pending = _InFlightFetch() + self._in_flight[cache_key] = pending + + if not is_leader: + pending.done.wait() + if pending.error is not None: + raise pending.error + assert pending.result is not None + return pending.result + + try: + data = self._fetch_prompt_from_api(name, version, label) + + # An older PostHog server ignores the label param and returns the latest + # version with no label field — surface that instead of failing silently. + if label is not None and data.get("label") != label: + log.warning( + "[PostHog Prompts] Requested label %r for prompt %r but the server " + "resolved %r. It may not support prompt labels yet and returned the " + "latest version instead.", + label, + name, + data.get("label"), + ) + + fetched = CachedPrompt( + prompt=data["prompt"], + fetched_at=time.time(), + name=data["name"], + version=data["version"], + label=data.get("label"), + config=_extract_config(data), + ) + # Cached before the in-flight entry is released, so a caller that + # arrives in between finds the entry instead of starting a request. + self._cache[cache_key] = fetched + pending.result = fetched + return fetched + except BaseException as error: + pending.error = error + raise + finally: + with self._in_flight_lock: + self._in_flight.pop(cache_key, None) + pending.done.set() + # Reported after the waiters are released, so an error-tracking callback + # that fetches the same prompt cannot wait on this very request. + if isinstance(pending.error, Exception): + self._maybe_capture_error( + pending.error, name=name, version=version, label=label + ) + def compile(self, prompt: str, variables: PromptVariables) -> str: """ Replace {{variableName}} placeholders with values. diff --git a/posthog/test/ai/test_prompts.py b/posthog/test/ai/test_prompts.py index 4e124f16d..7072bc90c 100644 --- a/posthog/test/ai/test_prompts.py +++ b/posthog/test/ai/test_prompts.py @@ -1,3 +1,4 @@ +import threading import unittest import warnings from unittest.mock import MagicMock, patch @@ -382,6 +383,125 @@ def test_hold_a_cooldown_after_each_failed_refetch( self.assertEqual(result.source, "api") self.assertEqual(mock_get.call_count, 4) + def _run_concurrent_gets(self, prompts, calls): + results = [None] * len(calls) + start_together = threading.Barrier(len(calls)) + + def run(index, kwargs): + start_together.wait(5) + results[index] = prompts.get("test-prompt", with_metadata=True, **kwargs) + + threads = [ + threading.Thread(target=run, args=(index, kwargs)) + for index, kwargs in enumerate(calls) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=5) + return results + + @patch("posthog.ai.prompts._get_session") + def test_share_one_request_between_concurrent_callers(self, mock_get_session): + mock_get = mock_get_session.return_value.get + release = threading.Event() + + def slow_get(url, **kwargs): + release.wait(5) + version = 2 if "version=2" in url else 1 + return MockResponse( + json_data={**self.mock_prompt_response, "version": version} + ) + + mock_get.side_effect = slow_get + prompts = Prompts(self.create_mock_posthog()) + + calls = [{}] * 10 + [{"version": 2}] + + # Release the network once every caller is waiting on it. + timer = threading.Timer(0.2, release.set) + timer.start() + results = self._run_concurrent_gets(prompts, calls) + timer.join() + + self.assertEqual(mock_get.call_count, 2) + self.assertEqual({r.source for r in results}, {"api"}) + self.assertEqual([r.version for r in results[:10]], [1] * 10) + self.assertEqual(results[10].version, 2) + + @patch("posthog.ai.prompts._get_session") + def test_concurrent_callers_share_one_failure_and_capture_it_once( + self, mock_get_session + ): + mock_get = mock_get_session.return_value.get + release = threading.Event() + + def failing_get(url, **kwargs): + release.wait(5) + raise Exception("Network error") + + mock_get.side_effect = failing_get + posthog = self.create_mock_posthog() + prompts = Prompts(posthog, capture_errors=True) + + timer = threading.Timer(0.2, release.set) + timer.start() + results = self._run_concurrent_gets( + prompts, [{"fallback": "fallback prompt"}] * 10 + ) + timer.join() + + self.assertEqual(mock_get.call_count, 1) + self.assertEqual({r.source for r in results}, {"code_fallback"}) + posthog.capture_exception.assert_called_once() + + @patch("posthog.ai.prompts._get_session") + def test_fork_hook_clears_inherited_in_flight_entries(self, mock_get_session): + from posthog.ai.prompts import ( + _InFlightFetch, + _cache_key, + _reset_in_flight_after_fork, + ) + + mock_get = mock_get_session.return_value.get + mock_get.return_value = MockResponse(json_data=self.mock_prompt_response) + prompts = Prompts(self.create_mock_posthog()) + # What a forked child inherits: an entry whose fetching thread no longer exists. + prompts._in_flight[_cache_key("test-prompt", None, None)] = _InFlightFetch() + + _reset_in_flight_after_fork() + result = prompts.get("test-prompt", with_metadata=True) + + self.assertEqual(result.source, "api") + self.assertEqual(mock_get.call_count, 1) + + @patch("posthog.ai.prompts._get_session") + def test_error_callback_that_fetches_the_same_prompt_does_not_deadlock( + self, mock_get_session + ): + mock_get = mock_get_session.return_value.get + mock_get.side_effect = Exception("Network error") + posthog = self.create_mock_posthog() + prompts = Prompts(posthog, capture_errors=True) + nested_results = [] + entered = threading.Event() + + def capture_exception(*args, **kwargs): + # The nested fetch fails too and reports again; only the first report fetches. + if entered.is_set(): + return + entered.set() + nested_results.append( + prompts.get("test-prompt", with_metadata=True, fallback="nested") + ) + + posthog.capture_exception.side_effect = capture_exception + + result = prompts.get("test-prompt", with_metadata=True, fallback="outer") + + self.assertEqual(result.prompt, "outer") + self.assertEqual([r.prompt for r in nested_results], ["nested"]) + @patch("posthog.ai.prompts._get_session") @patch("posthog.ai.prompts.log") def test_use_fallback_when_no_cache_and_fetch_fails_with_warning(