import logging import os import uuid from datetime import UTC, datetime, timedelta from typing import override import httpx from langfuse import __version__ as langfuse_version from langfuse.api import ( CreateGenerationBody, CreateSpanBody, IngestionEvent_GenerationCreate, IngestionEvent_SpanCreate, IngestionEvent_TraceCreate, LangfuseAPI, TraceBody, ) from langfuse.api.commons.types.usage import Usage from sqlalchemy.orm import sessionmaker from core.ops.base_trace_instance import BaseTraceInstance from core.ops.entities.trace_entity import ( BaseTraceInfo, DatasetRetrievalTraceInfo, GenerateNameTraceInfo, MessageTraceInfo, ModerationTraceInfo, SuggestedQuestionTraceInfo, ToolTraceInfo, TraceTaskName, WorkflowTraceInfo, ) from core.ops.utils import filter_none_values from core.repositories import DifyCoreRepositoryFactory from dify_trace_langfuse.config import LangfuseConfig from dify_trace_langfuse.entities.langfuse_trace_entity import ( GenerationUsage, LangfuseGeneration, LangfuseSpan, LangfuseTrace, LevelEnum, UnitEnum, ) from extensions.ext_database import db from graphon.enums import BuiltinNodeTypes from models import EndUser, WorkflowNodeExecutionTriggeredFrom from models.enums import MessageStatus logger = logging.getLogger(__name__) class LangFuseDataTrace(BaseTraceInstance): def __init__( self, langfuse_config: LangfuseConfig, ): super().__init__(langfuse_config) timeout = int(os.environ.get("LANGFUSE_TIMEOUT", 5)) self._http_client: httpx.Client | None = httpx.Client(timeout=timeout) self.langfuse_client = LangfuseAPI( base_url=langfuse_config.host, username=langfuse_config.public_key, password=langfuse_config.secret_key, x_langfuse_sdk_name="python", x_langfuse_sdk_version=langfuse_version, x_langfuse_public_key=langfuse_config.public_key, timeout=timeout, httpx_client=self._http_client, ) self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") def close(self) -> None: client = getattr(self, "_http_client", None) if client is None: return self._http_client = None try: client.close() except Exception: logger.debug("Failed to close Langfuse HTTP client", exc_info=True) def __del__(self) -> None: self.close() @staticmethod def _get_completion_start_time( start_time: datetime | None, time_to_first_token: float | int | None ) -> datetime | None: """Convert a relative TTFT value in seconds into Langfuse's absolute completion start time.""" if start_time is None or time_to_first_token is None: return None try: ttft_seconds = float(time_to_first_token) except (TypeError, ValueError): return None if ttft_seconds < 0: return None return start_time + timedelta(seconds=ttft_seconds) @override def trace(self, trace_info: BaseTraceInfo): match trace_info: case WorkflowTraceInfo(): self.workflow_trace(trace_info) case MessageTraceInfo(): self.message_trace(trace_info) case ModerationTraceInfo(): self.moderation_trace(trace_info) 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(): self.generate_name_trace(trace_info) case _: pass def workflow_trace(self, trace_info: WorkflowTraceInfo): trace_id = trace_info.trace_id or trace_info.workflow_run_id user_id = trace_info.metadata.get("user_id") metadata = trace_info.metadata metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id if trace_info.message_id: trace_id = trace_info.trace_id or trace_info.message_id name = TraceTaskName.MESSAGE_TRACE trace_data = LangfuseTrace( id=trace_id, user_id=user_id, name=name, input=dict(trace_info.workflow_run_inputs), output=dict(trace_info.workflow_run_outputs), metadata=metadata, session_id=trace_info.conversation_id, tags=["message", "workflow"], version=trace_info.workflow_run_version, ) self.add_trace(langfuse_trace_data=trace_data) workflow_span_data = LangfuseSpan( id=trace_info.workflow_run_id, name=TraceTaskName.WORKFLOW_TRACE, input=dict(trace_info.workflow_run_inputs), output=dict(trace_info.workflow_run_outputs), trace_id=trace_id, start_time=trace_info.start_time, end_time=trace_info.end_time, metadata=metadata, level=LevelEnum.DEFAULT if trace_info.error == "" else LevelEnum.ERROR, status_message=trace_info.error or "", ) self.add_span(langfuse_span_data=workflow_span_data) else: trace_data = LangfuseTrace( id=trace_id, user_id=user_id, name=TraceTaskName.WORKFLOW_TRACE, input=dict(trace_info.workflow_run_inputs), output=dict(trace_info.workflow_run_outputs), metadata=metadata, session_id=trace_info.conversation_id, tags=["workflow"], version=trace_info.workflow_run_version, ) self.add_trace(langfuse_trace_data=trace_data) # through workflow_run_id get all_nodes_execution using repository session_factory = sessionmaker(bind=db.engine) # Find the app's creator account app_id = trace_info.metadata.get("app_id") if not app_id: raise ValueError("No app_id found in trace_info metadata") service_account = self.get_service_account_with_tenant(app_id) workflow_node_execution_repository = DifyCoreRepositoryFactory.create_workflow_node_execution_repository( session_factory=session_factory, tenant_id=trace_info.tenant_id, user=service_account, app_id=app_id, triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, ) # Get all executions for this workflow run workflow_node_executions = workflow_node_execution_repository.get_by_workflow_execution( workflow_execution_id=trace_info.workflow_run_id ) for node_execution in workflow_node_executions: node_execution_id = node_execution.id tenant_id = trace_info.tenant_id # Use from trace_info instead app_id = trace_info.metadata.get("app_id") # Use from trace_info instead node_name = node_execution.title node_type = node_execution.node_type status = node_execution.status if node_type == BuiltinNodeTypes.LLM: inputs = node_execution.process_data.get("prompts", {}) if node_execution.process_data else {} else: inputs = node_execution.inputs or {} outputs = node_execution.outputs or {} created_at = node_execution.created_at or datetime.now() elapsed_time = node_execution.elapsed_time finished_at = created_at + timedelta(seconds=elapsed_time) execution_metadata = node_execution.metadata or {} metadata = {str(k): v for k, v in execution_metadata.items()} metadata.update( { "workflow_run_id": trace_info.workflow_run_id, "node_execution_id": node_execution_id, "tenant_id": tenant_id, "app_id": app_id, "node_name": node_name, "node_type": node_type, "status": status, } ) process_data = node_execution.process_data or {} model_provider = process_data.get("model_provider", None) model_name = process_data.get("model_name", None) if model_provider is not None and model_name is not None: metadata.update( { "model_provider": model_provider, "model_name": model_name, } ) # add generation span if process_data and process_data.get("model_mode") == "chat": total_token = metadata.get("total_tokens", 0) prompt_tokens = 0 completion_tokens = 0 completion_start_time = None try: usage_data = process_data.get("usage") if not isinstance(usage_data, dict): usage_data = outputs.get("usage") if not isinstance(usage_data, dict): usage_data = {} prompt_tokens = usage_data.get("prompt_tokens", 0) completion_tokens = usage_data.get("completion_tokens", 0) completion_start_time = self._get_completion_start_time( created_at, usage_data.get("time_to_first_token") ) except Exception: logger.error("Failed to extract usage", exc_info=True) # add generation generation_usage = GenerationUsage( input=prompt_tokens, output=completion_tokens, total=total_token, unit=UnitEnum.TOKENS, ) node_generation_data = LangfuseGeneration( id=node_execution_id, name=node_name, trace_id=trace_id, model=process_data.get("model_name"), start_time=created_at, completion_start_time=completion_start_time, end_time=finished_at, input=inputs, output=outputs, metadata=metadata, level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR), status_message=trace_info.error or "", parent_observation_id=trace_info.workflow_run_id if trace_info.message_id else None, usage=generation_usage, ) self.add_generation(langfuse_generation_data=node_generation_data) # add normal span else: span_data = LangfuseSpan( id=node_execution_id, name=node_name, input=inputs, output=outputs, trace_id=trace_id, start_time=created_at, end_time=finished_at, metadata=metadata, level=(LevelEnum.DEFAULT if status == "succeeded" else LevelEnum.ERROR), status_message=trace_info.error or "", parent_observation_id=trace_info.workflow_run_id if trace_info.message_id else None, ) self.add_span(langfuse_span_data=span_data) def message_trace(self, trace_info: MessageTraceInfo, **kwargs): # get message file data file_list = trace_info.file_list metadata = trace_info.metadata message_data = trace_info.message_data if message_data is None: return message_id = message_data.id user_id = message_data.from_account_id if message_data.from_end_user_id: end_user_data: EndUser | None = db.session.get(EndUser, message_data.from_end_user_id) if end_user_data is not None: user_id = end_user_data.session_id metadata["user_id"] = user_id trace_id = trace_info.trace_id or message_id trace_data = LangfuseTrace( id=trace_id, user_id=user_id, name=TraceTaskName.MESSAGE_TRACE, input={ "message": trace_info.inputs, "files": file_list, "message_tokens": trace_info.message_tokens, "answer_tokens": trace_info.answer_tokens, "total_tokens": trace_info.total_tokens, "error": trace_info.error, "provider_response_latency": message_data.provider_response_latency, "created_at": trace_info.start_time, }, output=trace_info.outputs, metadata=metadata, session_id=message_data.conversation_id, tags=["message", str(trace_info.conversation_mode)], version=None, release=None, public=None, ) self.add_trace(langfuse_trace_data=trace_data) # add generation generation_usage = GenerationUsage( input=trace_info.message_tokens, output=trace_info.answer_tokens, total=trace_info.total_tokens, unit=UnitEnum.TOKENS, totalCost=message_data.total_price, ) completion_start_time = self._get_completion_start_time( trace_info.start_time, trace_info.gen_ai_server_time_to_first_token, ) langfuse_generation_data = LangfuseGeneration( name="llm", trace_id=trace_id, start_time=trace_info.start_time, completion_start_time=completion_start_time, end_time=trace_info.end_time, model=message_data.model_id, input=trace_info.inputs, output=message_data.answer, metadata=metadata, level=(LevelEnum.DEFAULT if message_data.status != MessageStatus.ERROR else LevelEnum.ERROR), status_message=message_data.error or "", usage=generation_usage, ) try: langfuse_generation_data.id = str(message_id) except Exception: pass self.add_generation(langfuse_generation_data) # AGENT_TOOL_SPANS_V2 try: from sqlalchemy import select from core.db.session_factory import session_factory from models.model import MessageAgentThought with session_factory.create_session() as session: thoughts = session.scalars( select(MessageAgentThought) .where(MessageAgentThought.message_id == str(message_id)) .order_by(MessageAgentThought.position.asc()) ).all() logger.info("agent tool spans count=%s message_id=%s", len(thoughts), message_id) for at in thoughts: for tn in at.tools: if not tn: continue tool_inp = at.tool_inputs_dict.get(tn) outputs = at.tool_outputs_dict if isinstance(at.tool_outputs_dict, dict) else {} tool_output = outputs.get(tn) if tool_output in (None, "", {}): tool_output = at.observation or None meta = at.tool_meta.get(tn, {}) if isinstance(at.tool_meta, dict) else {} time_cost_f = float(getattr(at, "latency", 0) or 0) if not time_cost_f: try: time_cost_f = float((meta or {}).get("time_cost") or 0) except (TypeError, ValueError): time_cost_f = 0.0 if not time_cost_f: time_cost_f = 0.2 start_t = at.created_at or trace_info.start_time if start_t is not None and getattr(start_t, "tzinfo", None) is None: start_t = start_t.replace(tzinfo=UTC) end_t = start_t + timedelta(seconds=time_cost_f) if start_t else None err = (meta or {}).get("error") self.add_span( LangfuseSpan( name=tn, input=tool_inp, output=tool_output, trace_id=trace_id, parent_observation_id=str(message_id), start_time=start_t, end_time=end_t, metadata={**metadata, "tool_name": tn, "time_cost": time_cost_f}, level=(LevelEnum.DEFAULT if not err else LevelEnum.ERROR), status_message=str(err or ""), ) ) except Exception: logger.exception("agent tool spans patch failed") def moderation_trace(self, trace_info: ModerationTraceInfo): if trace_info.message_data is None: return span_data = LangfuseSpan( name=TraceTaskName.MODERATION_TRACE, input=trace_info.inputs, output={ "action": trace_info.action, "flagged": trace_info.flagged, "preset_response": trace_info.preset_response, "inputs": trace_info.inputs, }, trace_id=trace_info.trace_id or trace_info.message_id, start_time=trace_info.start_time or trace_info.message_data.created_at, end_time=trace_info.end_time or trace_info.message_data.created_at, metadata=trace_info.metadata, ) self.add_span(langfuse_span_data=span_data) def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): message_data = trace_info.message_data if message_data is None: return generation_usage = GenerationUsage( total=len(str(trace_info.suggested_question)), input=len(trace_info.inputs) if trace_info.inputs else 0, output=len(trace_info.suggested_question), unit=UnitEnum.CHARACTERS, ) generation_data = LangfuseGeneration( name=TraceTaskName.SUGGESTED_QUESTION_TRACE, input=trace_info.inputs, output=str(trace_info.suggested_question), trace_id=trace_info.trace_id or trace_info.message_id, start_time=trace_info.start_time, end_time=trace_info.end_time, metadata=trace_info.metadata, level=(LevelEnum.DEFAULT if message_data.status != MessageStatus.ERROR else LevelEnum.ERROR), status_message=message_data.error or "", usage=generation_usage, ) self.add_generation(langfuse_generation_data=generation_data) def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): if trace_info.message_data is None: return dataset_retrieval_span_data = LangfuseSpan( name=TraceTaskName.DATASET_RETRIEVAL_TRACE, input=trace_info.inputs, output={"documents": trace_info.documents}, trace_id=trace_info.trace_id or trace_info.message_id, start_time=trace_info.start_time or trace_info.message_data.created_at, end_time=trace_info.end_time or trace_info.message_data.updated_at, metadata=trace_info.metadata, ) self.add_span(langfuse_span_data=dataset_retrieval_span_data) def tool_trace(self, trace_info: ToolTraceInfo): tool_span_data = LangfuseSpan( name=trace_info.tool_name, input=trace_info.tool_inputs, output=trace_info.tool_outputs, trace_id=trace_info.trace_id or trace_info.message_id, start_time=trace_info.start_time, end_time=trace_info.end_time, metadata=trace_info.metadata, level=(LevelEnum.DEFAULT if trace_info.error == "" or trace_info.error is None else LevelEnum.ERROR), status_message=trace_info.error, ) self.add_span(langfuse_span_data=tool_span_data) def generate_name_trace(self, trace_info: GenerateNameTraceInfo): name_generation_trace_data = LangfuseTrace( name=TraceTaskName.GENERATE_NAME_TRACE, input=trace_info.inputs, output=trace_info.outputs, user_id=trace_info.tenant_id, metadata=trace_info.metadata, session_id=trace_info.conversation_id, ) self.add_trace(langfuse_trace_data=name_generation_trace_data) name_generation_span_data = LangfuseSpan( name=TraceTaskName.GENERATE_NAME_TRACE, input=trace_info.inputs, output=trace_info.outputs, trace_id=trace_info.conversation_id, start_time=trace_info.start_time, end_time=trace_info.end_time, metadata=trace_info.metadata, ) self.add_span(langfuse_span_data=name_generation_span_data) def _make_event_id(self) -> str: return str(uuid.uuid4()) def _now_iso(self) -> str: return datetime.now(UTC).isoformat() def add_trace(self, langfuse_trace_data: LangfuseTrace | None = None): data = filter_none_values(langfuse_trace_data.model_dump()) if langfuse_trace_data else {} try: body = TraceBody( id=data.get("id") or self._make_event_id(), name=data.get("name"), user_id=data.get("user_id"), input=data.get("input"), output=data.get("output"), metadata=data.get("metadata"), session_id=data.get("session_id"), version=data.get("version"), release=data.get("release"), tags=data.get("tags"), public=data.get("public"), ) event = IngestionEvent_TraceCreate( body=body, id=self._make_event_id(), timestamp=self._now_iso(), ) self.langfuse_client.ingestion.batch(batch=[event]) logger.debug("LangFuse Trace created successfully") except Exception as e: raise ValueError(f"LangFuse Failed to create trace: {str(e)}") def add_span(self, langfuse_span_data: LangfuseSpan | None = None): data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {} try: body = CreateSpanBody( id=data.get("id") or self._make_event_id(), trace_id=data.get("trace_id"), name=data.get("name"), start_time=data.get("start_time"), end_time=data.get("end_time"), input=data.get("input"), output=data.get("output"), metadata=data.get("metadata"), level=data.get("level"), status_message=data.get("status_message"), parent_observation_id=data.get("parent_observation_id"), version=data.get("version"), ) event = IngestionEvent_SpanCreate( body=body, id=self._make_event_id(), timestamp=self._now_iso(), ) self.langfuse_client.ingestion.batch(batch=[event]) logger.debug("LangFuse Span created successfully") except Exception as e: raise ValueError(f"LangFuse Failed to create span: {str(e)}") def update_span(self, span, langfuse_span_data: LangfuseSpan | None = None): format_span_data = filter_none_values(langfuse_span_data.model_dump()) if langfuse_span_data else {} span.end(**format_span_data) def add_generation(self, langfuse_generation_data: LangfuseGeneration | None = None): data = filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {} try: usage_data = data.pop("usage", None) usage = None if usage_data: usage = Usage( input=usage_data.get("input", 0) or 0, output=usage_data.get("output", 0) or 0, total=usage_data.get("total", 0) or 0, unit=usage_data.get("unit"), input_cost=usage_data.get("inputCost"), output_cost=usage_data.get("outputCost"), total_cost=usage_data.get("totalCost"), ) body = CreateGenerationBody( id=data.get("id") or self._make_event_id(), trace_id=data.get("trace_id"), name=data.get("name"), start_time=data.get("start_time"), end_time=data.get("end_time"), model=data.get("model"), model_parameters=data.get("model_parameters"), input=data.get("input"), output=data.get("output"), usage=usage, metadata=data.get("metadata"), level=data.get("level"), status_message=data.get("status_message"), parent_observation_id=data.get("parent_observation_id"), version=data.get("version"), completion_start_time=data.get("completion_start_time"), ) event = IngestionEvent_GenerationCreate( body=body, id=self._make_event_id(), timestamp=self._now_iso(), ) self.langfuse_client.ingestion.batch(batch=[event]) logger.debug("LangFuse Generation created successfully") except Exception as e: raise ValueError(f"LangFuse Failed to create generation: {str(e)}") def update_generation(self, generation, langfuse_generation_data: LangfuseGeneration | None = None): format_generation_data = ( filter_none_values(langfuse_generation_data.model_dump()) if langfuse_generation_data else {} ) generation.end(**format_generation_data) def api_check(self): try: projects = self.langfuse_client.projects.get() except Exception as e: logger.debug("LangFuse API check failed", exc_info=True) raise ValueError(f"LangFuse API check failed: {str(e)}") if not projects.data: raise ValueError("LangFuse API check failed: no project found for the provided credentials") return True def get_project_key(self): try: projects = self.langfuse_client.projects.get() return projects.data[0].id except Exception as e: logger.debug("LangFuse get project key failed", exc_info=True) raise ValueError(f"LangFuse get project key failed: {str(e)}")