"""Tencent APM tracing with idempotent client cleanup.""" import inspect import logging from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from core.ops.base_trace_instance import BaseTraceInstance from core.ops.entities.trace_entity import ( BaseTraceInfo, DatasetRetrievalTraceInfo, GenerateNameTraceInfo, MessageTraceInfo, ModerationTraceInfo, SuggestedQuestionTraceInfo, ToolTraceInfo, WorkflowTraceInfo, ) from core.repositories import SQLAlchemyWorkflowNodeExecutionRepository from dify_trace_tencent.client import TencentTraceClient from dify_trace_tencent.config import TencentConfig from dify_trace_tencent.entities.tencent_trace_entity import SpanData from dify_trace_tencent.span_builder import TencentSpanBuilder from dify_trace_tencent.utils import TencentTraceUtils from extensions.ext_database import db from graphon.entities.workflow_node_execution import ( WorkflowNodeExecution, ) from graphon.nodes import BuiltinNodeTypes from models import Account, App, TenantAccountJoin, WorkflowNodeExecutionTriggeredFrom logger = logging.getLogger(__name__) class TencentDataTrace(BaseTraceInstance): """ Tencent APM trace implementation with single responsibility principle. Acts as a coordinator that delegates specific tasks to specialized classes. The instance owns a long-lived ``TencentTraceClient``. Cleanup may happen explicitly in tests or implicitly during garbage collection, so shutdown must be safe to call multiple times. """ trace_client: TencentTraceClient _closed: bool def __init__(self, tencent_config: TencentConfig): super().__init__(tencent_config) self._closed = False self.trace_client = TencentTraceClient( service_name=tencent_config.service_name, endpoint=tencent_config.endpoint, token=tencent_config.token, metrics_export_interval_sec=5, ) def trace(self, trace_info: BaseTraceInfo) -> None: """Main tracing entry point - coordinates different trace types.""" match trace_info: case WorkflowTraceInfo(): self.workflow_trace(trace_info) case MessageTraceInfo(): self.message_trace(trace_info) case ModerationTraceInfo(): pass case SuggestedQuestionTraceInfo(): self.suggested_question_trace(trace_info) case DatasetRetrievalTraceInfo(): self.dataset_retrieval_trace(trace_info) case ToolTraceInfo(): self.tool_trace(trace_info) case GenerateNameTraceInfo(): pass def api_check(self) -> bool: return self.trace_client.api_check() def get_project_url(self) -> str: return self.trace_client.get_project_url() def workflow_trace(self, trace_info: WorkflowTraceInfo) -> None: """Handle workflow tracing by coordinating data retrieval and span construction.""" try: trace_id = TencentTraceUtils.convert_to_trace_id(trace_info.workflow_run_id) links = [] if trace_info.trace_id: links.append(TencentTraceUtils.create_link(trace_info.trace_id)) user_id = self._get_user_id(trace_info) workflow_spans = TencentSpanBuilder.build_workflow_spans(trace_info, trace_id, str(user_id), links) for span in workflow_spans: self.trace_client.add_span(span) self._process_workflow_nodes(trace_info, trace_id) # Record trace duration for entry span self._record_workflow_trace_duration(trace_info) except Exception: logger.exception("[Tencent APM] Failed to process workflow trace") def message_trace(self, trace_info: MessageTraceInfo) -> None: """Handle message tracing.""" try: trace_id = TencentTraceUtils.convert_to_trace_id(trace_info.message_id) user_id = self._get_user_id(trace_info) links = [] if trace_info.trace_id: links.append(TencentTraceUtils.create_link(trace_info.trace_id)) message_span = TencentSpanBuilder.build_message_span(trace_info, trace_id, str(user_id), links) self.trace_client.add_span(message_span) self._record_message_llm_metrics(trace_info) # Record trace duration for entry span self._record_message_trace_duration(trace_info) except Exception: logger.exception("[Tencent APM] Failed to process message trace") def tool_trace(self, trace_info: ToolTraceInfo) -> None: """Handle tool tracing.""" try: parent_span_id = None trace_root_id = None if trace_info.message_id: parent_span_id = TencentTraceUtils.convert_to_span_id(trace_info.message_id, "message") trace_root_id = trace_info.message_id if parent_span_id and trace_root_id: trace_id = TencentTraceUtils.convert_to_trace_id(trace_root_id) tool_span = TencentSpanBuilder.build_tool_span(trace_info, trace_id, parent_span_id) self.trace_client.add_span(tool_span) except Exception: logger.exception("[Tencent APM] Failed to process tool trace") def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo) -> None: """Handle dataset retrieval tracing.""" try: parent_span_id = None trace_root_id = None if trace_info.message_id: parent_span_id = TencentTraceUtils.convert_to_span_id(trace_info.message_id, "message") trace_root_id = trace_info.message_id if parent_span_id and trace_root_id: trace_id = TencentTraceUtils.convert_to_trace_id(trace_root_id) retrieval_span = TencentSpanBuilder.build_retrieval_span(trace_info, trace_id, parent_span_id) self.trace_client.add_span(retrieval_span) except Exception: logger.exception("[Tencent APM] Failed to process dataset retrieval trace") def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo) -> None: """Handle suggested question tracing""" try: logger.info("[Tencent APM] Processing suggested question trace") except Exception: logger.exception("[Tencent APM] Failed to process suggested question trace") def _process_workflow_nodes(self, trace_info: WorkflowTraceInfo, trace_id: int) -> None: """Process workflow node executions.""" try: workflow_span_id = TencentTraceUtils.convert_to_span_id(trace_info.workflow_run_id, "workflow") node_executions = self._get_workflow_node_executions(trace_info) for node_execution in node_executions: try: node_span = self._build_workflow_node_span(node_execution, trace_id, trace_info, workflow_span_id) if node_span: self.trace_client.add_span(node_span) if node_execution.node_type == BuiltinNodeTypes.LLM: self._record_llm_metrics(node_execution) except Exception: logger.exception("[Tencent APM] Failed to process node execution: %s", node_execution.id) except Exception: logger.exception("[Tencent APM] Failed to process workflow nodes") def _build_workflow_node_span( self, node_execution: WorkflowNodeExecution, trace_id: int, trace_info: WorkflowTraceInfo, workflow_span_id: int ) -> SpanData | None: """Build span for different node types""" try: if node_execution.node_type == BuiltinNodeTypes.LLM: return TencentSpanBuilder.build_workflow_llm_span( trace_id, workflow_span_id, trace_info, node_execution ) elif node_execution.node_type == BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL: return TencentSpanBuilder.build_workflow_retrieval_span( trace_id, workflow_span_id, trace_info, node_execution ) elif node_execution.node_type == BuiltinNodeTypes.TOOL: return TencentSpanBuilder.build_workflow_tool_span( trace_id, workflow_span_id, trace_info, node_execution ) else: # Handle all other node types as generic tasks return TencentSpanBuilder.build_workflow_task_span( trace_id, workflow_span_id, trace_info, node_execution ) except Exception: logger.debug( "[Tencent APM] Error building span for node %s: %s", node_execution.id, node_execution.node_type, exc_info=True, ) return None def _get_workflow_node_executions(self, trace_info: WorkflowTraceInfo) -> list[WorkflowNodeExecution]: """Retrieve workflow node executions from database.""" try: session_maker = sessionmaker(bind=db.engine) with Session(db.engine, expire_on_commit=False) as session: app_id = trace_info.metadata.get("app_id") if not app_id: raise ValueError("No app_id found in trace_info metadata") app_stmt = select(App).where(App.id == app_id) app = session.scalar(app_stmt) if not app: raise ValueError(f"App with id {app_id} not found") if not app.created_by: raise ValueError(f"App with id {app_id} has no creator") account_stmt = select(Account).where(Account.id == app.created_by) service_account = session.scalar(account_stmt) if not service_account: raise ValueError(f"Creator account not found for app {app_id}") current_tenant = session.scalar( select(TenantAccountJoin) .where(TenantAccountJoin.account_id == service_account.id, TenantAccountJoin.current.is_(True)) .limit(1) ) if not current_tenant: raise ValueError(f"Current tenant not found for account {service_account.id}") service_account.set_tenant_id(current_tenant.tenant_id) repository = SQLAlchemyWorkflowNodeExecutionRepository( session_factory=session_maker, user=service_account, app_id=trace_info.metadata.get("app_id"), triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) executions = repository.get_by_workflow_execution(workflow_execution_id=trace_info.workflow_run_id) return list(executions) except Exception: logger.exception("[Tencent APM] Failed to get workflow node executions") return [] def _get_user_id(self, trace_info: BaseTraceInfo) -> str: """Get user ID from trace info.""" try: tenant_id = None user_id = None if isinstance(trace_info, (WorkflowTraceInfo, GenerateNameTraceInfo)): tenant_id = trace_info.tenant_id if hasattr(trace_info, "metadata") and trace_info.metadata: user_id = trace_info.metadata.get("user_id") if user_id and tenant_id: stmt = ( select(Account.name) .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) .where(Account.id == user_id, TenantAccountJoin.tenant_id == tenant_id) ) session_maker = sessionmaker(bind=db.engine) with session_maker() as session: account_name = session.scalar(stmt) return account_name or str(user_id) elif user_id: return str(user_id) return "anonymous" except Exception: logger.exception("[Tencent APM] Failed to get user ID") return "unknown" def _record_llm_metrics(self, node_execution: WorkflowNodeExecution) -> None: """Record LLM performance metrics""" try: process_data = node_execution.process_data or {} outputs = node_execution.outputs or {} usage = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {}) model_provider = process_data.get("model_provider", "unknown") model_name = process_data.get("model_name", "unknown") model_mode = process_data.get("model_mode", "chat") # Record LLM duration if hasattr(self.trace_client, "record_llm_duration"): latency_s = float(usage.get("latency", 0.0)) if latency_s > 0: # Determine if streaming from usage metrics is_streaming = usage.get("time_to_first_token") is not None attributes = { "gen_ai.system": model_provider, "gen_ai.response.model": model_name, "gen_ai.operation.name": model_mode, "stream": "true" if is_streaming else "false", } self.trace_client.record_llm_duration(latency_s, attributes) # Record streaming metrics from usage time_to_first_token = usage.get("time_to_first_token") if time_to_first_token is not None and hasattr(self.trace_client, "record_time_to_first_token"): ttft_seconds = float(time_to_first_token) if ttft_seconds > 0: self.trace_client.record_time_to_first_token( ttft_seconds=ttft_seconds, provider=model_provider, model=model_name, operation_name=model_mode ) time_to_generate = usage.get("time_to_generate") if time_to_generate is not None and hasattr(self.trace_client, "record_time_to_generate"): ttg_seconds = float(time_to_generate) if ttg_seconds > 0: self.trace_client.record_time_to_generate( ttg_seconds=ttg_seconds, provider=model_provider, model=model_name, operation_name=model_mode ) # Record token usage if hasattr(self.trace_client, "record_token_usage"): # Extract token counts input_tokens = int(usage.get("prompt_tokens", 0)) output_tokens = int(usage.get("completion_tokens", 0)) if input_tokens > 0 or output_tokens > 0: server_address = f"{model_provider}" # Record input tokens if input_tokens > 0: self.trace_client.record_token_usage( token_count=input_tokens, token_type="input", operation_name=model_mode, request_model=model_name, response_model=model_name, server_address=server_address, provider=model_provider, ) # Record output tokens if output_tokens > 0: self.trace_client.record_token_usage( token_count=output_tokens, token_type="output", operation_name=model_mode, request_model=model_name, response_model=model_name, server_address=server_address, provider=model_provider, ) except Exception: logger.debug("[Tencent APM] Failed to record LLM metrics") def _record_message_llm_metrics(self, trace_info: MessageTraceInfo) -> None: """Record LLM metrics for message traces""" try: trace_metadata = trace_info.metadata or {} message_data = trace_info.message_data or {} provider_latency = 0.0 if isinstance(message_data, dict): provider_latency = float(message_data.get("provider_response_latency", 0.0) or 0.0) else: provider_latency = float(getattr(message_data, "provider_response_latency", 0.0) or 0.0) model_provider = trace_metadata.get("ls_provider") or ( message_data.get("model_provider", "") if isinstance(message_data, dict) else "" ) model_name = trace_metadata.get("ls_model_name") or ( message_data.get("model_id", "") if isinstance(message_data, dict) else "" ) # Record LLM duration if provider_latency > 0 and hasattr(self.trace_client, "record_llm_duration"): is_streaming = trace_info.is_streaming_request duration_attributes = { "gen_ai.system": model_provider, "gen_ai.response.model": model_name, "gen_ai.operation.name": "chat", # Message traces are always chat "stream": "true" if is_streaming else "false", } self.trace_client.record_llm_duration(provider_latency, duration_attributes) # Record streaming metrics for message traces if trace_info.is_streaming_request: # Record time to first token if trace_info.gen_ai_server_time_to_first_token is not None and hasattr( self.trace_client, "record_time_to_first_token" ): ttft_seconds = float(trace_info.gen_ai_server_time_to_first_token) if ttft_seconds > 0: self.trace_client.record_time_to_first_token( ttft_seconds=ttft_seconds, provider=str(model_provider or ""), model=str(model_name or "") ) # Record time to generate if trace_info.llm_streaming_time_to_generate is not None and hasattr( self.trace_client, "record_time_to_generate" ): ttg_seconds = float(trace_info.llm_streaming_time_to_generate) if ttg_seconds > 0: self.trace_client.record_time_to_generate( ttg_seconds=ttg_seconds, provider=str(model_provider or ""), model=str(model_name or "") ) # Record token usage if hasattr(self.trace_client, "record_token_usage"): input_tokens = int(trace_info.message_tokens or 0) output_tokens = int(trace_info.answer_tokens or 0) if input_tokens > 0: self.trace_client.record_token_usage( token_count=input_tokens, token_type="input", operation_name="chat", request_model=str(model_name or ""), response_model=str(model_name or ""), server_address=str(model_provider or ""), provider=str(model_provider or ""), ) if output_tokens > 0: self.trace_client.record_token_usage( token_count=output_tokens, token_type="output", operation_name="chat", request_model=str(model_name or ""), response_model=str(model_name or ""), server_address=str(model_provider or ""), provider=str(model_provider or ""), ) except Exception: logger.debug("[Tencent APM] Failed to record message LLM metrics") def _record_workflow_trace_duration(self, trace_info: WorkflowTraceInfo) -> None: """Record end-to-end workflow trace duration.""" try: if not hasattr(self.trace_client, "record_trace_duration"): return # Calculate duration from start_time and end_time to match span duration if trace_info.start_time and trace_info.end_time: duration_s = (trace_info.end_time - trace_info.start_time).total_seconds() else: # Fallback to workflow_run_elapsed_time if timestamps not available duration_s = float(trace_info.workflow_run_elapsed_time) if duration_s > 0: attributes = { "conversation_mode": "workflow", "workflow_status": trace_info.workflow_run_status, } # Add conversation_id if available if trace_info.conversation_id: attributes["has_conversation"] = "true" else: attributes["has_conversation"] = "false" self.trace_client.record_trace_duration(duration_s, attributes) except Exception: logger.debug("[Tencent APM] Failed to record workflow trace duration") def _record_message_trace_duration(self, trace_info: MessageTraceInfo) -> None: """Record end-to-end message trace duration.""" try: if not hasattr(self.trace_client, "record_trace_duration"): return # Calculate duration from start_time and end_time if trace_info.start_time and trace_info.end_time: duration = (trace_info.end_time - trace_info.start_time).total_seconds() if duration > 0: attributes = { "conversation_mode": trace_info.conversation_mode, } # Add streaming flag if available if hasattr(trace_info, "is_streaming_request"): attributes["stream"] = "true" if trace_info.is_streaming_request else "false" self.trace_client.record_trace_duration(duration, attributes) except Exception: logger.debug("[Tencent APM] Failed to record message trace duration") def close(self) -> None: """Synchronously and idempotently shutdown the underlying trace client.""" if getattr(self, "_closed", False): return self._closed = True trace_client = getattr(self, "trace_client", None) if trace_client is None: return try: shutdown_result = trace_client.shutdown() if inspect.isawaitable(shutdown_result): close_awaitable = getattr(shutdown_result, "close", None) if callable(close_awaitable): close_awaitable() except Exception: logger.exception("[Tencent APM] Failed to shutdown trace client during cleanup") def __del__(self): """Ensure best-effort cleanup on garbage collection without retrying shutdown.""" self.close()