blob: 2938a83eb0964d7bb7e045d5966f324f56d9dfca (
about) (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
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
|
