From c1cf8fe564aab52b8257be13a2c92a9ce26e6062 Mon Sep 17 00:00:00 2001
From: italojohnny <italojohnnydosanjos@gmail.com>
Date: Wed, 29 May 2024 16:01:54 -0300
Subject: [PATCH] change return to dict for token cases

---
 .../base/langflow/base/models/model.py        | 27 ++++++++++++-------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/src/backend/base/langflow/base/models/model.py b/src/backend/base/langflow/base/models/model.py
index b38d275f9..690dc01ba 100644
--- a/src/backend/base/langflow/base/models/model.py
+++ b/src/backend/base/langflow/base/models/model.py
@@ -53,19 +53,28 @@ class LCModelComponent(CustomComponent):
                 key in response_metadata["token_usage"] for key in inner_openai_keys
             ):
                 token_usage = response_metadata["token_usage"]
-                completion_tokens = token_usage["completion_tokens"]
-                prompt_tokens = token_usage["prompt_tokens"]
-                total_tokens = token_usage["total_tokens"]
-                finish_reason = response_metadata["finish_reason"]
-                status_message = f"Tokens:\nInput: {prompt_tokens}\nOutput: {completion_tokens}\nTotal Tokens: {total_tokens}\nStop Reason: {finish_reason}\nResponse: {content}"
+                status_message = {
+                    "tokens": {
+                        "input": token_usage["prompt_tokens"],
+                        "output": token_usage["completion_tokens"],
+                        "total": token_usage["total_tokens"],
+                        "stop_reason": response_metadata["finish_reason"],
+                        "response": content,
+                    }
+                }
+
             elif all(key in response_metadata for key in anthropic_keys) and all(
                 key in response_metadata["usage"] for key in inner_anthropic_keys
             ):
                 usage = response_metadata["usage"]
-                input_tokens = usage["input_tokens"]
-                output_tokens = usage["output_tokens"]
-                stop_reason = response_metadata["stop_reason"]
-                status_message = f"Tokens:\nInput: {input_tokens}\nOutput: {output_tokens}\nStop Reason: {stop_reason}\nResponse: {content}"
+                status_message = {
+                    "tokens": {
+                        "input": usage["input_tokens"],
+                        "output": usage["output_tokens"],
+                        "stop_reason": response_metadata["stop_reason"],
+                        "response": content,
+                    }
+                }
             else:
                 status_message = f"Response: {content}"
         else: