Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0bb6342472 | ||
|
|
d76bfe6e5b |
@@ -1,3 +1,11 @@
|
||||
# 6.7.5 - 2025-09-16
|
||||
|
||||
- feat: Django middleware now supports async request handling.
|
||||
|
||||
# 6.7.4 - 2025-09-05
|
||||
|
||||
- fix: Missing system prompts for some providers
|
||||
|
||||
# 6.7.3 - 2025-09-04
|
||||
|
||||
- fix: missing usage tokens in Gemini
|
||||
|
||||
@@ -3,7 +3,8 @@ import time
|
||||
import uuid
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from posthog.ai.types import TokenUsage
|
||||
from posthog.ai.types import TokenUsage, StreamingEventData
|
||||
from posthog.ai.utils import merge_system_prompt
|
||||
|
||||
try:
|
||||
from google import genai
|
||||
@@ -19,7 +20,6 @@ from posthog.ai.utils import (
|
||||
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,
|
||||
@@ -356,10 +356,8 @@ class Models:
|
||||
latency: float,
|
||||
output: Any,
|
||||
):
|
||||
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 +379,12 @@ class Models:
|
||||
# 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"""
|
||||
|
||||
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,
|
||||
|
||||
@@ -220,6 +220,30 @@ def format_gemini_response(response: Any) -> List[FormattedMessage]:
|
||||
return output
|
||||
|
||||
|
||||
def extract_gemini_system_instruction(config: Any) -> Optional[str]:
|
||||
"""
|
||||
Extract system instruction from Gemini config parameter.
|
||||
|
||||
Args:
|
||||
config: Config object or dict that may contain system instruction
|
||||
|
||||
Returns:
|
||||
System instruction string if present, None otherwise
|
||||
"""
|
||||
if config is None:
|
||||
return None
|
||||
|
||||
# Handle different config formats
|
||||
if hasattr(config, "system_instruction"):
|
||||
return config.system_instruction
|
||||
elif isinstance(config, dict) and "system_instruction" in config:
|
||||
return config["system_instruction"]
|
||||
elif isinstance(config, dict) and "systemInstruction" in config:
|
||||
return config["systemInstruction"]
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_gemini_tools(kwargs: Dict[str, Any]) -> Optional[Any]:
|
||||
"""
|
||||
Extract tool definitions from Gemini API kwargs.
|
||||
@@ -237,6 +261,38 @@ def extract_gemini_tools(kwargs: Dict[str, Any]) -> Optional[Any]:
|
||||
return None
|
||||
|
||||
|
||||
def format_gemini_input_with_system(
|
||||
contents: Any, config: Any = None
|
||||
) -> List[FormattedMessage]:
|
||||
"""
|
||||
Format Gemini input contents into standardized message format, including system instruction handling.
|
||||
|
||||
Args:
|
||||
contents: Input contents in various possible formats
|
||||
config: Config object or dict that may contain system instruction
|
||||
|
||||
Returns:
|
||||
List of formatted messages with role and content fields, with system message prepended if needed
|
||||
"""
|
||||
formatted_messages = format_gemini_input(contents)
|
||||
|
||||
# Check if system instruction is provided in config parameter
|
||||
system_instruction = extract_gemini_system_instruction(config)
|
||||
|
||||
if system_instruction is not None:
|
||||
has_system = any(msg.get("role") == "system" for msg in formatted_messages)
|
||||
if not has_system:
|
||||
from posthog.ai.types import FormattedMessage
|
||||
|
||||
system_message: FormattedMessage = {
|
||||
"role": "system",
|
||||
"content": system_instruction,
|
||||
}
|
||||
formatted_messages = [system_message] + list(formatted_messages)
|
||||
|
||||
return formatted_messages
|
||||
|
||||
|
||||
def format_gemini_input(contents: Any) -> List[FormattedMessage]:
|
||||
"""
|
||||
Format Gemini input contents into standardized message format for PostHog tracking.
|
||||
|
||||
@@ -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")
|
||||
|
||||
+18
-11
@@ -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,9 @@ 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.
|
||||
"""
|
||||
@@ -169,10 +171,11 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
|
||||
system = kwargs.get("system")
|
||||
return format_anthropic_input(messages, system)
|
||||
elif provider == "gemini":
|
||||
from posthog.ai.gemini.gemini_converter import format_gemini_input
|
||||
from posthog.ai.gemini.gemini_converter import format_gemini_input_with_system
|
||||
|
||||
contents = kwargs.get("contents", [])
|
||||
return format_gemini_input(contents)
|
||||
config = kwargs.get("config")
|
||||
return format_gemini_input_with_system(contents, config)
|
||||
elif provider == "openai":
|
||||
from posthog.ai.openai.openai_converter import format_openai_input
|
||||
|
||||
@@ -187,9 +190,11 @@ 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 +212,11 @@ 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
|
||||
|
||||
|
||||
@@ -1,10 +1,18 @@
|
||||
from typing import TYPE_CHECKING, cast
|
||||
from posthog import contexts, capture_exception
|
||||
from posthog import contexts
|
||||
from posthog.client import Client
|
||||
|
||||
try:
|
||||
from asgiref.sync import iscoroutinefunction
|
||||
except ImportError:
|
||||
# Fallback for older Django versions
|
||||
import asyncio
|
||||
|
||||
iscoroutinefunction = asyncio.iscoroutinefunction
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from django.http import HttpRequest, HttpResponse # noqa: F401
|
||||
from typing import Callable, Dict, Any, Optional # noqa: F401
|
||||
from typing import Callable, Dict, Any, Optional, Union, Awaitable # noqa: F401
|
||||
|
||||
|
||||
class PosthogContextMiddleware:
|
||||
@@ -33,9 +41,24 @@ class PosthogContextMiddleware:
|
||||
frontend. See the documentation for `set_context_session` and `identify_context` for more details.
|
||||
"""
|
||||
|
||||
# Django middleware capability flags
|
||||
sync_capable = True
|
||||
async_capable = True
|
||||
|
||||
def __init__(self, get_response):
|
||||
# type: (Callable[[HttpRequest], HttpResponse]) -> None
|
||||
self.get_response = get_response
|
||||
# type: (Union[Callable[[HttpRequest], HttpResponse], Callable[[HttpRequest], Awaitable[HttpResponse]]]) -> None
|
||||
self._is_coroutine = iscoroutinefunction(get_response)
|
||||
self._async_get_response = None # type: Optional[Callable[[HttpRequest], Awaitable[HttpResponse]]]
|
||||
self._sync_get_response = None # type: Optional[Callable[[HttpRequest], HttpResponse]]
|
||||
|
||||
if self._is_coroutine:
|
||||
self._async_get_response = cast(
|
||||
"Callable[[HttpRequest], Awaitable[HttpResponse]]", get_response
|
||||
)
|
||||
else:
|
||||
self._sync_get_response = cast(
|
||||
"Callable[[HttpRequest], HttpResponse]", get_response
|
||||
)
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
@@ -159,23 +182,39 @@ class PosthogContextMiddleware:
|
||||
|
||||
def __call__(self, request):
|
||||
# type: (HttpRequest) -> HttpResponse
|
||||
# Purely defensive around django's internal sync/async handling - this should be unreachable, but if it's reached, we may
|
||||
# as well return something semi-meaningful
|
||||
if self._is_coroutine:
|
||||
raise RuntimeError(
|
||||
"PosthogContextMiddleware received sync call but get_response is async"
|
||||
)
|
||||
|
||||
if self.request_filter and not self.request_filter(request):
|
||||
return self.get_response(request)
|
||||
assert self._sync_get_response is not None
|
||||
return self._sync_get_response(request)
|
||||
|
||||
with contexts.new_context(self.capture_exceptions, client=self.client):
|
||||
for k, v in self.extract_tags(request).items():
|
||||
contexts.tag(k, v)
|
||||
|
||||
return self.get_response(request)
|
||||
assert self._sync_get_response is not None
|
||||
return self._sync_get_response(request)
|
||||
|
||||
def process_exception(self, request, exception):
|
||||
async def __acall__(self, request):
|
||||
# type: (HttpRequest) -> HttpResponse
|
||||
if self.request_filter and not self.request_filter(request):
|
||||
return
|
||||
if self._async_get_response is not None:
|
||||
return await self._async_get_response(request)
|
||||
else:
|
||||
assert self._sync_get_response is not None
|
||||
return self._sync_get_response(request)
|
||||
|
||||
if not self.capture_exceptions:
|
||||
return
|
||||
with contexts.new_context(self.capture_exceptions, client=self.client):
|
||||
for k, v in self.extract_tags(request).items():
|
||||
contexts.tag(k, v)
|
||||
|
||||
if self.client:
|
||||
self.client.capture_exception(exception)
|
||||
else:
|
||||
capture_exception(exception)
|
||||
if self._async_get_response is not None:
|
||||
return await self._async_get_response(request)
|
||||
else:
|
||||
assert self._sync_get_response is not None
|
||||
return self._sync_get_response(request)
|
||||
|
||||
@@ -616,6 +616,8 @@ def test_tool_use_response(mock_client, mock_google_genai_client, mock_gemini_re
|
||||
|
||||
mock_config = MagicMock()
|
||||
mock_config.tools = [mock_tool]
|
||||
# Explicitly specify this config doesn't have system_instruction
|
||||
del mock_config.system_instruction
|
||||
|
||||
response = client.models.generate_content(
|
||||
model="gemini-2.5-flash",
|
||||
|
||||
@@ -0,0 +1,354 @@
|
||||
"""
|
||||
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 in config 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}]
|
||||
config = {"system_instruction": self.test_system_prompt}
|
||||
|
||||
client.models.generate_content(
|
||||
model="gemini-2.0-flash",
|
||||
contents=contents,
|
||||
config=config,
|
||||
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"])
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
VERSION = "6.7.3"
|
||||
VERSION = "6.7.5"
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(VERSION, end="") # noqa: T201
|
||||
|
||||
Reference in New Issue
Block a user