Merge branch 'master' into distributed-embedding-table

Author: oliverb123

ee/api/session_summaries.py

@@ -106,7 +106,7 @@ def _if_video_validation_enabled(self, user: User) -> bool | None:
     @staticmethod
     async def _get_summary_from_progress_stream(
         session_ids: list[str],
-        user_id: int,
+        user: User,
         team: Team,
         min_timestamp: datetime,
         max_timestamp: datetime,
@@ -117,7 +117,7 @@ async def _get_summary_from_progress_stream(
         results: list[tuple[SessionSummaryStreamUpdate, tuple[EnrichedSessionGroupSummaryPatternsList, str] | str]] = []
         async for update in execute_summarize_session_group(
             session_ids=session_ids,
-            user_id=user_id,
+            user=user,
             team=team,
             min_timestamp=min_timestamp,
             max_timestamp=max_timestamp,
@@ -171,7 +171,7 @@ def create_session_summaries(self, request: Request, **kwargs) -> Response:
         try:
             summary = async_to_sync(self._get_summary_from_progress_stream)(
                 session_ids=session_ids,
-                user_id=user.id,
+                user=user,
                 team=self.team,
                 min_timestamp=min_timestamp,
                 max_timestamp=max_timestamp,
@@ -217,15 +217,15 @@ def create_session_summaries(self, request: Request, **kwargs) -> Response:
     @staticmethod
     async def _summarize_session(
         session_id: str,
-        user_id: int,
+        user: User,
         team: Team,
         video_validation_enabled: bool | None,
         extra_summary_context: ExtraSummaryContext | None = None,
     ) -> SessionSummarySerializer | Exception:
         try:
             summary_raw = await execute_summarize_session(
                 session_id=session_id,
-                user_id=user_id,
+                user=user,
                 team=team,
                 video_validation_enabled=video_validation_enabled,
                 extra_summary_context=extra_summary_context,
@@ -240,7 +240,7 @@ async def _summarize_session(
     async def _get_individual_summaries(
         self,
         session_ids: list[str],
-        user_id: int,
+        user: User,
         team: Team,
         video_validation_enabled: bool | None,
         extra_summary_context: ExtraSummaryContext | None = None,
@@ -251,7 +251,7 @@ async def _get_individual_summaries(
                 tasks[session_id] = tg.create_task(
                     self._summarize_session(
                         session_id=session_id,
-                        user_id=user_id,
+                        user=user,
                         team=team,
                         video_validation_enabled=video_validation_enabled,
                         extra_summary_context=extra_summary_context,
@@ -262,9 +262,9 @@ async def _get_individual_summaries(
             res: SessionSummarySerializer | Exception = task.result()
             if isinstance(res, Exception):
                 logger.exception(
-                    f"Failed to generate individual session summary for session {session_id} from team {team.pk} by user {user_id}: {res}",
+                    f"Failed to generate individual session summary for session {session_id} from team {team.pk} by user {user.id}: {res}",
                     team_id=team.pk,
-                    user_id=user_id,
+                    user_id=user.id,
                 )
             else:
                 # Return only successful summaries
@@ -296,7 +296,7 @@ def create_session_summaries_individually(self, request: Request, **kwargs) -> R
         try:
             summaries = async_to_sync(self._get_individual_summaries)(
                 session_ids=session_ids,
-                user_id=user.id,
+                user=user,
                 team=self.team,
                 video_validation_enabled=video_validation_enabled,
                 extra_summary_context=extra_summary_context,

ee/api/test/test_session_summaries.py

@@ -123,7 +123,7 @@ def test_create_summaries_success(
         # Verify execute_summarize_session_group was called correctly
         mock_execute.assert_called_once_with(
             session_ids=["session1", "session2"],
-            user_id=self.user.pk,
+            user=self.user,
             team=self.team,
             min_timestamp=datetime(2024, 1, 1, 10, 0, 0),
             max_timestamp=datetime(2024, 1, 1, 11, 0, 0),

ee/hogai/chat_agent/session_summaries/nodes.py

@@ -488,7 +488,7 @@ async def _summarize(session_id: str) -> dict[str, Any]:
             nonlocal completed
             result = await execute_summarize_session(
                 session_id=session_id,
-                user_id=self._node._user.id,
+                user=self._node._user,
                 team=self._node._team,
                 model_to_use=SESSION_SUMMARIES_SYNC_MODEL,
                 video_validation_enabled=video_validation_enabled,
@@ -521,7 +521,7 @@ async def _summarize_sessions_as_group(
 
         async for update_type, data in execute_summarize_session_group(
             session_ids=session_ids,
-            user_id=self._node._user.id,
+            user=self._node._user,
             team=self._node._team,
             min_timestamp=min_timestamp,
             max_timestamp=max_timestamp,

ee/hogai/session_summaries/session/stream.py

@@ -1,6 +1,7 @@
 from collections.abc import Generator
 
 from posthog.models.team.team import Team
+from posthog.models.user import User
 from posthog.settings import SERVER_GATEWAY_INTERFACE
 from posthog.temporal.ai.session_summary.summarize_session import execute_summarize_session_stream
 
@@ -10,22 +11,22 @@
 
 def stream_recording_summary(
     session_id: str,
-    user_id: int,
+    user: User,
     team: Team,
     extra_summary_context: ExtraSummaryContext | None = None,
     local_reads_prod: bool = False,
 ) -> SyncIterableToAsync | Generator[str, None, None]:
     if SERVER_GATEWAY_INTERFACE == "ASGI":
         return _astream(
             session_id=session_id,
-            user_id=user_id,
+            user=user,
             team=team,
             extra_summary_context=extra_summary_context,
             local_reads_prod=local_reads_prod,
         )
     return execute_summarize_session_stream(
         session_id=session_id,
-        user_id=user_id,
+        user=user,
         team=team,
         extra_summary_context=extra_summary_context,
         local_reads_prod=local_reads_prod,
@@ -34,15 +35,15 @@ def stream_recording_summary(
 
 def _astream(
     session_id: str,
-    user_id: int,
+    user: User,
     team: Team,
     extra_summary_context: ExtraSummaryContext | None = None,
     local_reads_prod: bool = False,
 ) -> SyncIterableToAsync:
     return SyncIterableToAsync(
         execute_summarize_session_stream(
             session_id=session_id,
-            user_id=user_id,
+            user=user,
             team=team,
             extra_summary_context=extra_summary_context,
             local_reads_prod=local_reads_prod,

ee/hogai/session_summaries/session/summarize_session.py

@@ -63,6 +63,7 @@ class SingleSessionSummaryLlmInputs:
 
     session_id: str
     user_id: int
+    user_distinct_id_to_log: str | None = None
     summary_prompt: str
     system_prompt: str
     simplified_events_mapping: dict[str, list[str | int | None | list[str]]]
@@ -218,7 +219,12 @@ async def prepare_data_for_single_session_summary(
 
 
 def prepare_single_session_summary_input(
-    session_id: str, user_id: int, summary_data: SingleSessionSummaryData, model_to_use: str
+    session_id: str,
+    user_id: int,
+    summary_data: SingleSessionSummaryData,
+    model_to_use: str,
+    *,
+    user_distinct_id_to_log: str | None = None,
 ) -> SingleSessionSummaryLlmInputs:
     # Checking here instead of in the preparation function to keep mypy happy
     if summary_data.prompt_data is None:
@@ -241,6 +247,7 @@ def prepare_single_session_summary_input(
     input_data = SingleSessionSummaryLlmInputs(
         session_id=session_id,
         user_id=user_id,
+        user_distinct_id_to_log=user_distinct_id_to_log,
         summary_prompt=summary_data.prompt.summary_prompt,
         system_prompt=summary_data.prompt.system_prompt,
         simplified_events_mapping=summary_data.prompt_data.simplified_events_mapping,

ee/hogai/session_summaries/tests/test_summarize_session.py

@@ -48,7 +48,7 @@ def test_execute_summarize_session_stream_success(
             # Get the generator (stream simulation)
             empty_context = ExtraSummaryContext()
             result_generator = execute_summarize_session_stream(
-                session_id=mock_session_id, user_id=mock_user.id, team=mock_team, extra_summary_context=empty_context
+                session_id=mock_session_id, user=mock_user, team=mock_team, extra_summary_context=empty_context
             )
             # Get all results from generator (consume the stream fully)
             results = list(result_generator)
@@ -124,14 +124,14 @@ async def test_stream_recording_summary_asgi(
                 return_value=iter([ready_summary]),
             ) as mock_execute,
         ):
-            async_gen = stream_recording_summary(session_id=mock_session_id, user_id=mock_user.id, team=mock_team)
+            async_gen = stream_recording_summary(session_id=mock_session_id, user=mock_user, team=mock_team)
             assert isinstance(async_gen, SyncIterableToAsync)
             results = [chunk async for chunk in async_gen]
             assert len(results) == 1
             assert results[0] == ready_summary
             mock_execute.assert_called_once_with(
                 session_id=mock_session_id,
-                user_id=mock_user.id,
+                user=mock_user,
                 team=mock_team,
                 extra_summary_context=None,
                 local_reads_prod=False,
@@ -155,13 +155,13 @@ def test_stream_recording_summary_wsgi(
                 return_value=iter([ready_summary]),
             ) as mock_execute,
         ):
-            result_gen = stream_recording_summary(session_id=mock_session_id, user_id=mock_user.id, team=mock_team)
+            result_gen = stream_recording_summary(session_id=mock_session_id, user=mock_user, team=mock_team)
             results = list(result_gen)  # type: ignore[arg-type]
             assert len(results) == 1
             assert results[0] == ready_summary
             mock_execute.assert_called_once_with(
                 session_id=mock_session_id,
-                user_id=mock_user.id,
+                user=mock_user,
                 team=mock_team,
                 extra_summary_context=None,
                 local_reads_prod=False,

ee/hogai/tools/replay/summarize_sessions.py

@@ -459,7 +459,7 @@ async def _summarize(session_id: str) -> dict[str, Any]:
             nonlocal completed
             result = await execute_summarize_session(
                 session_id=session_id,
-                user_id=self._user.id,
+                user=self._user,
                 team=self._team,
                 model_to_use=SESSION_SUMMARIES_SYNC_MODEL,
                 video_validation_enabled=video_validation_enabled,
@@ -496,7 +496,7 @@ async def _summarize_sessions_as_group(
 
         async for update_type, data in execute_summarize_session_group(
             session_ids=session_ids,
-            user_id=self._user.id,
+            user=self._user,
             team=self._team,
             min_timestamp=min_timestamp,
             max_timestamp=max_timestamp,

ee/tasks/subscriptions/__init__.py

@@ -23,7 +23,7 @@
     send_slack_message_with_integration_async,
     send_slack_subscription_report,
 )
-from ee.tasks.subscriptions.subscription_utils import generate_assets, generate_assets_async
+from ee.tasks.subscriptions.subscription_utils import _has_asset_failed, generate_assets, generate_assets_async
 
 logger = structlog.get_logger(__name__)
 
@@ -141,6 +141,9 @@ async def deliver_subscription_report_async(
         "deliver_subscription_report_async.assets_generated", subscription_id=subscription_id, asset_count=len(assets)
     )
 
+    if any(_has_asset_failed(asset) for asset in assets):
+        get_subscription_failure_metric(subscription.target_type, "temporal", failure_type="asset_generation").add(1)
+
     if not assets:
         logger.warning("deliver_subscription_report_async.no_assets", subscription_id=subscription_id)
         capture_exception(Exception("No assets are in this subscription"), {"subscription_id": subscription.id})

posthog/session_recordings/session_recording_api.py

@@ -1258,7 +1258,7 @@ def summarize(self, request: request.Request, **kwargs):
         # If you want to test sessions locally - override `session_id` and `self.team.pk`
         # with session/team ids of your choice and set `local_reads_prod` to True
         return StreamingHttpResponse(
-            stream_recording_summary(session_id=session_id, user_id=user.pk, team=self.team),
+            stream_recording_summary(session_id=session_id, user=user, team=self.team),
             content_type=ServerSentEventRenderer.media_type,
         )
 

posthog/tasks/exports/image_exporter.py

@@ -205,30 +205,30 @@ def _screenshot_asset(
 
         try:
             WebDriverWait(driver, timeout).until(lambda x: x.find_element(By.CSS_SELECTOR, wait_for_css_selector))
-        except TimeoutException:
+        except TimeoutException as e:
             with posthoganalytics.new_context():
                 posthoganalytics.tag("stage", "image_exporter.page_load_timeout")
                 try:
                     driver.save_screenshot(image_path)
                     posthoganalytics.tag("image_path", image_path)
                 except Exception:
                     pass
-                capture_exception()
+                capture_exception(e)
 
             raise Exception(f"Timeout while waiting for the page to load")
 
         try:
             # Also wait until nothing is loading
             WebDriverWait(driver, 20).until_not(lambda x: x.find_element(By.CLASS_NAME, "Spinner"))
-        except TimeoutException:
+        except TimeoutException as e:
             with posthoganalytics.new_context():
                 posthoganalytics.tag("stage", "image_exporter.wait_for_spinner_timeout")
                 try:
                     driver.save_screenshot(image_path)
                     posthoganalytics.tag("image_path", image_path)
                 except Exception:
                     pass
-                capture_exception()
+                capture_exception(e)
 
         # Get the height of the visualization container specifically
         height = driver.execute_script(