ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/tests/context/test_context_compressor.py

@@ -0,0 +1,110 @@
+import os
+import sys
+from pathlib import Path
+import unittest
+
+# Add the project root to Python path
+project_root = Path(__file__).parent.parent.parent
+sys.path.insert(0, str(project_root))
+
+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
+
+
+class TestPromptCompressor(unittest.TestCase):
+    """Test cases for PromptCompressor.compress_batch function"""
+
+    def test_compress_batch_basic(self):
+
+        compressor = LLMCompressor(
+            llm_config=ModelConfig(
+                llm_model_name=os.environ["LLM_MODEL_NAME"],
+                llm_base_url=os.environ["LLM_BASE_URL"],
+                llm_api_key=os.environ["LLM_API_KEY"],
+            )
+        )
+
+        # 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=ModelConfig(
+                    llm_model_name=os.environ["LLM_MODEL_NAME"],
+                    llm_base_url=os.environ["LLM_BASE_URL"],
+                    llm_api_key=os.environ["LLM_API_KEY"],
+                )
+            )
+        )
+
+        # Create prompt processor with context_rule and model_config
+        processor = PromptProcessor(context_rule=context_rule, model_config=ModelConfig(
+            llm_model_name=os.environ["LLM_MODEL_NAME"],
+            llm_base_url=os.environ["LLM_BASE_URL"],
+            llm_api_key=os.environ["LLM_API_KEY"],
+        ))
+
+        # 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__":
+    unittest.main()

chapter9/gaia-experience/AWorld/tests/context/test_context_management.py

@@ -0,0 +1,170 @@
+import asyncio
+import os
+import sys
+from pathlib import Path
+import unittest
+
+# Add the project root to Python path
+project_root = Path(__file__).parent.parent.parent
+sys.path.insert(0, str(project_root))
+
+from aworld.core.context.session import Session
+from aworld.core.agent.swarm import Swarm
+from tests.base_test import assertEqual, assertIn, assertIsInstance, assertIsNotNone, assertTrue, run_multi_agent_as_team, run_task
+from aworld.runners.hook.hook_factory import HookFactory
+from aworld.core.context.base import Context
+from aworld.config.conf import AgentConfig, ContextRuleConfig, ModelConfig, OptimizationConfig, LlmCompressionConfig
+from aworld.agents.llm_agent import Agent
+from aworld.core.task import Task
+from tests.base_test import init_agent, run_agent, run_multi_agent_as_team
+
+class TestContextManagement(unittest.TestCase):
+
+    # def test_save_and_reload(self):
+    #     context = Context()
+    #     context.context_info.set("hello", "world")
+    #     task = Task(input="""What is an agent.""",
+    #                 swarm=Swarm(init_agent("1"), max_steps=1), context=context)
+    #     task.session_id = "1"
+    #     context.session = Session(session_id="1")
+    #     context.set_task(task)
+    #     context_manager = ContextManager()
+    #     checkpoint = asyncio.run(context_manager.save(context))
+
+    #     session_id = context.session_id
+    #     context = asyncio.run(context_manager.reload(session_id))
+    #     assertEqual(context.context_info.get("hello"), "world")
+        
+    def test_default_context_configuration(self):
+        mock_agent = init_agent("1")
+        response = run_agent(
+            input="""What is an agent. describe within 20 words""", agent=mock_agent)
+
+        assertIsNotNone(response.answer)
+        assertEqual(
+            mock_agent.conf.llm_config.llm_model_name, os.environ["LLM_MODEL_NAME"])
+
+        # Test default context rule behavior
+        assertIsNotNone(mock_agent.context_rule)
+        assertIsNotNone(
+            mock_agent.context_rule.optimization_config)
+
+    def test_custom_context_configuration(self):
+        """Test custom context configuration (README Configuration example)"""
+        # Create custom context rules
+        mock_agent = 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=os.environ["LLM_MODEL_NAME"],
+                    llm_base_url=os.environ["LLM_BASE_URL"],
+                    llm_api_key=os.environ["LLM_API_KEY"],
+                )
+            )
+        ))
+
+        response = run_agent(
+            input="""describe What is an agent in details""", agent=mock_agent)
+        assertIsNotNone(response.answer)
+
+        # Test configuration values
+        assertTrue(
+            mock_agent.context_rule.optimization_config.enabled)
+        assertTrue(
+            mock_agent.context_rule.llm_compression_config.enabled)
+
+
+    def test_multi_agent_state_trace(self):
+        class StateModifyAgent(Agent):
+            async def async_policy(self, observation, info=None, **kwargs):
+                result = await super().async_policy(observation, info, **kwargs)
+                self.context.context_info.set('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.context_info.get('policy_executed', True)
+                return result
+
+        # Create custom agent instance
+        custom_agent = StateModifyAgent(
+            conf=AgentConfig(
+                llm_model_name=os.environ["LLM_MODEL_NAME"],
+                llm_base_url=os.environ["LLM_BASE_URL"],
+                llm_api_key=os.environ["LLM_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=os.environ["LLM_MODEL_NAME"],
+                llm_base_url=os.environ["LLM_BASE_URL"],
+                llm_api_key=os.environ["LLM_API_KEY"]
+            ),
+            name="state_tracking_agent",
+            system_prompt="You are a helpful assistant.",
+            agent_prompt="You are a helpful assistant.",
+        )
+
+        response = run_multi_agent_as_team(
+            input="What is an agent. describe within 20 words",
+            agent1=custom_agent,
+            agent2=second_agent
+        )
+        assertIsNotNone(response.answer)
+
+        # Verify state changes after execution
+        assertTrue(custom_agent.context.context_info.get('policy_executed', True))
+
+    def test_multi_task_state_trace(self):
+        context = Context()
+        task = Task(input="What is an agent.", context=context)
+        new_context = task.context.deep_copy()
+        new_context.context_info.update({"hello": "world"})
+        run_task(context=new_context, agent=init_agent("1"))
+        assertEqual(new_context.context_info.get("hello"), "world")
+
+        task.context.merge_context(new_context)
+        assertEqual(task.context.context_info.get("hello"), "world")
+
+
+    def test_hook_registration(self):
+        from tests.runners.hook.llm_hook import TestPreLLMHook, TestPostLLMHook
+        """Test hook registration and retrieval"""
+        # Test that hooks are registered in _cls attribute
+        assertIn("TestPreLLMHook", HookFactory._cls)
+        assertIn("TestPostLLMHook", HookFactory._cls)
+
+        # Test hook creation using __call__ method
+        pre_hook = HookFactory("TestPreLLMHook")
+        post_hook = HookFactory("TestPostLLMHook")
+
+        assertIsInstance(pre_hook, TestPreLLMHook)
+        assertIsInstance(post_hook, TestPostLLMHook)
+
+    def test_hook_execution(self):
+        mock_agent = init_agent("1")
+        response = run_agent(
+            input="""What is an agent. describe within 20 words""", agent=mock_agent)
+        assertIsNotNone(response.answer)
+
+    def test_task_context_transfer(self):
+        mock_agent = init_agent("1")
+        context = Context()
+        context.context_info.update({"task": "What is an agent."})
+        run_task(context=context, agent=mock_agent)
+
+
+
+if __name__ == '__main__':
+    unittest.main()

chapter9/gaia-experience/AWorld/tests/context/test_prompt_template.py

@@ -0,0 +1,125 @@
+# Add the project root to Python path
+from pathlib import Path
+import sys
+import unittest
+
+project_root = Path(__file__).parent.parent.parent
+print(project_root)
+sys.path.insert(0, str(project_root))
+
+from aworld.agents.llm_agent import Agent
+from tests.base_test import assertEqual, init_agent, run_agent, run_task
+from aworld.core.context.base import Context
+from aworld.core.context.prompts.dynamic_variables import create_simple_field_getter, format_ordered_dict_json, \
+    get_field_values_from_list, get_value_by_path
+from aworld.core.context.prompts.string_prompt_template import StringPromptTemplate
+
+
+class TestPromptTemplate(unittest.TestCase):
+    def test_dynamic_variables(self):
+        context = Context()
+        context.context_info.update({"task": "chat"})
+
+        # Test dot separator
+        value_dot = get_value_by_path(context, "context_info.task")
+        assert "chat" == value_dot
+
+        # Test slash separator
+        value_slash = get_value_by_path(context, "context_info/task")
+        assert "chat" == value_slash
+
+    def test_formatted_field_getter(self):
+        context = Context()
+        value = {"steps": [1, 2, 3]}
+        context.trajectories.update(value)
+
+        getter = create_simple_field_getter(field_path="trajectories", default="default_value")
+        result = getter(context=context)
+        assert "steps" in value
+        # test default format function
+        assert "OrderedDict" not in result
+
+        # Test formatted field getter with processor
+        getter = create_simple_field_getter(field_path="trajectories", default="default_value",
+                                            processor=format_ordered_dict_json)
+        result = getter(context=context)
+        assert "steps" in result
+
+    def test_multiple_field_getters(self):
+        context = Context()
+        context.context_info.update({"task": "chat"})
+        context.trajectories.update({"steps": [1, 2, 3]})
+
+        field_paths = ["context_info.task", "trajectories.steps"]
+        result = get_field_values_from_list(context=context, field_paths=field_paths)
+        assert result["context_info_task"] == "chat"
+        assert result["trajectories_steps"] == "[1, 2, 3]"
+
+    def test_string_prompt_template(self):
+        # Use proper dot notation for nested field access
+        template = StringPromptTemplate.from_template(
+            "Hello {{name}}, welcome to {{place}}! Task: {{task}} Age: {{age}}",
+            partial_variables={"age": "1"})
+        assert "name" in template.input_variables
+        assert "place" in template.input_variables
+        assert "task" in template.input_variables
+
+        context = Context()
+        context.context_info.update({"task": "chat"})
+
+        # Pass task as a direct parameter since template expects it
+        result = template.format(context=context, name="Alice", place="AWorld", task="chat")
+        assert result == "Hello Alice, welcome to AWorld! Task: chat Age: 1"
+
+    def test_enhanced_field_values_basic(self):
+        context = Context()
+        context.context_info.update({"task": "chat"})
+
+        # Test retrieving both time variables and context fields
+        result = get_field_values_from_list(
+            context=context,
+            field_paths=["current_time", "context_info.task"],
+            default="not_found"
+        )
+
+        # Verify context field retrieved
+        assert result["context_info_task"] == "chat"
+
+        # Verify time variable retrieved (should be in HH:MM:SS format)
+        assert ":" in result["current_time"]
+        assert len(result["current_time"].split(":")) == 3
+
+    def test_undefined_system_prompt_template(self):
+        agent = init_agent()
+        agent._log_messages = lambda messages: assertEqual(messages[0]['content'], "You are a helpful assistant.")
+        result = run_task(
+            input="What is the weather in Beijing?",
+            agent=agent
+        )
+        assert result is not None
+
+    def test_custom_system_prompt_template(self):
+        context = Context()
+        context.context_info.set("name", "Qwen")
+        context.context_info.set("plan", [{"input": "query weather in Beijing"}])
+
+        system_prompt_template = StringPromptTemplate.from_template(
+            "Hello {{context_info.name}}, you are a {{role}}, {{context_info.plan}}",
+            partial_variables={"role": "assistant", "context_info.plan": lambda ob: "please " + ob[0]['input']})
+        
+        agent = init_agent(
+            system_prompt_template=system_prompt_template,
+        )
+
+        agent._log_messages = lambda messages: assertEqual(messages[0]['content'], "Hello Qwen, you are a assistant, please query weather in Beijing")
+
+        result = run_task(
+            input="What is the weather in Beijing?",
+            agent=agent,
+            context=context
+        )
+        assert result is not None
+
+
+if __name__ == "__main__":
+    unittest.main()