1 files changed, 56 insertions, 0 deletions
diff --git a/R2R/r2r/telemetry/telemetry_decorator.py b/R2R/r2r/telemetry/telemetry_decorator.py
new file mode 100755
index 00000000..2938a83e
--- /dev/null
+++ b/R2R/r2r/telemetry/telemetry_decorator.py
@@ -0,0 +1,56 @@
+import asyncio
+import logging
+from functools import wraps
+
+from r2r.telemetry.events import ErrorEvent, FeatureUsageEvent
+from r2r.telemetry.posthog import telemetry_client
+
+logger = logging.getLogger(__name__)
+
+
+def telemetry_event(event_name):
+    def decorator(func):
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            user_id = kwargs.get("user_id", "unknown_user")
+            try:
+                result = await func(*args, **kwargs)
+                try:
+                    telemetry_client.capture(
+                        FeatureUsageEvent(user_id=user_id, feature=event_name)
+                    )
+                except Exception as e:
+                    logger.error(f"Error in telemetry event logging: {str(e)}")
+                return result
+            except Exception as e:
+                try:
+                    telemetry_client.capture(
+                        ErrorEvent(
+                            user_id=user_id,
+                            endpoint=event_name,
+                            error_message=str(e),
+                        )
+                    )
+                except Exception as e:
+                    logger.error(f"Error in telemetry event logging: {str(e)}")
+
+                raise
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                future = asyncio.run_coroutine_threadsafe(
+                    async_wrapper(*args, **kwargs), loop
+                )
+                return future.result()
+            else:
+                return loop.run_until_complete(async_wrapper(*args, **kwargs))
+
+        return (
+            async_wrapper
+            if asyncio.iscoroutinefunction(func)
+            else sync_wrapper
+        )
+
+    return decorator