diff --git a/api/core/ops/langfuse_trace/langfuse_trace.py b/api/core/ops/langfuse_trace/langfuse_trace.py index b9ba068b19..f67e270ab1 100644 --- a/api/core/ops/langfuse_trace/langfuse_trace.py +++ b/api/core/ops/langfuse_trace/langfuse_trace.py @@ -213,9 +213,24 @@ class LangFuseDataTrace(BaseTraceInstance): if process_data and process_data.get("model_mode") == "chat": total_token = metadata.get("total_tokens", 0) + prompt_tokens = 0 + completion_tokens = 0 + try: + if outputs.get("usage"): + prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0) + completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0) + else: + prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0) + completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0) + 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( diff --git a/api/core/ops/langsmith_trace/langsmith_trace.py b/api/core/ops/langsmith_trace/langsmith_trace.py index 4ffd888bdd..e3494e2f23 100644 --- a/api/core/ops/langsmith_trace/langsmith_trace.py +++ b/api/core/ops/langsmith_trace/langsmith_trace.py @@ -199,6 +199,7 @@ class LangSmithDataTrace(BaseTraceInstance): ) 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( @@ -212,9 +213,23 @@ class LangSmithDataTrace(BaseTraceInstance): else: run_type = LangSmithRunType.tool + prompt_tokens = 0 + completion_tokens = 0 + try: + if outputs.get("usage"): + prompt_tokens = outputs.get("usage", {}).get("prompt_tokens", 0) + completion_tokens = outputs.get("usage", {}).get("completion_tokens", 0) + else: + prompt_tokens = process_data.get("usage", {}).get("prompt_tokens", 0) + completion_tokens = process_data.get("usage", {}).get("completion_tokens", 0) + except Exception: + logger.error("Failed to extract usage", exc_info=True) + node_dotted_order = generate_dotted_order(node_execution_id, created_at, workflow_dotted_order) langsmith_run = LangSmithRunModel( total_tokens=node_total_tokens, + input_tokens=prompt_tokens, + output_tokens=completion_tokens, name=node_type, inputs=inputs, run_type=run_type,