|
|
import os |
|
|
|
|
|
import pytest |
|
|
from injector import Injector |
|
|
|
|
|
from taskweaver.code_interpreter.code_generator.plugin_selection import PluginSelector |
|
|
from taskweaver.config.config_mgt import AppConfigSource |
|
|
from taskweaver.memory.plugin import PluginModule |
|
|
|
|
|
IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" |
|
|
|
|
|
|
|
|
@pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Test doesn't work in Github Actions.") |
|
|
def test_plugin_selector(): |
|
|
app_injector = Injector([PluginModule]) |
|
|
app_config = AppConfigSource( |
|
|
config={ |
|
|
"plugin.base_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "data/plugins"), |
|
|
"llm.embedding_api_type": "sentence_transformer", |
|
|
"llm.embedding_model": "all-mpnet-base-v2", |
|
|
"llm.api_key": "test_key", |
|
|
"code_generator.enable_auto_plugin_selection": True, |
|
|
}, |
|
|
) |
|
|
app_injector.binder.bind(AppConfigSource, to=app_config) |
|
|
plugin_selector = app_injector.get(PluginSelector) |
|
|
plugin_selector.generate_plugin_embeddings() |
|
|
|
|
|
query1 = "detect abnormal data points in ./data.csv." |
|
|
selected_plugins = plugin_selector.plugin_select(query1, top_k=3) |
|
|
assert any([p.name == "anomaly_detection" for p in selected_plugins]) |
|
|
assert len(selected_plugins) == 3 |
|
|
assert selected_plugins[0].name == "anomaly_detection" |
|
|
|
|
|
query2 = "summarize ./paper.pdf." |
|
|
selected_plugins = plugin_selector.plugin_select(query2, top_k=3) |
|
|
|
|
|
assert any([p.name == "paper_summary" for p in selected_plugins]) |
|
|
|
