You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
41 lines
1.5 KiB
41 lines
1.5 KiB
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
|
|
|
|
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
|
|
)
|
|
|
|
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)
|
|
|
|
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'], {})
|
|
|