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