import json
import logging
import os
import queue
import threading
import time
from datetime import timedelta
from typing import Any, Optional, Union
from uuid import UUID
from flask import current_app
from core.helper.encrypter import decrypt_token, encrypt_token, obfuscated_token
from core.ops.entities.config_entity import (
LangfuseConfig,
LangSmithConfig,
TracingProviderEnum,
)
from core.ops.entities.trace_entity import (
DatasetRetrievalTraceInfo,
GenerateNameTraceInfo,
MessageTraceInfo,
ModerationTraceInfo,
SuggestedQuestionTraceInfo,
ToolTraceInfo,
TraceTaskName,
WorkflowTraceInfo,
)
from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace
from core.ops.langsmith_trace.langsmith_trace import LangSmithDataTrace
from core.ops.utils import get_message_data
from extensions.ext_database import db
from models.model import App, AppModelConfig, Conversation, Message, MessageAgentThought, MessageFile, TraceAppConfig
from models.workflow import WorkflowAppLog, WorkflowRun
from tasks.ops_trace_task import process_trace_tasks
provider_config_map = {
TracingProviderEnum.LANGFUSE.value: {
"config_class": LangfuseConfig,
"secret_keys": ["public_key", "secret_key"],
"other_keys": ["host", "project_key"],
"trace_instance": LangFuseDataTrace,
},
TracingProviderEnum.LANGSMITH.value: {
"config_class": LangSmithConfig,
"secret_keys": ["api_key"],
"other_keys": ["project", "endpoint"],
"trace_instance": LangSmithDataTrace,
},
}
class OpsTraceManager:
@classmethod
def encrypt_tracing_config(
cls, tenant_id: str, tracing_provider: str, tracing_config: dict, current_trace_config=None
):
"""
Encrypt tracing config.
:param tenant_id: tenant id
:param tracing_provider: tracing provider
:param tracing_config: tracing config dictionary to be encrypted
:param current_trace_config: current tracing configuration for keeping existing values
:return: encrypted tracing configuration
"""
# Get the configuration class and the keys that require encryption
config_class, secret_keys, other_keys = (
provider_config_map[tracing_provider]["config_class"],
provider_config_map[tracing_provider]["secret_keys"],
provider_config_map[tracing_provider]["other_keys"],
)
new_config = {}
# Encrypt necessary keys
for key in secret_keys:
if key in tracing_config:
if "*" in tracing_config[key]:
# If the key contains '*', retain the original value from the current config
new_config[key] = current_trace_config.get(key, tracing_config[key])
else:
# Otherwise, encrypt the key
new_config[key] = encrypt_token(tenant_id, tracing_config[key])
for key in other_keys:
new_config[key] = tracing_config.get(key, "")
# Create a new instance of the config class with the new configuration
encrypted_config = config_class(**new_config)
return encrypted_config.model_dump()
@classmethod
def decrypt_tracing_config(cls, tenant_id: str, tracing_provider: str, tracing_config: dict):
"""
Decrypt tracing config
:param tenant_id: tenant id
:param tracing_provider: tracing provider
:param tracing_config: tracing config
:return:
"""
config_class, secret_keys, other_keys = (
provider_config_map[tracing_provider]["config_class"],
provider_config_map[tracing_provider]["secret_keys"],
provider_config_map[tracing_provider]["other_keys"],
)
new_config = {}
for key in secret_keys:
if key in tracing_config:
new_config[key] = decrypt_token(tenant_id, tracing_config[key])
for key in other_keys:
new_config[key] = tracing_config.get(key, "")
return config_class(**new_config).model_dump()
@classmethod
def obfuscated_decrypt_token(cls, tracing_provider: str, decrypt_tracing_config: dict):
"""
Decrypt tracing config
:param tracing_provider: tracing provider
:param decrypt_tracing_config: tracing config
:return:
"""
config_class, secret_keys, other_keys = (
provider_config_map[tracing_provider]["config_class"],
provider_config_map[tracing_provider]["secret_keys"],
provider_config_map[tracing_provider]["other_keys"],
)
new_config = {}
for key in secret_keys:
if key in decrypt_tracing_config:
new_config[key] = obfuscated_token(decrypt_tracing_config[key])
for key in other_keys:
new_config[key] = decrypt_tracing_config.get(key, "")
return config_class(**new_config).model_dump()
@classmethod
def get_decrypted_tracing_config(cls, app_id: str, tracing_provider: str):
"""
Get decrypted tracing config
:param app_id: app id
:param tracing_provider: tracing provider
:return:
"""
trace_config_data: TraceAppConfig = (
db.session.query(TraceAppConfig)
.filter(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider)
.first()
)
if not trace_config_data:
return None
# decrypt_token
tenant_id = db.session.query(App).filter(App.id == app_id).first().tenant_id
decrypt_tracing_config = cls.decrypt_tracing_config(
tenant_id, tracing_provider, trace_config_data.tracing_config
)
return decrypt_tracing_config
@classmethod
def get_ops_trace_instance(
cls,
app_id: Optional[Union[UUID, str]] = None,
):
"""
Get ops trace through model config
:param app_id: app_id
:return:
"""
if isinstance(app_id, UUID):
app_id = str(app_id)
if app_id is None:
return None
app: App = db.session.query(App).filter(App.id == app_id).first()
if app is None:
return None
app_ops_trace_config = json.loads(app.tracing) if app.tracing else None
if app_ops_trace_config is None:
return None
tracing_provider = app_ops_trace_config.get("tracing_provider")
if tracing_provider is None or tracing_provider not in provider_config_map:
return None
# decrypt_token
decrypt_trace_config = cls.get_decrypted_tracing_config(app_id, tracing_provider)
if app_ops_trace_config.get("enabled"):
trace_instance, config_class = (
provider_config_map[tracing_provider]["trace_instance"],
provider_config_map[tracing_provider]["config_class"],
)
tracing_instance = trace_instance(config_class(**decrypt_trace_config))
return tracing_instance
return None
@classmethod
def get_app_config_through_message_id(cls, message_id: str):
app_model_config = None
message_data = db.session.query(Message).filter(Message.id == message_id).first()
conversation_id = message_data.conversation_id
conversation_data = db.session.query(Conversation).filter(Conversation.id == conversation_id).first()
if conversation_data.app_model_config_id:
app_model_config = (
db.session.query(AppModelConfig)
.filter(AppModelConfig.id == conversation_data.app_model_config_id)
.first()
)
elif conversation_data.app_model_config_id is None and conversation_data.override_model_configs:
app_model_config = conversation_data.override_model_configs
return app_model_config
@classmethod
def update_app_tracing_config(cls, app_id: str, enabled: bool, tracing_provider: str):
"""
Update app tracing config
:param app_id: app id
:param enabled: enabled
:param tracing_provider: tracing provider
:return:
"""
# auth check
if tracing_provider not in provider_config_map and tracing_provider is not None:
raise ValueError(f"Invalid tracing provider: {tracing_provider}")
app_config: App = db.session.query(App).filter(App.id == app_id).first()
app_config.tracing = json.dumps(
{
"enabled": enabled,
"tracing_provider": tracing_provider,
}
)
db.session.commit()
@classmethod
def get_app_tracing_config(cls, app_id: str):
"""
Get app tracing config
:param app_id: app id
:return:
"""
app: App = db.session.query(App).filter(App.id == app_id).first()
if not app.tracing:
return {"enabled": False, "tracing_provider": None}
app_trace_config = json.loads(app.tracing)
return app_trace_config
@staticmethod
def check_trace_config_is_effective(tracing_config: dict, tracing_provider: str):
"""
Check trace config is effective
:param tracing_config: tracing config
:param tracing_provider: tracing provider
:return:
"""
config_type, trace_instance = (
provider_config_map[tracing_provider]["config_class"],
provider_config_map[tracing_provider]["trace_instance"],
)
tracing_config = config_type(**tracing_config)
return trace_instance(tracing_config).api_check()
@staticmethod
def get_trace_config_project_key(tracing_config: dict, tracing_provider: str):
"""
get trace config is project key
:param tracing_config: tracing config
:param tracing_provider: tracing provider
:return:
"""
config_type, trace_instance = (
provider_config_map[tracing_provider]["config_class"],
provider_config_map[tracing_provider]["trace_instance"],
)
tracing_config = config_type(**tracing_config)
return trace_instance(tracing_config).get_project_key()
@staticmethod
def get_trace_config_project_url(tracing_config: dict, tracing_provider: str):
"""
get trace config is project key
:param tracing_config: tracing config
:param tracing_provider: tracing provider
:return:
"""
config_type, trace_instance = (
provider_config_map[tracing_provider]["config_class"],
provider_config_map[tracing_provider]["trace_instance"],
)
tracing_config = config_type(**tracing_config)
return trace_instance(tracing_config).get_project_url()
class TraceTask:
def __init__(
self,
trace_type: Any,
message_id: Optional[str] = None,
workflow_run: Optional[WorkflowRun] = None,
conversation_id: Optional[str] = None,
user_id: Optional[str] = None,
timer: Optional[Any] = None,
**kwargs,
):
self.trace_type = trace_type
self.message_id = message_id
self.workflow_run = workflow_run
self.conversation_id = conversation_id
self.user_id = user_id
self.timer = timer
self.kwargs = kwargs
self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001")
self.app_id = None
def execute(self):
return self.preprocess()
def preprocess(self):
preprocess_map = {
TraceTaskName.CONVERSATION_TRACE: lambda: self.conversation_trace(**self.kwargs),
TraceTaskName.WORKFLOW_TRACE: lambda: self.workflow_trace(
self.workflow_run, self.conversation_id, self.user_id
),
TraceTaskName.MESSAGE_TRACE: lambda: self.message_trace(self.message_id),
TraceTaskName.MODERATION_TRACE: lambda: self.moderation_trace(self.message_id, self.timer, **self.kwargs),
TraceTaskName.SUGGESTED_QUESTION_TRACE: lambda: self.suggested_question_trace(
self.message_id, self.timer, **self.kwargs
),
TraceTaskName.DATASET_RETRIEVAL_TRACE: lambda: self.dataset_retrieval_trace(
self.message_id, self.timer, **self.kwargs
),
TraceTaskName.TOOL_TRACE: lambda: self.tool_trace(self.message_id, self.timer, **self.kwargs),
TraceTaskName.GENERATE_NAME_TRACE: lambda: self.generate_name_trace(
self.conversation_id, self.timer, **self.kwargs
),
}
return preprocess_map.get(self.trace_type, lambda: None)()
# process methods for different trace types
def conversation_trace(self, **kwargs):
return kwargs
def workflow_trace(self, workflow_run: WorkflowRun, conversation_id, user_id):
workflow_id = workflow_run.workflow_id
tenant_id = workflow_run.tenant_id
workflow_run_id = workflow_run.id
workflow_run_elapsed_time = workflow_run.elapsed_time
workflow_run_status = workflow_run.status
workflow_run_inputs = json.loads(workflow_run.inputs) if workflow_run.inputs else {}
workflow_run_outputs = json.loads(workflow_run.outputs) if workflow_run.outputs else {}
workflow_run_version = workflow_run.version
error = workflow_run.error or ""
total_tokens = workflow_run.total_tokens
file_list = workflow_run_inputs.get("sys.file") or []
query = workflow_run_inputs.get("query") or workflow_run_inputs.get("sys.query") or ""
# get workflow_app_log_id
workflow_app_log_data = (
db.session.query(WorkflowAppLog)
.filter_by(tenant_id=tenant_id, app_id=workflow_run.app_id, workflow_run_id=workflow_run.id)
.first()
)
workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None
# get message_id
message_data = (
db.session.query(Message.id)
.filter_by(conversation_id=conversation_id, workflow_run_id=workflow_run_id)
.first()
)
message_id = str(message_data.id) if message_data else None
metadata = {
"workflow_id": workflow_id,
"conversation_id": conversation_id,
"workflow_run_id": workflow_run_id,
"tenant_id": tenant_id,
"elapsed_time": workflow_run_elapsed_time,
"status": workflow_run_status,
"version": workflow_run_version,
"total_tokens": total_tokens,
"file_list": file_list,
"triggered_form": workflow_run.triggered_from,
"user_id": user_id,
}
workflow_trace_info = WorkflowTraceInfo(
workflow_data=workflow_run.to_dict(),
conversation_id=conversation_id,
workflow_id=workflow_id,
tenant_id=tenant_id,
workflow_run_id=workflow_run_id,
workflow_run_elapsed_time=workflow_run_elapsed_time,
workflow_run_status=workflow_run_status,
workflow_run_inputs=workflow_run_inputs,
workflow_run_outputs=workflow_run_outputs,
workflow_run_version=workflow_run_version,
error=error,
total_tokens=total_tokens,
file_list=file_list,
query=query,
metadata=metadata,
workflow_app_log_id=workflow_app_log_id,
message_id=message_id,
start_time=workflow_run.created_at,
end_time=workflow_run.finished_at,
)
return workflow_trace_info
def message_trace(self, message_id):
message_data = get_message_data(message_id)
if not message_data:
return {}
conversation_mode = db.session.query(Conversation.mode).filter_by(id=message_data.conversation_id).first()
conversation_mode = conversation_mode[0]
created_at = message_data.created_at
inputs = message_data.message
# get message file data
message_file_data = db.session.query(MessageFile).filter_by(message_id=message_id).first()
file_list = []
if message_file_data and message_file_data.url is not None:
file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else ""
file_list.append(file_url)
metadata = {
"conversation_id": message_data.conversation_id,
"ls_provider": message_data.model_provider,
"ls_model_name": message_data.model_id,
"status": message_data.status,
"from_end_user_id": message_data.from_account_id,
"from_account_id": message_data.from_account_id,
"agent_based": message_data.agent_based,
"workflow_run_id": message_data.workflow_run_id,
"from_source": message_data.from_source,
"message_id": message_id,
}
message_tokens = message_data.message_tokens
message_trace_info = MessageTraceInfo(
message_id=message_id,
message_data=message_data.to_dict(),
conversation_model=conversation_mode,
message_tokens=message_tokens,
answer_tokens=message_data.answer_tokens,
total_tokens=message_tokens + message_data.answer_tokens,
error=message_data.error or "",
inputs=inputs,
outputs=message_data.answer,
file_list=file_list,
start_time=created_at,
end_time=created_at + timedelta(seconds=message_data.provider_response_latency),
metadata=metadata,
message_file_data=message_file_data,
conversation_mode=conversation_mode,
)
return message_trace_info
def moderation_trace(self, message_id, timer, **kwargs):
moderation_result = kwargs.get("moderation_result")
inputs = kwargs.get("inputs")
message_data = get_message_data(message_id)
if not message_data:
return {}
metadata = {
"message_id": message_id,
"action": moderation_result.action,
"preset_response": moderation_result.preset_response,
"query": moderation_result.query,
}
# get workflow_app_log_id
workflow_app_log_id = None
if message_data.workflow_run_id:
workflow_app_log_data = (
db.session.query(WorkflowAppLog).filter_by(workflow_run_id=message_data.workflow_run_id).first()
)
workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None
moderation_trace_info = ModerationTraceInfo(
message_id=workflow_app_log_id or message_id,
inputs=inputs,
message_data=message_data.to_dict(),
flagged=moderation_result.flagged,
action=moderation_result.action,
preset_response=moderation_result.preset_response,
query=moderation_result.query,
start_time=timer.get("start"),
end_time=timer.get("end"),
metadata=metadata,
)
return moderation_trace_info
def suggested_question_trace(self, message_id, timer, **kwargs):
suggested_question = kwargs.get("suggested_question")
message_data = get_message_data(message_id)
if not message_data:
return {}
metadata = {
"message_id": message_id,
"ls_provider": message_data.model_provider,
"ls_model_name": message_data.model_id,
"status": message_data.status,
"from_end_user_id": message_data.from_account_id,
"from_account_id": message_data.from_account_id,
"agent_based": message_data.agent_based,
"workflow_run_id": message_data.workflow_run_id,
"from_source": message_data.from_source,
}
# get workflow_app_log_id
workflow_app_log_id = None
if message_data.workflow_run_id:
workflow_app_log_data = (
db.session.query(WorkflowAppLog).filter_by(workflow_run_id=message_data.workflow_run_id).first()
)
workflow_app_log_id = str(workflow_app_log_data.id) if workflow_app_log_data else None
suggested_question_trace_info = SuggestedQuestionTraceInfo(
message_id=workflow_app_log_id or message_id,
message_data=message_data.to_dict(),
inputs=message_data.message,
outputs=message_data.answer,
start_time=timer.get("start"),
end_time=timer.get("end"),
metadata=metadata,
total_tokens=message_data.message_tokens + message_data.answer_tokens,
status=message_data.status,
error=message_data.error,
from_account_id=message_data.from_account_id,
agent_based=message_data.agent_based,
from_source=message_data.from_source,
model_provider=message_data.model_provider,
model_id=message_data.model_id,
suggested_question=suggested_question,
level=message_data.status,
status_message=message_data.error,
)
return suggested_question_trace_info
def dataset_retrieval_trace(self, message_id, timer, **kwargs):
documents = kwargs.get("documents")
message_data = get_message_data(message_id)
if not message_data:
return {}
metadata = {
"message_id": message_id,
"ls_provider": message_data.model_provider,
"ls_model_name": message_data.model_id,
"status": message_data.status,
"from_end_user_id": message_data.from_account_id,
"from_account_id": message_data.from_account_id,
"agent_based": message_data.agent_based,
"workflow_run_id": message_data.workflow_run_id,
"from_source": message_data.from_source,
}
dataset_retrieval_trace_info = DatasetRetrievalTraceInfo(
message_id=message_id,
inputs=message_data.query or message_data.inputs,
documents=[doc.model_dump() for doc in documents],
start_time=timer.get("start"),
end_time=timer.get("end"),
metadata=metadata,
message_data=message_data.to_dict(),
)
return dataset_retrieval_trace_info
def tool_trace(self, message_id, timer, **kwargs):
tool_name = kwargs.get("tool_name")
tool_inputs = kwargs.get("tool_inputs")
tool_outputs = kwargs.get("tool_outputs")
message_data = get_message_data(message_id)
if not message_data:
return {}
tool_config = {}
time_cost = 0
error = None
tool_parameters = {}
created_time = message_data.created_at
end_time = message_data.updated_at
agent_thoughts: list[MessageAgentThought] = message_data.agent_thoughts
for agent_thought in agent_thoughts:
if tool_name in agent_thought.tools:
created_time = agent_thought.created_at
tool_meta_data = agent_thought.tool_meta.get(tool_name, {})
tool_config = tool_meta_data.get("tool_config", {})
time_cost = tool_meta_data.get("time_cost", 0)
end_time = created_time + timedelta(seconds=time_cost)
error = tool_meta_data.get("error", "")
tool_parameters = tool_meta_data.get("tool_parameters", {})
metadata = {
"message_id": message_id,
"tool_name": tool_name,
"tool_inputs": tool_inputs,
"tool_outputs": tool_outputs,
"tool_config": tool_config,
"time_cost": time_cost,
"error": error,
"tool_parameters": tool_parameters,
}
file_url = ""
message_file_data = db.session.query(MessageFile).filter_by(message_id=message_id).first()
if message_file_data:
message_file_id = message_file_data.id if message_file_data else None
type = message_file_data.type
created_by_role = message_file_data.created_by_role
created_user_id = message_file_data.created_by
file_url = f"{self.file_base_url}/{message_file_data.url}"
metadata.update(
{
"message_file_id": message_file_id,
"created_by_role": created_by_role,
"created_user_id": created_user_id,
"type": type,
}
)
tool_trace_info = ToolTraceInfo(
message_id=message_id,
message_data=message_data.to_dict(),
tool_name=tool_name,
start_time=timer.get("start") if timer else created_time,
end_time=timer.get("end") if timer else end_time,
tool_inputs=tool_inputs,
tool_outputs=tool_outputs,
metadata=metadata,
message_file_data=message_file_data,
error=error,
inputs=message_data.message,
outputs=message_data.answer,
tool_config=tool_config,
time_cost=time_cost,
tool_parameters=tool_parameters,
file_url=file_url,
)
return tool_trace_info
def generate_name_trace(self, conversation_id, timer, **kwargs):
generate_conversation_name = kwargs.get("generate_conversation_name")
inputs = kwargs.get("inputs")
tenant_id = kwargs.get("tenant_id")
start_time = timer.get("start")
end_time = timer.get("end")
metadata = {
"conversation_id": conversation_id,
"tenant_id": tenant_id,
}
generate_name_trace_info = GenerateNameTraceInfo(
conversation_id=conversation_id,
inputs=inputs,
outputs=generate_conversation_name,
start_time=start_time,
end_time=end_time,
metadata=metadata,
tenant_id=tenant_id,
)
return generate_name_trace_info
trace_manager_timer = None
trace_manager_queue = queue.Queue()
trace_manager_interval = int(os.getenv("TRACE_QUEUE_MANAGER_INTERVAL", 5))
trace_manager_batch_size = int(os.getenv("TRACE_QUEUE_MANAGER_BATCH_SIZE", 100))
class TraceQueueManager:
def __init__(self, app_id=None, user_id=None):
global trace_manager_timer
self.app_id = app_id
self.user_id = user_id
self.trace_instance = OpsTraceManager.get_ops_trace_instance(app_id)
self.flask_app = current_app._get_current_object()
if trace_manager_timer is None:
self.start_timer()
def add_trace_task(self, trace_task: TraceTask):
global trace_manager_timer, trace_manager_queue
try:
if self.trace_instance:
trace_task.app_id = self.app_id
trace_manager_queue.put(trace_task)
except Exception as e:
logging.error(f"Error adding trace task: {e}")
finally:
self.start_timer()
def collect_tasks(self):
global trace_manager_queue
tasks = []
while len(tasks) < trace_manager_batch_size and not trace_manager_queue.empty():
task = trace_manager_queue.get_nowait()
tasks.append(task)
trace_manager_queue.task_done()
return tasks
def run(self):
try:
tasks = self.collect_tasks()
if tasks:
self.send_to_celery(tasks)
except Exception as e:
logging.error(f"Error processing trace tasks: {e}")
def start_timer(self):
global trace_manager_timer
if trace_manager_timer is None or not trace_manager_timer.is_alive():
trace_manager_timer = threading.Timer(trace_manager_interval, self.run)
trace_manager_timer.name = f"trace_manager_timer_{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}"
trace_manager_timer.daemon = False
trace_manager_timer.start()
def send_to_celery(self, tasks: list[TraceTask]):
with self.flask_app.app_context():
for task in tasks:
trace_info = task.execute()
task_data = {
"app_id": task.app_id,
"trace_info_type": type(trace_info).__name__,
"trace_info": trace_info.model_dump() if trace_info else {},
}
process_trace_tasks.delay(task_data)