嵌入搜索提供程序#

NeMo Guardrails 利用嵌入搜索(也称为向量数据库)来实现 guardrails 流程 以及 知识库 功能。

为了提高嵌入搜索过程的效率,NeMo Guardrails 可以采用嵌入缓存机制。此机制存储已计算的嵌入,从而减少重复计算的需要并加速搜索过程。默认情况下,缓存机制处于禁用状态。

默认嵌入搜索使用 FastEmbed 计算嵌入(all-MiniLM-L6-v2 模型),并使用 Annoy 执行搜索。默认配置如下

core:
  embedding_search_provider:
    name: default
    parameters:
      embedding_engine: FastEmbed
      embedding_model: all-MiniLM-L6-v2
      use_batching: False
      max_batch_size: 10
      max_batch_hold: 0.01
      search_threshold: None
    cache:
      enabled: False
      key_generator: md5
      store: filesystem
      store_config: {}

knowledge_base:
  embedding_search_provider:
    name: default
    parameters:
      embedding_engine: FastEmbed
      embedding_model: all-MiniLM-L6-v2
      use_batching: False
      max_batch_size: 10
      max_batch_hold: 0.01
      search_threshold: None
    cache:
      enabled: False
      key_generator: md5
      store: filesystem
      store_config: {}

默认嵌入搜索提供程序也可以与 OpenAI 嵌入一起使用

core:
  embedding_search_provider:
    name: default
    parameters:
      embedding_engine: openai
      embedding_model: text-embedding-ada-002
    cache:
      enabled: False
      key_generator: md5
      store: filesystem
      store_config: {}

knowledge_base:
  embedding_search_provider:
    name: default
    parameters:
      embedding_engine: openai
      embedding_model: text-embedding-ada-002
    cache:
      enabled: False
      key_generator: md5
      store: filesystem
      store_config: {}

默认实现还旨在支持嵌入计算过程的异步执行,从而提高搜索功能的效率。

cache 配置是可选的。如果启用,它会使用指定的 key_generatorstore 来缓存嵌入。store_config 可用于提供存储所需的其他配置选项。默认的 cache 配置使用 md5 密钥生成器和 filesystem 存储。默认情况下,缓存处于禁用状态。

批量实现#

默认嵌入提供程序包含批量处理功能,旨在优化嵌入生成过程。此功能旨在在预定义的 10 毫秒延迟后启动嵌入生成过程。

自定义嵌入搜索提供程序#

您可以通过子类化 EmbeddingsIndex 来实现您自己的自定义嵌入搜索提供程序。为了快速参考,完整接口包含在下方

class EmbeddingsIndex:
    """The embeddings index is responsible for computing and searching a set of embeddings."""

    @property
    def embedding_size(self):
        raise NotImplementedError

    @property
    def cache_config(self):
      raise NotImplementedError

    async def _get_embeddings(self, texts: List[str]):
        raise NotImplementedError

    async def add_item(self, item: IndexItem):
        """Adds a new item to the index."""
        raise NotImplementedError()

    async def add_items(self, items: List[IndexItem]):
        """Adds multiple items to the index."""
        raise NotImplementedError()

    async def build(self):
        """Build the index, after the items are added.

        This is optional, might not be needed for all implementations."""
        pass

    async def search(self, text: str, max_results: int) -> List[IndexItem]:
        """Searches the index for the closest matches to the provided text."""
        raise NotImplementedError()

@dataclass
class IndexItem:
    text: str
    meta: Dict = field(default_factory=dict)

为了使用您的自定义嵌入搜索提供程序,您必须在您的 config.py 中注册它

def init(app: LLMRails):
    app.register_embedding_search_provider("simple", SimpleEmbeddingSearchProvider)

有关完整示例,请查看此测试配置