File size: 1,951 Bytes
3d3d712 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
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", # pragma: allowlist secret
"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
|