| from typing import Dict, List |
|
|
| import numpy as np |
| from injector import inject |
| from sklearn.metrics.pairwise import cosine_similarity |
|
|
| from taskweaver.llm import LLMApi |
| from taskweaver.memory.plugin import PluginEntry, PluginRegistry |
|
|
|
|
| class SelectedPluginPool: |
| def __init__(self): |
| self.selected_plugin_pool = [] |
| self._previous_used_plugin_cache = [] |
|
|
| def add_selected_plugins(self, external_plugin_pool: List[PluginEntry]): |
| """ |
| Add selected plugins to the pool |
| """ |
| self.selected_plugin_pool = self.merge_plugin_pool(self.selected_plugin_pool, external_plugin_pool) |
|
|
| def __len__(self) -> int: |
| return len(self.selected_plugin_pool) |
|
|
| def filter_unused_plugins(self, code: str): |
| """ |
| Filter out plugins that are not used in the code generated by LLM |
| """ |
| plugins_used_in_code = [p for p in self.selected_plugin_pool if p.name in code] |
| self._previous_used_plugin_cache = self.merge_plugin_pool( |
| self._previous_used_plugin_cache, |
| plugins_used_in_code, |
| ) |
| self.selected_plugin_pool = self._previous_used_plugin_cache |
|
|
| def get_plugins(self) -> List[PluginEntry]: |
| return self.selected_plugin_pool |
|
|
| @staticmethod |
| def merge_plugin_pool(pool1: List[PluginEntry], pool2: List[PluginEntry]) -> List[PluginEntry]: |
| """ |
| Merge two plugin pools and remove duplicates |
| """ |
| merged_list = pool1 + pool2 |
| result = [] |
|
|
| for item in merged_list: |
| is_duplicate = False |
| for existing_item in result: |
| if item.name == existing_item.name: |
| is_duplicate = True |
| break |
| if not is_duplicate: |
| result.append(item) |
| return result |
|
|
|
|
| class PluginSelector: |
| @inject |
| def __init__( |
| self, |
| plugin_registry: PluginRegistry, |
| llm_api: LLMApi, |
| plugin_only: bool = False, |
| ): |
| if plugin_only: |
| self.available_plugins = [p for p in plugin_registry.get_list() if p.plugin_only is True] |
| else: |
| self.available_plugins = plugin_registry.get_list() |
| self.llm_api = llm_api |
| self.plugin_embedding_dict: Dict[str, List[float]] = {} |
|
|
| def generate_plugin_embeddings(self): |
| plugin_intro_text_list: List[str] = [] |
| for p in self.available_plugins: |
| plugin_intro_text_list.append(p.name + ": " + p.spec.description) |
| plugin_embeddings = self.llm_api.get_embedding_list(plugin_intro_text_list) |
| for i, p in enumerate(self.available_plugins): |
| self.plugin_embedding_dict[p.name] = plugin_embeddings[i] |
|
|
| def plugin_select(self, user_query: str, top_k: int = 5) -> List[PluginEntry]: |
| user_query_embedding = np.array(self.llm_api.get_embedding(user_query)) |
|
|
| similarities = [] |
|
|
| if top_k >= len(self.available_plugins): |
| return self.available_plugins |
|
|
| for p in self.available_plugins: |
| similarity = cosine_similarity( |
| user_query_embedding.reshape( |
| 1, |
| -1, |
| ), |
| np.array(self.plugin_embedding_dict[p.name]).reshape(1, -1), |
| ) |
| similarities.append((p, similarity)) |
|
|
| plugins_rank = sorted( |
| similarities, |
| key=lambda x: x[1], |
| reverse=True, |
| )[:top_k] |
|
|
| selected_plugins = [p for p, sim in plugins_rank] |
|
|
| return selected_plugins |
|
|
