From 38db41b08cb36456cd7ce9a2171491748146c786 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 18:33:17 +0530
Subject: [PATCH 01/22] adding asdict

- added asdict to turn things into json compatible output
---
 swarms/structs/agent.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 30073fab..1471ff3e 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -57,6 +57,8 @@ from clusterops import (
     execute_with_cpu_cores,
 )
 from swarms.agents.ape_agent import auto_generate_prompt
+from dataclasses import asdict
+
 
 
 # Utils

From 0d2d7fd43911d81be4bab548d68687ad4b9fd250 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 18:38:27 +0530
Subject: [PATCH 02/22] add log methods to run

- adding updated log methods to run
---
 swarms/structs/agent.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 1471ff3e..39bf238b 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -783,6 +783,8 @@ class Agent:
                 or loop_count < self.max_loops
             ):
                 loop_count += 1
+                # Log step start
+                current_step_id = f"step_{loop_count}_{uuid.uuid4().hex}"
                 self.loop_count_print(loop_count, self.max_loops)
                 print("\n")
 
@@ -836,6 +838,9 @@ class Agent:
                         
                         # Convert to a str if the response is not a str
                         response = self.llm_output_parser(response)
+
+                        # Log step metadata
+                        step_meta = self.log_step_metadata(loop_count, task_prompt, response)
                         
                         # Print
                         if self.streaming_on is True:

From e8b8d48f24b5af6b9cc39d915c51a20d74bea1aa Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 18:41:23 +0530
Subject: [PATCH 03/22] add section for tools update

- make the run update for tools
---
 swarms/structs/agent.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 39bf238b..33f63e8b 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -826,7 +826,19 @@ class Agent:
                         # Check and execute tools
                         if self.tools is not None:
                             print(f"self.tools is not None: {response}")
+                            tool_result = self.parse_and_execute_tools(response)
                             self.parse_and_execute_tools(response)
+                            if tool_result:
+                                self.update_tool_usage(
+                                    step_meta["step_id"],
+                                    tool_result["tool"],
+                                    tool_result["args"],
+                                    tool_result["response"]
+                                )
+
+
+                        # Update agent output history
+                        self.agent_output.full_history = self.short_memory.return_history_as_string()
                         
                         # Log the step metadata
                         logged = self.log_step_metadata(

From 8b552badba3a6ab6f82671c7ded8ddb236659597 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 18:43:38 +0530
Subject: [PATCH 04/22] edit run return

- edit "run" return type with condition on how to return if output type is JSON .
---
 swarms/structs/agent.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 33f63e8b..c70fe201 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -998,7 +998,10 @@ class Agent:
 
             # return self.agent_output_type(all_responses)
 
-            return concat_strings(all_responses)
+            if self.output_type == "json":
+                return asdict(self.agent_output)
+            else:
+                return concat_strings(all_responses)
 
         except Exception as error:
             logger.info(

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 05/22] 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

From 2a8cf252a7be81784ec647523a3e4353012cafb2 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 18:53:50 +0530
Subject: [PATCH 06/22] Add Update_tool_usage

- update tool usage fn to add and account for tool usage metadata.
---
 swarms/structs/agent.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 6c8e5fea..4d800673 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -1948,6 +1948,17 @@ class Agent:
         # Add step to agent output tracking
         return self.step_pool.append(step_log)
 
+    def update_tool_usage(self, step_id: str, tool_name: str, tool_args: dict, tool_response: Any):
+        """Update tool usage information for a specific step."""
+        for step in self.agent_output.steps:
+            if step.step_id == step_id:
+                step.response.tool_calls.append({
+                    "tool": tool_name,
+                    "arguments": tool_args,
+                    "response": str(tool_response)
+                })
+                break
+                
     def _serialize_callable(
         self, attr_value: Callable
     ) -> Dict[str, Any]:

From b40c76ee5dcbeb657c9fd6a430e71ad74c7d9e6d Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 18:56:31 +0530
Subject: [PATCH 07/22] minor refactor

- minor refactor
---
 swarms/structs/agent.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 4d800673..8426905a 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -61,6 +61,8 @@ from dataclasses import asdict
 
 
 
+
+
 # Utils
 # Custom stopping condition
 def stop_when_repeats(response: str) -> bool:
@@ -812,6 +814,9 @@ class Agent:
                         )
                         response = self.call_llm(*response_args, **kwargs)
                         
+                        # Log step metadata
+                        step_meta = self.log_step_metadata(loop_count, task_prompt, response)
+
                         # Check if response is a dictionary and has 'choices' key
                         if isinstance(response, dict) and 'choices' in response:
                             response = response['choices'][0]['message']['content']
@@ -825,9 +830,7 @@ class Agent:
                         
                         # Check and execute tools
                         if self.tools is not None:
-                            print(f"self.tools is not None: {response}")
                             tool_result = self.parse_and_execute_tools(response)
-                            self.parse_and_execute_tools(response)
                             if tool_result:
                                 self.update_tool_usage(
                                     step_meta["step_id"],
@@ -836,7 +839,7 @@ class Agent:
                                     tool_result["response"]
                                 )
 
-
+                           
                         # Update agent output history
                         self.agent_output.full_history = self.short_memory.return_history_as_string()
                         
@@ -850,9 +853,6 @@ class Agent:
                         
                         # Convert to a str if the response is not a str
                         response = self.llm_output_parser(response)
-
-                        # Log step metadata
-                        step_meta = self.log_step_metadata(loop_count, task_prompt, response)
                         
                         # Print
                         if self.streaming_on is True:
@@ -1003,6 +1003,7 @@ class Agent:
             else:
                 return concat_strings(all_responses)
 
+
         except Exception as error:
             logger.info(
                 f"Error running agent: {error} optimize your input parameters"
@@ -1948,6 +1949,7 @@ class Agent:
         # Add step to agent output tracking
         return self.step_pool.append(step_log)
 
+
     def update_tool_usage(self, step_id: str, tool_name: str, tool_args: dict, tool_response: Any):
         """Update tool usage information for a specific step."""
         for step in self.agent_output.steps:

From 92912b0397a2c18281a2bbbfb0b2414ef5a3b9a0 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:02:12 +0530
Subject: [PATCH 08/22] Push libraries

- update correct libraries & dependencies .
---
 tests/agents/test_agent_logging.py | 7 +++++++
 1 file changed, 7 insertions(+)
 create mode 100644 tests/agents/test_agent_logging.py

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
new file mode 100644
index 00000000..e7fd464b
--- /dev/null
+++ b/tests/agents/test_agent_logging.py
@@ -0,0 +1,7 @@
+from unittest.mock import Mock, MagicMock
+from dataclasses import dataclass, field, asdict
+from typing import List, Dict, Any
+from datetime import datetime
+import unittest
+from swarms.schemas.agent_step_schemas import ManySteps, Step
+from swarms.structs.agent import Agent

From ef58dd4582ede4d6192a6a62f5513b2f99090f01 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:04:32 +0530
Subject: [PATCH 09/22] Create TestAgentLogging

- create class to execute modular unittests
- def setup for modular setup
- objective to keep setup minimal , so that tests aren't bloated and fast to run
- Since most param have a set default , init of necessary condition is a valid and supportive op wrt test speed .
---
 tests/agents/test_agent_logging.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index e7fd464b..13675858 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -5,3 +5,20 @@ from datetime import datetime
 import unittest
 from swarms.schemas.agent_step_schemas import ManySteps, Step
 from swarms.structs.agent import Agent
+
+class TestAgentLogging(unittest.TestCase):
+    def setUp(self):
+        self.mock_tokenizer = MagicMock()
+        self.mock_tokenizer.count_tokens.return_value = 100
+        
+        self.mock_short_memory = MagicMock()
+        self.mock_short_memory.get_memory_stats.return_value = {"message_count": 2}
+        
+        self.mock_long_memory = MagicMock()
+        self.mock_long_memory.get_memory_stats.return_value = {"item_count": 5}
+        
+        self.agent = Agent(
+            tokenizer=self.mock_tokenizer,
+            short_memory=self.mock_short_memory,
+            long_term_memory=self.mock_long_memory
+        )

From 4d300fccf94c2e65f88945bcd2b4c325eb4f4f83 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:08:25 +0530
Subject: [PATCH 10/22] Add basic test for log_step_metadata method

- Implemented test_log_step_metadata_basic to verify the correct logging of step metadata including step_id, timestamp, tokens, and memory usage.
- Confirmed that the token counts for total are accurately logged.
---
 tests/agents/test_agent_logging.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index 13675858..c68e9aa0 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -22,3 +22,15 @@ class TestAgentLogging(unittest.TestCase):
             short_memory=self.mock_short_memory,
             long_term_memory=self.mock_long_memory
         )
+
+    def test_log_step_metadata_basic(self):
+        log_result = self.agent.log_step_metadata(1, "Test prompt", "Test response")
+        
+        self.assertIn('step_id', log_result)
+        self.assertIn('timestamp', log_result)
+        self.assertIn('tokens', log_result)
+        self.assertIn('memory_usage', log_result)
+        
+        self.assertEqual(log_result['tokens']['total'], 200)
+
+    

From 672f3eeab449564f3986a80f971eb5b38b984ba1 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:09:35 +0530
Subject: [PATCH 11/22] Add test for log_step_metadata with no long-term memory

- Implemented test_log_step_metadata_no_long_term_memory to ensure that when long-term memory is None, the memory_usage for long_term is an empty dictionary in the log result.
---
 tests/agents/test_agent_logging.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index c68e9aa0..cbfb8b80 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -33,4 +33,8 @@ class TestAgentLogging(unittest.TestCase):
         
         self.assertEqual(log_result['tokens']['total'], 200)
 
+    def test_log_step_metadata_no_long_term_memory(self):
+        self.agent.long_term_memory = None
+        log_result = self.agent.log_step_metadata(1, "prompt", "response")
+        self.assertEqual(log_result['memory_usage']['long_term'], {})
     

From e80d1820927810708d1575d782f80fb108db2b55 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:10:37 +0530
Subject: [PATCH 12/22] Add test for timestamp in log_step_metadata

- Implemented test_log_step_metadata_timestamp to verify that a timestamp is included in the log result when logging step metadata.
---
 tests/agents/test_agent_logging.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index cbfb8b80..33e59843 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -38,3 +38,6 @@ class TestAgentLogging(unittest.TestCase):
         log_result = self.agent.log_step_metadata(1, "prompt", "response")
         self.assertEqual(log_result['memory_usage']['long_term'], {})
     
+    def test_log_step_metadata_timestamp(self):
+        log_result = self.agent.log_step_metadata(1, "prompt", "response")
+        self.assertIn('timestamp', log_result)

From 8900d645c8e03bfd0c585e411c5e462051dddba4 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:12:14 +0530
Subject: [PATCH 13/22] Add integration test for token counting in
 log_step_metadata

- Implemented test_token_counting_integration to verify the correct total token count when using a mocked tokenizer, ensuring that prompt and response token counts are accurately aggregated.
---
 tests/agents/test_agent_logging.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index 33e59843..98b0f8a5 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -41,3 +41,9 @@ class TestAgentLogging(unittest.TestCase):
     def test_log_step_metadata_timestamp(self):
         log_result = self.agent.log_step_metadata(1, "prompt", "response")
         self.assertIn('timestamp', log_result)
+
+    def test_token_counting_integration(self):
+        self.mock_tokenizer.count_tokens.side_effect = [150, 250]
+        log_result = self.agent.log_step_metadata(1, "prompt", "response")
+        
+        self.assertEqual(log_result['tokens']['total'], 400)

From d72d2e74ff047daa6bf359b952f3414655d3affe Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:14:23 +0530
Subject: [PATCH 14/22] Add test for agent output updating in log_step_metadata

- Implemented test_agent_output_updating to verify that logging step metadata correctly updates the total token count and ensures that the agent's output steps are properly tracked confirming only one step is recorded.
---
 tests/agents/test_agent_logging.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index 98b0f8a5..2fe13ea9 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -47,3 +47,14 @@ class TestAgentLogging(unittest.TestCase):
         log_result = self.agent.log_step_metadata(1, "prompt", "response")
         
         self.assertEqual(log_result['tokens']['total'], 400)
+
+    def test_agent_output_updating(self):
+            initial_total_tokens = sum(step['tokens']['total'] for step in self.agent.agent_output.steps)
+            self.agent.log_step_metadata(1, "prompt", "response")
+            
+            final_total_tokens = sum(step['tokens']['total'] for step in self.agent.agent_output.steps)
+            self.assertEqual(
+                final_total_tokens - initial_total_tokens,
+                200
+            )
+            self.assertEqual(len(self.agent.agent_output.steps), 1)

From 67a8493340f8db55c966ea0fe7310942f3bae82b Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:16:27 +0530
Subject: [PATCH 15/22] Integration test

- apart of modular and small unittests , its time to put to test for larger and more wholesome integration test
- an end to end test , to check if this change in code causes any breaks in flow and normal functioning .
---
 tests/agents/test_agent_logging.py | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index 2fe13ea9..dcecb1bb 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -58,3 +58,29 @@ class TestAgentLogging(unittest.TestCase):
                 200
             )
             self.assertEqual(len(self.agent.agent_output.steps), 1)
+
+class TestAgentLoggingIntegration(unittest.TestCase):
+    def setUp(self):
+        self.agent = Agent(agent_name="test-agent")
+
+    def test_full_logging_cycle(self):
+        task = "Test task"
+        max_loops = 1
+        
+        result = self.agent._run(task, max_loops=max_loops)
+        
+        self.assertIsInstance(result, dict)
+        self.assertIn('steps', result)
+        self.assertIsInstance(result['steps'], list)
+        self.assertEqual(len(result['steps']), max_loops)
+        
+        if result['steps']:
+            step = result['steps'][0]
+            self.assertIn('step_id', step)
+            self.assertIn('timestamp', step)
+            self.assertIn('task', step)
+            self.assertIn('response', step)
+            self.assertEqual(step['task'], task)
+            self.assertEqual(step['response'], f"Response for loop 1")
+        
+        self.assertTrue(len(self.agent.agent_output.steps) > 0)

From a3ef4508a040a953df9935ba9836361ae5bfe383 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:17:07 +0530
Subject: [PATCH 16/22] adding main

- add main to make this run
---
 tests/agents/test_agent_logging.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index dcecb1bb..884216e5 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -84,3 +84,6 @@ class TestAgentLoggingIntegration(unittest.TestCase):
             self.assertEqual(step['response'], f"Response for loop 1")
         
         self.assertTrue(len(self.agent.agent_output.steps) > 0)
+
+if __name__ == '__main__':
+    unittest.main()

From ad38711e2344ffaaa08692a737ce8ae93ea196d0 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 19:21:57 +0530
Subject: [PATCH 17/22] Update Docs

- Update the docs to return describe changes
- Add lines specific to how to toggle between output types .
---
 docs/swarms/agents/third_party.md | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/docs/swarms/agents/third_party.md b/docs/swarms/agents/third_party.md
index 3ab8700f..c8f9a496 100644
--- a/docs/swarms/agents/third_party.md
+++ b/docs/swarms/agents/third_party.md
@@ -509,6 +509,13 @@ logging.basicConfig(level=logging.INFO)
 monitored_agent = MonitoredAgent("MonitoredGriptapeAgent")
 result = monitored_agent.run("Summarize the latest AI research papers")
 ```
+Additionally the Agent class now includes built-in logging functionality and the ability to switch between JSON and string output. 
+
+To switch between JSON and string output:
+- Use `output_type="str"` for string output (default)
+- Use `output_type="json"` for JSON output
+
+The `output_type` parameter determines the format of the final result returned by the `run` method. When set to "str", it returns a string representation of the agent's response. When set to "json", it returns a JSON object containing detailed information about the agent's run, including all steps and metadata.
 
 ## 6. Best Practices for Custom Agent Development
 
@@ -614,4 +621,4 @@ The ability to seamlessly integrate agents from libraries like Griptape, Langcha
 
 As you embark on your journey with the swarms framework, remember that the field of AI and agent-based systems is rapidly evolving. Stay curious, keep experimenting, and don't hesitate to push the boundaries of what's possible with custom agents and integrated libraries.
 
-By embracing the power of the swarms framework and the ecosystem of agent libraries it supports, you're well-positioned to create the next generation of intelligent, adaptive, and collaborative AI systems. Happy agent building!
\ No newline at end of file
+By embracing the power of the swarms framework and the ecosystem of agent libraries it supports, you're well-positioned to create the next generation of intelligent, adaptive, and collaborative AI systems. Happy agent building!

From e100b468bcf34a17ff9bdb0029db8afe4d8d7885 Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Sun, 27 Oct 2024 23:22:00 +0530
Subject: [PATCH 18/22] slight refactor

---
 swarms/structs/agent.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 8426905a..1a4d8769 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -1947,7 +1947,7 @@ class Agent:
 
 
         # Add step to agent output tracking
-        return self.step_pool.append(step_log)
+        self.step_pool.append(step_log)
 
 
     def update_tool_usage(self, step_id: str, tool_name: str, tool_args: dict, tool_response: Any):

From a63541c61953446a7868138b4ee9d477ea30262e Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Mon, 28 Oct 2024 22:42:17 +0530
Subject: [PATCH 19/22] Reverting to puled agent.py with new changes

---
 swarms/structs/agent.py | 221 +++++++++++++++++++++++-----------------
 1 file changed, 129 insertions(+), 92 deletions(-)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 1a4d8769..2d07f106 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -57,11 +57,7 @@ from clusterops import (
     execute_with_cpu_cores,
 )
 from swarms.agents.ape_agent import auto_generate_prompt
-from dataclasses import asdict
-
-
-
-
+import yaml
 
 # Utils
 # Custom stopping condition
@@ -785,8 +781,6 @@ class Agent:
                 or loop_count < self.max_loops
             ):
                 loop_count += 1
-                # Log step start
-                current_step_id = f"step_{loop_count}_{uuid.uuid4().hex}"
                 self.loop_count_print(loop_count, self.max_loops)
                 print("\n")
 
@@ -805,21 +799,29 @@ class Agent:
                 while attempt < self.retry_attempts and not success:
                     try:
                         if self.long_term_memory is not None:
-                            logger.info("Querying long term memory...")
+                            logger.info(
+                                "Querying long term memory..."
+                            )
                             self.memory_query(task_prompt)
-                        
+
                         # Generate response using LLM
                         response_args = (
-                            (task_prompt, *args) if img is None else (task_prompt, img, *args)
+                            (task_prompt, *args)
+                            if img is None
+                            else (task_prompt, img, *args)
+                        )
+                        response = self.call_llm(
+                            *response_args, **kwargs
                         )
-                        response = self.call_llm(*response_args, **kwargs)
-                        
-                        # Log step metadata
-                        step_meta = self.log_step_metadata(loop_count, task_prompt, response)
 
                         # Check if response is a dictionary and has 'choices' key
-                        if isinstance(response, dict) and 'choices' in response:
-                            response = response['choices'][0]['message']['content']
+                        if (
+                            isinstance(response, dict)
+                            and "choices" in response
+                        ):
+                            response = response["choices"][0][
+                                "message"
+                            ]["content"]
                         elif isinstance(response, str):
                             # If response is already a string, use it as is
                             pass
@@ -827,48 +829,37 @@ class Agent:
                             raise ValueError(
                                 f"Unexpected response format: {type(response)}"
                             )
-                        
+
                         # Check and execute tools
                         if self.tools is not None:
-                            tool_result = self.parse_and_execute_tools(response)
-                            if tool_result:
-                                self.update_tool_usage(
-                                    step_meta["step_id"],
-                                    tool_result["tool"],
-                                    tool_result["args"],
-                                    tool_result["response"]
-                                )
-
-                           
-                        # Update agent output history
-                        self.agent_output.full_history = self.short_memory.return_history_as_string()
-                        
+                            print(
+                                f"self.tools is not None: {response}"
+                            )
+                            self.parse_and_execute_tools(response)
+
                         # Log the step metadata
                         logged = self.log_step_metadata(
-                            loop_count, 
-                            task_prompt, 
-                            response
+                            loop_count, task_prompt, response
                         )
                         logger.info(logged)
-                        
+
                         # Convert to a str if the response is not a str
                         response = self.llm_output_parser(response)
-                        
+
                         # Print
                         if self.streaming_on is True:
                             self.stream_response(response)
                         else:
                             print(response)
-                        
+
                         # Add the response to the memory
                         self.short_memory.add(
-                            role=self.agent_name,
-                            content=response
+                            role=self.agent_name, content=response
                         )
-                        
+
                         # Add to all responses
                         all_responses.append(response)
-                        
+
                         # TODO: Implement reliability check
                         if self.tools is not None:
                             # self.parse_function_call_and_execute(response)
@@ -997,12 +988,25 @@ class Agent:
             ]
 
             # return self.agent_output_type(all_responses)
-
-            if self.output_type == "json":
-                return asdict(self.agent_output)
-            else:
+            # More flexible output types
+            if self.output_type == "string":
                 return concat_strings(all_responses)
-
+            elif self.output_type == "list":
+                return all_responses
+            elif self.output_type == "json":
+                return self.agent_output.model_dump_json(indent=4)
+            elif self.output_type == "csv":
+                return self.dict_to_csv(
+                    self.agent_output.model_dump()
+                )
+            elif self.output_type == "dict":
+                return self.agent_output.model_dump()
+            elif self.output_type == "yaml":
+                return yaml.safe_dump(self.agent_output.model_dump(), sort_keys=False)
+            else:
+                raise ValueError(
+                    f"Invalid output type: {self.output_type}"
+                )
 
         except Exception as error:
             logger.info(
@@ -1011,20 +1015,79 @@ class Agent:
             raise error
 
     def __call__(
-        self, task: str = None, img: str = None, *args, **kwargs
-    ):
+        self,
+        task: Optional[str] = None,
+        img: Optional[str] = None,
+        is_last: bool = False,
+        device: str = "cpu",  # gpu
+        device_id: int = 0,
+        all_cores: bool = True,
+        *args,
+        **kwargs,
+    ) -> Any:
         """Call the agent
 
         Args:
-            task (str): _description_
-            img (str, optional): _description_. Defaults to None.
+            task (Optional[str]): The task to be performed. Defaults to None.
+            img (Optional[str]): The image to be processed. Defaults to None.
+            is_last (bool): Indicates if this is the last task. Defaults to False.
+            device (str): The device to use for execution. Defaults to "cpu".
+            device_id (int): The ID of the GPU to use if device is set to "gpu". Defaults to 0.
+            all_cores (bool): If True, uses all available CPU cores. Defaults to True.
         """
         try:
-            return self.run(task, img, *args, **kwargs)
+            if task is not None:
+                return self.run(
+                    task=task,
+                    is_last=is_last,
+                    device=device,
+                    device_id=device_id,
+                    all_cores=all_cores,
+                    *args,
+                    **kwargs,
+                )
+            elif img is not None:
+                return self.run(
+                    img=img,
+                    is_last=is_last,
+                    device=device,
+                    device_id=device_id,
+                    all_cores=all_cores,
+                    *args,
+                    **kwargs,
+                )
+            else:
+                raise ValueError(
+                    "Either 'task' or 'img' must be provided."
+                )
         except Exception as error:
             logger.error(f"Error calling agent: {error}")
             raise error
 
+    def dict_to_csv(self, data: dict) -> str:
+        """
+        Convert a dictionary to a CSV string.
+
+        Args:
+            data (dict): The dictionary to convert.
+
+        Returns:
+            str: The CSV string representation of the dictionary.
+        """
+        import csv
+        import io
+
+        output = io.StringIO()
+        writer = csv.writer(output)
+
+        # Write header
+        writer.writerow(data.keys())
+
+        # Write values
+        writer.writerow(data.values())
+
+        return output.getvalue()
+
     def parse_and_execute_tools(self, response: str, *args, **kwargs):
         # Extract json from markdown
         # response = extract_code_from_markdown(response)
@@ -1883,42 +1946,36 @@ class Agent:
         """Parse the output from the LLM"""
         try:
             if isinstance(response, dict):
-                if 'choices' in response:
-                    return response['choices'][0]['message']['content']
+                if "choices" in response:
+                    return response["choices"][0]["message"][
+                        "content"
+                    ]
                 else:
-                    return json.dumps(response)  # Convert dict to string
+                    return json.dumps(
+                        response
+                    )  # Convert dict to string
             elif isinstance(response, str):
                 return response
             else:
-                return str(response)  # Convert any other type to string
+                return str(
+                    response
+                )  # Convert any other type to string
         except Exception as e:
             logger.error(f"Error parsing LLM output: {e}")
-            return str(response)  # Return string representation as fallback
+            return str(
+                response
+            )  # Return string representation as fallback
 
     def log_step_metadata(
         self, loop: int, task: str, response: str
     ) -> Step:
-        """Log metadata for each step of agent execution."""
-        # Generate unique step ID
-        step_id = f"step_{loop}_{uuid.uuid4().hex}"
-        
-        # Calculate token usage
+        # # # Step Metadata
         # full_memory = self.short_memory.return_history_as_string()
         # prompt_tokens = self.tokenizer.count_tokens(full_memory)
         # completion_tokens = self.tokenizer.count_tokens(response)
-        # 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
-        }
+        # self.tokenizer.count_tokens(prompt_tokens + completion_tokens)
 
         step_log = Step(
-            step_id=step_id,
-            time=time.time(),
-            tokens = total_tokens,
             response=AgentChatCompletionResponse(
                 id=self.agent_id,
                 agent_name=self.agent_name,
@@ -1933,34 +1990,14 @@ class Agent:
                 ),
                 # usage=UsageInfo(
                 #     prompt_tokens=prompt_tokens,
-                #     completion_tokens=completion_tokens,
                 #     total_tokens=total_tokens,
+                #     completion_tokens=completion_tokens,
                 # ),
-                tool_calls=[],
-                memory_usage=memory_usage
             ),
         )
 
-        # 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
         self.step_pool.append(step_log)
 
-
-    def update_tool_usage(self, step_id: str, tool_name: str, tool_args: dict, tool_response: Any):
-        """Update tool usage information for a specific step."""
-        for step in self.agent_output.steps:
-            if step.step_id == step_id:
-                step.response.tool_calls.append({
-                    "tool": tool_name,
-                    "arguments": tool_args,
-                    "response": str(tool_response)
-                })
-                break
-                
     def _serialize_callable(
         self, attr_value: Callable
     ) -> Dict[str, Any]:

From 062176fd8adfcffc837dc13cae869d144da8528a Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Mon, 28 Oct 2024 23:04:44 +0530
Subject: [PATCH 20/22] Update agent.py

---
 swarms/structs/agent.py | 63 +++++++++++++++++++++++++++++++++++------
 1 file changed, 54 insertions(+), 9 deletions(-)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 2d07f106..737e7a65 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -781,6 +781,8 @@ class Agent:
                 or loop_count < self.max_loops
             ):
                 loop_count += 1
+                # Log step start
+                current_step_id = f"step_{loop_count}_{uuid.uuid4().hex}"
                 self.loop_count_print(loop_count, self.max_loops)
                 print("\n")
 
@@ -814,6 +816,8 @@ class Agent:
                             *response_args, **kwargs
                         )
 
+                        # Log step metadata
+                        step_meta = self.log_step_metadata(loop_count, task_prompt, response)
                         # Check if response is a dictionary and has 'choices' key
                         if (
                             isinstance(response, dict)
@@ -832,10 +836,18 @@ class Agent:
 
                         # Check and execute tools
                         if self.tools is not None:
-                            print(
-                                f"self.tools is not None: {response}"
-                            )
-                            self.parse_and_execute_tools(response)
+                            tool_result = self.parse_and_execute_tools(response)
+                            if tool_result:
+                                self.update_tool_usage(
+                                    step_meta["step_id"],
+                                    tool_result["tool"],
+                                    tool_result["args"],
+                                    tool_result["response"]
+                                )
+ 
+ 
+                        # Update agent output history
+                        self.agent_output.full_history = self.short_memory.return_history_as_string()
 
                         # Log the step metadata
                         logged = self.log_step_metadata(
@@ -1969,13 +1981,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,
@@ -1990,14 +2016,33 @@ 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
             ),
         )
-
+ 
+        # 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
         self.step_pool.append(step_log)
 
+    def update_tool_usage(self, step_id: str, tool_name: str, tool_args: dict, tool_response: Any):
+        """Update tool usage information for a specific step."""
+        for step in self.agent_output.steps:
+            if step.step_id == step_id:
+                step.response.tool_calls.append({
+                    "tool": tool_name,
+                    "arguments": tool_args,
+                    "response": str(tool_response)
+                })
+                break
+
     def _serialize_callable(
         self, attr_value: Callable
     ) -> Dict[str, Any]:

From fe39cd5393cbfe802bced41b86959b5b17d8943f Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Tue, 29 Oct 2024 01:22:21 +0530
Subject: [PATCH 21/22] added tools and memory responses to the json output.

---
 swarms/structs/agent.py | 31 ++++++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index 737e7a65..e9ddf8b0 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -1991,11 +1991,32 @@ class Agent:
         # completion_tokens = self.tokenizer.count_tokens(response)
         # total_tokens = prompt_tokens + completion_tokens
         total_tokens=self.tokenizer.count_tokens(task) + self.tokenizer.count_tokens(response),
- 
+
+        # Get memory responses
+        memory_responses = {
+            "short_term": self.short_memory.return_history_as_string() if self.short_memory else None,
+            "long_term": self.long_term_memory.query(task) if self.long_term_memory else None
+        }
+
+        # Get tool responses if tool was used
+        tool_response = None
+        if self.tools:
+            try:
+                tool_call_output = parse_and_execute_json(self.tools, response, parse_md=True)
+                if tool_call_output:
+                    tool_response = {
+                        "tool_name": tool_call_output.get("tool_name", "unknown"),
+                        "tool_args": tool_call_output.get("args", {}),
+                        "tool_output": str(tool_call_output.get("output", ""))
+                    }
+            except Exception as e:
+                logger.debug(f"No tool call detected in response: {e}")
+                
         # 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
+            "short_term": len(self.short_memory.messages) if self.short_memory else 0,
+            "long_term": self.long_term_memory.count if self.long_term_memory else 0,
+            "responses": memory_responses
         }
  
         step_log = Step(
@@ -2019,14 +2040,14 @@ class Agent:
                 #     completion_tokens=completion_tokens,
                 #     total_tokens=total_tokens,
                 # ),
-                tool_calls=[],
+                tool_calls=[] if tool_response is None else [tool_response],
                 memory_usage=memory_usage
             ),
         )
  
         # Update total tokens if agent_output exists
         if hasattr(self, 'agent_output'):
-            self.agent_output.total_tokens += step.response.total_tokens
+            self.agent_output.total_tokens += self.response.total_tokens
  
  
         # Add step to agent output tracking

From 750f9a1e14fb4762fb7c7cb07768d19d44191ecb Mon Sep 17 00:00:00 2001
From: Sambhav Dixit <94298612+sambhavnoobcoder@users.noreply.github.com>
Date: Tue, 29 Oct 2024 01:22:57 +0530
Subject: [PATCH 22/22] Updated test , especially for tool gen mocking

---
 tests/agents/test_agent_logging.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tests/agents/test_agent_logging.py b/tests/agents/test_agent_logging.py
index 884216e5..b2106dd6 100644
--- a/tests/agents/test_agent_logging.py
+++ b/tests/agents/test_agent_logging.py
@@ -5,6 +5,15 @@ from datetime import datetime
 import unittest
 from swarms.schemas.agent_step_schemas import ManySteps, Step
 from swarms.structs.agent import Agent
+from swarms.tools.tool_parse_exec import parse_and_execute_json
+
+# Mock parse_and_execute_json for testing
+parse_and_execute_json = MagicMock()
+parse_and_execute_json.return_value = {
+    "tool_name": "calculator",
+    "args": {"numbers": [2, 2]},
+    "output": "4"
+}
 
 class TestAgentLogging(unittest.TestCase):
     def setUp(self):