Skip to content

Commit eb23c3d

Browse files
carlos-marchal-phcarlos-marchal-ph
committed
chore: lint
1 parent 6d18346 commit eb23c3d

File tree

3 files changed

+161
-71
lines changed

3 files changed

+161
-71
lines changed

‎posthog/ai/openai/openai_converter.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -607,5 +607,5 @@ def format_openai_streaming_input(
607607
Formatted input ready for PostHog tracking
608608
"""
609609
from posthog.ai.utils import merge_system_prompt
610-
610+
611611
return merge_system_prompt(kwargs, "openai")

‎posthog/ai/utils.py‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,9 @@ def extract_available_tool_calls(provider: str, kwargs: Dict[str, Any]):
158158
return None
159159

160160

161-
def merge_system_prompt(kwargs: Dict[str, Any], provider: str) -> List[FormattedMessage]:
161+
def merge_system_prompt(
162+
kwargs: Dict[str, Any], provider: str
163+
) -> List[FormattedMessage]:
162164
"""
163165
Merge system prompts and format messages for the given provider.
164166
"""
@@ -173,11 +175,11 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str) -> List[Formatted
173175

174176
contents = kwargs.get("contents", [])
175177
formatted_messages = format_gemini_input(contents)
176-
178+
177179
# Check if system instruction is provided in config parameter
178180
config = kwargs.get("config")
179181
system_instruction = None
180-
182+
181183
if config is not None:
182184
# Handle different config formats
183185
if hasattr(config, "system_instruction"):
@@ -186,16 +188,15 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str) -> List[Formatted
186188
system_instruction = config["system_instruction"]
187189
elif isinstance(config, dict) and "systemInstruction" in config:
188190
system_instruction = config["systemInstruction"]
189-
191+
190192
if system_instruction is not None:
191193
has_system = any(msg.get("role") == "system" for msg in formatted_messages)
192194
if not has_system:
193-
system_message = cast(FormattedMessage, {
194-
"role": "system",
195-
"content": system_instruction
196-
})
195+
system_message = cast(
196+
FormattedMessage, {"role": "system", "content": system_instruction}
197+
)
197198
formatted_messages = [system_message] + list(formatted_messages)
198-
199+
199200
return formatted_messages
200201
elif provider == "openai":
201202
from posthog.ai.openai.openai_converter import format_openai_input
@@ -211,9 +212,10 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str) -> List[Formatted
211212
if kwargs.get("system") is not None:
212213
has_system = any(msg.get("role") == "system" for msg in messages)
213214
if not has_system:
214-
system_msg = cast(FormattedMessage, {
215-
"role": "system", "content": kwargs.get("system")
216-
})
215+
system_msg = cast(
216+
FormattedMessage,
217+
{"role": "system", "content": kwargs.get("system")},
218+
)
217219
messages = [system_msg] + messages
218220

219221
# For Responses API, add instructions to the system prompt if provided
@@ -232,9 +234,10 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str) -> List[Formatted
232234
)
233235
else:
234236
# Create a new system message with instructions
235-
instruction_msg = cast(FormattedMessage, {
236-
"role": "system", "content": kwargs.get("instructions")
237-
})
237+
instruction_msg = cast(
238+
FormattedMessage,
239+
{"role": "system", "content": kwargs.get("instructions")},
240+
)
238241
messages = [instruction_msg] + messages
239242

240243
return messages

‎posthog/test/ai/test_system_prompts.py‎

Lines changed: 142 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -22,18 +22,32 @@ def setUp(self):
2222
self.test_system_prompt = "You are a helpful AI assistant."
2323
self.test_user_message = "Hello, how are you?"
2424
self.test_response = "I'm doing well, thank you!"
25-
25+
2626
# Create mock PostHog client
2727
self.client = MagicMock()
2828
self.client.privacy_mode = False
2929

3030
def _assert_system_prompt_captured(self, captured_input):
3131
"""Helper to assert system prompt is correctly captured."""
32-
self.assertEqual(len(captured_input), 2, "Should have 2 messages (system + user)")
33-
self.assertEqual(captured_input[0]["role"], "system", "First message should be system")
34-
self.assertEqual(captured_input[0]["content"], self.test_system_prompt, "System content should match")
35-
self.assertEqual(captured_input[1]["role"], "user", "Second message should be user")
36-
self.assertEqual(captured_input[1]["content"], self.test_user_message, "User content should match")
32+
self.assertEqual(
33+
len(captured_input), 2, "Should have 2 messages (system + user)"
34+
)
35+
self.assertEqual(
36+
captured_input[0]["role"], "system", "First message should be system"
37+
)
38+
self.assertEqual(
39+
captured_input[0]["content"],
40+
self.test_system_prompt,
41+
"System content should match",
42+
)
43+
self.assertEqual(
44+
captured_input[1]["role"], "user", "Second message should be user"
45+
)
46+
self.assertEqual(
47+
captured_input[1]["content"],
48+
self.test_user_message,
49+
"User content should match",
50+
)
3751

3852
# OpenAI Tests
3953
def test_openai_messages_array_system_prompt(self):
@@ -47,21 +61,38 @@ def test_openai_messages_array_system_prompt(self):
4761
self.skipTest("OpenAI package not available")
4862

4963
mock_response = ChatCompletion(
50-
id="test", model="gpt-4", object="chat.completion", created=int(time.time()),
51-
choices=[Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(
52-
content=self.test_response, role="assistant"))],
53-
usage=CompletionUsage(completion_tokens=10, prompt_tokens=20, total_tokens=30),
64+
id="test",
65+
model="gpt-4",
66+
object="chat.completion",
67+
created=int(time.time()),
68+
choices=[
69+
Choice(
70+
finish_reason="stop",
71+
index=0,
72+
message=ChatCompletionMessage(
73+
content=self.test_response, role="assistant"
74+
),
75+
)
76+
],
77+
usage=CompletionUsage(
78+
completion_tokens=10, prompt_tokens=20, total_tokens=30
79+
),
5480
)
5581

56-
with patch("openai.resources.chat.completions.Completions.create", return_value=mock_response):
82+
with patch(
83+
"openai.resources.chat.completions.Completions.create",
84+
return_value=mock_response,
85+
):
5786
client = OpenAI(posthog_client=self.client, api_key="test")
58-
87+
5988
messages = [
6089
{"role": "system", "content": self.test_system_prompt},
61-
{"role": "user", "content": self.test_user_message}
90+
{"role": "user", "content": self.test_user_message},
6291
]
63-
64-
client.chat.completions.create(model="gpt-4", messages=messages, posthog_distinct_id="test-user")
92+
93+
client.chat.completions.create(
94+
model="gpt-4", messages=messages, posthog_distinct_id="test-user"
95+
)
6596

6697
self.assertEqual(len(self.client.capture.call_args_list), 1)
6798
properties = self.client.capture.call_args_list[0][1]["properties"]
@@ -78,56 +109,99 @@ def test_openai_separate_system_parameter(self):
78109
self.skipTest("OpenAI package not available")
79110

80111
mock_response = ChatCompletion(
81-
id="test", model="gpt-4", object="chat.completion", created=int(time.time()),
82-
choices=[Choice(finish_reason="stop", index=0, message=ChatCompletionMessage(
83-
content=self.test_response, role="assistant"))],
84-
usage=CompletionUsage(completion_tokens=10, prompt_tokens=20, total_tokens=30),
112+
id="test",
113+
model="gpt-4",
114+
object="chat.completion",
115+
created=int(time.time()),
116+
choices=[
117+
Choice(
118+
finish_reason="stop",
119+
index=0,
120+
message=ChatCompletionMessage(
121+
content=self.test_response, role="assistant"
122+
),
123+
)
124+
],
125+
usage=CompletionUsage(
126+
completion_tokens=10, prompt_tokens=20, total_tokens=30
127+
),
85128
)
86129

87-
with patch("openai.resources.chat.completions.Completions.create", return_value=mock_response):
130+
with patch(
131+
"openai.resources.chat.completions.Completions.create",
132+
return_value=mock_response,
133+
):
88134
client = OpenAI(posthog_client=self.client, api_key="test")
89-
135+
90136
messages = [{"role": "user", "content": self.test_user_message}]
91-
137+
92138
client.chat.completions.create(
93-
model="gpt-4", messages=messages, system=self.test_system_prompt, posthog_distinct_id="test-user"
139+
model="gpt-4",
140+
messages=messages,
141+
system=self.test_system_prompt,
142+
posthog_distinct_id="test-user",
94143
)
95144

96145
self.assertEqual(len(self.client.capture.call_args_list), 1)
97146
properties = self.client.capture.call_args_list[0][1]["properties"]
98147
self._assert_system_prompt_captured(properties["$ai_input"])
99148

100-
101149
def test_openai_streaming_system_parameter(self):
102150
"""Test OpenAI streaming with system parameter."""
103151
try:
104152
from posthog.ai.openai import OpenAI
105153
from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
106-
from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
154+
from openai.types.chat.chat_completion_chunk import Choice as ChoiceChunk
107155
from openai.types.chat.chat_completion_chunk import ChoiceDelta
108156
from openai.types.completion_usage import CompletionUsage
109157
except ImportError:
110158
self.skipTest("OpenAI package not available")
111159

112160
chunk1 = ChatCompletionChunk(
113-
id="test", model="gpt-4", object="chat.completion.chunk", created=int(time.time()),
114-
choices=[ChoiceChunk(finish_reason=None, index=0, delta=ChoiceDelta(content="Hello", role="assistant"))]
161+
id="test",
162+
model="gpt-4",
163+
object="chat.completion.chunk",
164+
created=int(time.time()),
165+
choices=[
166+
ChoiceChunk(
167+
finish_reason=None,
168+
index=0,
169+
delta=ChoiceDelta(content="Hello", role="assistant"),
170+
)
171+
],
115172
)
116-
173+
117174
chunk2 = ChatCompletionChunk(
118-
id="test", model="gpt-4", object="chat.completion.chunk", created=int(time.time()),
119-
choices=[ChoiceChunk(finish_reason="stop", index=0, delta=ChoiceDelta(content=" there!", role=None))],
120-
usage=CompletionUsage(completion_tokens=10, prompt_tokens=20, total_tokens=30)
175+
id="test",
176+
model="gpt-4",
177+
object="chat.completion.chunk",
178+
created=int(time.time()),
179+
choices=[
180+
ChoiceChunk(
181+
finish_reason="stop",
182+
index=0,
183+
delta=ChoiceDelta(content=" there!", role=None),
184+
)
185+
],
186+
usage=CompletionUsage(
187+
completion_tokens=10, prompt_tokens=20, total_tokens=30
188+
),
121189
)
122190

123-
with patch("openai.resources.chat.completions.Completions.create", return_value=[chunk1, chunk2]):
191+
with patch(
192+
"openai.resources.chat.completions.Completions.create",
193+
return_value=[chunk1, chunk2],
194+
):
124195
client = OpenAI(posthog_client=self.client, api_key="test")
125-
196+
126197
messages = [{"role": "user", "content": self.test_user_message}]
127-
198+
128199
response_generator = client.chat.completions.create(
129-
model="gpt-4", messages=messages, system=self.test_system_prompt,
130-
stream=True, posthog_distinct_id="test-user"
200+
model="gpt-4",
201+
messages=messages,
202+
system=self.test_system_prompt,
203+
stream=True,
204+
posthog_distinct_id="test-user",
131205
)
132206

133207
list(response_generator) # Consume generator
@@ -151,15 +225,19 @@ def test_anthropic_messages_array_system_prompt(self):
151225
mock_response.usage.cache_read_input_tokens = None
152226
mock_response.usage.cache_creation_input_tokens = None
153227
mock_create.return_value = mock_response
154-
228+
155229
client = Anthropic(posthog_client=self.client, api_key="test")
156-
230+
157231
messages = [
158232
{"role": "system", "content": self.test_system_prompt},
159-
{"role": "user", "content": self.test_user_message}
233+
{"role": "user", "content": self.test_user_message},
160234
]
161-
162-
client.messages.create(model="claude-3-5-sonnet-20241022", messages=messages, posthog_distinct_id="test-user")
235+
236+
client.messages.create(
237+
model="claude-3-5-sonnet-20241022",
238+
messages=messages,
239+
posthog_distinct_id="test-user",
240+
)
163241

164242
self.assertEqual(len(self.client.capture.call_args_list), 1)
165243
properties = self.client.capture.call_args_list[0][1]["properties"]
@@ -179,21 +257,23 @@ def test_anthropic_separate_system_parameter(self):
179257
mock_response.usage.cache_read_input_tokens = None
180258
mock_response.usage.cache_creation_input_tokens = None
181259
mock_create.return_value = mock_response
182-
260+
183261
client = Anthropic(posthog_client=self.client, api_key="test")
184-
262+
185263
messages = [{"role": "user", "content": self.test_user_message}]
186-
264+
187265
client.messages.create(
188-
model="claude-3-5-sonnet-20241022", messages=messages,
189-
system=self.test_system_prompt, posthog_distinct_id="test-user"
266+
model="claude-3-5-sonnet-20241022",
267+
messages=messages,
268+
system=self.test_system_prompt,
269+
posthog_distinct_id="test-user",
190270
)
191271

192272
self.assertEqual(len(self.client.capture.call_args_list), 1)
193273
properties = self.client.capture.call_args_list[0][1]["properties"]
194274
self._assert_system_prompt_captured(properties["$ai_input"])
195275

196-
# Gemini Tests
276+
# Gemini Tests
197277
def test_gemini_contents_array_system_prompt(self):
198278
"""Test Gemini with system prompt in contents array."""
199279
try:
@@ -218,13 +298,17 @@ def test_gemini_contents_array_system_prompt(self):
218298
mock_genai_class.return_value = mock_client_instance
219299

220300
client = Client(posthog_client=self.client, api_key="test")
221-
301+
222302
contents = [
223303
{"role": "system", "content": self.test_system_prompt},
224-
{"role": "user", "content": self.test_user_message}
304+
{"role": "user", "content": self.test_user_message},
225305
]
226-
227-
client.models.generate_content(model="gemini-2.0-flash", contents=contents, posthog_distinct_id="test-user")
306+
307+
client.models.generate_content(
308+
model="gemini-2.0-flash",
309+
contents=contents,
310+
posthog_distinct_id="test-user",
311+
)
228312

229313
self.assertEqual(len(self.client.capture.call_args_list), 1)
230314
properties = self.client.capture.call_args_list[0][1]["properties"]
@@ -254,14 +338,17 @@ def test_gemini_system_instruction_parameter(self):
254338
mock_genai_class.return_value = mock_client_instance
255339

256340
client = Client(posthog_client=self.client, api_key="test")
257-
341+
258342
contents = [{"role": "user", "content": self.test_user_message}]
259343
config = {"system_instruction": self.test_system_prompt}
260-
344+
261345
client.models.generate_content(
262-
model="gemini-2.0-flash", contents=contents, config=config, posthog_distinct_id="test-user"
346+
model="gemini-2.0-flash",
347+
contents=contents,
348+
config=config,
349+
posthog_distinct_id="test-user",
263350
)
264351

265352
self.assertEqual(len(self.client.capture.call_args_list), 1)
266353
properties = self.client.capture.call_args_list[0][1]["properties"]
267-
self._assert_system_prompt_captured(properties["$ai_input"])
354+
self._assert_system_prompt_captured(properties["$ai_input"])

0 commit comments

Comments
 (0)