chore: add PostHog telemetry tracing for LLM assistant operations (baserow#4299)

* Add PostHog telemetry tracing for LLM assistant operations

This commit adds comprehensive PostHog telemetry integration to track
and analyze LLM assistant operations including:

- uDSPy module execution (ChainOfThought, ReAct, Predict)
- LLM API calls with automatic generation tracking
- Tool invocations
- Performance metrics (latency, token usage)
- Error tracking with exception messages

The telemetry system:
- Captures hierarchical span relationships for nested operations
- Patches OpenAI client for automatic generation event tracking
- Handles errors gracefully without breaking assistant execution
- Respects POSTHOG_ENABLED setting
- Adds workspace and user context to all events

Run the telemetry test suite:
```bash
just pytest -k test_telemetry
```

All 8 tests should pass, covering:
- Trace context manager (success and exception paths)
- Module execution tracing (ChainOfThought, ReAct, Predict)
- LM callback enrichment
- Tool execution tracing
- Exception handling (module and tool failures)
- Telemetry disabled state

To test manually:
1. Ensure POSTHOG_ENABLED=True in settings
2. Set up PostHog project credentials
3. Use the assistant and verify events appear in PostHog
4. Check for $ai_trace, $ai_span, and $ai_generation events
5. Verify span hierarchy and context propagation

* Fix issue when POSTHOG_PROJECT_API_KEY is not set

* Fix tests; capture all spans; add trace input and output

* Bump udspy version

Author: silvestrid

backend/requirements/base.in

@@ -64,7 +64,7 @@ Brotli==1.1.0
 loguru==0.7.2
 django-cachalot==2.6.2
 celery-singleton==0.3.1
-posthog==3.5.0
+posthog==7.0.1
 https://github.com/fellowapp/prosemirror-py/archive/refs/tags/v0.3.5.zip
 rich==13.7.1
 tzdata==2025.2
@@ -88,4 +88,4 @@ httpcore==1.0.9 # Pinned to address vulnerability.
 genson==1.3.0
 pyotp==2.9.0
 qrcode==8.2
-udspy==0.1.7
+udspy==0.1.8

backend/requirements/base.txt

@@ -125,6 +125,7 @@ distro==1.9.0
     # via
     #   anthropic
     #   openai
+    #   posthog
 dj-database-url==2.1.0
     # via -r base.in
 django==5.0.14
@@ -311,8 +312,6 @@ mdurl==0.1.2
     # via markdown-it-py
 mistralai==1.1.0
     # via -r base.in
-monotonic==1.6
-    # via posthog
 msgpack==1.1.0
     # via channels-redis
 mypy-extensions==1.0.0
@@ -454,7 +453,7 @@ pgvector==0.4.1
     # via -r base.in
 pillow==10.3.0
     # via -r base.in
-posthog==3.5.0
+posthog==7.0.1
     # via -r base.in
 prometheus-client==0.21.1
     # via flower
@@ -665,6 +664,7 @@ typing-extensions==4.11.0
     #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-sdk
     #   opentelemetry-semantic-conventions
+    #   posthog
     #   prosemirror
     #   pydantic
     #   pydantic-core
@@ -679,7 +679,7 @@ tzdata==2025.2
     #   -r base.in
     #   django-celery-beat
     #   kombu
-udspy==0.1.7
+udspy==0.1.8
     # via -r base.in
 unicodecsv==0.14.1
     # via -r base.in

backend/src/baserow/config/settings/base.py

@@ -11,7 +11,6 @@
 from django.core.exceptions import ImproperlyConfigured
 
 import dj_database_url
-import posthog
 import sentry_sdk
 from corsheaders.defaults import default_headers
 from sentry_sdk.integrations.django import DjangoIntegration
@@ -1243,13 +1242,8 @@ def __setitem__(self, key, value):
 )
 
 POSTHOG_PROJECT_API_KEY = os.getenv("POSTHOG_PROJECT_API_KEY", "")
-POSTHOG_HOST = os.getenv("POSTHOG_HOST", "")
-POSTHOG_ENABLED = POSTHOG_PROJECT_API_KEY and POSTHOG_HOST
-if POSTHOG_ENABLED:
-    posthog.project_api_key = POSTHOG_PROJECT_API_KEY
-    posthog.host = POSTHOG_HOST
-else:
-    posthog.disabled = True
+POSTHOG_HOST = os.getenv("POSTHOG_HOST") or None
+POSTHOG_ENABLED = bool(POSTHOG_PROJECT_API_KEY)
 
 BASEROW_BUILDER_DOMAINS = os.getenv("BASEROW_BUILDER_DOMAINS", None)
 BASEROW_BUILDER_DOMAINS = (

backend/src/baserow/core/posthog.py

@@ -6,13 +6,20 @@
 from django.contrib.auth.models import AbstractUser
 from django.dispatch import receiver
 
-import posthog
 from loguru import logger
+from posthog import Posthog
 
 from baserow.core.action.signals import ActionCommandType, action_done
 from baserow.core.models import Workspace
 from baserow.core.utils import exception_capturer
 
+posthog_client = Posthog(
+    settings.POSTHOG_PROJECT_API_KEY,
+    settings.POSTHOG_HOST,
+    # disabled=True will automatically avoid sending any data, even if capture is called
+    disabled=not settings.POSTHOG_ENABLED,
+)
+
 
 def capture_event(distinct_id: str, event: str, properties: dict):
     """
@@ -28,7 +35,7 @@ def capture_event(distinct_id: str, event: str, properties: dict):
         return
 
     try:
-        posthog.capture(
+        posthog_client.capture(
             distinct_id=distinct_id,
             event=event,
             properties=properties,

backend/tests/baserow/core/test_posthog.py

@@ -26,7 +26,7 @@ def scope(cls, *args, **kwargs):
 
 @pytest.mark.django_db
 @override_settings(POSTHOG_ENABLED=False)
-@patch("baserow.core.posthog.posthog")
+@patch("baserow.core.posthog.posthog_client")
 def test_not_capture_event_if_not_enabled(mock_posthog, data_fixture):
     user = data_fixture.create_user()
     capture_user_event(user, "test", {})
@@ -35,7 +35,7 @@ def test_not_capture_event_if_not_enabled(mock_posthog, data_fixture):
 
 @pytest.mark.django_db
 @override_settings(POSTHOG_ENABLED=True)
-@patch("baserow.core.posthog.posthog")
+@patch("baserow.core.posthog.posthog_client")
 def test_capture_event_if_enabled(mock_posthog, data_fixture):
     user = data_fixture.create_user()
     workspace = data_fixture.create_workspace()

enterprise/backend/src/baserow_enterprise/assistant/assistant.py

@@ -15,6 +15,7 @@
     AssistantMessageCancelled,
     AssistantModelNotSupportedError,
 )
+from baserow_enterprise.assistant.telemetry import PosthogTracingCallback
 from baserow_enterprise.assistant.tools.navigation.types import AnyNavigationRequestType
 from baserow_enterprise.assistant.tools.navigation.utils import unsafe_navigate_to
 from baserow_enterprise.assistant.tools.registries import assistant_tool_registry
@@ -190,22 +191,42 @@ def _init_assistant(self):
                 self._user, self._workspace, self.tool_helpers
             )
         ]
-        self.callbacks = AssistantCallbacks(self.tool_helpers)
+
+        self._assistant_callbacks = AssistantCallbacks(self.tool_helpers)
+        self._telemetry_callbacks = PosthogTracingCallback()
+        self._callbacks = [self._assistant_callbacks, self._telemetry_callbacks]
 
         module_kwargs = {
             "temperature": settings.BASEROW_ENTERPRISE_ASSISTANT_LLM_TEMPERATURE,
             "response_format": {"type": "json_object"},
         }
-
-        self.search_user_docs_tool = next(
-            (tool for tool in tools if tool.name == "search_user_docs"), None
-        )
+        self.search_user_docs_tool = self._get_search_user_docs_tool(tools)
         self.agent_tools = tools
         self._request_router = udspy.ChainOfThought(RequestRouter, **module_kwargs)
         self._assistant = udspy.ReAct(
             ChatSignature, tools=self.agent_tools, max_iters=20, **module_kwargs
         )
 
+    def _get_search_user_docs_tool(
+        self, tools: list[udspy.Tool | Callable]
+    ) -> udspy.Tool | None:
+        """
+        Retrieves the search_user_docs tool from the list of tools if available.
+
+        :param tools: The list of tools to search through.
+        :return: The search_user_docs as udspy.Tool or None if not found.
+        """
+
+        search_user_docs_tool = next(
+            (tool for tool in tools if tool.name == "search_user_docs"), None
+        )
+        if search_user_docs_tool is None or isinstance(
+            search_user_docs_tool, udspy.Tool
+        ):
+            return search_user_docs_tool
+
+        return udspy.Tool(search_user_docs_tool)
+
     async def acreate_chat_message(
         self,
         role: AssistantChatMessage.Role,
@@ -360,7 +381,7 @@ async def _acreate_ai_message_response(
         :return: The created AiMessage instance to return to the user.
         """
 
-        sources = self.callbacks.sources
+        sources = self._assistant_callbacks.sources
         ai_msg = await self.acreate_chat_message(
             AssistantChatMessage.Role.AI,
             prediction.answer,
@@ -449,7 +470,7 @@ async def _process_router_stream(
                 messages.append(
                     AiMessageChunk(
                         content=event.content,
-                        sources=self.callbacks.sources,
+                        sources=self._assistant_callbacks.sources,
                     )
                 )
 
@@ -472,7 +493,6 @@ async def _process_router_stream(
                                 "the local knowledge base. \n\n"
                                 "You can find setup instructions at: https://baserow.io/user-docs"
                             ),
-                            sources=[],
                         )
                     )
             elif getattr(event, "answer", None):
@@ -510,7 +530,7 @@ async def _process_agent_stream(
                 messages.append(
                     AiMessageChunk(
                         content=event.content,
-                        sources=self.callbacks.sources,
+                        sources=self._assistant_callbacks.sources,
                     )
                 )
 
@@ -586,11 +606,12 @@ async def astream_messages(
             AssistantChatMessage.Role.HUMAN,
             message.content,
         )
+        default_callbacks = udspy.settings.callbacks
 
         with udspy.settings.context(
             lm=self._lm_client,
-            callbacks=[*udspy.settings.callbacks, self.callbacks],
-        ):
+            callbacks=[*default_callbacks, *self._callbacks],
+        ), self._telemetry_callbacks.trace(self._chat, human_msg.content):
             message_id = str(human_msg.id)
             yield AiStartedMessage(message_id=message_id)