diff --git a/.sampo/changesets/prompts-get-all-by-label.md b/.sampo/changesets/prompts-get-all-by-label.md new file mode 100644 index 000000000..16736bce2 --- /dev/null +++ b/.sampo/changesets/prompts-get-all-by-label.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +`Prompts.get_all(label="production")` fetches every prompt that carries a label in one request and stores them in the prompt cache, so later `get(name, label=...)` calls are cache hits. Apps with many prompts no longer need one request per prompt per cache cycle. Against a PostHog server that does not support labels on the prompt list endpoint yet, the call fails with a clear error instead of caching wrong versions. diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index 48700b656..1786df28a 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -11,7 +11,7 @@ import urllib.parse import warnings from dataclasses import dataclass -from typing import Any, Dict, Literal, Optional, Union, overload +from typing import Any, Dict, List, Literal, Optional, Union, overload from posthog.request import USER_AGENT, _get_session from posthog.utils import remove_trailing_slash @@ -20,6 +20,10 @@ APP_ENDPOINT = "https://us.posthog.com" DEFAULT_CACHE_TTL_SECONDS = 300 # 5 minutes +# Backstop against a server whose pagination never terminates. 100 pages of the +# default page size covers 10,000 prompts; past that get_all raises rather than +# returning a truncated result. +_MAX_PROMPT_LIST_PAGES = 100 PromptVariables = Dict[str, Union[str, int, float, bool]] PromptCacheKey = tuple[str, Optional[int], Optional[str]] @@ -108,6 +112,48 @@ def _is_prompt_api_response(data: Any) -> bool: ) +def _row_resolves_label(row: Dict[str, Any], label: str) -> bool: + """Check that the server resolved this list row through the requested label. + + An older server ignores the label param on the list endpoint and returns the + latest version of every prompt. A row that was resolved through a label + carries a matching name and version entry in its all_labels field. + """ + all_labels = row.get("all_labels") + if not isinstance(all_labels, list): + return False + return any( + isinstance(entry, dict) + and entry.get("name") == label + and entry.get("version") == row.get("version") + for entry in all_labels + ) + + +def _is_same_origin(url: str, host: str) -> bool: + parsed = urllib.parse.urlsplit(url) + expected = urllib.parse.urlsplit(host) + return (parsed.scheme, parsed.netloc) == (expected.scheme, expected.netloc) + + +def _authentication_error(reference: str) -> Exception: + return Exception( + f"[PostHog Prompts] Authentication failed for {reference}. " + "The key may be missing, expired, or the wrong type. Prompt fetches " + "require a personal API key (starts with 'phx_'); a project secret " + "key is not accepted. Pass this key as personal_api_key when you " + "construct Prompts directly, or as secret_key when you configure the " + "PostHog client." + ) + + +def _access_denied_error(reference: str) -> Exception: + return Exception( + f"[PostHog Prompts] Access denied for {reference}. " + "Check that your personal_api_key has the correct permissions and the LLM prompts feature is enabled." + ) + + class Prompts: """ Fetch and compile LLM prompts from PostHog. @@ -142,6 +188,9 @@ class Prompts: # Fetch the version a label currently points to prod_prompt = prompts.get('support-system-prompt', label='production') + # Fetch all prompts at a label in one request and warm the cache + prod_prompts = prompts.get_all(label='production') + # Compile with variables system_prompt = prompts.compile(template, { 'company': 'Acme Corp', @@ -302,6 +351,85 @@ def get( return fallback raise + def get_all(self, *, label: str) -> Dict[str, PromptResult]: + """ + Fetch every prompt that carries a label, in one batch. + + Returns a dict mapping prompt name to :class:`PromptResult`, with each + prompt at the version the label points to. Prompts without the label + are not included. + + Each fetched prompt is stored in the cache, so later + ``get(name, label=...)`` calls are served from cache within the TTL. + An app with many prompts can call this once per cache cycle instead of + making one ``get()`` request per prompt. + + Args: + label: The label to resolve, e.g. 'production'. + + Returns: + Dict of prompt name to PromptResult. + + Raises: + Exception: If the request fails, or the server does not support + fetching prompts by label on the list endpoint (PostHog + releases from before September 2026). + """ + try: + rows = self._fetch_prompt_list_from_api(label) + except Exception as error: + self._maybe_capture_error(error, name="*", version=None, label=label) + raise + + now = time.time() + results: Dict[str, PromptResult] = {} + skipped: List[str] = [] + for row in rows: + if not _is_prompt_api_response(row) or not _row_resolves_label(row, label): + skipped.append(str(row.get("name")) if isinstance(row, dict) else "?") + continue + + config = _extract_config(row) + self._cache[_cache_key(row["name"], None, label)] = CachedPrompt( + prompt=row["prompt"], + fetched_at=now, + name=row["name"], + version=row["version"], + label=label, + config=config, + ) + results[row["name"]] = PromptResult( + source="api", + prompt=row["prompt"], + name=row["name"], + version=row["version"], + label=label, + config=copy.deepcopy(config), + ) + + if rows and not results: + # Nothing resolved the label, so the server most likely ignored the + # label param and served latest versions. Caching those under the + # label would be the silent wrong-version failure labels exist to + # prevent, so fail loudly instead. + compat_error = Exception( + f'[PostHog Prompts] The server returned prompts, but none resolve label "{label}". ' + "It may not support fetching prompts by label on the list endpoint yet. " + "Upgrade PostHog, or fetch prompts one by one with get()." + ) + self._maybe_capture_error(compat_error, name="*", version=None, label=label) + raise compat_error + + if skipped: + log.warning( + "[PostHog Prompts] Skipped %d prompt(s) that did not resolve label %r: %s", + len(skipped), + label, + ", ".join(skipped), + ) + + return results + def _get_internal( self, name: str, @@ -477,6 +605,96 @@ def _maybe_capture_error( except Exception: log.debug("[PostHog Prompts] Failed to capture exception to error tracking") + def _require_credentials(self) -> None: + if not self._personal_api_key: + raise Exception( + "[PostHog Prompts] personal_api_key is required to fetch prompts. " + "Please provide it when initializing the Prompts instance." + ) + if not self._project_api_key: + raise Exception( + "[PostHog Prompts] project_api_key is required to fetch prompts. " + "Please provide it when initializing the Prompts instance." + ) + + def _fetch_prompt_list_from_api(self, label: str) -> List[Dict[str, Any]]: + """ + Fetch all prompts at a label from the paginated list endpoint. + + Endpoint: + {host}/api/environments/@current/llm_prompts/ + ?token={encoded_project_api_key}&label={label}&content=full + Auth: Bearer {personal_api_key} + + Follows pagination links until the last page. Returns the raw rows. + """ + self._require_credentials() + + query = urllib.parse.urlencode( + {"token": self._project_api_key, "label": label, "content": "full"} + ) + url: Optional[str] = ( + f"{self._host}/api/environments/@current/llm_prompts/?{query}" + ) + reference = f'prompts with label "{label}"' + headers = { + "Authorization": f"Bearer {self._personal_api_key}", + "User-Agent": USER_AGENT, + } + + rows: List[Dict[str, Any]] = [] + pages = 0 + while url: + if pages >= _MAX_PROMPT_LIST_PAGES: + # A truncated result must not look complete: callers would cache + # a partial prompt set and treat missing prompts as unlabeled. + raise Exception( + f"[PostHog Prompts] {reference} spans more than " + f"{_MAX_PROMPT_LIST_PAGES} pages. Refusing to return an " + "incomplete result." + ) + + response = _get_session().get(url, headers=headers, timeout=10) + + if not response.ok: + if response.status_code == 401: + 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}" + ) + + try: + data = response.json() + except Exception: + raise Exception( + f"[PostHog Prompts] Invalid response format for {reference}" + ) + + if not isinstance(data, dict) or not isinstance(data.get("results"), list): + raise Exception( + f"[PostHog Prompts] Invalid response format for {reference}" + ) + + rows.extend(data["results"]) + + # The Authorization header goes to every followed link, so a link + # off the configured host must never be requested. + next_url = data.get("next") + if next_url is not None and ( + not isinstance(next_url, str) + or not _is_same_origin(next_url, self._host) + ): + raise Exception( + f"[PostHog Prompts] Refusing to follow a pagination link off the " + f"configured host while fetching {reference}." + ) + url = next_url + pages += 1 + + return rows + def _fetch_prompt_from_api( self, name: str, version: Optional[int] = None, label: Optional[str] = None ) -> Dict[str, Any]: @@ -501,16 +719,7 @@ def _fetch_prompt_from_api( Raises: Exception: If the prompt cannot be fetched """ - if not self._personal_api_key: - raise Exception( - "[PostHog Prompts] personal_api_key is required to fetch prompts. " - "Please provide it when initializing the Prompts instance." - ) - if not self._project_api_key: - raise Exception( - "[PostHog Prompts] project_api_key is required to fetch prompts. " - "Please provide it when initializing the Prompts instance." - ) + self._require_credentials() encoded_name = urllib.parse.quote(name, safe="") query_params: Dict[str, Union[str, int]] = {"token": self._project_api_key} @@ -535,20 +744,10 @@ def _fetch_prompt_from_api( raise Exception(f"[PostHog Prompts] {prompt_title} not found") if response.status_code == 401: - raise Exception( - f"[PostHog Prompts] Authentication failed for {prompt_reference}. " - "The key may be missing, expired, or the wrong type. Prompt fetches " - "require a personal API key (starts with 'phx_'); a project secret " - "key is not accepted. Pass this key as personal_api_key when you " - "construct Prompts directly, or as secret_key when you configure the " - "PostHog client." - ) + raise _authentication_error(prompt_reference) if response.status_code == 403: - raise Exception( - f"[PostHog Prompts] Access denied for {prompt_reference}. " - "Check that your personal_api_key has the correct permissions and the LLM prompts feature is enabled." - ) + raise _access_denied_error(prompt_reference) raise Exception( f"[PostHog Prompts] Failed to fetch {prompt_title}: HTTP {response.status_code}" diff --git a/posthog/test/ai/test_prompts.py b/posthog/test/ai/test_prompts.py index 02d7ea5d2..c255a079a 100644 --- a/posthog/test/ai/test_prompts.py +++ b/posthog/test/ai/test_prompts.py @@ -1359,5 +1359,131 @@ def test_clear_all_prompts_from_cache(self, mock_get_session): self.assertEqual(mock_get.call_count, 4) +class TestPromptsGetAll(TestPrompts): + """Tests for fetching all prompts at a label in one request.""" + + def labeled_row(self, name, version=1, label="production", config=None): + return { + "id": f"id-{name}", + "name": name, + "prompt": f"Prompt for {name}", + "version": version, + "all_labels": [{"name": label, "version": version}], + "config": config, + } + + def list_response(self, rows, next_url=None): + return MockResponse( + json_data={ + "count": len(rows), + "next": next_url, + "previous": None, + "results": rows, + } + ) + + @patch("posthog.ai.prompts._get_session") + def test_fetches_all_pages_and_seeds_the_cache(self, mock_get_session): + mock_get = mock_get_session.return_value.get + next_url = "https://us.posthog.com/api/environments/@current/llm_prompts/?token=phc_test_key&label=production&content=full&limit=100&offset=100" + mock_get.side_effect = [ + self.list_response( + [self.labeled_row("prompt-a", config={"temperature": 0})], + next_url=next_url, + ), + self.list_response([self.labeled_row("prompt-b", version=3)]), + ] + + prompts = Prompts(self.create_mock_posthog()) + results = prompts.get_all(label="production") + + self.assertEqual(mock_get.call_count, 2) + self.assertEqual(mock_get.call_args_list[1].args[0], next_url) + self.assertEqual( + results, + { + "prompt-a": PromptResult( + source="api", + prompt="Prompt for prompt-a", + name="prompt-a", + version=1, + label="production", + config={"temperature": 0}, + ), + "prompt-b": PromptResult( + source="api", + prompt="Prompt for prompt-b", + name="prompt-b", + version=3, + label="production", + ), + }, + ) + + # Later labeled get() calls are cache hits, not new requests. + cached = prompts.get("prompt-b", label="production", with_metadata=True) + self.assertEqual(mock_get.call_count, 2) + self.assertEqual(cached.source, "cache") + self.assertEqual(cached.version, 3) + + @patch("posthog.ai.prompts._get_session") + def test_raises_when_the_server_ignores_the_label(self, mock_get_session): + # An old server ignores ?label= and returns latest versions; none of the + # rows resolve the label, and caching them would serve wrong versions. + mock_get = mock_get_session.return_value.get + row = self.labeled_row("prompt-a") + row["all_labels"] = [] + mock_get.return_value = self.list_response([row]) + + prompts = Prompts(self.create_mock_posthog()) + + with self.assertRaises(Exception) as ctx: + prompts.get_all(label="production") + self.assertIn("none resolve label", str(ctx.exception)) + self.assertEqual(prompts._cache, {}) + + @patch("posthog.ai.prompts._get_session") + def test_skips_a_row_whose_label_moved_and_keeps_the_rest(self, mock_get_session): + # A label that moved between the query and the response leaves one stale + # row; the rest of the batch must still be returned and cached. + mock_get = mock_get_session.return_value.get + moved = self.labeled_row("prompt-a") + moved["all_labels"] = [{"name": "production", "version": 2}] + mock_get.return_value = self.list_response( + [moved, self.labeled_row("prompt-b")] + ) + + prompts = Prompts(self.create_mock_posthog()) + results = prompts.get_all(label="production") + + self.assertEqual(list(results), ["prompt-b"]) + + @patch("posthog.ai.prompts._get_session") + def test_refuses_a_pagination_link_off_the_configured_host(self, mock_get_session): + mock_get = mock_get_session.return_value.get + mock_get.return_value = self.list_response( + [self.labeled_row("prompt-a")], + next_url="https://attacker.example.com/collect", + ) + + prompts = Prompts(self.create_mock_posthog()) + + with self.assertRaises(Exception) as ctx: + prompts.get_all(label="production") + self.assertIn("off the configured host", str(ctx.exception)) + self.assertEqual(mock_get.call_count, 1) + + @patch("posthog.ai.prompts._get_session") + def test_raises_on_http_error(self, mock_get_session): + mock_get = mock_get_session.return_value.get + mock_get.return_value = MockResponse(status_code=500, ok=False) + + prompts = Prompts(self.create_mock_posthog()) + + with self.assertRaises(Exception) as ctx: + prompts.get_all(label="production") + self.assertIn("HTTP 500", str(ctx.exception)) + + if __name__ == "__main__": unittest.main() diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 2186387d2..503e324b2 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -1295,6 +1295,7 @@ method posthog.ai.otel.processor.PostHogSpanProcessor.shutdown() -> None method posthog.ai.prompts.Prompts.clear_cache(name: Optional[str] = None, *, version: Optional[int] = None) -> None method posthog.ai.prompts.Prompts.compile(prompt: str, variables: PromptVariables) -> str method posthog.ai.prompts.Prompts.get(name: str, *, with_metadata: Optional[bool] = None, cache_ttl_seconds: Optional[int] = None, fallback: Optional[str] = None, version: Optional[int] = None, label: Optional[str] = None) -> Union[str, PromptResult] +method posthog.ai.prompts.Prompts.get_all(*, label: str) -> Dict[str, PromptResult] method posthog.ai.stream.AsyncStreamWrapper.aclose() -> None method posthog.ai.stream.AsyncStreamWrapper.close() -> None method posthog.async_client.AsyncClient.alias(previous_id: ID_TYPES, distinct_id: Optional[str], timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None) -> Optional[str]