跳到主要内容

如何在 Wren AI 中自定义 LLM、嵌入器或文档存储

决定您想添加的 LLM、嵌入器或文档存储​

在 Wren AI 内部，我们使用 Haystack 和 LiteLLM 来提供 LLM、嵌入模型和文档存储功能。您可以在此处找到支持的LLM和文档存储列表。

Haystack 提供了广泛的 LLM 和文档存储支持，并且拥有简单的 API 和出色的开发者体验，我们可以轻松地向 Wren AI 添加自定义组件。

对于嵌入器，请确保您选择的文档存储支持它。例如，这些是Qdrant 支持的嵌入模型。此外，您可以参考 Haystack 支持的LLM来检查是否支持相应的嵌入器。

在 llm、embedder 或 document_store 包下创建提供者定义文件（在此查看）​

文件结构应如下所示

src
|__ providers
|   |__ llm
|   |__ embedder
|   |__ document_store

例如，如果您想添加Mistral作为新的 LLM 提供者，您可以在llm包下添加一个名为mistral.py的新文件。

创建一个继承自 LLMProvider、EmbedderProvider 或 DocumentStoreProvider 的类​

以下是OpenAILLMProvider 实现的示例，其中有几点需要考虑

  1. 该类应继承自 LLMProvider 并实现必要的方法
  2. 我们使用生成器类的异步版本；否则可能会有性能问题。
  3. 请确保提供者名称与文件名相同，并以_llm作为后缀。
  4. 请确保您为提供者定义了几个默认变量，例如GENERATION_MODEL、GENERATION_MODEL_KWARGS等，这些变量也应在 env 文件中定义。
OPENAI_API_BASE = "https://api.openai.com/v1"
GENERATION_MODEL_NAME = "gpt-4o-mini"
GENERATION_MODEL_KWARGS = {
    "temperature": 0,
    "n": 1,
    "max_tokens": 4096,
    "response_format": {"type": "json_object"},
}

@provider("openai_llm")
class OpenAILLMProvider(LLMProvider):
    def __init__(
        ...
    ):
        ...

    def get_generator(
        ...
    ):
        return AsyncGenerator(
            ...
        )

其他提供者（如EmbedderProvider和DocumentStoreProvider）应遵循类似的模式。您可以在此处查看官方实现作为参考。

配置您的提供者​

创建自定义提供者类后，您需要在wren-ai-service/config.yaml文件中对其进行配置。此配置文件对于告知 Wren AI 如何与您的提供者及其模型交互至关重要。

对于 LLM 提供者，添加一个具有以下结构的配置块。这定义了如何初始化和使用您的自定义 LLM 提供者

type: llm
provider: custom_llm_name
models:
  - model: model_name
    kwargs:
      temperature: 0
      max_tokens: 4096
      # other model-specific parameters
api_base: api_endpoint
# other provider-specific configurations

请注意，提供者特定的配置是可选的，取决于您的实现。配置中的参数名称必须与您提供者的构造函数中的参数名称匹配。例如，如果您的构造函数接受project_id和organization_id，您需要在 YAML 文件中配置相同的名称。

对于嵌入器提供者，请使用此结构进行配置

type: embedder
provider: custom_embedder_name
models:
  - model: model_name
    dimension: 1536 # specify your embedding dimension
api_base: api_endpoint
timeout: 30 # optional timeout in seconds

最后，在 pipeline 部分配置您的自定义提供者。此部分定义了 LLM 和嵌入器等不同组件如何协同工作

type: pipeline
pipes:
  - name: pipeline_name
    llm: custom_llm_name.model_name
    embedder: custom_embedder_name.model_name
    # other pipeline configurations

要更深入地了解如何配置自定义提供者，包括实际示例和最佳实践，请查阅我们详细的配置文档。此外，有关全面的配置示例，您可以参考配置示例文件。