diff --git a/CHANGELOG.md b/CHANGELOG.md index 0659316..d5826b5 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/docs/index.rst b/docs/index.rst index 4595a1e..bbb8366 100644 --- a/docs/index.rst +++ b/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 ----------- diff --git a/simplemind/__init__.py b/simplemind/__init__.py index 8438abe..bfa7045 100644 --- a/simplemind/__init__.py +++ b/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", diff --git a/simplemind/models.py b/simplemind/models.py index c6325ed..a031d31 100644 --- a/simplemind/models.py +++ b/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"" - 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]: