| import os |
|
|
| from injector import Injector |
|
|
| from taskweaver.code_interpreter.code_generator.plugin_selection import SelectedPluginPool |
| from taskweaver.config.config_mgt import AppConfigSource |
| from taskweaver.logging import LoggingModule |
| from taskweaver.memory.plugin import PluginModule, PluginRegistry |
|
|
|
|
| def test_plugin_pool(): |
| app_injector = Injector( |
| [PluginModule, LoggingModule], |
| ) |
| app_config = AppConfigSource( |
| config={ |
| "app_dir": os.path.dirname(os.path.abspath(__file__)), |
| "llm.api_key": "this_is_not_a_real_key", |
| "plugin.base_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "data/plugins"), |
| }, |
| ) |
| app_injector.binder.bind(AppConfigSource, to=app_config) |
|
|
| plugin_registry = app_injector.get(PluginRegistry) |
|
|
| plugins = plugin_registry.get_list() |
|
|
| selected_plugin_pool = SelectedPluginPool() |
|
|
| selected_plugin_pool.add_selected_plugins(plugins[:1]) |
| assert len(selected_plugin_pool) == 1 |
|
|
| selected_plugin_pool.add_selected_plugins(plugins[:1]) |
| assert len(selected_plugin_pool) == 1 |
|
|
| selected_plugin_pool.add_selected_plugins(plugins[1:3]) |
| assert len(selected_plugin_pool) == 3 |
|
|
| selected_plugin_pool.add_selected_plugins(plugins[2:4]) |
| assert len(selected_plugin_pool) == 4 |
|
|
| selected_plugin_pool.filter_unused_plugins("xcxcxc anomaly_detection() ababab") |
| assert len(selected_plugin_pool) == 1 |
| assert selected_plugin_pool.get_plugins()[0].name == "anomaly_detection" |
|
|
| selected_plugin_pool.filter_unused_plugins("") |
| assert len(selected_plugin_pool) == 1 |
|
|
| selected_plugin_pool.add_selected_plugins(plugins[1:4]) |
| assert len(selected_plugin_pool) == 4 |
|
|
| selected_plugin_pool.filter_unused_plugins("abc sql_pull_data def") |
| assert len(selected_plugin_pool) == 2 |
|
|
| selected_plugin_pool.filter_unused_plugins("") |
| assert len(selected_plugin_pool) == 2 |
|
|
