Expose LLM usage in workflows (#21766)

Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com>

api/core/ops/aliyun_trace/aliyun_trace.py

@@ -372,6 +372,7 @@ class AliyunDataTrace(BaseTraceInstance):
     ) -> SpanData:
         process_data = node_execution.process_data or {}
         outputs = node_execution.outputs or {}
+        usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
         return SpanData(
             trace_id=trace_id,
             parent_span_id=workflow_span_id,
@@ -385,9 +386,9 @@ class AliyunDataTrace(BaseTraceInstance):
                 GEN_AI_FRAMEWORK: "dify",
                 GEN_AI_MODEL_NAME: process_data.get("model_name", ""),
                 GEN_AI_SYSTEM: process_data.get("model_provider", ""),
-                GEN_AI_USAGE_INPUT_TOKENS: str(outputs.get("usage", {}).get("prompt_tokens", 0)),
-                GEN_AI_USAGE_OUTPUT_TOKENS: str(outputs.get("usage", {}).get("completion_tokens", 0)),
-                GEN_AI_USAGE_TOTAL_TOKENS: str(outputs.get("usage", {}).get("total_tokens", 0)),
+                GEN_AI_USAGE_INPUT_TOKENS: str(usage_data.get("prompt_tokens", 0)),
+                GEN_AI_USAGE_OUTPUT_TOKENS: str(usage_data.get("completion_tokens", 0)),
+                GEN_AI_USAGE_TOTAL_TOKENS: str(usage_data.get("total_tokens", 0)),
                 GEN_AI_PROMPT: json.dumps(process_data.get("prompts", []), ensure_ascii=False),
                 GEN_AI_COMPLETION: str(outputs.get("text", "")),
                 GEN_AI_RESPONSE_FINISH_REASON: outputs.get("finish_reason", ""),

api/core/ops/arize_phoenix_trace/arize_phoenix_trace.py

@@ -213,11 +213,12 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
                     if model:
                         node_metadata["ls_model_name"] = model
 
-                    usage = json.loads(node_execution.outputs).get("usage", {}) if node_execution.outputs else {}
-                    if usage:
-                        node_metadata["total_tokens"] = usage.get("total_tokens", 0)
-                        node_metadata["prompt_tokens"] = usage.get("prompt_tokens", 0)
-                        node_metadata["completion_tokens"] = usage.get("completion_tokens", 0)
+                    outputs = json.loads(node_execution.outputs).get("usage", {})
+                    usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                    if usage_data:
+                        node_metadata["total_tokens"] = usage_data.get("total_tokens", 0)
+                        node_metadata["prompt_tokens"] = usage_data.get("prompt_tokens", 0)
+                        node_metadata["completion_tokens"] = usage_data.get("completion_tokens", 0)
                 elif node_execution.node_type == "dataset_retrieval":
                     span_kind = OpenInferenceSpanKindValues.RETRIEVER.value
                 elif node_execution.node_type == "tool":
@@ -246,14 +247,19 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
                         if model:
                             node_span.set_attribute(SpanAttributes.LLM_MODEL_NAME, model)
 
-                        usage = json.loads(node_execution.outputs).get("usage", {}) if node_execution.outputs else {}
-                        if usage:
-                            node_span.set_attribute(SpanAttributes.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens", 0))
+                        outputs = json.loads(node_execution.outputs).get("usage", {})
+                        usage_data = (
+                            process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                        )
+                        if usage_data:
                             node_span.set_attribute(
-                                SpanAttributes.LLM_TOKEN_COUNT_PROMPT, usage.get("prompt_tokens", 0)
+                                SpanAttributes.LLM_TOKEN_COUNT_TOTAL, usage_data.get("total_tokens", 0)
                             )
                             node_span.set_attribute(
-                                SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, usage.get("completion_tokens", 0)
+                                SpanAttributes.LLM_TOKEN_COUNT_PROMPT, usage_data.get("prompt_tokens", 0)
+                            )
+                            node_span.set_attribute(
+                                SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, usage_data.get("completion_tokens", 0)
                             )
                 finally:
                     node_span.end(end_time=datetime_to_nanos(finished_at))

api/core/ops/langfuse_trace/langfuse_trace.py

@@ -181,12 +181,9 @@ class LangFuseDataTrace(BaseTraceInstance):
                 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)
+                    usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                    prompt_tokens = usage_data.get("prompt_tokens", 0)
+                    completion_tokens = usage_data.get("completion_tokens", 0)
                 except Exception:
                     logger.error("Failed to extract usage", exc_info=True)
 

api/core/ops/langsmith_trace/langsmith_trace.py

@@ -206,12 +206,9 @@ class LangSmithDataTrace(BaseTraceInstance):
             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)
+                usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                prompt_tokens = usage_data.get("prompt_tokens", 0)
+                completion_tokens = usage_data.get("completion_tokens", 0)
             except Exception:
                 logger.error("Failed to extract usage", exc_info=True)
 

api/core/ops/opik_trace/opik_trace.py

@@ -222,10 +222,10 @@ class OpikDataTrace(BaseTraceInstance):
                 )
 
                 try:
-                    if outputs.get("usage"):
-                        total_tokens = outputs["usage"].get("total_tokens", 0)
-                        prompt_tokens = outputs["usage"].get("prompt_tokens", 0)
-                        completion_tokens = outputs["usage"].get("completion_tokens", 0)
+                    usage_data = process_data.get("usage", {}) if "usage" in process_data else outputs.get("usage", {})
+                    total_tokens = usage_data.get("total_tokens", 0)
+                    prompt_tokens = usage_data.get("prompt_tokens", 0)
+                    completion_tokens = usage_data.get("completion_tokens", 0)
                 except Exception:
                     logger.error("Failed to extract usage", exc_info=True)
 

api/core/workflow/nodes/llm/node.py

@@ -221,15 +221,6 @@ class LLMNode(BaseNode[LLMNodeData]):
                 jinja2_variables=self.node_data.prompt_config.jinja2_variables,
             )
 
-            process_data = {
-                "model_mode": model_config.mode,
-                "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
-                    model_mode=model_config.mode, prompt_messages=prompt_messages
-                ),
-                "model_provider": model_config.provider,
-                "model_name": model_config.model,
-            }
-
             # handle invoke result
             generator = self._invoke_llm(
                 node_data_model=self.node_data.model,
@@ -253,6 +244,17 @@ class LLMNode(BaseNode[LLMNodeData]):
                 elif isinstance(event, LLMStructuredOutput):
                     structured_output = event
 
+            process_data = {
+                "model_mode": model_config.mode,
+                "prompts": PromptMessageUtil.prompt_messages_to_prompt_for_saving(
+                    model_mode=model_config.mode, prompt_messages=prompt_messages
+                ),
+                "usage": jsonable_encoder(usage),
+                "finish_reason": finish_reason,
+                "model_provider": model_config.provider,
+                "model_name": model_config.model,
+            }
+
             outputs = {"text": result_text, "usage": jsonable_encoder(usage), "finish_reason": finish_reason}
             if structured_output:
                 outputs["structured_output"] = structured_output.structured_output

api/core/workflow/nodes/parameter_extractor/parameter_extractor_node.py

@@ -253,7 +253,12 @@ class ParameterExtractorNode(BaseNode):
             status=WorkflowNodeExecutionStatus.SUCCEEDED,
             inputs=inputs,
             process_data=process_data,
-            outputs={"__is_success": 1 if not error else 0, "__reason": error, **result},
+            outputs={
+                "__is_success": 1 if not error else 0,
+                "__reason": error,
+                "__usage": jsonable_encoder(usage),
+                **result,
+            },
             metadata={
                 WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: usage.total_tokens,
                 WorkflowNodeExecutionMetadataKey.TOTAL_PRICE: usage.total_price,

api/core/workflow/nodes/question_classifier/question_classifier_node.py

@@ -145,7 +145,11 @@ class QuestionClassifierNode(LLMNode):
                 "model_provider": model_config.provider,
                 "model_name": model_config.model,
             }
-            outputs = {"class_name": category_name, "class_id": category_id}
+            outputs = {
+                "class_name": category_name,
+                "class_id": category_id,
+                "usage": jsonable_encoder(usage),
+            }
 
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.SUCCEEDED,

web/app/components/workflow/constants.ts

@@ -480,6 +480,10 @@ export const LLM_OUTPUT_STRUCT: Var[] = [
     variable: 'text',
     type: VarType.string,
   },
+  {
+    variable: 'usage',
+    type: VarType.object,
+  },
 ]
 
 export const KNOWLEDGE_RETRIEVAL_OUTPUT_STRUCT: Var[] = [
@@ -501,6 +505,10 @@ export const QUESTION_CLASSIFIER_OUTPUT_STRUCT = [
     variable: 'class_name',
     type: VarType.string,
   },
+  {
+    variable: 'usage',
+    type: VarType.object,
+  },
 ]
 
 export const HTTP_REQUEST_OUTPUT_STRUCT: Var[] = [
@@ -546,6 +554,10 @@ export const PARAMETER_EXTRACTOR_COMMON_STRUCT: Var[] = [
     variable: '__reason',
     type: VarType.string,
   },
+  {
+    variable: '__usage',
+    type: VarType.object,
+  },
 ]
 
 export const FILE_STRUCT: Var[] = [

web/app/components/workflow/nodes/llm/panel.tsx

@@ -282,6 +282,11 @@ const Panel: FC<NodePanelProps<LLMNodeType>> = ({
             type='string'
             description={t(`${i18nPrefix}.outputVars.output`)}
           />
+          <VarItem
+            name='usage'
+            type='object'
+            description={t(`${i18nPrefix}.outputVars.usage`)}
+          />
           {inputs.structured_output_enabled && (
             <>
               <Split className='mt-3' />

web/app/components/workflow/nodes/parameter-extractor/panel.tsx

@@ -190,12 +190,17 @@ const Panel: FC<NodePanelProps<ParameterExtractorNodeType>> = ({
               <VarItem
                 name='__is_success'
                 type={VarType.number}
-                description={t(`${i18nPrefix}.isSuccess`)}
+                description={t(`${i18nPrefix}.outputVars.isSuccess`)}
               />
               <VarItem
                 name='__reason'
                 type={VarType.string}
-                description={t(`${i18nPrefix}.errorReason`)}
+                description={t(`${i18nPrefix}.outputVars.errorReason`)}
+              />
+              <VarItem
+                name='__usage'
+                type='object'
+                description={t(`${i18nPrefix}.outputVars.usage`)}
               />
             </>
           </OutputVars>

web/app/components/workflow/nodes/question-classifier/panel.tsx

@@ -129,6 +129,11 @@ const Panel: FC<NodePanelProps<QuestionClassifierNodeType>> = ({
               type='string'
               description={t(`${i18nPrefix}.outputVars.className`)}
             />
+            <VarItem
+              name='usage'
+              type='object'
+              description={t(`${i18nPrefix}.outputVars.usage`)}
+            />
           </>
         </OutputVars>
       </div>

web/i18n/de-DE/workflow.ts

@@ -669,6 +669,7 @@ const translation = {
       inputVars: 'Eingabevariablen',
       outputVars: {
         className: 'Klassennamen',
+        usage: 'Nutzungsinformationen des Modells',
       },
       class: 'Klasse',
       classNamePlaceholder: 'Geben Sie Ihren Klassennamen ein',
@@ -682,6 +683,11 @@ const translation = {
     },
     parameterExtractor: {
       inputVar: 'Eingabevariable',
+      outputVars: {
+        isSuccess: 'Ist Erfolg. Bei Erfolg beträgt der Wert 1, bei Misserfolg beträgt der Wert 0.',
+        errorReason: 'Fehlergrund',
+        usage: 'Nutzungsinformationen des Modells',
+      },
       extractParameters: 'Parameter extrahieren',
       importFromTool: 'Aus Tools importieren',
       addExtractParameter: 'Extraktionsparameter hinzufügen',
@@ -701,8 +707,6 @@ const translation = {
       advancedSetting: 'Erweiterte Einstellung',
       reasoningMode: 'Schlussfolgerungsmodus',
       reasoningModeTip: 'Sie können den entsprechenden Schlussfolgerungsmodus basierend auf der Fähigkeit des Modells wählen, auf Anweisungen zur Funktionsaufruf- oder Eingabeaufforderungen zu reagieren.',
-      isSuccess: 'Ist Erfolg. Bei Erfolg beträgt der Wert 1, bei Misserfolg beträgt der Wert 0.',
-      errorReason: 'Fehlergrund',
     },
     iteration: {
       deleteTitle: 'Iterationsknoten löschen?',

web/i18n/en-US/workflow.ts

@@ -673,6 +673,7 @@ const translation = {
       inputVars: 'Input Variables',
       outputVars: {
         className: 'Class Name',
+        usage: 'Model Usage Information',
       },
       class: 'Class',
       classNamePlaceholder: 'Write your class name',
@@ -686,6 +687,11 @@ const translation = {
     },
     parameterExtractor: {
       inputVar: 'Input Variable',
+      outputVars: {
+        isSuccess: 'Is Success.On success the value is 1, on failure the value is 0.',
+        errorReason: 'Error Reason',
+        usage: 'Model Usage Information',
+      },
       extractParameters: 'Extract Parameters',
       importFromTool: 'Import from tools',
       addExtractParameter: 'Add Extract Parameter',
@@ -705,8 +711,6 @@ const translation = {
       advancedSetting: 'Advanced Setting',
       reasoningMode: 'Reasoning Mode',
       reasoningModeTip: 'You can choose the appropriate reasoning mode based on the model\'s ability to respond to instructions for function calling or prompts.',
-      isSuccess: 'Is Success.On success the value is 1, on failure the value is 0.',
-      errorReason: 'Error Reason',
     },
     iteration: {
       deleteTitle: 'Delete Iteration Node?',

web/i18n/es-ES/workflow.ts

@@ -667,6 +667,7 @@ const translation = {
       inputVars: 'Variables de entrada',
       outputVars: {
         className: 'Nombre de la clase',
+        usage: 'Información de uso del modelo',
       },
       class: 'Clase',
       classNamePlaceholder: 'Escribe el nombre de tu clase',
@@ -680,6 +681,11 @@ const translation = {
     },
     parameterExtractor: {
       inputVar: 'Variable de entrada',
+      outputVars: {
+        isSuccess: 'Es éxito. En caso de éxito el valor es 1, en caso de fallo el valor es 0.',
+        errorReason: 'Motivo del error',
+        usage: 'Información de uso del modelo',
+      },
       extractParameters: 'Extraer parámetros',
       importFromTool: 'Importar desde herramientas',
       addExtractParameter: 'Agregar parámetro de extracción',
@@ -699,8 +705,6 @@ const translation = {
       advancedSetting: 'Configuración avanzada',
       reasoningMode: 'Modo de razonamiento',
       reasoningModeTip: 'Puede elegir el modo de razonamiento apropiado basado en la capacidad del modelo para responder a instrucciones para llamadas de funciones o indicaciones.',
-      isSuccess: 'Es éxito. En caso de éxito el valor es 1, en caso de fallo el valor es 0.',
-      errorReason: 'Motivo del error',
     },
     iteration: {
       deleteTitle: '¿Eliminar nodo de iteración?',

web/i18n/fa-IR/workflow.ts

@@ -669,6 +669,7 @@ const translation = {
       inputVars: 'متغیرهای ورودی',
       outputVars: {
         className: 'نام کلاس',
+        usage: 'اطلاعات استفاده از مدل',
       },
       class: 'کلاس',
       classNamePlaceholder: 'نام کلاس خود را بنویسید',
@@ -682,6 +683,11 @@ const translation = {
     },
     parameterExtractor: {
       inputVar: 'متغیر ورودی',
+      outputVars: {
+        isSuccess: 'موفقیت‌آمیز است. در صورت موفقیت مقدار 1 و در صورت شکست مقدار 0 است.',
+        errorReason: 'دلیل خطا',
+        usage: 'اطلاعات استفاده از مدل',
+      },
       extractParameters: 'استخراج پارامترها',
       importFromTool: 'وارد کردن از ابزارها',
       addExtractParameter: 'افزودن پارامتر استخراج شده',
@@ -701,8 +707,6 @@ const translation = {
       advancedSetting: 'تنظیمات پیشرفته',
       reasoningMode: 'حالت استدلال',
       reasoningModeTip: 'می‌توانید حالت استدلال مناسب را بر اساس توانایی مدل برای پاسخ به دستورات برای فراخوانی عملکردها یا پیشنهادات انتخاب کنید.',
-      isSuccess: 'موفقیت‌آمیز است. در صورت موفقیت مقدار 1 و در صورت شکست مقدار 0 است.',
-      errorReason: 'دلیل خطا',
     },
     iteration: {
       deleteTitle: 'حذف نود تکرار؟',

web/i18n/fr-FR/workflow.ts

@@ -669,6 +669,7 @@ const translation = {
       inputVars: 'Variables de saisie',
       outputVars: {
         className: 'Nom de la classe',
+        usage: 'Informations sur l\'utilisation du modèle',
       },
       class: 'Classe',
       classNamePlaceholder: 'Écrivez le nom de votre classe',
@@ -682,6 +683,11 @@ const translation = {
     },
     parameterExtractor: {
       inputVar: 'Variable de saisie',
+      outputVars: {
+        isSuccess: 'Est réussi. En cas de succès, la valeur est 1, en cas d\'échec, la valeur est 0.',
+        errorReason: 'Raison de l\'erreur',
+        usage: 'Informations sur l\'utilisation du modèle',
+      },
       extractParameters: 'Extraire des paramètres',
       importFromTool: 'Importer des outils',
       addExtractParameter: 'Ajouter un paramètre d\'extraction',
@@ -701,8 +707,6 @@ const translation = {
       advancedSetting: 'Paramètre avancé',
       reasoningMode: 'Mode de raisonnement',
       reasoningModeTip: 'Vous pouvez choisir le mode de raisonnement approprié en fonction de la capacité du modèle à répondre aux instructions pour les appels de fonction ou les invites.',
-      isSuccess: 'Est réussi. En cas de succès, la valeur est 1, en cas d\'échec, la valeur est 0.',
-      errorReason: 'Raison de l\'erreur',
     },
     iteration: {
       deleteTitle: 'Supprimer le nœud d\'itération?',

web/i18n/hi-IN/workflow.ts

@@ -685,6 +685,7 @@ const translation = {
       inputVars: 'इनपुट वेरिएबल्स',
       outputVars: {
         className: 'क्लास नाम',
+        usage: 'मॉडल उपयोग जानकारी',
       },
       class: 'क्लास',
       classNamePlaceholder: 'अपना क्लास नाम लिखें',
@@ -699,6 +700,11 @@ const translation = {
     },
     parameterExtractor: {
       inputVar: 'इनपुट वेरिएबल',
+      outputVars: {
+        isSuccess: 'सफलता है। सफलता पर मान 1 है, असफलता पर मान 0 है।',
+        errorReason: 'त्रुटि का कारण',
+        usage: 'मॉडल उपयोग जानकारी',
+      },
       extractParameters: 'पैरामीटर्स निकालें',
       importFromTool: 'उपकरणों से आयात करें',
       addExtractParameter: 'एक्सट्रेक्ट पैरामीटर जोड़ें',
@@ -721,8 +727,6 @@ const translation = {
       reasoningMode: 'रीज़निंग मोड',
       reasoningModeTip:
         'फ़ंक्शन कॉलिंग या प्रॉम्प्ट्स के लिए निर्देशों का जवाब देने की मॉडल की क्षमता के आधार पर उपयुक्त रीज़निंग मोड चुन सकते हैं।',
-      isSuccess: 'सफलता है। सफलता पर मान 1 है, असफलता पर मान 0 है।',
-      errorReason: 'त्रुटि का कारण',
     },
     iteration: {
       deleteTitle: 'इटरेशन नोड हटाएं?',

web/i18n/it-IT/workflow.ts

@@ -688,6 +688,7 @@ const translation = {
       inputVars: 'Variabili di Input',
       outputVars: {
         className: 'Nome Classe',
+        usage: 'Informazioni sull\'utilizzo del modello',
       },
       class: 'Classe',
       classNamePlaceholder: 'Scrivi il nome della tua classe',
@@ -702,6 +703,11 @@ const translation = {
     },
     parameterExtractor: {
       inputVar: 'Variabile di Input',
+      outputVars: {
+        isSuccess: 'È successo. In caso di successo il valore è 1, in caso di fallimento il valore è 0.',
+        errorReason: 'Motivo dell\'errore',
+        usage: 'Informazioni sull\'utilizzo del modello',
+      },
       extractParameters: 'Estrai Parametri',
       importFromTool: 'Importa dagli strumenti',
       addExtractParameter: 'Aggiungi Parametro Estratto',
@@ -724,9 +730,6 @@ const translation = {
       reasoningMode: 'Modalità di ragionamento',
       reasoningModeTip:
         'Puoi scegliere la modalità di ragionamento appropriata in base alla capacità del modello di rispondere alle istruzioni per la chiamata delle funzioni o i prompt.',
-      isSuccess:
-        'È successo. In caso di successo il valore è 1, in caso di fallimento il valore è 0.',
-      errorReason: 'Motivo dell\'errore',
     },
     iteration: {
       deleteTitle: 'Eliminare Nodo Iterazione?',

web/i18n/ja-JP/workflow.ts

@@ -676,6 +676,7 @@ const translation = {
       inputVars: '入力変数',
       outputVars: {
         className: 'クラス名',
+        usage: 'モデル使用量',
       },
       class: 'クラス',
       classNamePlaceholder: 'クラス名を入力してください',
@@ -689,6 +690,11 @@ const translation = {
     },
     parameterExtractor: {
       inputVar: '入力変数',
+      outputVars: {
+        isSuccess: '成功。成功した場合の値は 1、失敗した場合の値は 0 です。',
+        errorReason: 'エラーの理由',
+        usage: 'モデル使用量',
+      },
       extractParameters: 'パラメーターを抽出',
       importFromTool: 'ツールからインポート',
       addExtractParameter: '抽出パラメーターを追加',
@@ -708,8 +714,6 @@ const translation = {
       advancedSetting: '高度な設定',
       reasoningMode: '推論モード',
       reasoningModeTip: '関数呼び出しやプロンプトの指示に応答するモデルの能力に基づいて、適切な推論モードを選択できます。',
-      isSuccess: '成功。成功した場合の値は 1、失敗した場合の値は 0 です。',
-      errorReason: 'エラーの理由',
     },
     iteration: {
       deleteTitle: 'イテレーションノードを削除しますか？',