chore(llma): apply prompt SDK review cleanups (#456)

* chore(llma): apply prompt SDK review cleanups

* chore(llma): add review cleanup changeset

Author: Radu-Raicea

.sampo/changesets/prompt-sdk-review-cleanups.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+chore(llma): clean up prompt SDK review follow-ups

posthog/ai/prompts.py

@@ -35,9 +35,12 @@ def _cache_key(name: str, version: Optional[int]) -> PromptCacheKey:
     return (name, version)
 
 
-def _prompt_reference(name: str, version: Optional[int]) -> str:
+def _prompt_reference(
+    name: str, version: Optional[int], *, capitalize: bool = False
+) -> str:
     """Format a prompt reference for logs and errors."""
-    label = f'prompt "{name}"'
+    prefix = "Prompt" if capitalize else "prompt"
+    label = f'{prefix} "{name}"'
     if version is not None:
         return f"{label} version {version}"
     return label
@@ -291,7 +294,7 @@ def _fetch_prompt_from_api(self, name: str, version: Optional[int] = None) -> st
         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:]
+        prompt_label = _prompt_reference(name, version, capitalize=True)
 
         headers = {
             "Authorization": f"Bearer {self._personal_api_key}",

posthog/test/ai/test_prompts.py

@@ -1,6 +1,8 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
+from parameterized import parameterized
+
 from posthog.ai.prompts import Prompts
 
 
@@ -256,36 +258,31 @@ def test_throw_when_no_cache_no_fallback_and_fetch_fails(self, mock_get_session)
 
         self.assertIn("Network error", str(context.exception))
 
+    @parameterized.expand(
+        [
+            ("latest", {}, 'Prompt "nonexistent-prompt" not found'),
+            (
+                "versioned",
+                {"version": 3},
+                'Prompt "nonexistent-prompt" version 3 not found',
+            ),
+        ]
+    )
     @patch("posthog.ai.prompts._get_session")
-    def test_handle_404_response(self, mock_get_session):
-        """Should handle 404 response."""
-        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")
-
-        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."""
+    def test_handle_404_response(
+        self, _scenario, get_kwargs, expected_message, mock_get_session
+    ):
+        """Should handle 404 responses for latest and versioned prompts."""
         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)
+            prompts.get("nonexistent-prompt", **get_kwargs)
 
-        self.assertIn(
-            'Prompt "nonexistent-prompt" version 3 not found',
-            str(context.exception),
-        )
+        self.assertIn(expected_message, str(context.exception))
 
     @patch("posthog.ai.prompts._get_session")
     def test_handle_403_response(self, mock_get_session):
@@ -618,6 +615,28 @@ def test_handle_variables_with_dots(self):
 class TestPromptsClearCache(TestPrompts):
     """Tests for the Prompts.clear_cache() method."""
 
+    def _populate_versioned_cache(self, prompts, mock_get):
+        """Populate cache with latest and versioned entries for the same prompt."""
+        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),
+        ]
+
+        prompts.get("test-prompt")
+        prompts.get("test-prompt", version=1)
+
+        return latest_prompt_response, versioned_prompt_response
+
     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()
@@ -663,30 +682,14 @@ def test_clear_a_specific_prompt_from_cache(self, mock_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)
+        _, versioned_prompt_response = self._populate_versioned_cache(prompts, mock_get)
         self.assertEqual(mock_get.call_count, 2)
 
+        mock_get.side_effect = [MockResponse(json_data=versioned_prompt_response)]
         prompts.clear_cache("test-prompt", version=1)
 
         prompts.get("test-prompt")
@@ -699,31 +702,19 @@ def test_clear_a_specific_prompt_version_from_cache(self, mock_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)
+        latest_prompt_response, versioned_prompt_response = (
+            self._populate_versioned_cache(prompts, mock_get)
+        )
         self.assertEqual(mock_get.call_count, 2)
 
+        mock_get.side_effect = [
+            MockResponse(json_data=latest_prompt_response),
+            MockResponse(json_data=versioned_prompt_response),
+        ]
         prompts.clear_cache("test-prompt")
 
         prompts.get("test-prompt")