import json import logging import os import uuid from datetime import datetime, timedelta from typing import Optional, cast from langsmith import Client from langsmith.schemas import RunBase from core.ops.base_trace_instance import BaseTraceInstance from core.ops.entities.config_entity import LangSmithConfig from core.ops.entities.trace_entity import ( BaseTraceInfo, DatasetRetrievalTraceInfo, GenerateNameTraceInfo, MessageTraceInfo, ModerationTraceInfo, SuggestedQuestionTraceInfo, ToolTraceInfo, TraceTaskName, WorkflowTraceInfo, ) from core.ops.langsmith_trace.entities.langsmith_trace_entity import ( LangSmithRunModel, LangSmithRunType, LangSmithRunUpdateModel, ) from core.ops.utils import filter_none_values, generate_dotted_order from extensions.ext_database import db from models.model import EndUser, MessageFile from models.workflow import WorkflowNodeExecution logger = logging.getLogger(__name__) class LangSmithDataTrace(BaseTraceInstance): def __init__( self, langsmith_config: LangSmithConfig, ): super().__init__(langsmith_config) self.langsmith_key = langsmith_config.api_key self.project_name = langsmith_config.project self.project_id = None self.langsmith_client = Client(api_key=langsmith_config.api_key, api_url=langsmith_config.endpoint) self.file_base_url = os.getenv("FILES_URL", "http://127.0.0.1:5001") def trace(self, trace_info: BaseTraceInfo): if isinstance(trace_info, WorkflowTraceInfo): self.workflow_trace(trace_info) if isinstance(trace_info, MessageTraceInfo): self.message_trace(trace_info) if isinstance(trace_info, ModerationTraceInfo): self.moderation_trace(trace_info) if isinstance(trace_info, SuggestedQuestionTraceInfo): self.suggested_question_trace(trace_info) if isinstance(trace_info, DatasetRetrievalTraceInfo): self.dataset_retrieval_trace(trace_info) if isinstance(trace_info, ToolTraceInfo): self.tool_trace(trace_info) if isinstance(trace_info, GenerateNameTraceInfo): self.generate_name_trace(trace_info) def workflow_trace(self, trace_info: WorkflowTraceInfo): trace_id = trace_info.message_id or trace_info.workflow_run_id if trace_info.start_time is None: trace_info.start_time = datetime.now() message_dotted_order = ( generate_dotted_order(trace_info.message_id, trace_info.start_time) if trace_info.message_id else None ) workflow_dotted_order = generate_dotted_order( trace_info.workflow_run_id, trace_info.workflow_data.created_at, message_dotted_order, ) metadata = trace_info.metadata metadata["workflow_app_log_id"] = trace_info.workflow_app_log_id if trace_info.message_id: message_run = LangSmithRunModel( id=trace_info.message_id, name=TraceTaskName.MESSAGE_TRACE.value, inputs=dict(trace_info.workflow_run_inputs), outputs=dict(trace_info.workflow_run_outputs), run_type=LangSmithRunType.chain, start_time=trace_info.start_time, end_time=trace_info.end_time, extra={ "metadata": metadata, }, tags=["message", "workflow"], error=trace_info.error, trace_id=trace_id, dotted_order=message_dotted_order, file_list=[], serialized=None, parent_run_id=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, ) self.add_run(message_run) langsmith_run = LangSmithRunModel( file_list=trace_info.file_list, total_tokens=trace_info.total_tokens, id=trace_info.workflow_run_id, name=TraceTaskName.WORKFLOW_TRACE.value, inputs=dict(trace_info.workflow_run_inputs), run_type=LangSmithRunType.tool, start_time=trace_info.workflow_data.created_at, end_time=trace_info.workflow_data.finished_at, outputs=dict(trace_info.workflow_run_outputs), extra={ "metadata": metadata, }, error=trace_info.error, tags=["workflow"], parent_run_id=trace_info.message_id or None, trace_id=trace_id, dotted_order=workflow_dotted_order, serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, ) self.add_run(langsmith_run) # through workflow_run_id get all_nodes_execution workflow_nodes_execution_id_records = ( db.session.query(WorkflowNodeExecution.id) .filter(WorkflowNodeExecution.workflow_run_id == trace_info.workflow_run_id) .all() ) for node_execution_id_record in workflow_nodes_execution_id_records: node_execution = ( db.session.query( WorkflowNodeExecution.id, WorkflowNodeExecution.tenant_id, WorkflowNodeExecution.app_id, WorkflowNodeExecution.title, WorkflowNodeExecution.node_type, WorkflowNodeExecution.status, WorkflowNodeExecution.inputs, WorkflowNodeExecution.outputs, WorkflowNodeExecution.created_at, WorkflowNodeExecution.elapsed_time, WorkflowNodeExecution.process_data, WorkflowNodeExecution.execution_metadata, ) .filter(WorkflowNodeExecution.id == node_execution_id_record.id) .first() ) if not node_execution: continue node_execution_id = node_execution.id tenant_id = node_execution.tenant_id app_id = node_execution.app_id node_name = node_execution.title node_type = node_execution.node_type status = node_execution.status if node_type == "llm": inputs = ( json.loads(node_execution.process_data).get("prompts", {}) if node_execution.process_data else {} ) else: inputs = json.loads(node_execution.inputs) if node_execution.inputs else {} outputs = json.loads(node_execution.outputs) if node_execution.outputs else {} 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 = ( json.loads(node_execution.execution_metadata) if node_execution.execution_metadata else {} ) node_total_tokens = execution_metadata.get("total_tokens", 0) metadata = execution_metadata.copy() metadata.update( { "workflow_run_id": trace_info.workflow_run_id, "node_execution_id": node_execution_id, "tenant_id": tenant_id, "app_id": app_id, "app_name": node_name, "node_type": node_type, "status": status, } ) process_data = json.loads(node_execution.process_data) if node_execution.process_data else {} if process_data and process_data.get("model_mode") == "chat": run_type = LangSmithRunType.llm metadata.update( { "ls_provider": process_data.get("model_provider", ""), "ls_model_name": process_data.get("model_name", ""), } ) elif node_type == "knowledge-retrieval": run_type = LangSmithRunType.retriever else: run_type = LangSmithRunType.tool node_dotted_order = generate_dotted_order(node_execution_id, created_at, workflow_dotted_order) langsmith_run = LangSmithRunModel( total_tokens=node_total_tokens, name=node_type, inputs=inputs, run_type=run_type, start_time=created_at, end_time=finished_at, outputs=outputs, file_list=trace_info.file_list, extra={ "metadata": metadata, }, parent_run_id=trace_info.workflow_run_id, tags=["node_execution"], id=node_execution_id, trace_id=trace_id, dotted_order=node_dotted_order, error="", serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, ) self.add_run(langsmith_run) def message_trace(self, trace_info: MessageTraceInfo): # get message file data file_list = cast(list[str], trace_info.file_list) or [] message_file_data: Optional[MessageFile] = trace_info.message_file_data file_url = f"{self.file_base_url}/{message_file_data.url}" if message_file_data else "" file_list.append(file_url) 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 metadata["user_id"] = user_id if message_data.from_end_user_id: end_user_data: Optional[EndUser] = ( db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() ) if end_user_data is not None: end_user_id = end_user_data.session_id metadata["end_user_id"] = end_user_id message_run = LangSmithRunModel( input_tokens=trace_info.message_tokens, output_tokens=trace_info.answer_tokens, total_tokens=trace_info.total_tokens, id=message_id, name=TraceTaskName.MESSAGE_TRACE.value, inputs=trace_info.inputs, run_type=LangSmithRunType.chain, start_time=trace_info.start_time, end_time=trace_info.end_time, outputs=message_data.answer, extra={"metadata": metadata}, tags=["message", str(trace_info.conversation_mode)], error=trace_info.error, file_list=file_list, serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, trace_id=None, dotted_order=None, parent_run_id=None, ) self.add_run(message_run) # create llm run parented to message run llm_run = LangSmithRunModel( input_tokens=trace_info.message_tokens, output_tokens=trace_info.answer_tokens, total_tokens=trace_info.total_tokens, name="llm", inputs=trace_info.inputs, run_type=LangSmithRunType.llm, start_time=trace_info.start_time, end_time=trace_info.end_time, outputs=message_data.answer, extra={"metadata": metadata}, parent_run_id=message_id, tags=["llm", str(trace_info.conversation_mode)], error=trace_info.error, file_list=file_list, serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, trace_id=None, dotted_order=None, id=str(uuid.uuid4()), ) self.add_run(llm_run) def moderation_trace(self, trace_info: ModerationTraceInfo): if trace_info.message_data is None: return langsmith_run = LangSmithRunModel( name=TraceTaskName.MODERATION_TRACE.value, inputs=trace_info.inputs, outputs={ "action": trace_info.action, "flagged": trace_info.flagged, "preset_response": trace_info.preset_response, "inputs": trace_info.inputs, }, run_type=LangSmithRunType.tool, extra={"metadata": trace_info.metadata}, tags=["moderation"], parent_run_id=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, id=str(uuid.uuid4()), serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, trace_id=None, dotted_order=None, error="", file_list=[], ) self.add_run(langsmith_run) def suggested_question_trace(self, trace_info: SuggestedQuestionTraceInfo): message_data = trace_info.message_data if message_data is None: return suggested_question_run = LangSmithRunModel( name=TraceTaskName.SUGGESTED_QUESTION_TRACE.value, inputs=trace_info.inputs, outputs=trace_info.suggested_question, run_type=LangSmithRunType.tool, extra={"metadata": trace_info.metadata}, tags=["suggested_question"], parent_run_id=trace_info.message_id, start_time=trace_info.start_time or message_data.created_at, end_time=trace_info.end_time or message_data.updated_at, id=str(uuid.uuid4()), serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, trace_id=None, dotted_order=None, error="", file_list=[], ) self.add_run(suggested_question_run) def dataset_retrieval_trace(self, trace_info: DatasetRetrievalTraceInfo): if trace_info.message_data is None: return dataset_retrieval_run = LangSmithRunModel( name=TraceTaskName.DATASET_RETRIEVAL_TRACE.value, inputs=trace_info.inputs, outputs={"documents": trace_info.documents}, run_type=LangSmithRunType.retriever, extra={"metadata": trace_info.metadata}, tags=["dataset_retrieval"], parent_run_id=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, id=str(uuid.uuid4()), serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, trace_id=None, dotted_order=None, error="", file_list=[], ) self.add_run(dataset_retrieval_run) def tool_trace(self, trace_info: ToolTraceInfo): tool_run = LangSmithRunModel( name=trace_info.tool_name, inputs=trace_info.tool_inputs, outputs=trace_info.tool_outputs, run_type=LangSmithRunType.tool, extra={ "metadata": trace_info.metadata, }, tags=["tool", trace_info.tool_name], parent_run_id=trace_info.message_id, start_time=trace_info.start_time, end_time=trace_info.end_time, file_list=[cast(str, trace_info.file_url)], id=str(uuid.uuid4()), serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, trace_id=None, dotted_order=None, error=trace_info.error or "", ) self.add_run(tool_run) def generate_name_trace(self, trace_info: GenerateNameTraceInfo): name_run = LangSmithRunModel( name=TraceTaskName.GENERATE_NAME_TRACE.value, inputs=trace_info.inputs, outputs=trace_info.outputs, run_type=LangSmithRunType.tool, extra={"metadata": trace_info.metadata}, tags=["generate_name"], start_time=trace_info.start_time or datetime.now(), end_time=trace_info.end_time or datetime.now(), id=str(uuid.uuid4()), serialized=None, events=[], session_id=None, session_name=None, reference_example_id=None, input_attachments={}, output_attachments={}, trace_id=None, dotted_order=None, error="", file_list=[], parent_run_id=None, ) self.add_run(name_run) def add_run(self, run_data: LangSmithRunModel): data = run_data.model_dump() if self.project_id: data["session_id"] = self.project_id elif self.project_name: data["session_name"] = self.project_name data = filter_none_values(data) try: self.langsmith_client.create_run(**data) logger.debug("LangSmith Run created successfully.") except Exception as e: raise ValueError(f"LangSmith Failed to create run: {str(e)}") def update_run(self, update_run_data: LangSmithRunUpdateModel): data = update_run_data.model_dump() data = filter_none_values(data) try: self.langsmith_client.update_run(**data) logger.debug("LangSmith Run updated successfully.") except Exception as e: raise ValueError(f"LangSmith Failed to update run: {str(e)}") def api_check(self): try: random_project_name = f"test_project_{datetime.now().strftime('%Y%m%d%H%M%S')}" self.langsmith_client.create_project(project_name=random_project_name) self.langsmith_client.delete_project(project_name=random_project_name) return True except Exception as e: logger.debug(f"LangSmith API check failed: {str(e)}") raise ValueError(f"LangSmith API check failed: {str(e)}") def get_project_url(self): try: run_data = RunBase( id=uuid.uuid4(), name="tool", inputs={"input": "test"}, outputs={"output": "test"}, run_type=LangSmithRunType.tool, start_time=datetime.now(), ) project_url = self.langsmith_client.get_run_url( run=run_data, project_id=self.project_id, project_name=self.project_name ) return project_url.split("/r/")[0] except Exception as e: logger.debug(f"LangSmith get run url failed: {str(e)}") raise ValueError(f"LangSmith get run url failed: {str(e)}")