From 68806eb0e8c0978c3f2a79857fc203fdd57a4b35 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 18:52:10 +0530
Subject: [PATCH] update the log_step_metadata

- add newer methods to the log_step_metadata for logging
---
 swarms/structs/agent.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index c70fe201..6c8e5fea 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -1897,13 +1897,27 @@ class Agent:
     def log_step_metadata(
         self, loop: int, task: str, response: str
     ) -> Step:
-        # # # Step Metadata
+        """Log metadata for each step of agent execution."""
+        # Generate unique step ID
+        step_id = f"step_{loop}_{uuid.uuid4().hex}"
+        
+        # Calculate token usage
         # full_memory = self.short_memory.return_history_as_string()
         # prompt_tokens = self.tokenizer.count_tokens(full_memory)
         # completion_tokens = self.tokenizer.count_tokens(response)
-        # self.tokenizer.count_tokens(prompt_tokens + completion_tokens)
+        # total_tokens = prompt_tokens + completion_tokens
+        total_tokens=self.tokenizer.count_tokens(task) + self.tokenizer.count_tokens(response),
+
+        # Create memory usage tracking
+        memory_usage = {
+            "short_term": len(self.short_memory.messages),
+            "long_term": self.long_term_memory.count if hasattr(self, 'long_term_memory') else 0
+        }
 
         step_log = Step(
+            step_id=step_id,
+            time=time.time(),
+            tokens = total_tokens,
             response=AgentChatCompletionResponse(
                 id=self.agent_id,
                 agent_name=self.agent_name,
@@ -1918,13 +1932,21 @@ class Agent:
                 ),
                 # usage=UsageInfo(
                 #     prompt_tokens=prompt_tokens,
-                #     total_tokens=total_tokens,
                 #     completion_tokens=completion_tokens,
+                #     total_tokens=total_tokens,
                 # ),
+                tool_calls=[],
+                memory_usage=memory_usage
             ),
         )
 
-        self.step_pool.append(step_log)
+        # Update total tokens if agent_output exists
+        if hasattr(self, 'agent_output'):
+            self.agent_output.total_tokens += step.response.total_tokens
+
+
+        # Add step to agent output tracking
+        return self.step_pool.append(step_log)
 
     def _serialize_callable(
         self, attr_value: Callable