feat(llma): support prompt versions in prompts sdk (#454)

* feat(llma): support prompt versions in prompts sdk

* fix(llma): enforce clear_cache version requires name

.sampo/changesets/prompt-sdk-versions.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+feat(llma): support fetching versioned prompts from the prompts sdk

posthog/ai/prompts.py

@@ -19,6 +19,7 @@ APP_ENDPOINT = "https://us.posthog.com"
 DEFAULT_CACHE_TTL_SECONDS = 300  # 5 minutes
 
 PromptVariables = Dict[str, Union[str, int, float, bool]]
+PromptCacheKey = tuple[str, Optional[int]]
 
 
 class CachedPrompt:
@@ -29,6 +30,19 @@ class CachedPrompt:
         self.fetched_at = fetched_at
 
 
+def _cache_key(name: str, version: Optional[int]) -> PromptCacheKey:
+    """Build a cache key for latest or versioned prompt fetches."""
+    return (name, version)
+
+
+def _prompt_reference(name: str, version: Optional[int]) -> str:
+    """Format a prompt reference for logs and errors."""
+    label = f'prompt "{name}"'
+    if version is not None:
+        return f"{label} version {version}"
+    return label
+
+
 def _is_prompt_api_response(data: Any) -> bool:
     """Check if the response is a valid prompt API response."""
     return (
@@ -63,6 +77,9 @@ class Prompts:
         # Fetch with caching and fallback
         template = prompts.get('support-system-prompt', fallback='You are a helpful assistant.')
 
+        # Fetch a specific published version
+        prompt_v1 = prompts.get('support-system-prompt', version=1)
+
         # Compile with variables
         system_prompt = prompts.compile(template, {
             'company': 'Acme Corp',
@@ -93,7 +110,7 @@ class Prompts:
         self._default_cache_ttl_seconds = (
             default_cache_ttl_seconds or DEFAULT_CACHE_TTL_SECONDS
         )
-        self._cache: Dict[str, CachedPrompt] = {}
+        self._cache: Dict[PromptCacheKey, CachedPrompt] = {}
 
         if posthog is not None:
             self._personal_api_key = getattr(posthog, "personal_api_key", None) or ""
@@ -112,6 +129,7 @@ class Prompts:
         *,
         cache_ttl_seconds: Optional[int] = None,
         fallback: Optional[str] = None,
+        version: Optional[int] = None,
     ) -> str:
         """
         Fetch a prompt by name from the PostHog API.
@@ -126,6 +144,8 @@ class Prompts:
             name: The name of the prompt to fetch
             cache_ttl_seconds: Cache TTL in seconds (defaults to instance default)
             fallback: Fallback prompt to use if fetch fails and no cache available
+            version: Specific prompt version to fetch. If None, fetches the latest
+                version
 
         Returns:
             The prompt string
@@ -138,9 +158,10 @@ class Prompts:
             if cache_ttl_seconds is not None
             else self._default_cache_ttl_seconds
         )
+        cache_key = _cache_key(name, version)
 
         # Check cache first
-        cached = self._cache.get(name)
+        cached = self._cache.get(cache_key)
         now = time.time()
 
         if cached is not None:
@@ -151,21 +172,22 @@ class Prompts:
 
         # Try to fetch from API
         try:
-            prompt = self._fetch_prompt_from_api(name)
+            prompt = self._fetch_prompt_from_api(name, version)
             fetched_at = time.time()
 
             # Update cache
-            self._cache[name] = CachedPrompt(prompt=prompt, fetched_at=fetched_at)
+            self._cache[cache_key] = CachedPrompt(prompt=prompt, fetched_at=fetched_at)
 
             return prompt
 
         except Exception as error:
+            prompt_reference = _prompt_reference(name, version)
             # Fallback order:
             # 1. Return stale cache (with warning)
             if cached is not None:
                 log.warning(
-                    '[PostHog Prompts] Failed to fetch prompt "%s", using stale cache: %s',
+                    "[PostHog Prompts] Failed to fetch %s, using stale cache: %s",
-                    name,
+                    prompt_reference,
                     error,
                 )
                 return cached.prompt
@@ -173,8 +195,8 @@ class Prompts:
             # 2. Return fallback (with warning)
             if fallback is not None:
                 log.warning(
-                    '[PostHog Prompts] Failed to fetch prompt "%s", using fallback: %s',
+                    "[PostHog Prompts] Failed to fetch %s, using fallback: %s",
-                    name,
+                    prompt_reference,
                     error,
                 )
                 return fallback
@@ -207,27 +229,43 @@ class Prompts:
 
         return re.sub(r"\{\{([\w.-]+)\}\}", replace_variable, prompt)
 
-    def clear_cache(self, name: Optional[str] = None) -> None:
+    def clear_cache(
+        self, name: Optional[str] = None, *, version: Optional[int] = None
+    ) -> None:
         """
         Clear cached prompts.
 
         Args:
-            name: Specific prompt to clear. If None, clears all cached prompts.
+            name: Specific prompt name to clear. If None, clears all cached prompts.
+            version: Specific prompt version to clear. Requires name.
         """
-        if name is not None:
+        if version is not None and name is None:
-            self._cache.pop(name, None)
+            raise ValueError("'version' requires 'name' to be provided")
-        else:
-            self._cache.clear()
 
-    def _fetch_prompt_from_api(self, name: str) -> str:
+        if name is None:
+            self._cache.clear()
+            return
+
+        if version is not None:
+            self._cache.pop(_cache_key(name, version), None)
+            return
+
+        keys_to_clear = [key for key in self._cache if key[0] == name]
+        for key in keys_to_clear:
+            self._cache.pop(key, None)
+
+    def _fetch_prompt_from_api(self, name: str, version: Optional[int] = None) -> str:
         """
         Fetch prompt from PostHog API.
 
-        Endpoint: {host}/api/environments/@current/llm_prompts/name/{encoded_name}/?token={encoded_project_api_key}
+        Endpoint:
+            {host}/api/environments/@current/llm_prompts/name/{encoded_name}/
+            ?token={encoded_project_api_key}[&version={version}]
         Auth: Bearer {personal_api_key}
 
         Args:
             name: The name of the prompt to fetch
+            version: Specific prompt version to fetch. If None, fetches the latest
 
         Returns:
             The prompt string
@@ -247,8 +285,13 @@ class Prompts:
             )
 
         encoded_name = urllib.parse.quote(name, safe="")
-        encoded_project_api_key = urllib.parse.quote(self._project_api_key, safe="")
+        query_params: Dict[str, Union[str, int]] = {"token": self._project_api_key}
-        url = f"{self._host}/api/environments/@current/llm_prompts/name/{encoded_name}/?token={encoded_project_api_key}"
+        if version is not None:
+            query_params["version"] = version
+        encoded_query = urllib.parse.urlencode(query_params)
+        url = f"{self._host}/api/environments/@current/llm_prompts/name/{encoded_name}/?{encoded_query}"
+        prompt_reference = _prompt_reference(name, version)
+        prompt_label = prompt_reference[:1].upper() + prompt_reference[1:]
 
         headers = {
             "Authorization": f"Bearer {self._personal_api_key}",
@@ -259,28 +302,28 @@ class Prompts:
 
         if not response.ok:
             if response.status_code == 404:
-                raise Exception(f'[PostHog Prompts] Prompt "{name}" not found')
+                raise Exception(f"[PostHog Prompts] {prompt_label} not found")
 
             if response.status_code == 403:
                 raise Exception(
-                    f'[PostHog Prompts] Access denied for prompt "{name}". '
+                    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 Exception(
-                f'[PostHog Prompts] Failed to fetch prompt "{name}": HTTP {response.status_code}'
+                f"[PostHog Prompts] Failed to fetch {prompt_label}: HTTP {response.status_code}"
             )
 
         try:
             data = response.json()
         except Exception:
             raise Exception(
-                f'[PostHog Prompts] Invalid response format for prompt "{name}"'
+                f"[PostHog Prompts] Invalid response format for {prompt_label}"
             )
 
         if not _is_prompt_api_response(data):
             raise Exception(
-                f'[PostHog Prompts] Invalid response format for prompt "{name}"'
+                f"[PostHog Prompts] Invalid response format for {prompt_label}"
             )
 
         return data["prompt"]

posthog/test/ai/test_prompts.py

@@ -72,6 +72,30 @@ class TestPromptsGet(TestPrompts):
             call_args[1]["headers"]["Authorization"], "Bearer phx_test_key"
         )
 
+    @patch("posthog.ai.prompts._get_session")
+    def test_successfully_fetch_a_specific_prompt_version(self, mock_get_session):
+        """Should successfully fetch a specific prompt version."""
+        mock_get = mock_get_session.return_value.get
+        versioned_prompt_response = {
+            **self.mock_prompt_response,
+            "prompt": "Prompt version 1",
+            "version": 1,
+        }
+        mock_get.return_value = MockResponse(json_data=versioned_prompt_response)
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        result = prompts.get("test-prompt", version=1)
+
+        self.assertEqual(result, versioned_prompt_response["prompt"])
+        mock_get.assert_called_once()
+        call_args = mock_get.call_args
+        self.assertEqual(
+            call_args[0][0],
+            "https://us.posthog.com/api/environments/@current/llm_prompts/name/test-prompt/?token=phc_test_key&version=1",
+        )
+
     @patch("posthog.ai.prompts._get_session")
     @patch("posthog.ai.prompts.time.time")
     def test_return_cached_prompt_when_fresh(self, mock_time, mock_get_session):
@@ -96,6 +120,41 @@ class TestPromptsGet(TestPrompts):
         self.assertEqual(result2, self.mock_prompt_response["prompt"])
         self.assertEqual(mock_get.call_count, 1)  # No additional fetch
 
+    @patch("posthog.ai.prompts._get_session")
+    def test_cache_latest_and_versioned_prompts_separately(self, mock_get_session):
+        """Should cache latest and historical prompt versions separately."""
+        mock_get = mock_get_session.return_value.get
+        latest_prompt_response = {
+            **self.mock_prompt_response,
+            "prompt": "Latest prompt",
+            "version": 2,
+        }
+        versioned_prompt_response = {
+            **self.mock_prompt_response,
+            "prompt": "Prompt version 1",
+            "version": 1,
+        }
+
+        mock_get.side_effect = [
+            MockResponse(json_data=latest_prompt_response),
+            MockResponse(json_data=versioned_prompt_response),
+        ]
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        self.assertEqual(prompts.get("test-prompt"), latest_prompt_response["prompt"])
+        self.assertEqual(
+            prompts.get("test-prompt", version=1),
+            versioned_prompt_response["prompt"],
+        )
+        self.assertEqual(prompts.get("test-prompt"), latest_prompt_response["prompt"])
+        self.assertEqual(
+            prompts.get("test-prompt", version=1),
+            versioned_prompt_response["prompt"],
+        )
+        self.assertEqual(mock_get.call_count, 2)
+
     @patch("posthog.ai.prompts._get_session")
     @patch("posthog.ai.prompts.time.time")
     def test_refetch_when_cache_is_stale(self, mock_time, mock_get_session):
@@ -211,6 +270,23 @@ class TestPromptsGet(TestPrompts):
 
         self.assertIn('Prompt "nonexistent-prompt" not found', str(context.exception))
 
+    @patch("posthog.ai.prompts._get_session")
+    def test_handle_404_response_for_specific_prompt_version(self, mock_get_session):
+        """Should handle 404 response for a specific prompt version."""
+        mock_get = mock_get_session.return_value.get
+        mock_get.return_value = MockResponse(status_code=404, ok=False)
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        with self.assertRaises(Exception) as context:
+            prompts.get("nonexistent-prompt", version=3)
+
+        self.assertIn(
+            'Prompt "nonexistent-prompt" version 3 not found',
+            str(context.exception),
+        )
+
     @patch("posthog.ai.prompts._get_session")
     def test_handle_403_response(self, mock_get_session):
         """Should handle 403 response."""
@@ -542,6 +618,16 @@ class TestPromptsCompile(TestPrompts):
 class TestPromptsClearCache(TestPrompts):
     """Tests for the Prompts.clear_cache() method."""
 
+    def test_clear_cache_with_version_and_no_name_raises_value_error(self):
+        """Should enforce that versioned cache clearing requires a prompt name."""
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        with self.assertRaises(ValueError) as context:
+            prompts.clear_cache(version=1)
+
+        self.assertIn("requires 'name'", str(context.exception))
+
     @patch("posthog.ai.prompts._get_session")
     def test_clear_a_specific_prompt_from_cache(self, mock_get_session):
         """Should clear a specific prompt from cache."""
@@ -573,6 +659,77 @@ class TestPromptsClearCache(TestPrompts):
         prompts.get("other-prompt")
         self.assertEqual(mock_get.call_count, 3)
 
+    @patch("posthog.ai.prompts._get_session")
+    def test_clear_a_specific_prompt_version_from_cache(self, mock_get_session):
+        """Should clear only the requested prompt version from cache."""
+        mock_get = mock_get_session.return_value.get
+        latest_prompt_response = {
+            **self.mock_prompt_response,
+            "prompt": "Latest prompt",
+            "version": 2,
+        }
+        versioned_prompt_response = {
+            **self.mock_prompt_response,
+            "prompt": "Prompt version 1",
+            "version": 1,
+        }
+
+        mock_get.side_effect = [
+            MockResponse(json_data=latest_prompt_response),
+            MockResponse(json_data=versioned_prompt_response),
+            MockResponse(json_data=versioned_prompt_response),
+        ]
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        prompts.get("test-prompt")
+        prompts.get("test-prompt", version=1)
+        self.assertEqual(mock_get.call_count, 2)
+
+        prompts.clear_cache("test-prompt", version=1)
+
+        prompts.get("test-prompt")
+        self.assertEqual(mock_get.call_count, 2)
+
+        prompts.get("test-prompt", version=1)
+        self.assertEqual(mock_get.call_count, 3)
+
+    @patch("posthog.ai.prompts._get_session")
+    def test_clear_a_prompt_name_clears_all_cached_versions(self, mock_get_session):
+        """Should clear latest and versioned cache entries for the same prompt name."""
+        mock_get = mock_get_session.return_value.get
+        latest_prompt_response = {
+            **self.mock_prompt_response,
+            "prompt": "Latest prompt",
+            "version": 2,
+        }
+        versioned_prompt_response = {
+            **self.mock_prompt_response,
+            "prompt": "Prompt version 1",
+            "version": 1,
+        }
+
+        mock_get.side_effect = [
+            MockResponse(json_data=latest_prompt_response),
+            MockResponse(json_data=versioned_prompt_response),
+            MockResponse(json_data=latest_prompt_response),
+            MockResponse(json_data=versioned_prompt_response),
+        ]
+
+        posthog = self.create_mock_posthog()
+        prompts = Prompts(posthog)
+
+        prompts.get("test-prompt")
+        prompts.get("test-prompt", version=1)
+        self.assertEqual(mock_get.call_count, 2)
+
+        prompts.clear_cache("test-prompt")
+
+        prompts.get("test-prompt")
+        prompts.get("test-prompt", version=1)
+        self.assertEqual(mock_get.call_count, 4)
+
     @patch("posthog.ai.prompts._get_session")
     def test_clear_all_prompts_from_cache(self, mock_get_session):
         """Should clear all prompts from cache when no name is provided."""