Refactor conversation plugin hooks and add plugin interface

CHANGELOG.md

@@ -1,6 +1,14 @@
 Release History
 ===============
 
+
+## 0.1.3 (2024-10-30)
+
+- Make Conversation a context manager.
+- Add more robust conversation plugin hooks.
+- Remove `send_hook` from `BaseProvider`. Replaced with `pre_send_hook` and `post_send_hook`.
+- Change plugin hooks to try/except NotImplementedError.
+
 ## 0.1.2 (2024-10-29)
 
 - Add ollama provider.

docs/index.rst

@@ -143,6 +143,73 @@ This example adds a simple custom memory plugin to the conversation.
     conversation.add_message("user", "Write a poem about the moon")
     print(conversation.send().text)
 
+Plugin Development
+~~~~~~~~~~~~~~~~~~
+
+Plugins in SimpleMind follow a simple hook-based architecture. The ``send_hook`` method shown above is just one of several hooks available. Here's a more detailed example showing the complete plugin interface:
+
+.. code-block:: python
+
+    from simplemind.plugins import BasePlugin
+
+    class CustomPlugin(BasePlugin):
+        def __init__(self):
+            self.conversation_history = []
+
+        def initialize_hook(self, conversation):
+            """Called when the plugin is first added to a conversation."""
+            print("Plugin initialized!")
+
+        def pre_send_hook(self, conversation):
+            """Called before the conversation is sent to the AI provider."""
+            # Add any system messages or modify the conversation
+            conversation.add_message("system", "Remember to be helpful.")
+
+        def send_hook(self, conversation):
+            """Called during the send process."""
+            # Add messages or modify the conversation
+            self.conversation_history.append(conversation.messages)
+
+        def post_send_hook(self, conversation, response):
+            """Called after receiving a response from the AI provider."""
+            # Process or modify the response
+            return response
+
+        def cleanup_hook(self):
+            """Called when the plugin is removed or the conversation ends."""
+            self.conversation_history.clear()
+
+All plugins should inherit from ``BasePlugin``, which provides default no-op implementations of these hooks. You only need to implement the hooks you want to use. Here's a simpler example:
+
+.. code-block:: python
+
+    from simplemind.plugins import BasePlugin
+
+    class LoggingPlugin(BasePlugin):
+        def pre_send_hook(self, conversation):
+            print(f"Sending conversation with {len(conversation.messages)} messages")
+
+        def post_send_hook(self, conversation, response):
+            print(f"Received response: {response.text[:50]}...")
+            return response
+
+    conversation = sm.create_conversation()
+    conversation.add_plugin(LoggingPlugin())
+    conversation.add_message("user", "Hello!")
+    response = conversation.send()
+
+Plugins can be used to implement features like:
+
+- Conversation logging
+- Memory management
+- Response filtering
+- Token counting
+- Custom prompt engineering
+- Analytics and monitoring
+
+Multiple plugins can be added to a single conversation, and they will be executed in the order they were added.
+
+
 Contributing
 -----------
 

simplemind/__init__.py

@@ -1,21 +1,34 @@
-from .models import Conversation
+from typing import List, Optional
+
+from .models import Conversation, BasePlugin
 from .utils import find_provider
 from .settings import settings
 
 
-def create_conversation(llm_model=None, llm_provider=None):
+def create_conversation(
+    llm_model=None, llm_provider=None, *, plugins: Optional[List[BasePlugin]] = None
+):
     """Create a new conversation."""
 
-    return Conversation(
+    # Create the conversation.
+    conversation = Conversation(
         llm_model=llm_model, llm_provider=llm_provider or settings.DEFAULT_LLM_PROVIDER
     )
 
+    # Add plugins to the conversation.
+    for plugin in plugins or []:
+        conversation.add_plugin(plugin)
+
+    return conversation
+
 
 def generate_data(prompt, *, llm_model=None, llm_provider=None, response_model=None):
     """Generate structured data from a given prompt."""
 
+    # Find the provider.
     provider = find_provider(llm_provider or settings.DEFAULT_LLM_PROVIDER)
 
+    # Generate the data.
     return provider.structured_response(
         prompt=prompt,
         llm_model=llm_model,
@@ -25,13 +38,15 @@ def generate_data(prompt, *, llm_model=None, llm_provider=None, response_model=N
 
 def generate_text(prompt, *, llm_model=None, llm_provider=None, **kwargs):
     """Generate text from a given prompt."""
+
+    # Find the provider.
     provider = find_provider(llm_provider or settings.DEFAULT_LLM_PROVIDER)
 
+    # Generate the text.
     return provider.generate_text(prompt=prompt, llm_model=llm_model, **kwargs)
 
 
 __all__ = [
-    "Conversation",
     "create_conversation",
     "find_provider",
     "generate_data",

simplemind/models.py

@@ -25,9 +25,29 @@ class SMBaseModel(BaseModel):
 class BasePlugin(ABC):
     """The base conversation plugin class."""
 
-    @abstractmethod
-    def send_hook(self, conversation: "Conversation"):
-        """Send a hook to the plugin."""
+    # @abstractmethod
+    def initialize_hook(self, conversation: "Conversation"):
+        """Initialize a hook for the plugin."""
+        raise NotImplementedError
+
+    # @abstractmethod
+    def cleanup_hook(self, conversation: "Conversation"):
+        """Cleanup a hook for the plugin."""
+        raise NotImplementedError
+
+    # @abstractmethod
+    def add_message_hook(self, conversation: "Conversation", message: "Message"):
+        """Add a message hook for the plugin."""
+        raise NotImplementedError
+
+    # @abstractmethod
+    def pre_send_hook(self, conversation: "Conversation"):
+        """Pre-send hook for the plugin."""
+        raise NotImplementedError
+
+    # @abstractmethod
+    def post_send_hook(self, conversation: "Conversation", response: "Message"):
+        """Post-send hook for the plugin."""
         raise NotImplementedError
 
 
@@ -60,28 +80,82 @@ class Conversation(SMBaseModel):
     def __str__(self):
         return f"<Conversation id={self.id!r}>"
 
-    def prepend_system_message(self, role: str, text: str, meta: Optional[Dict[str, Any]] = None):
+    def __enter__(self):
+        # Execute all initialize hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "initialize_hook"):
+                try:
+                    plugin.initialize_hook(self)
+                except NotImplementedError:
+                    pass
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Execute all cleanup hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "cleanup_hook"):
+                try:
+                    plugin.cleanup_hook(self)
+                except NotImplementedError:
+                    pass
+
+    def prepend_system_message(
+        self, role: str, text: str, meta: Optional[Dict[str, Any]] = None
+    ):
+        """Prepend a system message to the conversation."""
         self.messages = [Message(role=role, text=text, meta=meta or {})] + self.messages
 
     def add_message(
         self, role: MESSAGE_ROLE, text: str, meta: Optional[Dict[str, Any]] = None
     ):
         """Add a new message to the conversation."""
+
+        # Ensure meta is a dict.
         if meta is None:
             meta = {}
+
+        # Execute all add-message hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "add_message_hook"):
+                try:
+                    plugin.add_message_hook(
+                        self, Message(role=role, text=text, meta=meta)
+                    )
+                except NotImplementedError:
+                    pass
+
+        # Add the message to the conversation.
         self.messages.append(Message(role=role, text=text, meta=meta))
 
     def send(
         self, llm_model: Optional[str] = None, llm_provider: Optional[str] = None
     ) -> Message:
         """Send the conversation to the LLM."""
-        for plugin in self.plugins:
-            plugin.send_hook(self)
 
+        # Execute all pre send hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "pre_send_hook"):
+                try:
+                    plugin.pre_send_hook(self)
+                except NotImplementedError:
+                    pass
+
+        # Find the provider and send the conversation.
         provider = find_provider(llm_provider or self.llm_provider)
         response = provider.send_conversation(self)
 
+        # Execute all post-send hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "post_send_hook"):
+                try:
+                    plugin.post_send_hook(self, response)
+                except NotImplementedError:
+                    pass
+
+        # Add the response to the conversation.
         self.add_message(role="assistant", text=response.text, meta=response.meta)
+
         return response
 
     def get_last_message(self, role: MESSAGE_ROLE) -> Optional[Message]: