Embedding 搜索提供程序

NeMo Guardrails 利用 embedding 搜索（也称为向量数据库）来实现guardrails 流程以及知识库功能。

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

默认的 embedding 搜索使用 FastEmbed 计算 embeddings（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: sha256
      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: sha256
      store: filesystem
      store_config: {}

默认的 embedding 搜索提供程序也可以与 OpenAI embeddings 一起使用

core:
  embedding_search_provider:
    name: default
    parameters:
      embedding_engine: openai
      embedding_model: text-embedding-ada-002
    cache:
      enabled: False
      key_generator: sha256
      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: sha256
      store: filesystem
      store_config: {}

默认实现还设计用于支持 embedding 计算过程的异步执行，从而提高搜索功能的效率。

cache 配置是可选的。如果启用，它将使用指定的 key_generator 和 store 来缓存 embeddings。store_config 可用于提供存储所需的额外配置选项。默认的 cache 配置使用 sha256 键生成器和 filesystem 存储。默认情况下，缓存处于禁用状态。

批处理实现

默认的 embedding 提供程序包含一个批处理功能，旨在优化 embedding 生成过程。此功能设计用于在预定义的 10 毫秒延迟后启动 embedding 生成过程。

自定义 Embedding 搜索提供程序

您可以通过继承 EmbeddingsIndex 类来实现自己的自定义 embedding 搜索提供程序。为了快速参考，完整的接口包含在下面

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)

为了使用您的自定义 embedding 搜索提供程序，您需要在 config.py 中注册它

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

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