enh: use cached plugin if the config and module_path are unchanged (single plugin initialization call).

Signed-off-by: habeck <habeck@us.ibm.com>

Author: tedhabeck

cpex/framework/isolated/client.py

@@ -116,7 +116,7 @@ def _is_venv_cache_valid(self, venv_path: str, requirements_file: str) -> bool:
             # Compare hashes
             cached_hash = metadata.get("requirements_hash")
             if cached_hash != current_hash:
-                logger.info("Requirements changed. Cached hash: %s, Current hash: %s",cached_hash, current_hash)
+                logger.info("Requirements changed. Cached hash: %s, Current hash: %s", cached_hash, current_hash)
                 return False
 
             logger.info("Valid venv cache found for %s", venv_path)

cpex/framework/isolated/worker.py

@@ -9,6 +9,7 @@
 """
 
 import asyncio
+import hashlib
 import importlib.metadata
 import json
 import logging
@@ -28,6 +29,39 @@
 logger = logging.getLogger(__name__)
 
 
+class TaskProcessor:
+    """
+    A Caching task processor that only reloads the plugin if the config has changed.
+    """
+
+    config_hash: str
+    module_path_hash: str
+    hook_ref: HookRef | None
+    executor: PluginExecutor | None
+
+    def __init__(self) -> None:
+        """Initialize defaults."""
+        hasher = hashlib.sha256()
+        hasher.update(b"")
+        self.config_hash = hasher.hexdigest()
+        self.module_path_hash = self.config_hash
+        self.hook_ref = None
+        self.executor = None
+
+    def compute_hash(self, json_config_or_module_path: str):
+        """Compute the hash of the supplied string"""
+        hasher = hashlib.sha256()
+        hasher.update(json_config_or_module_path.encode())
+        return hasher.hexdigest()
+
+    def initialize(self, hook_ref: HookRef, executor: PluginExecutor, json_config: str, module_path: str):
+        """Assign locals, and compute hashes."""
+        self.hook_ref = hook_ref
+        self.executor = executor
+        self.config_hash = self.compute_hash(json_config_or_module_path=json_config)
+        self.module_path_hash = self.compute_hash(json_config_or_module_path=module_path)
+
+
 def get_environment_info():
     """Get information about current Python environment."""
     return {
@@ -55,7 +89,7 @@ def get_proper_config(name, module_path):
     return None
 
 
-async def process_task(task_data):
+async def process_task(task_data, tp: TaskProcessor):
     """Process the task received from parent."""
     task_type = task_data.get("task_type")
 
@@ -71,28 +105,33 @@ async def process_task(task_data):
         json_config = task_data.get("config")
         config_raw = json.loads(json_config)
         module_path: str = task_data.get("script_path")
-        sys.path.append(str(Path(module_path).resolve()))
-        config = get_proper_config(config_raw.get("name"), module_path)
-        hook_type = task_data.get(HOOK_TYPE)
-        cls_name: str = task_data.get("class_name")
-        mod_name, n_cls_name = parse_class_name(cls_name)
-        module: ModuleType = importlib.import_module(mod_name)
-        # cool, we found the module, and verified it implemented the hook type.
-        class_ = getattr(module, n_cls_name)
-        plugin_type = cast(Type[Plugin], class_)
-        plugin = plugin_type(config)
-        await plugin.initialize()
-        # now invoke the hook
-        plugin_ref = PluginRef(plugin)
-        hook_ref = HookRef(hook_type, plugin_ref)
-        executor = PluginExecutor(None, 30)
+        if tp.module_path_hash != tp.compute_hash(module_path) or tp.config_hash != tp.compute_hash(json_config):
+            sys.path.append(str(Path(module_path).resolve()))
+            config = get_proper_config(config_raw.get("name"), module_path)
+            hook_type = task_data.get(HOOK_TYPE)
+            cls_name: str = task_data.get("class_name")
+            mod_name, n_cls_name = parse_class_name(cls_name)
+            module: ModuleType = importlib.import_module(mod_name)
+            # cool, we found the module, and verified it implemented the hook type.
+            class_ = getattr(module, n_cls_name)
+            plugin_type = cast(Type[Plugin], class_)
+            plugin = plugin_type(config)
+            await plugin.initialize()
+            # now invoke the hook
+            plugin_ref = PluginRef(plugin)
+            hook_ref = HookRef(hook_type, plugin_ref)
+            executor = PluginExecutor(None, 30)
+            tp.initialize(hook_ref=hook_ref, executor=executor, json_config=json_config, module_path=module_path)
         # retrieve the context
         context = task_data.get("context")
         plugin_context = PluginContext(
             state=context.get("state"), global_context=context.get("global_context"), metadata=context.get("metadata")
         )
-        result = await executor.execute_plugin(
-            hook_ref, payload=task_data.get("payload"), local_context=plugin_context, violations_as_exceptions=False
+        result = await tp.executor.execute_plugin(
+            hookref=tp.hook_ref,
+            payload=task_data.get("payload"),
+            local_context=plugin_context,
+            violations_as_exceptions=False,
         )
         return result
 
@@ -102,6 +141,8 @@ async def main():
     logger.info("Worker process started, waiting for tasks...")
 
     try:
+        # Cache the plugin so that it only has to be initialized once
+        tp = TaskProcessor()
         # Continuously read and process tasks
         while True:
             try:
@@ -125,7 +166,7 @@ async def main():
                     break
 
                 # Process the task
-                response = await process_task(task_data)
+                response = await process_task(task_data, tp)
 
                 # Serialize response
                 if response:

tests/unit/cpex/framework/isolated/test_worker.py

@@ -16,7 +16,7 @@
 
 import pytest
 
-from cpex.framework.isolated.worker import get_environment_info, get_proper_config, main, process_task
+from cpex.framework.isolated.worker import TaskProcessor, get_environment_info, get_proper_config, main, process_task
 
 
 class TestWorkerFunctions:
@@ -81,9 +81,10 @@ def test_get_proper_config_no_plugins(self, mock_load_config):
     @pytest.mark.asyncio
     async def test_process_task_info(self):
         """Test processing info task."""
-        task_data = {"task_type": "info"}
-
-        result = await process_task(task_data)
+        config_dict = {"name": "test_plugin", "kind": "isolated_venv", "config": {}}
+        task_data = {"task_type": "info", "config": json.dumps(config_dict)}
+        tp = TaskProcessor()
+        result = await process_task(task_data, tp)
 
         assert result["status"] == "success"
         assert "environment" in result
@@ -132,8 +133,8 @@ async def test_process_task_load_and_run_hook_success(self, mock_executor_class,
             "payload": {"name": "test_tool", "args": {}},
             "context": {"state": {}, "global_context": {"request_id": "req-123"}, "metadata": {}},
         }
-
-        result = await process_task(task_data)
+        tp = TaskProcessor()
+        result = await process_task(task_data, tp=tp)
 
         assert result is not None
         mock_plugin_instance.initialize.assert_called_once()
@@ -155,10 +156,10 @@ async def test_process_task_load_and_run_hook_no_config(self, mock_get_config):
             "payload": {},
             "context": {"state": {}, "global_context": {}, "metadata": {}},
         }
-
+        tp = TaskProcessor()
         # Should raise an error or return None
         with pytest.raises((AttributeError, TypeError)):
-            await process_task(task_data)
+            await process_task(task_data, tp)
 
     @pytest.mark.asyncio
     @patch("cpex.framework.isolated.worker.get_proper_config")
@@ -180,9 +181,9 @@ async def test_process_task_load_and_run_hook_import_error(self, mock_import, mo
             "payload": {},
             "context": {"state": {}, "global_context": {}, "metadata": {}},
         }
-
+        tp = TaskProcessor()
         with pytest.raises(ImportError):
-            await process_task(task_data)
+            await process_task(task_data, tp)
 
     @pytest.mark.asyncio
     @patch("cpex.framework.isolated.worker.get_proper_config")
@@ -214,6 +215,7 @@ async def test_process_task_with_different_hook_types(self, mock_executor_class,
         mock_executor_class.return_value = mock_executor
 
         hook_types = ["tool_pre_invoke", "tool_post_invoke", "prompt_pre_fetch", "prompt_post_fetch"]
+        tp = TaskProcessor()
 
         for hook_type in hook_types:
             config_dict = {"name": "test_plugin", "kind": "isolated_venv"}
@@ -226,17 +228,16 @@ async def test_process_task_with_different_hook_types(self, mock_executor_class,
                 "payload": {},
                 "context": {"state": {}, "global_context": {"request_id": "req-123"}, "metadata": {}},
             }
-
-            result = await process_task(task_data)
+            result = await process_task(task_data, tp)
             assert result is not None
 
     @pytest.mark.asyncio
     async def test_process_task_unknown_task_type(self):
         """Test processing task with unknown task type."""
         task_data = {"task_type": "unknown_type"}
-
+        tp = TaskProcessor()
         # Should return None or handle gracefully
-        result = await process_task(task_data)
+        result = await process_task(task_data, tp)
         assert result is None
 
     @pytest.mark.asyncio
@@ -282,8 +283,9 @@ async def test_process_task_with_metadata(self, mock_executor_class, mock_import
                 "metadata": {"custom": "data"},
             },
         }
+        tp = TaskProcessor()
 
-        result = await process_task(task_data)
+        result = await process_task(task_data, tp)
 
         assert result is not None
         # Verify executor was called with proper context