unit tests

Author: heyitsaamir

packages/ai/src/microsoft/teams/ai/agent.py

@@ -9,7 +9,7 @@
 from .chat_prompt import ChatPrompt, ChatSendResult
 from .function import Function
 from .memory import ListMemory, Memory
-from .message import Message
+from .message import Message, SystemMessage
 
 
 class Agent(ChatPrompt):
@@ -26,7 +26,8 @@ async def send(
         self,
         input: str | Message,
         *,
+        system_message: SystemMessage | None = None,
         memory: Memory | None = None,
         on_chunk: Callable[[str], Awaitable[None]] | Callable[[str], None] | None = None,
     ) -> ChatSendResult:
-        return await super().send(input, memory=memory or self.memory, on_chunk=on_chunk)
+        return await super().send(input, memory=memory or self.memory, system_message=system_message, on_chunk=on_chunk)

packages/ai/tests/test_chat_prompt.py

@@ -14,12 +14,16 @@
     FunctionCall,
     ListMemory,
     Memory,
+    Message,
     ModelMessage,
     SystemMessage,
     UserMessage,
 )
+from microsoft.teams.ai.plugin import BaseAIPlugin
 from pydantic import BaseModel
 
+# pyright: basic
+
 
 class MockFunctionParams(BaseModel):
     value: str
@@ -61,6 +65,11 @@ async def generate_text(
             try:
                 params = function.parameter_schema(value="test_input")
                 result = function.handler(params)
+                # Handle both sync and async results
+                from inspect import isawaitable
+
+                if isawaitable(result):
+                    result = await result
                 # In real implementation, function result would be added to memory
                 # and conversation would continue recursively
                 content += f" | Function result: {result}"
@@ -277,3 +286,284 @@ async def test_different_message_types(self, mock_model: MockAIModel) -> None:
         model_msg = ModelMessage(content="Model message", function_calls=None)
         result3 = await prompt.send(model_msg)
         assert result3.response.content == "GENERATED - Model message"
+
+
+class MockPlugin(BaseAIPlugin):
+    """Mock plugin for testing that tracks all hook calls"""
+
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.before_send_called = False
+        self.after_send_called = False
+        self.before_function_called: list[tuple[str, BaseModel]] = []
+        self.after_function_called: list[tuple[str, BaseModel, str]] = []
+        self.build_functions_called = False
+        self.build_system_message_called = False
+        self.input_modifications: list[str] = []
+        self.response_modifications: list[str] = []
+        self.function_result_modifications: list[str] = []
+
+    async def on_before_send(self, input: Message) -> Message | None:
+        self.before_send_called = True
+        if self.input_modifications:
+            modification = self.input_modifications.pop(0)
+            if isinstance(input, UserMessage):
+                return UserMessage(content=f"{modification}: {input.content}")
+            elif isinstance(input, ModelMessage):
+                return ModelMessage(
+                    content=f"{modification}: {input.content}" if input.content else modification,
+                    function_calls=input.function_calls,
+                )
+        return input
+
+    async def on_after_send(self, response: ModelMessage) -> ModelMessage | None:
+        self.after_send_called = True
+        if self.response_modifications:
+            modification = self.response_modifications.pop(0)
+            return ModelMessage(
+                content=f"{modification}: {response.content}" if response.content else modification,
+                function_calls=response.function_calls,
+            )
+        return response
+
+    async def on_before_function_call(self, function_name: str, args: BaseModel) -> None:
+        self.before_function_called.append((function_name, args))
+
+    async def on_after_function_call(self, function_name: str, args: BaseModel, result: str) -> str | None:
+        self.after_function_called.append((function_name, args, result))
+        if self.function_result_modifications:
+            modification = self.function_result_modifications.pop(0)
+            return f"{modification}: {result}"
+        return result
+
+    async def on_build_functions(self, functions: list[Function[BaseModel]]) -> list[Function[BaseModel]] | None:
+        self.build_functions_called = True
+        return functions
+
+    async def on_build_system_message(self, system_message: SystemMessage | None) -> SystemMessage | None:
+        self.build_system_message_called = True
+        if system_message is None:
+            return SystemMessage(content="Plugin-generated system message")
+        return SystemMessage(content=f"Plugin-modified: {system_message.content}")
+
+
+class TestChatPromptPlugins:
+    """Test suite for plugin functionality in ChatPrompt"""
+
+    def test_plugin_initialization_and_registration(self, mock_model: MockAIModel) -> None:
+        """Test plugin initialization and registration"""
+        plugin1 = MockPlugin("plugin1")
+        plugin2 = MockPlugin("plugin2")
+
+        # Test initialization with plugins
+        prompt = ChatPrompt(mock_model, plugins=[plugin1])
+        assert len(prompt.plugins) == 1
+        assert prompt.plugins[0] is plugin1
+
+        # Test with_plugin method
+        result = prompt.with_plugin(plugin2)
+        assert result is prompt  # Should return self for chaining
+        assert len(prompt.plugins) == 2
+        assert prompt.plugins[1] is plugin2
+
+    @pytest.mark.asyncio
+    async def test_on_before_send_hook(self, mock_model: MockAIModel) -> None:
+        """Test that on_before_send hook can modify input messages"""
+        plugin = MockPlugin("test_plugin")
+        plugin.input_modifications = ["MODIFIED"]
+
+        prompt = ChatPrompt(mock_model, plugins=[plugin])
+        result = await prompt.send("Original message")
+
+        assert plugin.before_send_called
+        assert result.response.content == "GENERATED - MODIFIED: Original message"
+
+    @pytest.mark.asyncio
+    async def test_on_after_send_hook(self, mock_model: MockAIModel) -> None:
+        """Test that on_after_send hook can modify response messages"""
+        plugin = MockPlugin("test_plugin")
+        plugin.response_modifications = ["RESPONSE_MODIFIED"]
+
+        prompt = ChatPrompt(mock_model, plugins=[plugin])
+        result = await prompt.send("Test message")
+
+        assert plugin.after_send_called
+        assert result.response.content == "RESPONSE_MODIFIED: GENERATED - Test message"
+
+    @pytest.mark.asyncio
+    async def test_on_build_system_message_hook(self, mock_model: MockAIModel) -> None:
+        """Test that on_build_system_message hook is called"""
+        plugin = MockPlugin("test_plugin")
+
+        prompt = ChatPrompt(mock_model, plugins=[plugin])
+
+        # Test with None system message
+        await prompt.send("Test", system_message=None)
+        assert plugin.build_system_message_called
+
+        # Reset and test with existing system message
+        plugin.build_system_message_called = False
+        system_msg = SystemMessage(content="Original system")
+        await prompt.send("Test", system_message=system_msg)
+        assert plugin.build_system_message_called
+
+    @pytest.mark.asyncio
+    async def test_function_call_hooks(self, mock_function_handler: Mock) -> None:
+        """Test that function call hooks are properly executed"""
+        plugin = MockPlugin("test_plugin")
+        plugin.function_result_modifications = ["FUNCTION_MODIFIED"]
+
+        # Create a mock model that will call functions
+        mock_model = MockAIModel(should_call_function=True)
+
+        # Create function with mock handler
+        test_function = Function(
+            name="test_function",
+            description="A test function",
+            parameter_schema=MockFunctionParams,
+            handler=mock_function_handler,
+        )
+
+        prompt = ChatPrompt(mock_model, functions=[test_function], plugins=[plugin])
+        result = await prompt.send("Call the function")
+
+        # Verify before hook was called
+        assert len(plugin.before_function_called) == 1
+        assert plugin.before_function_called[0][0] == "test_function"
+        assert isinstance(plugin.before_function_called[0][1], MockFunctionParams)
+
+        # Verify after hook was called and modified result
+        assert len(plugin.after_function_called) == 1
+        assert plugin.after_function_called[0][0] == "test_function"
+        assert result.response.content is not None
+        assert "FUNCTION_MODIFIED: Function executed successfully" in result.response.content
+
+    @pytest.mark.asyncio
+    async def test_on_build_functions_hook(
+        self, mock_model: MockAIModel, test_function: Function[MockFunctionParams]
+    ) -> None:
+        """Test that on_build_functions hook is called when functions are present"""
+        plugin = MockPlugin("test_plugin")
+
+        prompt = ChatPrompt(mock_model, functions=[test_function], plugins=[plugin])
+        await prompt.send("Test message")
+
+        assert plugin.build_functions_called
+
+    @pytest.mark.asyncio
+    async def test_multiple_plugins_execution_order(self, mock_model: MockAIModel) -> None:
+        """Test that multiple plugins execute in correct order"""
+        plugin1 = MockPlugin("plugin1")
+        plugin1.input_modifications = ["FIRST"]
+        plugin1.response_modifications = ["FIRST_RESP"]
+
+        plugin2 = MockPlugin("plugin2")
+        plugin2.input_modifications = ["SECOND"]
+        plugin2.response_modifications = ["SECOND_RESP"]
+
+        prompt = ChatPrompt(mock_model, plugins=[plugin1, plugin2])
+        result = await prompt.send("Original")
+
+        # Both plugins should be called
+        assert plugin1.before_send_called
+        assert plugin2.before_send_called
+        assert plugin1.after_send_called
+        assert plugin2.after_send_called
+
+        # Input should be modified by both plugins in order
+        assert result.response.content == "SECOND_RESP: FIRST_RESP: GENERATED - SECOND: FIRST: Original"
+
+    @pytest.mark.asyncio
+    async def test_plugin_returns_none_preserves_original(self, mock_model: MockAIModel) -> None:
+        """Test that when plugin returns None, original values are preserved"""
+
+        class NoOpPlugin(BaseAIPlugin):
+            def __init__(self):
+                super().__init__("noop")
+
+            async def on_before_send(self, input: Message) -> Message | None:
+                return None  # Return None to preserve original
+
+            async def on_after_send(self, response: ModelMessage) -> ModelMessage | None:
+                return None  # Return None to preserve original
+
+        plugin = NoOpPlugin()
+        prompt = ChatPrompt(mock_model, plugins=[plugin])
+        result = await prompt.send("Test message")
+
+        # Should be unchanged since plugin returned None
+        assert result.response.content == "GENERATED - Test message"
+
+    @pytest.mark.asyncio
+    async def test_empty_plugin_list_maintains_compatibility(self, mock_model: MockAIModel) -> None:
+        """Test that ChatPrompt with no plugins behaves identically to original implementation"""
+        prompt_with_plugins = ChatPrompt(mock_model, plugins=[])
+        prompt_without_plugins = ChatPrompt(mock_model)
+
+        result_with = await prompt_with_plugins.send("Test message")
+        result_without = await prompt_without_plugins.send("Test message")
+
+        assert result_with.response.content == result_without.response.content
+
+    @pytest.mark.asyncio
+    async def test_plugin_with_async_function_handler(self, mock_function_handler: Mock) -> None:
+        """Test plugin hooks work correctly with async function handlers"""
+        plugin = MockPlugin("async_test")
+        plugin.function_result_modifications = ["ASYNC_MODIFIED"]
+
+        # Use the existing mock function handler (it's already set up correctly)
+        mock_model = MockAIModel(should_call_function=True)
+
+        test_function = Function(
+            name="test_function",  # Use same name as MockAIModel expects
+            description="A test function",
+            parameter_schema=MockFunctionParams,
+            handler=mock_function_handler,
+        )
+
+        prompt = ChatPrompt(mock_model, functions=[test_function], plugins=[plugin])
+        result = await prompt.send("Call the function")
+
+        # Verify function was called and result was modified by plugin
+        assert len(plugin.before_function_called) == 1
+        assert len(plugin.after_function_called) == 1
+        assert result.response.content is not None
+        assert "ASYNC_MODIFIED: Function executed successfully" in result.response.content
+
+    @pytest.mark.asyncio
+    async def test_plugin_error_handling(self, mock_model: MockAIModel) -> None:
+        """Test that plugin errors don't break the chat flow"""
+
+        class FaultyPlugin(BaseAIPlugin):
+            def __init__(self):
+                super().__init__("faulty")
+
+            async def on_before_send(self, input: Message) -> Message | None:
+                raise ValueError("Plugin error")
+
+        plugin = FaultyPlugin()
+        prompt = ChatPrompt(mock_model, plugins=[plugin])
+
+        # Plugin error should propagate (this is expected behavior)
+        with pytest.raises(ValueError, match="Plugin error"):
+            await prompt.send("Test message")
+
+    @pytest.mark.asyncio
+    async def test_base_plugin_default_implementations(self, mock_model: MockAIModel) -> None:
+        """Test that BaseAIPlugin provides working default implementations"""
+        base_plugin = BaseAIPlugin("base")
+        prompt = ChatPrompt(mock_model, plugins=[base_plugin])
+
+        # Should work without any issues using default implementations
+        result = await prompt.send("Test with base plugin")
+        assert result.response.content == "GENERATED - Test with base plugin"
+
+        # Test with functions too
+        def handler(params: MockFunctionParams) -> str:
+            return "Base plugin test"
+
+        test_function = Function("test", "test", MockFunctionParams, handler)
+        prompt_with_func = ChatPrompt(mock_model, functions=[test_function], plugins=[base_plugin])
+
+        result2 = await prompt_with_func.send("Test with function")
+        assert result2.response.content == "GENERATED - Test with function"