Upload 7 files

tests/base_test.py

@@ -0,0 +1,99 @@
+
+
+class BaseTest:
+
+    # Custom assert methods implementation
+    def assertIsNotNone(self, value, msg=None):
+        """Assert that value is not None"""
+        if value is None:
+            raise AssertionError(msg or f"Expected not None, but got None")
+    
+    def assertEqual(self, first, second, msg=None):
+        """Assert that first equals second"""
+        if first != second:
+            raise AssertionError(msg or f"Expected {first} == {second}, but {first} != {second}")
+    
+    def assertTrue(self, expr, msg=None):
+        """Assert that expr is True"""
+        if not expr:
+            raise AssertionError(msg or f"Expected True, but got {expr}")
+    
+    def assertFalse(self, expr, msg=None):
+        """Assert that expr is False"""
+        if expr:
+            raise AssertionError(msg or f"Expected False, but got {expr}")
+    
+    def assertAlmostEqual(self, first, second, places=7, msg=None):
+        """Assert that first and second are approximately equal"""
+        if round(abs(second - first), places) != 0:
+            raise AssertionError(msg or f"Expected {first} ~= {second} (within {places} decimal places)")
+    
+    def assertIs(self, first, second, msg=None):
+        """Assert that first is second (same object identity)"""
+        if first is not second:
+            raise AssertionError(msg or f"Expected {first} is {second}, but they are different objects")
+    
+    def assertIn(self, member, container, msg=None):
+        """Assert that member is in container"""
+        if member not in container:
+            raise AssertionError(msg or f"Expected {member} in {container}")
+    
+    def assertIsInstance(self, obj, cls, msg=None):
+        """Assert that obj is an instance of cls"""
+        if not isinstance(obj, cls):
+            raise AssertionError(msg or f"Expected {obj} to be instance of {cls}, but got {type(obj)}")
+    
+    def assertIsNone(self, value, msg=None):
+        """Assert that value is None"""
+        if value is not None:
+            raise AssertionError(msg or f"Expected None, but got {value}")
+    
+    def assertNotEqual(self, first, second, msg=None):
+        """Assert that first does not equal second"""
+        if first == second:
+            raise AssertionError(msg or f"Expected {first} != {second}, but they are equal")
+    
+    def assertGreater(self, first, second, msg=None):
+        """Assert that first is greater than second"""
+        if not first > second:
+            raise AssertionError(msg or f"Expected {first} > {second}")
+    
+    def assertLess(self, first, second, msg=None):
+        """Assert that first is less than second"""
+        if not first < second:
+            raise AssertionError(msg or f"Expected {first} < {second}")
+    
+    def assertGreaterEqual(self, first, second, msg=None):
+        """Assert that first is greater than or equal to second"""
+        if not first >= second:
+            raise AssertionError(msg or f"Expected {first} >= {second}")
+    
+    def assertLessEqual(self, first, second, msg=None):
+        """Assert that first is less than or equal to second"""
+        if not first <= second:
+            raise AssertionError(msg or f"Expected {first} <= {second}")
+    
+    def assertNotIn(self, member, container, msg=None):
+        """Assert that member is not in container"""
+        if member in container:
+            raise AssertionError(msg or f"Expected {member} not in {container}")
+    
+    def assertIsNot(self, first, second, msg=None):
+        """Assert that first is not second (different object identity)"""
+        if first is second:
+            raise AssertionError(msg or f"Expected {first} is not {second}, but they are the same object")
+    
+    def assertRaises(self, exception_class, callable_obj=None, *args, **kwargs):
+        """Assert that calling callable_obj raises exception_class"""
+        if callable_obj is None:
+            # Return a context manager for use with 'with' statement
+            return self._AssertRaisesContext(exception_class)
+        else:
+            try:
+                callable_obj(*args, **kwargs)
+                raise AssertionError(f"Expected {exception_class.__name__} to be raised, but no exception was raised")
+            except exception_class:
+                pass  # Expected exception was raised
+            except Exception as e:
+                raise AssertionError(f"Expected {exception_class.__name__} to be raised, but got {type(e).__name__}: {e}")
+    

tests/test.json

@@ -0,0 +1 @@
+{"test": "test content"}

tests/test_context_compressor.py

@@ -0,0 +1,133 @@
+import unittest
+import sys
+import os
+from unittest.mock import Mock, patch, MagicMock
+from pathlib import Path
+
+# Add the project root to Python path
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from tests.base_test import BaseTest
+
+from aworld.config.conf import AgentConfig, ModelConfig, ContextRuleConfig, OptimizationConfig, LlmCompressionConfig
+from aworld.core.context.processor import CompressionResult, CompressionType
+from aworld.core.context.processor.llm_compressor import LLMCompressor
+from aworld.core.context.processor.prompt_processor import PromptProcessor
+from aworld.core.context.base import AgentContext, ContextUsage
+
+
+class TestPromptCompressor(BaseTest):
+    """Test cases for PromptCompressor.compress_batch function"""
+
+    def __init__(self):
+        """Set up test fixtures"""
+        self.mock_model_name = "qwen/qwen3-1.7b"
+        self.mock_base_url = "http://localhost:1234/v1"
+        self.mock_api_key = "lm-studio"
+        self.mock_llm_config = ModelConfig(
+            llm_model_name=self.mock_model_name,
+            llm_base_url=self.mock_base_url,
+            llm_api_key=self.mock_api_key
+        )
+        os.environ["LLM_API_KEY"] = self.mock_api_key
+        os.environ["LLM_BASE_URL"] = self.mock_base_url
+        os.environ["LLM_MODEL_NAME"] = self.mock_model_name
+
+    def test_compress_batch_basic(self):
+        
+        compressor = LLMCompressor(
+            llm_config=self.mock_llm_config
+        )
+        
+        # Test data
+        contents = [
+            "[SYSTEM]You are a helpful assistant.\n[USER]This is the first long text content that needs compression. This is the first long text content that needs compression.",
+        ]
+        
+        # Execute compress_batch
+        results = compressor.compress_batch(contents)
+        
+        # Assertions
+        for result in results:
+            self.assertIsInstance(result, CompressionResult)
+            self.assertEqual(result.compression_type, CompressionType.LLM_BASED)
+            self.assertTrue('This is the first long text content that needs compression. This is the first long text content that needs compression.' not in result.compressed_content)
+
+    def test_compress_messages(self):
+        """Test compress_messages function from PromptProcessor"""
+        
+        # Create context rule with compression enabled
+        context_rule = ContextRuleConfig(
+            optimization_config=OptimizationConfig(
+                enabled=True,
+                max_token_budget_ratio=0.8
+            ),
+            llm_compression_config=LlmCompressionConfig(
+                enabled=True,
+                trigger_compress_token_length=10,  # Low threshold to trigger compression
+                compress_model=self.mock_llm_config
+            )
+        )
+        
+        # Create agent context
+        agent_context = AgentContext(
+            agent_id="test_agent",
+            agent_name="test_agent",
+            agent_desc="Test agent for compression",
+            system_prompt="You are a helpful assistant.",
+            agent_prompt="You are a helpful assistant.",
+            model_config=self.mock_llm_config,
+            context_rule=context_rule,
+            context_usage=ContextUsage(total_context_length=4096)
+        )
+        
+        # Create prompt processor
+        processor = PromptProcessor(agent_context)
+        
+        # Test messages with repeated content that needs compression
+        messages = [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant."
+            },
+            {
+                "role": "user", 
+                "content": "This is the first long text content that needs compression. This is the first long text content that needs compression."
+            },
+            {
+                "role": "assistant",
+                "content": "I understand you want me to help with compression."
+            }
+        ]
+        
+        # Execute compress_messages
+        compressed_messages = processor.compress_messages(messages)
+
+        # Assertions
+        self.assertIsInstance(compressed_messages, list)
+        self.assertEqual(len(compressed_messages), len(messages))
+        
+        # Find the user message and verify it was processed
+        user_message = None
+        for msg in compressed_messages:
+            if msg.get("role") == "user":
+                user_message = msg
+                break
+        
+        self.assertIsNotNone(user_message)
+        # The original repeated text should be compressed
+        original_content = "This is the first long text content that needs compression. This is the first long text content that needs compression."
+        self.assertNotEqual(user_message["content"], original_content)
+        # The compressed content should be shorter than original
+        self.assertLess(len(user_message["content"]), len(original_content))
+
+if __name__ == '__main__':
+    testPromptCompressor = TestPromptCompressor()
+    testPromptCompressor.test_compress_batch_basic()
+    testPromptCompressor = TestPromptCompressor()
+    testPromptCompressor.test_compress_messages()
+
+
+
+

tests/test_context_hook.py

@@ -0,0 +1,32 @@
+
+# Add the project root to Python path
+from pathlib import Path
+import sys
+
+from aworld.logs.util import color_log
+
+
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from aworld.core.task import Task
+from aworld.core.agent.base import AgentFactory
+from aworld.core.agent.swarm import Swarm
+from aworld.runner import Runners
+from aworld.agents.llm_agent import Agent
+from aworld.config.conf import AgentConfig, ContextRuleConfig, ModelConfig, OptimizationConfig, LlmCompressionConfig
+from aworld.core.context.base import Context
+from aworld.core.event.base import Message
+from aworld.runners.hook.hooks import PreLLMCallHook, PostLLMCallHook
+from aworld.runners.hook.hook_factory import HookFactory
+from aworld.utils.common import convert_to_snake
+from tests.base_test import BaseTest
+
+@HookFactory.register(name="CheckContextPreLLMHook", desc="Test pre-LLM hook")
+class CheckContextPreLLMHook(PreLLMCallHook):
+    def name(self):
+        return convert_to_snake("CheckContextPreLLMHook")
+    async def exec(self, message: Message, context: Context = None) -> Message:
+        assert context.state.get("task") == "What is an agent."
+        color_log("CheckContextPreLLMHook test ok", color="green")
+        return message

tests/test_context_management.py

@@ -0,0 +1,236 @@
+import os
+import random
+import sys
+from pathlib import Path
+
+# Add the project root to Python path
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from aworld.core.task import Task
+from aworld.core.agent.swarm import Swarm
+from aworld.runner import Runners
+from aworld.agents.llm_agent import Agent
+from aworld.config.conf import AgentConfig, ContextRuleConfig, ModelConfig, OptimizationConfig, LlmCompressionConfig
+from aworld.core.context.base import Context
+from aworld.runners.hook.hook_factory import HookFactory
+from tests.base_test import BaseTest
+
+
+class TestContextManagement(BaseTest):
+    """Test cases for Context Management system based on README examples"""
+
+    def __init__(self):
+        """Set up test fixtures"""
+        self.mock_model_name = "qwen/qwen3-1.7b"
+        self.mock_base_url = "http://localhost:1234/v1"
+        self.mock_api_key = "lm-studio"
+        os.environ["LLM_API_KEY"] = self.mock_api_key
+        os.environ["LLM_BASE_URL"] = self.mock_base_url
+        os.environ["LLM_MODEL_NAME"] = self.mock_model_name
+
+    def init_agent(self, config_type: str = "1", context_rule: ContextRuleConfig = None):
+        if config_type == "1":
+            conf = AgentConfig(
+                llm_model_name=self.mock_model_name,
+                llm_base_url=self.mock_base_url,
+                llm_api_key=self.mock_api_key
+            )
+        else:
+            conf = AgentConfig(
+                llm_config=ModelConfig(
+                    llm_model_name=self.mock_model_name,
+                    llm_base_url=self.mock_base_url,
+                    llm_api_key=self.mock_api_key
+                )
+            )
+        return Agent(
+            conf=conf,
+            name="my_agent" + str(random.randint(0, 1000000)),
+            system_prompt="You are a helpful assistant.",
+            agent_prompt="You are a helpful assistant.",
+            context_rule=context_rule
+        )
+
+    class _AssertRaisesContext:
+        """Context manager for assertRaises"""
+
+        def __init__(self, expected_exception):
+            self.expected_exception = expected_exception
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            if exc_type is None:
+                raise AssertionError(f"Expected {self.expected_exception.__name__} to be raised, but no exception was raised")
+            if not issubclass(exc_type, self.expected_exception):
+                raise AssertionError(f"Expected {self.expected_exception.__name__} to be raised, but got {exc_type.__name__}: {exc_value}")
+            return True  # Suppress the exception
+
+    def fail(self, msg=None):
+        """Fail immediately with the given message"""
+        raise AssertionError(msg or "Test failed")
+
+    def run_agent(self, input, agent: Agent):
+        swarm = Swarm(agent, max_steps=1)
+        return Runners.sync_run(
+            input=input,
+            swarm=swarm
+        )
+
+    def run_multi_agent(self, input, agent1: Agent, agent2: Agent):
+        swarm = Swarm(agent1, agent2, max_steps=1)
+        return Runners.sync_run(
+            input=input,
+            swarm=swarm
+        )
+
+    def run_task(self, context: Context, agent: Agent):
+        swarm = Swarm(agent, max_steps=1)
+        task = Task(input="""What is an agent.""", swarm=swarm, context=context)
+        return Runners.sync_run_task(task)
+
+    def test_default_context_configuration(self):
+
+        # No need to explicitly configure context_rule, system automatically uses default configuration
+        # Default configuration is equivalent to:
+        # context_rule=ContextRuleConfig(
+        #     optimization_config=OptimizationConfig(
+        #         enabled=True,
+        #         max_token_budget_ratio=1.0  # Use 100% of context window
+        #     ),
+        #     llm_compression_config=LlmCompressionConfig(
+        #         enabled=False  # Compression disabled by default
+        #     )
+        # )
+        mock_agent = self.init_agent("1")
+        response = self.run_agent(input="""What is an agent. describe within 20 words""", agent=mock_agent)
+
+        self.assertIsNotNone(response.answer)
+        self.assertEqual(mock_agent.agent_context.model_config.llm_model_name, self.mock_model_name)
+
+        # Test default context rule behavior
+        self.assertIsNotNone(mock_agent.agent_context.context_rule)
+        self.assertIsNotNone(mock_agent.agent_context.context_rule.optimization_config)
+
+    def test_custom_context_configuration(self):
+        """Test custom context configuration (README Configuration example)"""
+        # Create custom context rules
+        mock_agent = self.init_agent(context_rule=ContextRuleConfig(
+            optimization_config=OptimizationConfig(
+                enabled=True,
+                max_token_budget_ratio=0.00015
+            ),
+            llm_compression_config=LlmCompressionConfig(
+                enabled=True,
+                trigger_compress_token_length=100,
+                compress_model=ModelConfig(
+                    llm_model_name=self.mock_model_name,
+                    llm_base_url=self.mock_base_url,
+                    llm_api_key=self.mock_api_key,
+                )
+            )
+        ))
+
+        response = self.run_agent(input="""describe What is an agent in details""", agent=mock_agent)
+        self.assertIsNotNone(response.answer)
+
+        # Test configuration values
+        self.assertTrue(mock_agent.agent_context.context_rule.optimization_config.enabled)
+        self.assertTrue(mock_agent.agent_context.context_rule.llm_compression_config.enabled)
+
+    def test_state_management_and_recovery(self):
+        class StateModifyAgent(Agent):
+            async def async_policy(self, observation, info=None, **kwargs):
+                result = await super().async_policy(observation, info, **kwargs)
+                self.context.state['policy_executed'] = True
+                return result
+
+        class StateTrackingAgent(Agent):
+            async def async_policy(self, observation, info=None, **kwargs):
+                result = await super().async_policy(observation, info, **kwargs)
+                assert self.context.state['policy_executed'] == True
+                return result
+
+        # Create custom agent instance
+        custom_agent = StateModifyAgent(
+            conf=AgentConfig(
+                llm_model_name=self.mock_model_name,
+                llm_base_url=self.mock_base_url,
+                llm_api_key=self.mock_api_key
+            ),
+            name="state_modify_agent",
+            system_prompt="You are a Python expert who provides detailed and practical answers.",
+            agent_prompt="You are a Python expert who provides detailed and practical answers.",
+        )
+
+        # Create a second agent for multi-agent testing
+        second_agent = StateTrackingAgent(
+            conf=AgentConfig(
+                llm_model_name=self.mock_model_name,
+                llm_base_url=self.mock_base_url,
+                llm_api_key=self.mock_api_key
+            ),
+            name="state_tracking_agent",
+            system_prompt="You are a helpful assistant.",
+            agent_prompt="You are a helpful assistant.",
+        )
+
+        response = self.run_multi_agent(
+            input="What is an agent. describe within 20 words",
+            agent1=custom_agent,
+            agent2=second_agent
+        )
+        self.assertIsNotNone(response.answer)
+
+        # Verify state changes after execution
+        self.assertTrue(custom_agent.context.state.get('policy_executed', True))
+        self.assertTrue(second_agent.agent_context.state.get('policy_executed', True))
+
+
+class TestHookSystem(TestContextManagement):
+    def __init__(self):
+        super().__init__()
+
+    def test_hook_registration(self):
+        from tests.test_llm_hook import TestPreLLMHook, TestPostLLMHook
+        """Test hook registration and retrieval"""
+        # Test that hooks are registered in _cls attribute
+        self.assertIn("TestPreLLMHook", HookFactory._cls)
+        self.assertIn("TestPostLLMHook", HookFactory._cls)
+
+        # Test hook creation using __call__ method
+        pre_hook = HookFactory("TestPreLLMHook")
+        post_hook = HookFactory("TestPostLLMHook")
+
+        self.assertIsInstance(pre_hook, TestPreLLMHook)
+        self.assertIsInstance(post_hook, TestPostLLMHook)
+
+    def test_hook_execution(self):
+        from tests.test_llm_hook import TestPreLLMHook, TestPostLLMHook
+
+        mock_agent = self.init_agent("1")
+        response = self.run_agent(input="""What is an agent. describe within 20 words""", agent=mock_agent)
+        self.assertIsNotNone(response.answer)
+
+    def test_task_context_transfer(self):
+        from tests.test_context_hook import CheckContextPreLLMHook
+
+        mock_agent = self.init_agent("1")
+        context = Context.instance()
+        context.state.update({"task": "What is an agent."})
+        self.run_task(context=context, agent=mock_agent)
+
+
+if __name__ == '__main__':
+    testContextManagement = TestContextManagement()
+    testContextManagement.test_default_context_configuration()
+    testContextManagement.test_custom_context_configuration()
+    testContextManagement.test_state_management_and_recovery()
+    testHookSystem = TestHookSystem()
+    testHookSystem.test_hook_registration()
+    testHookSystem = TestHookSystem()
+    testHookSystem.test_hook_execution()
+    testHookSystem = TestHookSystem()
+    testHookSystem.test_task_context_transfer()

tests/test_demo.py

@@ -0,0 +1,11 @@
+import unittest
+
+
+
+class TestShellTool(unittest.TestCase):
+
+
+    def test_init(self):
+        """Test initialization"""
+        print("Test initialization")
+

tests/test_llm_hook.py

@@ -0,0 +1,49 @@
+
+# Add the project root to Python path
+from pathlib import Path
+import sys
+
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from aworld.core.task import Task
+from aworld.core.agent.base import AgentFactory
+from aworld.core.agent.swarm import Swarm
+from aworld.runner import Runners
+from aworld.agents.llm_agent import Agent
+from aworld.config.conf import AgentConfig, ContextRuleConfig, ModelConfig, OptimizationConfig, LlmCompressionConfig
+from aworld.core.context.base import Context
+from aworld.core.event.base import Message
+from aworld.runners.hook.hooks import PreLLMCallHook, PostLLMCallHook
+from aworld.runners.hook.hook_factory import HookFactory
+from aworld.utils.common import convert_to_snake
+from tests.base_test import BaseTest
+
+
+@HookFactory.register(name="TestPreLLMHook", desc="Test pre-LLM hook")
+class TestPreLLMHook(PreLLMCallHook):
+    def name(self):
+        return convert_to_snake("TestPreLLMHook")
+    async def exec(self, message: Message, context: Context = None) -> Message:
+        agent = AgentFactory.agent_instance(message.sender)
+        agent_context = agent.agent_context
+        if agent_context is not None:
+            agent_context.step = 1 
+        assert agent_context.step == 1 or agent_context.step == 2
+        return message
+
+@HookFactory.register(name="TestPostLLMHook", desc="Test post-LLM hook")
+class TestPostLLMHook(PostLLMCallHook):
+    def name(self):
+        return convert_to_snake("TestPostLLMHook")
+    async def exec(self, message: Message, context: Context = None) -> Message:
+        agent = AgentFactory.agent_instance(message.sender)
+        agent_context = agent.agent_context
+        if agent_context is not None and agent_context.llm_output is not None:
+            # Test dynamic prompt adjustment based on LLM output
+            if hasattr(agent_context.llm_output, 'content'):
+                content = agent_context.llm_output.content.lower()
+                if content is not None:
+                    agent_context.agent_prompt = "Success mode activated"
+        assert agent_context.agent_prompt == "Success mode activated"
+        return message