From b075d8cb49521136f9334e92f4575e73ef259d01 Mon Sep 17 00:00:00 2001
From: Asankhaya Sharma <codelion@users.noreply.github.com>
Date: Thu, 24 Oct 2024 16:13:05 +0800
Subject: [PATCH] Add plugins

Plugins are loaded in the following order
- from package plugins folder
- current directory/optillm/plugins folder
for same SLUG we will override with local plugins
---
 MANIFEST.in |  1 +
 optillm.py  | 65 +++++++++++++++++++++++++++++++++++++++++++----------
 setup.py    |  6 ++++-
 3 files changed, 59 insertions(+), 13 deletions(-)
 create mode 100644 MANIFEST.in

diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 0000000..6cae0ef
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1 @@
+include optillm/plugins/*.py
diff --git a/optillm.py b/optillm.py
index b298e52..12b196f 100644
--- a/optillm.py
+++ b/optillm.py
@@ -103,18 +103,59 @@ def get_config():
 plugin_approaches = {}
 
 def load_plugins():
-    plugin_dir = os.path.join(os.path.dirname(__file__), 'optillm/plugins')
-    plugin_files = glob.glob(os.path.join(plugin_dir, '*.py'))
-
-    for plugin_file in plugin_files:
-        module_name = os.path.basename(plugin_file)[:-3]  # Remove .py extension
-        spec = importlib.util.spec_from_file_location(module_name, plugin_file)
-        module = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(module)
-        
-        if hasattr(module, 'SLUG') and hasattr(module, 'run'):
-            plugin_approaches[module.SLUG] = module.run
-            logger.info(f"Loaded plugin: {module.SLUG}")
+   # Clear existing plugins first but modify the global dict in place
+   plugin_approaches.clear()
+   
+   # Get installed package plugins directory
+   import optillm
+   package_plugin_dir = os.path.join(os.path.dirname(optillm.__file__), 'plugins')
+   
+   # Get local project plugins directory
+   current_dir = os.getcwd()
+   local_plugin_dir = os.path.join(current_dir, 'optillm', 'plugins')
+   
+   plugin_dirs = []
+   
+   # Add package plugin dir
+   plugin_dirs.append((package_plugin_dir, "package"))
+   
+   # Add local plugin dir only if it's different from package dir
+   if local_plugin_dir != package_plugin_dir:
+       plugin_dirs.append((local_plugin_dir, "local"))
+   
+   for plugin_dir, source in plugin_dirs:
+       logger.info(f"Looking for {source} plugins in: {plugin_dir}")
+       
+       if not os.path.exists(plugin_dir):
+           logger.debug(f"{source.capitalize()} plugin directory not found: {plugin_dir}")
+           continue
+           
+       plugin_files = glob.glob(os.path.join(plugin_dir, '*.py'))
+       if not plugin_files:
+           logger.debug(f"No plugin files found in {source} directory: {plugin_dir}")
+           continue
+           
+       logger.info(f"Found {source} plugin files: {plugin_files}")
+       
+       for plugin_file in plugin_files:
+           try:
+               module_name = os.path.basename(plugin_file)[:-3]  # Remove .py extension
+               spec = importlib.util.spec_from_file_location(module_name, plugin_file)
+               module = importlib.util.module_from_spec(spec)
+               spec.loader.exec_module(module)
+               
+               if hasattr(module, 'SLUG') and hasattr(module, 'run'):
+                   if module.SLUG in plugin_approaches:
+                       logger.info(f"Overriding {source} plugin: {module.SLUG}")
+                   plugin_approaches[module.SLUG] = module.run
+                   logger.info(f"Loaded {source} plugin: {module.SLUG}")
+               else:
+                   logger.warning(f"Plugin {module_name} from {source} missing required attributes (SLUG and run)")
+           except Exception as e:
+               logger.error(f"Error loading {source} plugin {plugin_file}: {str(e)}")
+   
+   if not plugin_approaches:
+       logger.warning("No plugins loaded from any location")
 
 def parse_combined_approach(model: str, known_approaches: list, plugin_approaches: dict):
     if model == 'auto':
diff --git a/setup.py b/setup.py
index 32ab43a..ae76402 100644
--- a/setup.py
+++ b/setup.py
@@ -2,9 +2,13 @@
 
 setup(
     name="optillm",
-    version="0.0.5",
+    version="0.0.6",
     packages=find_packages(),
     py_modules=['optillm'],
+    package_data={
+        'optillm': ['plugins/*.py'],  # Include plugin files
+    },
+    include_package_data=True,  # This is important
     install_requires=[
         "numpy",
         "networkx",