import os import uuid from contextlib import contextmanager from datetime import datetime from typing import Dict, Any, Optional, List from datetime import datetime as dt, timezone as tz import logging from .business_event_context import BusinessEventContext from common.models.entitlements import BusinessEventLog from common.extensions import db from .celery_utils import current_celery 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.document_version_file_size = kwargs.get('document_version_file_size') 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 = [] self.llm_metrics = { 'total_tokens': 0, 'prompt_tokens': 0, 'completion_tokens': 0, 'total_time': 0, 'call_count': 0, 'interaction_type': None } self._log_buffer = [] 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}'") def update_llm_metrics(self, metrics: dict): self.llm_metrics['total_tokens'] += metrics['total_tokens'] self.llm_metrics['prompt_tokens'] += metrics['prompt_tokens'] self.llm_metrics['completion_tokens'] += metrics['completion_tokens'] self.llm_metrics['total_time'] += metrics['time_elapsed'] self.llm_metrics['call_count'] += 1 self.llm_metrics['interaction_type'] = metrics['interaction_type'] def reset_llm_metrics(self): self.llm_metrics['total_tokens'] = 0 self.llm_metrics['prompt_tokens'] = 0 self.llm_metrics['completion_tokens'] = 0 self.llm_metrics['total_time'] = 0 self.llm_metrics['call_count'] = 0 self.llm_metrics['interaction_type'] = None @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 = str(uuid.uuid4()) # 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 self.log(f"Start") try: yield finally: if self.llm_metrics['call_count'] > 0: self.log_final_metrics() self.reset_llm_metrics() self.log(f"End") # 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'): log_data = { '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, 'document_version_file_size': self.document_version_file_size, 'chat_session_id': self.chat_session_id, 'interaction_id': self.interaction_id, 'environment': self.environment, 'message': message, } self._log_buffer.append(log_data) def log_llm_metrics(self, metrics: dict, level: str = 'info'): self.update_llm_metrics(metrics) message = "LLM Metrics" logger = logging.getLogger('business_events') log_data = { '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, 'document_version_file_size': self.document_version_file_size, 'chat_session_id': self.chat_session_id, 'interaction_id': self.interaction_id, 'environment': self.environment, 'llm_metrics_total_tokens': metrics['total_tokens'], 'llm_metrics_prompt_tokens': metrics['prompt_tokens'], 'llm_metrics_completion_tokens': metrics['completion_tokens'], 'llm_metrics_total_time': metrics['time_elapsed'], 'llm_interaction_type': metrics['interaction_type'], 'message': message, } self._log_buffer.append(log_data) def log_final_metrics(self, level: str = 'info'): logger = logging.getLogger('business_events') message = "Final LLM Metrics" log_data = { '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, 'document_version_file_size': self.document_version_file_size, 'chat_session_id': self.chat_session_id, 'interaction_id': self.interaction_id, 'environment': self.environment, 'llm_metrics_total_tokens': self.llm_metrics['total_tokens'], 'llm_metrics_prompt_tokens': self.llm_metrics['prompt_tokens'], 'llm_metrics_completion_tokens': self.llm_metrics['completion_tokens'], 'llm_metrics_total_time': self.llm_metrics['total_time'], 'llm_metrics_call_count': self.llm_metrics['call_count'], 'llm_interaction_type': self.llm_metrics['interaction_type'], 'message': message, } self._log_buffer.append(log_data) @staticmethod def _direct_db_persist(log_entries: List[Dict[str, Any]]): """Fallback method to directly persist logs to DB if async fails""" try: db_entries = [] for entry in log_entries: event_log = BusinessEventLog( timestamp=entry.pop('timestamp'), event_type=entry.pop('event_type'), tenant_id=entry.pop('tenant_id'), trace_id=entry.pop('trace_id'), span_id=entry.pop('span_id', None), span_name=entry.pop('span_name', None), parent_span_id=entry.pop('parent_span_id', None), document_version_id=entry.pop('document_version_id', None), document_version_file_size=entry.pop('document_version_file_size', None), chat_session_id=entry.pop('chat_session_id', None), interaction_id=entry.pop('interaction_id', None), environment=entry.pop('environment', None), llm_metrics_total_tokens=entry.pop('llm_metrics_total_tokens', None), llm_metrics_prompt_tokens=entry.pop('llm_metrics_prompt_tokens', None), llm_metrics_completion_tokens=entry.pop('llm_metrics_completion_tokens', None), llm_metrics_total_time=entry.pop('llm_metrics_total_time', None), llm_metrics_call_count=entry.pop('llm_metrics_call_count', None), llm_interaction_type=entry.pop('llm_interaction_type', None), message=entry.pop('message', None) ) db_entries.append(event_log) # Bulk insert db.session.bulk_save_objects(db_entries) db.session.commit() except Exception as e: logger = logging.getLogger('business_events') logger.error(f"Failed to persist logs directly to DB: {e}") db.session.rollback() def _flush_log_buffer(self): """Flush the log buffer to the database via a Celery task""" if self._log_buffer: try: # Send to Celery task current_celery.send_task( 'persist_business_events', args=[self._log_buffer], queue='entitlements' # Or dedicated log queue ) except Exception as e: # Fallback to direct DB write in case of issues with Celery logger = logging.getLogger('business_events') logger.error(f"Failed to send logs to Celery. Falling back to direct DB: {e}") self._direct_db_persist(self._log_buffer) # Clear the buffer after sending self._log_buffer = [] def __enter__(self): self.log(f'Starting Trace for {self.event_type}') return BusinessEventContext(self).__enter__() def __exit__(self, exc_type, exc_val, exc_tb): if self.llm_metrics['call_count'] > 0: self.log_final_metrics() self.reset_llm_metrics() self.log(f'Ending Trace for {self.event_type}') self._flush_log_buffer() return BusinessEventContext(self).__exit__(exc_type, exc_val, exc_tb)