apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-07-26 23:51:14 +02:00
Code Issues Projects Releases Packages Wiki Activity

feat(backend): integrate PostHog analytics for enhanced observability

Browse source 
- Added PostHog configuration options to .env.example files for both Docker and Surfsense backend.
- Introduced PostHog dependency in pyproject.toml.
- Implemented analytics middleware to capture various events across the application, including user authentication, automation runs, and API requests.
- Enhanced existing routes and services to emit analytics events, providing insights into user interactions and system performance.
- Ensured graceful shutdown of analytics clients in worker processes and application lifecycles.
This commit is contained in:
DESKTOP-RTLN3BA\$punk 2026-07-22 22:16:28 -07:00
parent ca4f231577
commit dbedf0cfa5
47 changed files with 1618 additions and 513 deletions
Download patch file Download diff file Expand all files Collapse all files

83
surfsense_backend/app/tasks/celery_tasks/connector_tasks.py
View file

@ -1,5 +1,6 @@
"""Celery tasks for connector indexing."""
import contextlib
import logging
import time
import traceback
@ -8,6 +9,7 @@ from collections.abc import Awaitable, Callable
from celery import current_task
from app.celery_app import celery_app
from app.observability import analytics as ph_analytics
from app.observability import metrics as ot_metrics, otel as ot
from app.tasks.celery_tasks import (
get_celery_session_maker,
@ -17,8 +19,18 @@ from app.tasks.celery_tasks import (
logger = logging.getLogger(__name__)
def run_async_celery_task[T](coro_factory: Callable[[], Awaitable[T]]) -> T:
"""Run connector sync work and record aggregate connector metrics."""
def run_async_celery_task[T](
coro_factory: Callable[[], Awaitable[T]],
*,
workspace_id: int | None = None,
user_id: str | None = None,
) -> T:
"""Run connector sync work and record aggregate connector metrics.
When ``workspace_id``/``user_id`` are provided, also emits the PostHog
``connector_indexing_completed``/``_failed`` outcome (the frontend only
knows indexing *started*, never whether it worked). No-op without a key.
"""
task_name = getattr(current_task, "name", None) or "unknown"
t0 = time.perf_counter()
status = "failed"
@ -45,6 +57,29 @@ def run_async_celery_task[T](coro_factory: Callable[[], Awaitable[T]]) -> T:
status=status,
error_category=error_category,
)
if user_id and ph_analytics.is_enabled():
event = (
"connector_indexing_completed"
if status == "success"
else "connector_indexing_failed"
)
with contextlib.suppress(Exception):
ph_analytics.capture(
event,
distinct_id=str(user_id),
properties={
"workspace_id": workspace_id,
# ``connector_type`` is the Celery task name, e.g.
# index_notion_pages / index_github_repos.
"connector_type": task_name,
"status": status,
"error_category": error_category,
"duration_ms": int(elapsed_s * 1000),
},
groups={"workspace": str(workspace_id)}
if workspace_id is not None
else None,
)
def _handle_greenlet_error(e: Exception, task_name: str, connector_id: int) -> None:
@ -91,7 +126,9 @@ def index_notion_pages_task(
return run_async_celery_task(
lambda: _index_notion_pages(
connector_id, workspace_id, user_id, start_date, end_date
)
),
workspace_id=workspace_id,
user_id=user_id,
)
except Exception as e:
_handle_greenlet_error(e, "index_notion_pages", connector_id)
@ -129,7 +166,9 @@ def index_github_repos_task(
return run_async_celery_task(
lambda: _index_github_repos(
connector_id, workspace_id, user_id, start_date, end_date
)
),
workspace_id=workspace_id,
user_id=user_id,
)
@ -164,7 +203,9 @@ def index_confluence_pages_task(
return run_async_celery_task(
lambda: _index_confluence_pages(
connector_id, workspace_id, user_id, start_date, end_date
)
),
workspace_id=workspace_id,
user_id=user_id,
)
@ -200,7 +241,9 @@ def index_google_calendar_events_task(
return run_async_celery_task(
lambda: _index_google_calendar_events(
connector_id, workspace_id, user_id, start_date, end_date
)
),
workspace_id=workspace_id,
user_id=user_id,
)
except Exception as e:
_handle_greenlet_error(e, "index_google_calendar_events", connector_id)
@ -238,7 +281,9 @@ def index_google_gmail_messages_task(
return run_async_celery_task(
lambda: _index_google_gmail_messages(
connector_id, workspace_id, user_id, start_date, end_date
)
),
workspace_id=workspace_id,
user_id=user_id,
)
@ -275,7 +320,9 @@ def index_google_drive_files_task(
workspace_id,
user_id,
items_dict,
)
),
workspace_id=workspace_id,
user_id=user_id,
)
@ -315,7 +362,9 @@ def index_onedrive_files_task(
workspace_id,
user_id,
items_dict,
)
),
workspace_id=workspace_id,
user_id=user_id,
)
@ -355,7 +404,9 @@ def index_dropbox_files_task(
workspace_id,
user_id,
items_dict,
)
),
workspace_id=workspace_id,
user_id=user_id,
)
@ -393,7 +444,9 @@ def index_elasticsearch_documents_task(
return run_async_celery_task(
lambda: _index_elasticsearch_documents(
connector_id, workspace_id, user_id, start_date, end_date
)
),
workspace_id=workspace_id,
user_id=user_id,
)
@ -428,7 +481,9 @@ def index_bookstack_pages_task(
return run_async_celery_task(
lambda: _index_bookstack_pages(
connector_id, workspace_id, user_id, start_date, end_date
)
),
workspace_id=workspace_id,
user_id=user_id,
)
@ -463,7 +518,9 @@ def index_composio_connector_task(
return run_async_celery_task(
lambda: _index_composio_connector(
connector_id, workspace_id, user_id, start_date, end_date
)
),
workspace_id=workspace_id,
user_id=user_id,
)

98
surfsense_backend/app/tasks/celery_tasks/document_tasks.py
View file

@ -4,11 +4,13 @@ import asyncio
import contextlib
import logging
import os
import time
from uuid import UUID
from app.celery_app import celery_app
from app.config import config
from app.notifications.service import NotificationService
from app.observability import analytics as ph_analytics
from app.observability import metrics as ot_metrics
from app.services.task_logging_service import TaskLoggingService
from app.tasks.celery_tasks import get_celery_session_maker, run_async_celery_task
@ -22,6 +24,44 @@ from app.tasks.document_processors import (
logger = logging.getLogger(__name__)
def _capture_doc_processing(
status: str,
*,
user_id: str | None,
workspace_id: int,
doc_type: str,
file_size: int | None = None,
duration_ms: int | None = None,
) -> None:
"""Emit ``document_processing_completed``/``_failed`` from a Celery task.
The frontend only knows the upload POST succeeded, never whether ingestion
actually worked — this is the authoritative outcome. No-op when PostHog is
unconfigured. ``distinct_id`` is the owning user's id so it joins the same
person the web app identifies.
"""
if not ph_analytics.is_enabled() or not user_id:
return
event = (
"document_processing_completed"
if status == "success"
else "document_processing_failed"
)
ph_analytics.capture(
event,
distinct_id=str(user_id),
properties={
"workspace_id": workspace_id,
"doc_type": doc_type,
"file_size": file_size,
"duration_ms": duration_ms,
"status": status,
},
groups={"workspace": str(workspace_id)},
)
# ===== Redis heartbeat for document processing tasks =====
# Same mechanism as connector indexing heartbeats (search_source_connectors_routes.py).
# A background coroutine refreshes a Redis key every 60s with a 2-min TTL.
@ -268,11 +308,30 @@ def process_extension_document_task(
workspace_id: ID of the workspace
user_id: ID of the user
"""
return run_async_celery_task(
lambda: _process_extension_document(
individual_document_dict, workspace_id, user_id
_t0 = time.perf_counter()
try:
result = run_async_celery_task(
lambda: _process_extension_document(
individual_document_dict, workspace_id, user_id
)
)
except Exception:
_capture_doc_processing(
"failed",
user_id=user_id,
workspace_id=workspace_id,
doc_type="extension",
duration_ms=int((time.perf_counter() - _t0) * 1000),
)
raise
_capture_doc_processing(
"success",
user_id=user_id,
workspace_id=workspace_id,
doc_type="extension",
duration_ms=int((time.perf_counter() - _t0) * 1000),
)
return result
async def _process_extension_document(
@ -430,12 +489,14 @@ def process_file_upload_task(
)
return
file_size: int | None = None
try:
file_size = os.path.getsize(file_path)
logger.info(f"[process_file_upload] File size: {file_size} bytes")
except Exception as e:
logger.warning(f"[process_file_upload] Could not get file size: {e}")
_t0 = time.perf_counter()
try:
run_async_celery_task(
lambda: _process_file_upload(file_path, filename, workspace_id, user_id)
@ -448,7 +509,23 @@ def process_file_upload_task(
f"[process_file_upload] Task failed for {filename}: {e}\n"
f"Traceback:\n{traceback.format_exc()}"
)
_capture_doc_processing(
"failed",
user_id=user_id,
workspace_id=workspace_id,
doc_type="file_upload",
file_size=file_size,
duration_ms=int((time.perf_counter() - _t0) * 1000),
)
raise
_capture_doc_processing(
"success",
user_id=user_id,
workspace_id=workspace_id,
doc_type="file_upload",
file_size=file_size,
duration_ms=int((time.perf_counter() - _t0) * 1000),
)
async def _process_file_upload(
@ -682,6 +759,7 @@ def process_file_upload_with_document_task(
)
return
_t0 = time.perf_counter()
try:
run_async_celery_task(
lambda: _process_file_with_document(
@ -702,7 +780,21 @@ def process_file_upload_with_document_task(
f"[process_file_upload_with_document] Task failed for {filename}: {e}\n"
f"Traceback:\n{traceback.format_exc()}"
)
_capture_doc_processing(
"failed",
user_id=user_id,
workspace_id=workspace_id,
doc_type="file_upload_2phase",
duration_ms=int((time.perf_counter() - _t0) * 1000),
)
raise
_capture_doc_processing(
"success",
user_id=user_id,
workspace_id=workspace_id,
doc_type="file_upload_2phase",
duration_ms=int((time.perf_counter() - _t0) * 1000),
)
async def _mark_document_failed(document_id: int, reason: str):

15
surfsense_backend/app/tasks/celery_tasks/video_presentation_tasks.py
View file

@ -12,6 +12,7 @@ from app.agents.video_presentation.state import State as VideoPresentationState
from app.celery_app import celery_app
from app.config import config as app_config
from app.db import VideoPresentation, VideoPresentationStatus
from app.observability import analytics as ph_analytics
from app.services.billable_calls import (
BillingSettlementError,
QuotaInsufficientError,
@ -239,6 +240,20 @@ async def _generate_video_presentation(
logger.info(f"Successfully generated video presentation: {video_pres.id}")
# Credit-consuming deliverable — the frontend never confirms
# completion. Attributed to the workspace owner resolved above.
if owner_user_id:
ph_analytics.capture(
"video_presentation_generated",
distinct_id=str(owner_user_id),
properties={
"workspace_id": workspace_id,
"video_presentation_id": video_pres.id,
"slide_count": len(serializable_slides),
},
groups={"workspace": str(workspace_id)},
)
return {
"status": "ready",
"video_presentation_id": video_pres.id,

42
surfsense_backend/app/tasks/chat/streaming/flows/new_chat/orchestrator.py
View file

@ -97,6 +97,10 @@ from app.tasks.chat.streaming.flows.shared.rate_limit_recovery import (
log_rate_limit_recovered,
reroute_to_next_auto_pin,
)
from app.tasks.chat.streaming.flows.shared.analytics import (
build_llm_callback_handler,
capture_chat_turn_completed,
)
from app.tasks.chat.streaming.flows.shared.span import (
close_chat_request_span,
open_chat_request_span,
@ -487,6 +491,23 @@ async def stream_new_chat(
"recursion_limit": 10_000,
}
# PostHog LLM analytics: attach a callback so the full agent trace
# tree (LLM calls, tools, subagents) is captured per turn. No-op when
# PostHog is unconfigured.
_llm_handler = build_llm_callback_handler(
distinct_id=user_id,
trace_id=stream_result.turn_id,
properties={
"workspace_id": workspace_id,
"chat_id": chat_id,
"$ai_session_id": str(chat_id),
"flow": flow,
},
groups={"workspace": str(workspace_id)},
)
if _llm_handler is not None:
config["callbacks"] = [_llm_handler]
# --- Block 4: First SSE frames ---
for sse in iter_initial_frames(
@ -830,6 +851,27 @@ async def stream_new_chat(
log_prefix="stream_new_chat",
)
# Authoritative server-side product event. Inside the shield so it
# survives client-disconnect cancellation (abandoned tabs still
# produce a turn event), and while ``stream_result`` is still live
# (it's dropped to None below for GC).
capture_chat_turn_completed(
flow=flow,
outcome=chat_outcome,
error_category=chat_error_category,
workspace_id=workspace_id,
chat_id=chat_id,
user_id=user_id,
auth_context=auth_context,
agent_mode=chat_agent_mode,
client_platform=fs_platform,
filesystem_mode=fs_mode,
turn_id=getattr(stream_result, "turn_id", None),
request_id=request_id,
duration_ms=int((time.perf_counter() - _t_total) * 1000),
accumulator=accumulator,
)
# Persist any sandbox-produced files to local storage so they remain
# downloadable after the Daytona sandbox auto-deletes.
if stream_result and stream_result.sandbox_files:

39
surfsense_backend/app/tasks/chat/streaming/flows/resume_chat/orchestrator.py
View file

@ -76,6 +76,10 @@ from app.tasks.chat.streaming.flows.shared.rate_limit_recovery import (
log_rate_limit_recovered,
reroute_to_next_auto_pin,
)
from app.tasks.chat.streaming.flows.shared.analytics import (
build_llm_callback_handler,
capture_chat_turn_completed,
)
from app.tasks.chat.streaming.flows.shared.span import (
close_chat_request_span,
open_chat_request_span,
@ -387,6 +391,22 @@ async def stream_resume_chat(
"recursion_limit": 10_000,
}
# PostHog LLM analytics — same trace as the original turn's
# conversation via ``$ai_session_id``. No-op when unconfigured.
_llm_handler = build_llm_callback_handler(
distinct_id=user_id,
trace_id=stream_result.turn_id,
properties={
"workspace_id": workspace_id,
"chat_id": chat_id,
"$ai_session_id": str(chat_id),
"flow": "resume",
},
groups={"workspace": str(workspace_id)},
)
if _llm_handler is not None:
config["callbacks"] = [_llm_handler]
# --- First SSE frames ---
for sse in iter_initial_frames(
@ -599,6 +619,25 @@ async def stream_resume_chat(
log_prefix="stream_resume",
)
# Authoritative server-side product event (see new_chat
# orchestrator for rationale). ``flow="resume"``.
capture_chat_turn_completed(
flow="resume",
outcome=chat_outcome,
error_category=chat_error_category,
workspace_id=workspace_id,
chat_id=chat_id,
user_id=user_id,
auth_context=auth_context,
agent_mode=chat_agent_mode,
client_platform=fs_platform,
filesystem_mode=fs_mode,
turn_id=getattr(stream_result, "turn_id", None),
request_id=request_id,
duration_ms=int((time.perf_counter() - _t_total) * 1000),
accumulator=accumulator,
)
# Release the lock from the original interrupted turn or any
# re-interrupt/bailout. Skip on ``BusyError`` (lock not held here).
if not busy_error_raised:

117
surfsense_backend/app/tasks/chat/streaming/flows/shared/analytics.py Normal file
View file

@ -0,0 +1,117 @@
"""PostHog chat-turn analytics for streaming flows.
Emits a single authoritative ``chat_turn_completed`` product event per turn,
shared by the new-chat and resume orchestrators so every chat source (web,
desktop, PAT scripts, gateway, automations) is tracked identically —
including sources the frontend can never observe. No-op when PostHog is
unconfigured.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any
from app.config import config
from app.observability import analytics
if TYPE_CHECKING:
from app.auth.context import AuthContext
from app.services.token_tracking_service import TurnTokenAccumulator
logger = logging.getLogger(__name__)
def build_llm_callback_handler(
*,
distinct_id: str | None,
trace_id: str | None,
properties: dict[str, Any] | None = None,
groups: dict[str, str] | None = None,
) -> Any | None:
"""Build a PostHog LangChain ``CallbackHandler`` for a chat turn.
Attaching this to the LangGraph ``config["callbacks"]`` captures the full
agent trace tree ($ai_trace / $ai_span / $ai_generation) — every LLM call,
tool, subagent, and retriever — joined to the ``chat_turn_completed`` event
via ``trace_id`` (the turn id) and grouped per conversation via
``$ai_session_id`` (in ``properties``).
Returns ``None`` when PostHog is disabled or the package/handler is
unavailable, so callers can simply skip attaching callbacks. ``privacy_mode``
(default on) suppresses prompt/completion bodies — chat content includes
users' private documents.
"""
client_obj = analytics.get_client()
if client_obj is None or not distinct_id:
return None
try:
from posthog.ai.langchain import CallbackHandler
return CallbackHandler(
client=client_obj,
distinct_id=distinct_id,
trace_id=trace_id,
properties=properties or {},
privacy_mode=config.POSTHOG_AI_PRIVACY_MODE,
groups=groups,
)
except Exception:
logger.debug("PostHog LLM callback handler unavailable", exc_info=True)
return None
def capture_chat_turn_completed(
*,
flow: str,
outcome: str,
error_category: str | None,
workspace_id: int,
chat_id: int,
user_id: str | None,
auth_context: AuthContext | None,
agent_mode: str,
client_platform: str,
filesystem_mode: str,
turn_id: str | None,
request_id: str | None,
duration_ms: int,
accumulator: TurnTokenAccumulator,
) -> None:
"""Capture ``chat_turn_completed``. Best-effort; never raises."""
if not analytics.is_enabled() or not user_id:
return
props: dict[str, Any] = {
"flow": flow,
"outcome": outcome,
"error_category": error_category,
"workspace_id": workspace_id,
"chat_id": chat_id,
"agent_mode": agent_mode,
"client_platform": client_platform,
"filesystem_mode": filesystem_mode,
"turn_id": turn_id,
"request_id": request_id,
"duration_ms": duration_ms,
# Cost is micro-USD (integer), matching TurnTokenAccumulator; do not
# convert to float dollars.
"total_tokens": accumulator.grand_total,
"prompt_tokens": accumulator.total_prompt_tokens,
"completion_tokens": accumulator.total_completion_tokens,
"cost_micros": accumulator.total_cost_micros,
}
groups = {"workspace": str(workspace_id)}
if auth_context is not None:
analytics.capture_for(
auth_context, "chat_turn_completed", props, groups=groups
)
else:
analytics.capture(
"chat_turn_completed",
distinct_id=user_id,
properties=props,
groups=groups,
)