110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
import os
|
|
import uuid
|
|
from contextlib import contextmanager
|
|
from datetime import datetime
|
|
from typing import Dict, Any, Optional
|
|
from datetime import datetime as dt, timezone as tz
|
|
from portkey_ai import Portkey, Config
|
|
import logging
|
|
|
|
from .business_event_context import BusinessEventContext
|
|
from common.models.monitoring import BusinessEventLog
|
|
from common.extensions import db
|
|
|
|
|
|
class BusinessEvent:
|
|
# The BusinessEvent class itself is a context manager, but it doesn't use the @contextmanager decorator.
|
|
# Instead, it defines __enter__ and __exit__ methods explicitly. This is because we're doing something a bit more
|
|
# complex - we're interacting with the BusinessEventContext and the _business_event_stack.
|
|
|
|
def __init__(self, event_type: str, tenant_id: int, **kwargs):
|
|
self.event_type = event_type
|
|
self.tenant_id = tenant_id
|
|
self.trace_id = str(uuid.uuid4())
|
|
self.span_id = None
|
|
self.span_name = None
|
|
self.parent_span_id = None
|
|
self.document_version_id = kwargs.get('document_version_id')
|
|
self.chat_session_id = kwargs.get('chat_session_id')
|
|
self.interaction_id = kwargs.get('interaction_id')
|
|
self.environment = os.environ.get("FLASK_ENV", "development")
|
|
self.span_counter = 0
|
|
self.spans = []
|
|
|
|
def update_attribute(self, attribute: str, value: any):
|
|
if hasattr(self, attribute):
|
|
setattr(self, attribute, value)
|
|
else:
|
|
raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{attribute}'")
|
|
|
|
@contextmanager
|
|
def create_span(self, span_name: str):
|
|
# The create_span method is designed to be used as a context manager. We want to perform some actions when
|
|
# entering the span (like setting the span ID and name) and some actions when exiting the span (like removing
|
|
# these temporary attributes). The @contextmanager decorator allows us to write this method in a way that
|
|
# clearly separates the "entry" and "exit" logic, with the yield statement in between.
|
|
|
|
parent_span_id = self.span_id
|
|
self.span_counter += 1
|
|
new_span_id = f"{self.trace_id}-{self.span_counter}"
|
|
|
|
# Save the current span info
|
|
self.spans.append((self.span_id, self.span_name, self.parent_span_id))
|
|
|
|
# Set the new span info
|
|
self.span_id = new_span_id
|
|
self.span_name = span_name
|
|
self.parent_span_id = parent_span_id
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
# Restore the previous span info
|
|
if self.spans:
|
|
self.span_id, self.span_name, self.parent_span_id = self.spans.pop()
|
|
else:
|
|
self.span_id = None
|
|
self.span_name = None
|
|
self.parent_span_id = None
|
|
|
|
def log(self, message: str, level: str = 'info'):
|
|
logger = logging.getLogger('business_events')
|
|
log_data = {
|
|
'event_type': self.event_type,
|
|
'tenant_id': self.tenant_id,
|
|
'trace_id': self.trace_id,
|
|
'span_id': self.span_id,
|
|
'span_name': self.span_name,
|
|
'parent_span_id': self.parent_span_id,
|
|
'document_version_id': self.document_version_id,
|
|
'chat_session_id': self.chat_session_id,
|
|
'interaction_id': self.interaction_id,
|
|
'environment': self.environment
|
|
}
|
|
# log to Graylog
|
|
getattr(logger, level)(message, extra=log_data)
|
|
|
|
# Log to database
|
|
event_log = BusinessEventLog(
|
|
timestamp=dt.now(tz=tz.utc),
|
|
event_type=self.event_type,
|
|
tenant_id=self.tenant_id,
|
|
trace_id=self.trace_id,
|
|
span_id=self.span_id,
|
|
span_name=self.span_name,
|
|
parent_span_id=self.parent_span_id,
|
|
document_version_id=self.document_version_id,
|
|
chat_session_id=self.chat_session_id,
|
|
interaction_id=self.interaction_id,
|
|
environment=self.environment,
|
|
message=message
|
|
)
|
|
db.session.add(event_log)
|
|
db.session.commit()
|
|
|
|
def __enter__(self):
|
|
return BusinessEventContext(self).__enter__()
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
return BusinessEventContext(self).__exit__(exc_type, exc_val, exc_tb)
|