"""
Plugin Registry for OllamaTools.
This module provides the PluginRegistry class for managing OllamaTools plugins.
"""
import os
import sys
import importlib
import logging
import pkg_resources
from typing import Dict, Any, List, Optional, Set, Type, Callable, Union
from ollamatools.plugins import OllamaToolsPlugin
from ollamatools.plugins.errors import (
PluginError,
PluginInitializationError,
PluginDependencyError,
PluginNotFoundError,
handle_plugin_exception
)
# Configure logging
logger = logging.getLogger(__name__)
class PluginRegistry:
"""Registry for OllamaTools plugins."""
def __init__(self):
"""Initialize the plugin registry."""
self.plugins = {}
self.initialized_plugins = {}
def register_plugin(self, plugin_class, config=None):
"""
Register a plugin class.
Args:
plugin_class: The plugin class to register
config: Optional configuration for the plugin
Returns:
PluginRegistry: Self for method chaining
"""
plugin_name = plugin_class.__name__
self.plugins[plugin_name] = {
"class": plugin_class,
"config": config
}
logger.debug(f"Registered plugin: {plugin_name}")
return self
def discover_plugins(self):
"""
Discover plugins using entry points.
Returns:
PluginRegistry: Self for method chaining
"""
try:
for entry_point in pkg_resources.iter_entry_points('ollamatools.plugins'):
try:
plugin_class = entry_point.load()
self.register_plugin(plugin_class)
logger.info(f"Discovered plugin: {plugin_class.__name__} from entry point {entry_point.name}")
except Exception as e:
logger.error(f"Error loading plugin from entry point {entry_point.name}: {e}")
except Exception as e:
logger.error(f"Error discovering plugins: {e}")
return self
def initialize_plugin(self, plugin_name, context):
"""
Initialize a plugin and its dependencies.
Args:
plugin_name: The name of the plugin to initialize
context: Runtime context for the plugin
Returns:
OllamaToolsPlugin: The initialized plugin instance
Raises:
PluginNotFoundError: If the plugin is not found
PluginInitializationError: If the plugin fails to initialize
PluginDependencyError: If a dependency cannot be initialized
"""
# Return already initialized plugin
if plugin_name in self.initialized_plugins:
return self.initialized_plugins[plugin_name]
# Check if plugin exists
if plugin_name not in self.plugins:
raise PluginNotFoundError(f"Plugin {plugin_name} not found")
logger.debug(f"Initializing plugin: {plugin_name}")
plugin_info = self.plugins[plugin_name]
plugin_class = plugin_info["class"]
plugin_config = plugin_info["config"]
try:
# Create plugin instance
plugin_instance = plugin_class(plugin_config)
# Initialize dependencies first
for dependency in plugin_instance.dependencies:
try:
self.initialize_plugin(dependency, context)
except PluginError as e:
raise PluginDependencyError(
f"Failed to initialize dependency {dependency} for plugin {plugin_name}",
details={"plugin": plugin_name, "dependency": dependency, "original_error": str(e)}
) from e
# Initialize the plugin
plugin_instance.initialize(context)
self.initialized_plugins[plugin_name] = plugin_instance
logger.info(f"Initialized plugin: {plugin_name}")
return plugin_instance
except PluginError:
# Re-raise plugin errors
raise
except Exception as e:
# Wrap other exceptions
raise PluginInitializationError(
f"Failed to initialize plugin {plugin_name}: {str(e)}",
details={"plugin": plugin_name}
) from e
def initialize_all(self, context):
"""
Initialize all registered plugins.
Args:
context: Runtime context for the plugins
Returns:
Dict[str, OllamaToolsPlugin]: Dictionary of initialized plugins
"""
for plugin_name in list(self.plugins.keys()):
try:
self.initialize_plugin(plugin_name, context)
except PluginError as e:
logger.error(f"Error initializing plugin {plugin_name}: {e}")
handle_plugin_exception(e)
except Exception as e:
logger.error(f"Unexpected error initializing plugin {plugin_name}: {e}")
return self.initialized_plugins
def get_plugin(self, plugin_name):
"""
Get an initialized plugin.
Args:
plugin_name: The name of the plugin to get
Returns:
Optional[OllamaToolsPlugin]: The plugin instance or None if not found
"""
return self.initialized_plugins.get(plugin_name)
def get_plugins_by_type(self, plugin_type):
"""
Get all initialized plugins of a specific type.
Args:
plugin_type: The plugin type class
Returns:
Dict[str, OllamaToolsPlugin]: Dictionary of plugin names to instances
"""
return {
name: plugin for name, plugin in self.initialized_plugins.items()
if isinstance(plugin, plugin_type)
}
def shutdown(self):
"""
Shutdown all initialized plugins.
Returns:
PluginRegistry: Self for method chaining
"""
for plugin_name, plugin in list(self.initialized_plugins.items()):
try:
logger.debug(f"Shutting down plugin: {plugin_name}")
plugin.shutdown()
logger.info(f"Shut down plugin: {plugin_name}")
except Exception as e:
logger.error(f"Error shutting down plugin {plugin_name}: {e}")
self.initialized_plugins = {}
return self