fix: always capture system prompt

Author: carlos-marchal-ph

posthog/ai/gemini/gemini.py

@@ -19,7 +19,6 @@
     merge_usage_stats,
 )
 from posthog.ai.gemini.gemini_converter import (
-    format_gemini_input,
     extract_gemini_usage_from_chunk,
     extract_gemini_content_from_chunk,
     format_gemini_streaming_output,
@@ -359,7 +358,7 @@ def _capture_streaming_event(
         from posthog.ai.types import StreamingEventData
 
         # Prepare standardized event data
-        formatted_input = self._format_input(contents)
+        formatted_input = self._format_input(contents, **kwargs)
         sanitized_input = sanitize_gemini(formatted_input)
 
         event_data = StreamingEventData(
@@ -381,10 +380,13 @@ def _capture_streaming_event(
         # Use the common capture function
         capture_streaming_event(self._ph_client, event_data)
 
-    def _format_input(self, contents):
+    def _format_input(self, contents, **kwargs):
         """Format input contents for PostHog tracking"""
+        from posthog.ai.utils import merge_system_prompt
 
-        return format_gemini_input(contents)
+        # Create kwargs dict with contents for merge_system_prompt
+        input_kwargs = {"contents": contents, **kwargs}
+        return merge_system_prompt(input_kwargs, "gemini")
 
     def generate_content_stream(
         self,

posthog/ai/openai/openai_converter.py

@@ -606,7 +606,6 @@ def format_openai_streaming_input(
     Returns:
         Formatted input ready for PostHog tracking
     """
-    if api_type == "chat":
-        return kwargs.get("messages")
-    else:  # responses API
-        return kwargs.get("input")
+    from posthog.ai.utils import merge_system_prompt
+    
+    return merge_system_prompt(kwargs, "openai")

posthog/ai/utils.py

@@ -1,9 +1,9 @@
 import time
 import uuid
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, List, Optional, cast
 
 from posthog.client import Client as PostHogClient
-from posthog.ai.types import StreamingEventData, TokenUsage
+from posthog.ai.types import FormattedMessage, StreamingEventData, TokenUsage
 from posthog.ai.sanitization import (
     sanitize_openai,
     sanitize_anthropic,
@@ -158,7 +158,7 @@ def extract_available_tool_calls(provider: str, kwargs: Dict[str, Any]):
     return None
 
 
-def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
+def merge_system_prompt(kwargs: Dict[str, Any], provider: str) -> List[FormattedMessage]:
     """
     Merge system prompts and format messages for the given provider.
     """
@@ -172,7 +172,17 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
         from posthog.ai.gemini.gemini_converter import format_gemini_input
 
         contents = kwargs.get("contents", [])
-        return format_gemini_input(contents)
+        formatted_messages = format_gemini_input(contents)
+
+        if kwargs.get("system_instruction") is not None:
+            system_instruction = kwargs.get("system_instruction")
+            system_message = cast(FormattedMessage, {
+                "role": "system", 
+                "content": system_instruction
+            })
+            formatted_messages = [system_message] + list(formatted_messages)
+        
+        return formatted_messages
     elif provider == "openai":
         from posthog.ai.openai.openai_converter import format_openai_input
 
@@ -187,9 +197,10 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
         if kwargs.get("system") is not None:
             has_system = any(msg.get("role") == "system" for msg in messages)
             if not has_system:
-                messages = [
-                    {"role": "system", "content": kwargs.get("system")}
-                ] + messages
+                system_msg = cast(FormattedMessage, {
+                    "role": "system", "content": kwargs.get("system")
+                })
+                messages = [system_msg] + messages
 
         # For Responses API, add instructions to the system prompt if provided
         if kwargs.get("instructions") is not None:
@@ -207,9 +218,10 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
                 )
             else:
                 # Create a new system message with instructions
-                messages = [
-                    {"role": "system", "content": kwargs.get("instructions")}
-                ] + messages
+                instruction_msg = cast(FormattedMessage, {
+                    "role": "system", "content": kwargs.get("instructions")
+                })
+                messages = [instruction_msg] + messages
 
         return messages
 

posthog/test/ai/test_system_prompts.py

@@ -0,0 +1,267 @@
+"""
+Tests for system prompt capture across all LLM providers.
+
+This test suite ensures that system prompts are correctly captured in analytics
+regardless of how they're passed to the providers:
+- As first message in messages/contents array (standard format)
+- As separate system parameter (Anthropic, OpenAI)
+- As instructions parameter (OpenAI Responses API)
+- As system_instruction parameter (Gemini)
+"""
+
+import time
+import unittest
+from unittest.mock import patch, MagicMock
+
+
+class TestSystemPromptCapture(unittest.TestCase):
+    """Test system prompt capture for all providers."""
+
+    def setUp(self):
+        super().setUp()
+        self.test_system_prompt = "You are a helpful AI assistant."
+        self.test_user_message = "Hello, how are you?"
+        self.test_response = "I'm doing well, thank you!"
+        
+        # Create mock PostHog client
+        self.client = MagicMock()
+        self.client.privacy_mode = False
+
+    def _assert_system_prompt_captured(self, captured_input):
+        """Helper to assert system prompt is correctly captured."""
+        self.assertEqual(len(captured_input), 2, "Should have 2 messages (system + user)")
+        self.assertEqual(captured_input[0]["role"], "system", "First message should be system")
+        self.assertEqual(captured_input[0]["content"], self.test_system_prompt, "System content should match")
+        self.assertEqual(captured_input[1]["role"], "user", "Second message should be user")
+        self.assertEqual(captured_input[1]["content"], self.test_user_message, "User content should match")
+
+    # OpenAI Tests
+    def test_openai_messages_array_system_prompt(self):
+        """Test OpenAI with system prompt in messages array."""
+        try:
+            from posthog.ai.openai import OpenAI
+            from openai.types.chat import ChatCompletion, ChatCompletionMessage
+            from openai.types.chat.chat_completion import Choice
+            from openai.types.completion_usage import CompletionUsage
+        except ImportError:
+            self.skipTest("OpenAI package not available")
+
+        mock_response = ChatCompletion(
+            id="test", model="gpt-4", object="chat.completion", created=int(time.time()),
+            choices=[Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(
+                content=self.test_response, role="assistant"))],
+            usage=CompletionUsage(completion_tokens=10, prompt_tokens=20, total_tokens=30),
+        )
+
+        with patch("openai.resources.chat.completions.Completions.create", return_value=mock_response):
+            client = OpenAI(posthog_client=self.client, api_key="test")
+            
+            messages = [
+                {"role": "system", "content": self.test_system_prompt},
+                {"role": "user", "content": self.test_user_message}
+            ]
+            
+            client.chat.completions.create(model="gpt-4", messages=messages, posthog_distinct_id="test-user")
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    def test_openai_separate_system_parameter(self):
+        """Test OpenAI with system prompt as separate parameter."""
+        try:
+            from posthog.ai.openai import OpenAI
+            from openai.types.chat import ChatCompletion, ChatCompletionMessage
+            from openai.types.chat.chat_completion import Choice
+            from openai.types.completion_usage import CompletionUsage
+        except ImportError:
+            self.skipTest("OpenAI package not available")
+
+        mock_response = ChatCompletion(
+            id="test", model="gpt-4", object="chat.completion", created=int(time.time()),
+            choices=[Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(
+                content=self.test_response, role="assistant"))],
+            usage=CompletionUsage(completion_tokens=10, prompt_tokens=20, total_tokens=30),
+        )
+
+        with patch("openai.resources.chat.completions.Completions.create", return_value=mock_response):
+            client = OpenAI(posthog_client=self.client, api_key="test")
+            
+            messages = [{"role": "user", "content": self.test_user_message}]
+            
+            client.chat.completions.create(
+                model="gpt-4", messages=messages, system=self.test_system_prompt, posthog_distinct_id="test-user"
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+
+    def test_openai_streaming_system_parameter(self):
+        """Test OpenAI streaming with system parameter."""
+        try:
+            from posthog.ai.openai import OpenAI
+            from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+            from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk  
+            from openai.types.chat.chat_completion_chunk import ChoiceDelta
+            from openai.types.completion_usage import CompletionUsage
+        except ImportError:
+            self.skipTest("OpenAI package not available")
+
+        chunk1 = ChatCompletionChunk(
+            id="test", model="gpt-4", object="chat.completion.chunk", created=int(time.time()),
+            choices=[ChoiceChunk(finish_reason=None, index=0, delta=ChoiceDelta(content="Hello", role="assistant"))]
+        )
+        
+        chunk2 = ChatCompletionChunk(
+            id="test", model="gpt-4", object="chat.completion.chunk", created=int(time.time()),
+            choices=[ChoiceChunk(finish_reason="stop", index=0, delta=ChoiceDelta(content=" there!", role=None))],
+            usage=CompletionUsage(completion_tokens=10, prompt_tokens=20, total_tokens=30)
+        )
+
+        with patch("openai.resources.chat.completions.Completions.create", return_value=[chunk1, chunk2]):
+            client = OpenAI(posthog_client=self.client, api_key="test")
+            
+            messages = [{"role": "user", "content": self.test_user_message}]
+            
+            response_generator = client.chat.completions.create(
+                model="gpt-4", messages=messages, system=self.test_system_prompt,
+                stream=True, posthog_distinct_id="test-user"
+            )
+
+            list(response_generator)  # Consume generator
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    # Anthropic Tests
+    def test_anthropic_messages_array_system_prompt(self):
+        """Test Anthropic with system prompt in messages array."""
+        try:
+            from posthog.ai.anthropic import Anthropic
+        except ImportError:
+            self.skipTest("Anthropic package not available")
+
+        with patch("anthropic.resources.messages.Messages.create") as mock_create:
+            mock_response = MagicMock()
+            mock_response.usage.input_tokens = 20
+            mock_response.usage.output_tokens = 10
+            mock_response.usage.cache_read_input_tokens = None
+            mock_response.usage.cache_creation_input_tokens = None
+            mock_create.return_value = mock_response
+            
+            client = Anthropic(posthog_client=self.client, api_key="test")
+            
+            messages = [
+                {"role": "system", "content": self.test_system_prompt},
+                {"role": "user", "content": self.test_user_message}
+            ]
+            
+            client.messages.create(model="claude-3-5-sonnet-20241022", messages=messages, posthog_distinct_id="test-user")
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    def test_anthropic_separate_system_parameter(self):
+        """Test Anthropic with system prompt as separate parameter."""
+        try:
+            from posthog.ai.anthropic import Anthropic
+        except ImportError:
+            self.skipTest("Anthropic package not available")
+
+        with patch("anthropic.resources.messages.Messages.create") as mock_create:
+            mock_response = MagicMock()
+            mock_response.usage.input_tokens = 20
+            mock_response.usage.output_tokens = 10
+            mock_response.usage.cache_read_input_tokens = None
+            mock_response.usage.cache_creation_input_tokens = None
+            mock_create.return_value = mock_response
+            
+            client = Anthropic(posthog_client=self.client, api_key="test")
+            
+            messages = [{"role": "user", "content": self.test_user_message}]
+            
+            client.messages.create(
+                model="claude-3-5-sonnet-20241022", messages=messages,
+                system=self.test_system_prompt, posthog_distinct_id="test-user"
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    # Gemini Tests  
+    def test_gemini_contents_array_system_prompt(self):
+        """Test Gemini with system prompt in contents array."""
+        try:
+            from posthog.ai.gemini import Client
+        except ImportError:
+            self.skipTest("Gemini package not available")
+
+        with patch("google.genai.Client") as mock_genai_class:
+            mock_response = MagicMock()
+            mock_response.candidates = [MagicMock()]
+            mock_response.candidates[0].content.parts = [MagicMock()]
+            mock_response.candidates[0].content.parts[0].text = self.test_response
+            mock_response.usage_metadata.prompt_token_count = 20
+            mock_response.usage_metadata.candidates_token_count = 10
+            mock_response.usage_metadata.cached_content_token_count = None
+            mock_response.usage_metadata.thoughts_token_count = None
+
+            mock_client_instance = MagicMock()
+            mock_models_instance = MagicMock()
+            mock_models_instance.generate_content.return_value = mock_response
+            mock_client_instance.models = mock_models_instance
+            mock_genai_class.return_value = mock_client_instance
+
+            client = Client(posthog_client=self.client, api_key="test")
+            
+            contents = [
+                {"role": "system", "content": self.test_system_prompt},
+                {"role": "user", "content": self.test_user_message}
+            ]
+            
+            client.models.generate_content(model="gemini-2.0-flash", contents=contents, posthog_distinct_id="test-user")
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])
+
+    def test_gemini_system_instruction_parameter(self):
+        """Test Gemini with system_instruction parameter."""
+        try:
+            from posthog.ai.gemini import Client
+        except ImportError:
+            self.skipTest("Gemini package not available")
+
+        with patch("google.genai.Client") as mock_genai_class:
+            mock_response = MagicMock()
+            mock_response.candidates = [MagicMock()]
+            mock_response.candidates[0].content.parts = [MagicMock()]
+            mock_response.candidates[0].content.parts[0].text = self.test_response
+            mock_response.usage_metadata.prompt_token_count = 20
+            mock_response.usage_metadata.candidates_token_count = 10
+            mock_response.usage_metadata.cached_content_token_count = None
+            mock_response.usage_metadata.thoughts_token_count = None
+
+            mock_client_instance = MagicMock()
+            mock_models_instance = MagicMock()
+            mock_models_instance.generate_content.return_value = mock_response
+            mock_client_instance.models = mock_models_instance
+            mock_genai_class.return_value = mock_client_instance
+
+            client = Client(posthog_client=self.client, api_key="test")
+            
+            contents = [{"role": "user", "content": self.test_user_message}]
+            
+            client.models.generate_content(
+                model="gemini-2.0-flash", contents=contents,
+                system_instruction=self.test_system_prompt, posthog_distinct_id="test-user"
+            )
+
+            self.assertEqual(len(self.client.capture.call_args_list), 1)
+            properties = self.client.capture.call_args_list[0][1]["properties"]
+            self._assert_system_prompt_captured(properties["$ai_input"])