1807435339
- Introduction of dynamic Processors - Introduction of caching system - Introduction of a better template manager - Adaptation of ModelVariables to support dynamic Processors / Retrievers / Specialists - Start adaptation of chat client
246 lines
10 KiB
Python
246 lines
10 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
|
|
import logging
|
|
|
|
from .business_event_context import BusinessEventContext
|
|
from common.models.entitlements 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.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
|
|
}
|
|
|
|
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"Starting span {span_name}")
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
if self.llm_metrics['call_count'] > 0:
|
|
self.log_final_metrics()
|
|
self.reset_llm_metrics()
|
|
self.log(f"Ending span {span_name}")
|
|
# 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,
|
|
'document_version_file_size': self.document_version_file_size,
|
|
'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,
|
|
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
|
|
)
|
|
db.session.add(event_log)
|
|
db.session.commit()
|
|
|
|
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 = {
|
|
'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'],
|
|
}
|
|
# 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,
|
|
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
|
|
)
|
|
db.session.add(event_log)
|
|
db.session.commit()
|
|
|
|
def log_final_metrics(self, level: str = 'info'):
|
|
logger = logging.getLogger('business_events')
|
|
message = "Final LLM Metrics"
|
|
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,
|
|
'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'],
|
|
}
|
|
# 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,
|
|
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
|
|
)
|
|
db.session.add(event_log)
|
|
db.session.commit()
|
|
|
|
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}')
|
|
return BusinessEventContext(self).__exit__(exc_type, exc_val, exc_tb)
|