如何在Wren AI中使用您首选的LLM、嵌入器或文档存储
我们强烈推荐使用OpenAI GPT-4o或GPT-4o-mini与Wren AI结合。这些模型已经经过广泛测试,以确保最佳性能和兼容性。
虽然从技术上讲可以集成其他AI模型,但请注意,它们尚未与我们的系统完全测试。因此,使用替代模型需自行承担风险,可能会导致意外行为或性能不佳。
使用您的自定义LLM、嵌入器或文档存储运行Wren AI
要使用自定义的LLM、Embedder或Document Store设置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.yamlwget -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 # 重命名文件
- 更新您的配置
打开~/.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)- add the following configuration to your
config.yaml
under thelitellm_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- 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
- add the following configuration to your
-
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)
- Basically you need to fill in a new model configuration under the
-
-
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 yourconfig.yaml
under theollama_embedder
section, also make sure theembedding_model_dim
is set to the dimension of the embedding model under thedocument_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.
-
- 启动 Wren AI
- 通过运行启动器应用程序启动Wren AI,并从下拉菜单中选择“自定义”。
Ollama 集成:
- 在桌面应用程序中运行Ollama:
- 仅适用于Windows/MacOS用户。
- 从ollama.com安装Ollama。
- 启动Ollama桌面应用程序或在终端中运行
ollama serve
。 - 使用以下命令拉取所需的模型:
ollama pull
。 - 在
config.yaml
的ollama_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.yaml
的ollama_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
的新文件。
创建一个继承自LLMProvider
、EmbedderProvider
或DocumentStoreProvider
的类
以下是OpenAILLMProvider
实现的示例,有几件事你需要考虑:
- 该类应继承自
LLMProvider
并实现 必要的方法 - 我们使用生成器类的异步版本;否则会出现性能问题。
- 请确保提供者名称与文件名相同,并以
_llm
作为后缀。 - 请确保您为提供者定义了几个默认变量,例如
GENERATION_MODEL
、GENERATION_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(
...
)
其他提供者如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
最后,在管道部分配置您的自定义提供者。此部分定义了不同组件(如LLMs和嵌入器)如何协同工作:
type: pipeline
pipes:
- name: pipeline_name
llm: custom_llm_name.model_name
embedder: custom_embedder_name.model_name
# other pipeline configurations
要更深入地了解如何配置自定义提供程序,包括实际示例和最佳实践,请查看我们的详细配置文档。此外,您还可以参考配置示例文件以获取全面的配置示例。