Skip to main content

如何在Wren AI中使用您首选的LLM、嵌入器或文档存储

warning

我们强烈推荐使用OpenAI GPT-4o或GPT-4o-mini与Wren AI结合。这些模型已经经过广泛测试,以确保最佳性能和兼容性。

虽然从技术上讲可以集成其他AI模型,但请注意,它们尚未与我们的系统完全测试。因此,使用替代模型需自行承担风险,可能会导致意外行为或性能不佳。

使用您的自定义LLM、嵌入器或文档存储运行Wren AI

要使用自定义的LLM、Embedder或Document Store设置Wren AI,请按照以下步骤操作:

  1. 复制并重命名配置文件
    首先,您需要复制示例配置文件并重命名它。此文件将用于配置您的自定义提供程序。
  • 替换为您正在使用的Wren AI的版本号。

  • 对于MacOS或Linux用户:打开您的终端并运行以下命令:

    wget -O config.example.yaml https://raw.githubusercontent.com/canner/WrenAI//docker/config.example.yaml && \
    mkdir -p ~/.wrenai && cp config.example.yaml ~/.wrenai/config.yaml
    wget -O .env.example https://raw.githubusercontent.com/canner/WrenAI//docker/.env.example && \
    mkdir -p ~/.wrenai && cp .env.example ~/.wrenai/.env
  • 对于Windows用户: 打开PowerShell并执行以下命令:

    wget -O config.example.yaml https://raw.githubusercontent.com/canner/WrenAI//docker/config.example.yaml
    mkdir -p ~/.wrenai
    cp config.example.yaml ~/.wrenai/config.yaml
    notepad ~/.wrenai/config.yaml # 填写所需的配置
    wget -O .env.example https://raw.githubusercontent.com/canner/WrenAI//docker/.env.example
    mkdir -p ~/.wrenai
    cp .env.example ~/.wrenai/.env.example.txt
    notepad ~/.wrenai/.env.example.txt # 填写所需的配置
    mv ~/.wrenai/.env.example.txt ~/.wrenai/.env # 重命名文件
  1. 更新您的配置
    打开 ~/.wrenai/config.yaml 文件并更新它以匹配您的自定义 LLM、Embedder 或 Document Store 设置。您可以参考 部分 获取如何配置这些设置的指导。此外,您可能需要更新 .env 文件以填写所需的 API 密钥。
  • For custom LLM

    • We are now using LiteLLM to support LLMs, so basically you can use any LLMs supported by LiteLLM.

    • For example, if you want to use llama3.1:8b from LM Studio(For local LLMs, we recommend you use LM Studio for more consistent JSON output support)

      1. add the following configuration to your config.yaml under the litellm_llm section:
      type: llm
      provider: litellm_llm
      timeout: 120
      models:
      # omitted other model definitions
      - kwargs:
      n: 1
      temperature: 0
      response_format:
      type: json_object
      # please replace with your model name here, should be lm_studio/<MODEL_NAME>
      model: lm_studio/mlx-community/meta-llama-3.1-8b-instruct
      api_base: http://host.docker.internal:1234/v1
      api_key_name: LLM_LM_STUDIO_API_KEY
      1. add the following environment variable to the .env file in the ~/.wrenai directory:
      LLM_LM_STUDIO_API_KEY=random # just put a random string here, should not be empty
    • Please refer to the LiteLLM documentation for more details about each LLM's supported parameters.

      • Basically you need to fill in a new model configuration under the litellm_llm section. For example, what should be the model name, api_base, api_key_name, and keyword arguments, etc. accordding to the LiteLLM documentation. (Please refer to the example configuration above)
  • For custom Embedder

    • As of now, we only support embedding models from OpenAI, Azure OpenAI, OpenAI-compatible models, and Ollama.

    • For example, if you want to use nomic-embed-text from Ollama, add the following configuration to your config.yaml under the ollama_embedder section, also make sure the embedding_model_dim is set to the dimension of the embedding model under the document_store section:

      ---
      type: embedder
      provider: ollama_embedder
      models:
      - model: nomic-embed-text
      dimension: 768
      url: http://localhost:11434
      timeout: 120
      ---
      type: document_store
      provider: qdrant
      location: http://qdrant:6333
      embedding_model_dim: 768
      timeout: 120
      recreate_index: true
    • If you are using Ollama, please add EMBEDDER_OLLAMA_URL=http://host.docker.internal:11434 to the .env file in the ~/.wrenai directory.

  1. 启动 Wren AI
  • 通过运行启动器应用程序启动Wren AI,并从下拉菜单中选择“自定义”。
note

Ollama 集成:

  • 在桌面应用程序中运行Ollama:
    • 仅适用于Windows/MacOS用户。
    • ollama.com安装Ollama。
    • 启动Ollama桌面应用程序或在终端中运行ollama serve
    • 使用以下命令拉取所需的模型:ollama pull
    • config.yamlollama_embedder/ollama_llm部分设置url以指向您的Ollama服务器(默认:http://docker.host.internal:11434)。
  • 在Docker容器中运行Ollama:
    • 适用于Windows/MacOS/Linux用户。
    • 使用以下命令在Docker容器中运行Ollama:docker run -d --network wrenai_wren -v ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama
    • config.yamlollama_embedder/ollama_llm部分设置url以指向您的Ollama服务器(默认:http://ollama:11434)。

向Wren AI添加自定义LLM、嵌入器或文档存储

Wren AI 团队正在努力提升文本到 SQL 的用户体验和性能。然而,我们希望借助社区的力量来加快进展。我们热烈欢迎贡献者提供反馈、创建问题或提交拉取请求。

目前我们收到用户的反馈,他们希望使用他们偏好的LLM或文档存储。我们很高兴地宣布,我们已经使您能够将您偏好的LLM或文档存储添加到Wren AI中。

以下是您可以添加您偏好的LLM或文档存储并为此主题做出贡献的方法:

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

在Wren AI的背后,我们使用Haystack来提供LLM和文档存储功能。您可以找到支持的LLMs文档存储列表。

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

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

创建一个提供者定义文件 在 llm、embedder 或 document_store 包下

文件结构应如下所示:

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

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

创建一个继承自LLMProviderEmbedderProviderDocumentStoreProvider的类

以下是OpenAILLMProvider实现的示例,有几件事你需要考虑:

  1. 该类应继承自 LLMProvider 并实现 必要的方法
  2. 我们使用生成器类的异步版本;否则会出现性能问题。
  3. 请确保提供者名称与文件名相同,并以_llm作为后缀。
  4. 请确保您为提供者定义了几个默认变量,例如GENERATION_MODELGENERATION_MODEL_KWARGS等,并且这些变量也应在环境文件中定义。
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(
...
)

其他提供者如EmbedderProviderDocumentStoreProvider应遵循类似的模式。您可以查看官方实现作为参考这里

配置您的提供商

创建自定义提供者类后,您需要在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_idorganization_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

最后,在管道部分配置您的自定义提供者。此部分定义了不同组件(如LLMs和嵌入器)如何协同工作:

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

要更深入地了解如何配置自定义提供程序,包括实际示例和最佳实践,请查看我们的详细配置文档。此外,您还可以参考配置示例文件以获取全面的配置示例。