Add AI Plugins

Author: heyitsaamir

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

@@ -5,14 +5,16 @@
 
 import inspect
 from dataclasses import dataclass
+from inspect import isawaitable
 from typing import Any, Awaitable, Callable, TypeVar
 
 from pydantic import BaseModel
 
 from .ai_model import AIModel
-from .function import Function
+from .function import Function, FunctionHandler
 from .memory import Memory
-from .message import Message, ModelMessage, UserMessage
+from .message import Message, ModelMessage, SystemMessage, UserMessage
+from .plugin import AIPluginProtocol
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -23,24 +25,74 @@ class ChatSendResult:
 
 
 class ChatPrompt:
-    def __init__(self, model: AIModel, *, functions: list[Function[Any]] | None = None):
+    def __init__(
+        self,
+        model: AIModel,
+        *,
+        functions: list[Function[Any]] | None = None,
+        plugins: list[AIPluginProtocol] | None = None,
+    ):
         self.model = model
         self.functions: dict[str, Function[Any]] = {func.name: func for func in functions} if functions else {}
+        self.plugins: list[AIPluginProtocol] = plugins or []
 
     def with_function(self, function: Function[T]) -> "ChatPrompt":
         self.functions[function.name] = function
         return self
 
+    def with_plugin(self, plugin: AIPluginProtocol) -> "ChatPrompt":
+        """Add a plugin to the chat prompt."""
+        self.plugins.append(plugin)
+        return self
+
     async def send(
         self,
         input: str | Message,
         *,
         memory: Memory | None = None,
         on_chunk: Callable[[str], Awaitable[None]] | Callable[[str], None] | None = None,
+        system_message: SystemMessage | None = None,
     ) -> ChatSendResult:
         if isinstance(input, str):
             input = UserMessage(content=input)
 
+        # Allow plugins to modify the input before sending
+        current_input = input
+        for plugin in self.plugins:
+            plugin_result = await plugin.on_before_send(current_input)
+            if plugin_result is not None:
+                current_input = plugin_result
+
+        # Allow plugins to modify the system message
+        current_system_message = system_message
+        for plugin in self.plugins:
+            plugin_result = await plugin.on_build_system_message(current_system_message)
+            if plugin_result is not None:
+                current_system_message = plugin_result
+
+        # Wrap functions with plugin hooks
+        wrapped_functions: dict[str, Function[BaseModel]] | None = None
+        if self.functions:
+            wrapped_functions = {}
+            for name, func in self.functions.items():
+                wrapped_functions[name] = Function[BaseModel](
+                    name=func.name,
+                    description=func.description,
+                    parameter_schema=func.parameter_schema,
+                    handler=self._wrap_function_handler(func.handler, name),
+                )
+
+        # Allow plugins to modify the functions before sending to model
+        if wrapped_functions:
+            functions_list = list(wrapped_functions.values())
+            for plugin in self.plugins:
+                plugin_result = await plugin.on_build_functions(functions_list)
+                if plugin_result is not None:
+                    functions_list = plugin_result
+
+            # Convert back to dict for model
+            wrapped_functions = {func.name: func for func in functions_list}
+
         async def on_chunk_fn(chunk: str):
             if not on_chunk:
                 return
@@ -49,10 +101,40 @@ async def on_chunk_fn(chunk: str):
                 await res
 
         response = await self.model.generate_text(
-            input,
-            memory=memory,
-            functions=self.functions if self.functions else None,
-            on_chunk=on_chunk_fn if on_chunk else None,
+            current_input, memory=memory, functions=wrapped_functions, on_chunk=on_chunk_fn if on_chunk else None
         )
 
-        return ChatSendResult(response=response)
+        # Allow plugins to modify the response after receiving
+        current_response = response
+        for plugin in self.plugins:
+            plugin_result = await plugin.on_after_send(current_response)
+            if plugin_result is not None:
+                current_response = plugin_result
+
+        return ChatSendResult(response=current_response)
+
+    def _wrap_function_handler(
+        self, original_handler: FunctionHandler[BaseModel], function_name: str
+    ) -> FunctionHandler[BaseModel]:
+        """Wrap a function handler with plugin before/after hooks."""
+
+        async def wrapped_handler(params: BaseModel) -> str:
+            # Run before function call hooks
+            for plugin in self.plugins:
+                await plugin.on_before_function_call(function_name, params)
+
+            # Call the original function (could be sync or async)
+            result = original_handler(params)
+            if isawaitable(result):
+                result = await result
+
+            # Run after function call hooks
+            current_result = result
+            for plugin in self.plugins:
+                plugin_result = await plugin.on_after_function_call(function_name, params, current_result)
+                if plugin_result is not None:
+                    current_result = plugin_result
+
+            return current_result
+
+        return wrapped_handler

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

@@ -0,0 +1,85 @@
+"""
+Copyright (c) Microsoft Corporation. All rights reserved.
+Licensed under the MIT License.
+"""
+
+from abc import abstractmethod
+from typing import Protocol, TypeVar, runtime_checkable
+
+from pydantic import BaseModel
+
+from .function import Function
+from .message import Message, ModelMessage, SystemMessage
+
+T = TypeVar("T")
+
+
+@runtime_checkable
+class AIPluginProtocol(Protocol):
+    """Protocol defining the interface for AI plugins."""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Unique name of the plugin."""
+        ...
+
+    async def on_before_send(self, input: Message) -> Message | None:
+        """Modify input before sending to model."""
+        ...
+
+    async def on_after_send(self, response: ModelMessage) -> ModelMessage | None:
+        """Modify response after receiving from model."""
+        ...
+
+    async def on_before_function_call(self, function_name: str, args: BaseModel) -> None:
+        """Called before a function is executed."""
+        ...
+
+    async def on_after_function_call(self, function_name: str, args: BaseModel, result: str) -> str | None:
+        """Called after a function is executed."""
+        ...
+
+    async def on_build_functions(self, functions: list[Function[BaseModel]]) -> list[Function[BaseModel]] | None:
+        """Modify the functions array passed to the model."""
+        ...
+
+    async def on_build_system_message(self, system_message: SystemMessage | None) -> SystemMessage | None:
+        """Modify the system message before sending to model."""
+        ...
+
+
+class BaseAIPlugin:
+    """Base implementation of AIPlugin with no-op methods."""
+
+    def __init__(self, name: str):
+        self._name = name
+
+    @property
+    def name(self) -> str:
+        """Unique name of the plugin."""
+        return self._name
+
+    async def on_before_send(self, input: Message) -> Message | None:
+        """Modify input before sending to model."""
+        return input
+
+    async def on_after_send(self, response: ModelMessage) -> ModelMessage | None:
+        """Modify response after receiving from model."""
+        return response
+
+    async def on_before_function_call(self, function_name: str, args: BaseModel) -> None:
+        """Called before a function is executed."""
+        pass
+
+    async def on_after_function_call(self, function_name: str, args: BaseModel, result: str) -> str | None:
+        """Called after a function is executed."""
+        return result
+
+    async def on_build_functions(self, functions: list[Function[BaseModel]]) -> list[Function[BaseModel]] | None:
+        """Modify the functions array passed to the model."""
+        return functions
+
+    async def on_build_system_message(self, system_message: SystemMessage | None) -> SystemMessage | None:
+        """Modify the system message before sending to model."""
+        return system_message