feat: support gemini #237
Merged
feat: support gemini #237
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
There was a problem hiding this comment.
PR Summary
Added support for Google's Gemini AI models in the PostHog Python SDK, providing a drop-in replacement for the google-genai library with automatic usage tracking.
- Missing optional dependency: Add
google-genaitoextras_requirein both setup files to allow installation via pip extras syntax - Add version constraint for
google-genaidependency to prevent compatibility issues - Sync package paths between
setup.pyandsetup_analytics.py(posthog vs posthoganalytics) - Consider adding error handling tests for API failures and rate limits in
test_gemini.py - Add docstring examples in
gemini.pyshowing how to use privacy mode and group analytics features
6 file(s) reviewed, 4 comment(s)
Edit PR Review Bot Settings | Greptile
Comment on lines
+72
to
+81
| def __init__( | ||
| self, | ||
| api_key: Optional[str] = None, | ||
| posthog_client: Optional[PostHogClient] = None, | ||
| posthog_distinct_id: Optional[str] = None, | ||
| posthog_properties: Optional[Dict[str, Any]] = None, | ||
| posthog_privacy_mode: bool = False, | ||
| posthog_groups: Optional[Dict[str, Any]] = None, | ||
| **kwargs, | ||
| ): |
There was a problem hiding this comment.
logic: Missing validation check for posthog_client in Models.init() - could lead to NoneType errors later
Suggested change
| def __init__( | |
| self, | |
| api_key: Optional[str] = None, | |
| posthog_client: Optional[PostHogClient] = None, | |
| posthog_distinct_id: Optional[str] = None, | |
| posthog_properties: Optional[Dict[str, Any]] = None, | |
| posthog_privacy_mode: bool = False, | |
| posthog_groups: Optional[Dict[str, Any]] = None, | |
| **kwargs, | |
| ): | |
| def __init__( | |
| self, | |
| api_key: Optional[str] = None, | |
| posthog_client: Optional[PostHogClient] = None, | |
| posthog_distinct_id: Optional[str] = None, | |
| posthog_properties: Optional[Dict[str, Any]] = None, | |
| posthog_privacy_mode: bool = False, | |
| posthog_groups: Optional[Dict[str, Any]] = None, | |
| **kwargs, | |
| ): | |
| if posthog_client is None: | |
| raise ValueError("posthog_client is required for PostHog tracking") |
Comment on lines
+281
to
+287
| if hasattr(self._ph_client, "capture"): | ||
| self._ph_client.capture( | ||
| distinct_id=distinct_id, | ||
| event="$ai_generation", | ||
| properties=event_properties, | ||
| groups=groups, | ||
| ) |
There was a problem hiding this comment.
style: Check for capture method should happen in init to fail fast if PostHog client is incompatible
Comment on lines
+88
to
+139
| def test_new_client_streaming_with_generate_content_stream(mock_client, mock_google_genai_client): | ||
| """Test the new generate_content_stream method""" | ||
|
|
||
| def mock_streaming_response(): | ||
| mock_chunk1 = MagicMock() | ||
| mock_chunk1.text = "Hello " | ||
| mock_usage1 = MagicMock() | ||
| mock_usage1.prompt_token_count = 10 | ||
| mock_usage1.candidates_token_count = 5 | ||
| mock_chunk1.usage_metadata = mock_usage1 | ||
|
|
||
| mock_chunk2 = MagicMock() | ||
| mock_chunk2.text = "world!" | ||
| mock_usage2 = MagicMock() | ||
| mock_usage2.prompt_token_count = 10 | ||
| mock_usage2.candidates_token_count = 10 | ||
| mock_chunk2.usage_metadata = mock_usage2 | ||
| 10000 |
|
|
| yield mock_chunk1 | ||
| yield mock_chunk2 | ||
|
|
||
| # Mock the generate_content_stream method | ||
| mock_google_genai_client.models.generate_content_stream.return_value = mock_streaming_response() | ||
|
|
||
| client = Client(api_key="test-key", posthog_client=mock_client) | ||
|
|
||
| response = client.models.generate_content_stream( | ||
| model="gemini-2.0-flash", | ||
| contents=["Write a short story"], | ||
| posthog_distinct_id="test-id", | ||
| posthog_properties={"feature": "streaming"}, | ||
| ) | ||
|
|
||
| chunks = list(response) | ||
| assert len(chunks) == 2 | ||
| assert chunks[0].text == "Hello " | ||
| assert chunks[1].text == "world!" | ||
|
|
||
| # Check that the streaming event was captured | ||
| assert mock_client.capture.call_count == 1 | ||
| call_args = mock_client.capture.call_args[1] | ||
| props = call_args["properties"] | ||
|
|
||
| assert call_args["distinct_id"] == "test-id" | ||
| assert call_args["event"] == "$ai_generation" | ||
| assert props["$ai_provider"] == "gemini" | ||
| assert props["$ai_model"] == "gemini-2.0-flash" | ||
| assert props["$ai_input_tokens"] == 10 | ||
| assert props["$ai_output_tokens"] == 10 | ||
| assert props["feature"] == "streaming" | ||
| assert isinstance(props["$ai_latency"], float) | ||
|
|
There was a problem hiding this comment.
logic: Streaming test only verifies successful case. Should test error handling during streaming and verify token counts are accurate when stream is interrupted.
Comment on lines
+263
to
+297
| def test_new_client_override_defaults(mock_client, mock_google_genai_client, mock_gemini_response): | ||
| """Test overriding client defaults per call""" | ||
| mock_google_genai_client.models.generate_content.return_value = mock_gemini_response | ||
|
|
||
| client = Client( | ||
| api_key="test-key", | ||
| posthog_client=mock_client, | ||
| posthog_distinct_id="default_user", | ||
| posthog_properties={"team": "ai"}, | ||
| posthog_privacy_mode=False, | ||
| posthog_groups={"company": "acme_corp"}, | ||
| ) | ||
|
|
||
| # Override defaults in call | ||
| client.models.generate_content( | ||
| model="gemini-2.0-flash", | ||
| contents=["Hello"], | ||
| posthog_distinct_id="specific_user", | ||
| posthog_properties={"feature": "chat", "urgent": True}, | ||
| posthog_privacy_mode=True, | ||
| posthog_groups={"organization": "special_org"}, | ||
| ) | ||
|
|
||
| call_args = mock_client.capture.call_args[1] | ||
| props = call_args["properties"] | ||
|
|
||
| # Check overrides | ||
| assert call_args["distinct_id"] == "specific_user" | ||
| assert call_args["groups"] == {"organization": "special_org"} | ||
| assert props["$ai_input"] is None # privacy mode was overridden | ||
|
|
||
| # Check merged properties (defaults + call-specific) | ||
| assert props["team"] == "ai" # from defaults | ||
| assert props["feature"] == "chat" # from call | ||
| assert props["urgent"] is True # from call |
There was a problem hiding this comment.
style: Test should verify that all default properties are preserved when not explicitly overridden.
kappa90
approved these changes
May 23, 2025
posthog/ai/gemini/gemini.py
Outdated
| posthog_distinct_id, posthog_trace_id, posthog_properties, posthog_privacy_mode, posthog_groups | ||
| ) | ||
|
|
||
| if trace_id is None: |
There was a problem hiding this comment.
Nit: Why don't you put this in self._merge_posthog_params()? Seems cleaner to me.
| privacy_mode, | ||
| [{"content": output, "role": "assistant"}], | ||
| ), | ||
| "$ai_http_status": 200, |
There was a problem hiding this comment.
I guess we don't capture non-200 statuses in LLM Observability?
posthog/ai/gemini/gemini.py
Outdated
| posthog_distinct_id, posthog_trace_id, posthog_properties, posthog_privacy_mode, posthog_groups | ||
| ) | ||
|
|
||
| if trace_id is None: |
| @@ -198,6 +249,22 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str): | |||
| if kwargs.get("system") is None: | |||
| return messages | |||
| return [{"role": "system", "content": kwargs.get("system")}] + messages | |||
| elif provider == "gemini": | |||
| contents = kwargs.get("contents", []) | |||
| if isinstance(contents, str): | |||
There was a problem hiding this comment.
This whole bit is a repetition of format_input from gemini.py
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Added support for the google-genai gemini library withing posthogs ai sdk