mirror of
https://github.com/hanzoai/insights-python.git
synced 2026-10-30 18:33:17 +00:00
zeekay
3ca982343e
- Remove Posthog = Insights alias from __init__.py - Remove PosthogContextMiddleware alias from django.py - Remove PostHogTracingProcessor alias from processor.py - Change $lib from "posthog-python" to "insights-python" - Change ingestion URLs from posthog.com to insights.hanzo.ai - Rename all posthog_* kwargs to insights_* across AI wrappers - Rename __posthog_exception_captured to __insights_exception_captured - Rename posthog_context_stack contextvar to insights_context_stack - Rename posthog🎏 Redis prefix to insights🎏 - Rename $$_posthog_redacted_* sentinels to $$_insights_redacted_* - Remove POSTHOG_MW_* Django settings fallback, X-POSTHOG-* headers - Rename Prompts(posthog=) param to Prompts(client=) - Update APP_ENDPOINT to us.insights.hanzo.ai - Update all tests, examples, docs, mypy config
364 lines
14 KiB
Python
364 lines
14 KiB
Python
"""
|
|
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 MagicMock, patch
|
|
|
|
from hanzo_insights.client import Client
|
|
from hanzo_insights.test.test_utils import FAKE_TEST_API_KEY
|
|
|
|
|
|
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 Insights client
|
|
self.client = Client(FAKE_TEST_API_KEY)
|
|
self.client._enqueue = 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 openai.types.chat import ChatCompletion, ChatCompletionMessage
|
|
from openai.types.chat.chat_completion import Choice
|
|
from openai.types.completion_usage import CompletionUsage
|
|
|
|
from hanzo_insights.ai.openai import OpenAI
|
|
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(insights_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, insights_distinct_id="test-user"
|
|
)
|
|
|
|
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
|
|
properties = self.client._enqueue.call_args_list[0][0][0]["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 openai.types.chat import ChatCompletion, ChatCompletionMessage
|
|
from openai.types.chat.chat_completion import Choice
|
|
from openai.types.completion_usage import CompletionUsage
|
|
|
|
from hanzo_insights.ai.openai import OpenAI
|
|
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(insights_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,
|
|
insights_distinct_id="test-user",
|
|
)
|
|
|
|
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
|
|
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
|
|
self._assert_system_prompt_captured(properties["$ai_input"])
|
|
|
|
def test_openai_streaming_system_parameter(self):
|
|
"""Test OpenAI streaming with system parameter."""
|
|
try:
|
|
from openai.types.chat.chat_completion_chunk import (
|
|
ChatCompletionChunk,
|
|
ChoiceDelta,
|
|
)
|
|
from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
|
|
from openai.types.completion_usage import CompletionUsage
|
|
|
|
from hanzo_insights.ai.openai import OpenAI
|
|
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(insights_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,
|
|
insights_distinct_id="test-user",
|
|
)
|
|
|
|
list(response_generator) # Consume generator
|
|
|
|
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
|
|
properties = self.client._enqueue.call_args_list[0][0][0]["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 hanzo_insights.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(insights_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,
|
|
insights_distinct_id="test-user",
|
|
)
|
|
|
|
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
|
|
properties = self.client._enqueue.call_args_list[0][0][0]["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 hanzo_insights.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(insights_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,
|
|
insights_distinct_id="test-user",
|
|
)
|
|
|
|
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
|
|
properties = self.client._enqueue.call_args_list[0][0][0]["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 hanzo_insights.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(insights_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,
|
|
insights_distinct_id="test-user",
|
|
)
|
|
|
|
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
|
|
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
|
|
self._assert_system_prompt_captured(properties["$ai_input"])
|
|
|
|
def test_gemini_system_instruction_parameter(self):
|
|
"""Test Gemini with system_instruction in config parameter."""
|
|
try:
|
|
from hanzo_insights.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(insights_client=self.client, api_key="test")
|
|
|
|
contents = [{"role": "user", "content": self.test_user_message}]
|
|
config = {"system_instruction": self.test_system_prompt}
|
|
|
|
client.models.generate_content(
|
|
model="gemini-2.0-flash",
|
|
contents=contents,
|
|
config=config,
|
|
insights_distinct_id="test-user",
|
|
)
|
|
|
|
self.assertEqual(len(self.client._enqueue.call_args_list), 1)
|
|
properties = self.client._enqueue.call_args_list[0][0][0]["properties"]
|
|
self._assert_system_prompt_captured(properties["$ai_input"])
|