How to setup Wren AI using your Custom LLM or Embedder

warning

We highly recommend using OpenAI o3-mini, GPT-4o or GPT-4o-mini with Wren AI. These models have been extensively tested to ensure optimal performance and compatibility.

While it is technically possible to integrate other AI models, please note that they have not been fully tested with our system. Therefore, using alternative models is at your own risk and may result in unexpected behavior or suboptimal performance.

To set up Wren AI with your custom LLM, Embedder or Document Store, follow these steps:

Check if Wren AI already supports your preferred LLM or embedding models

Please first check if Wren AI already supports your preferred LLM or embedding models by checking LiteLLM documentation. Also you could check if Wren AI already supports the configuration examples here.

• If Wren AI already supports your preferred LLM or embedding models and provides the configuration example, you could try setting up the configuration by referring to the example and go to step 4 to launch Wren AI.
• If Wren AI already supports your preferred LLM or embedding models but does not provide the configuration example, please follow the steps below to set up the configuration. We definitely need your help to improve the documentation and add the configuration example. Welcome to contribute to Wren AI!
• If Wren AI does not support your preferred LLM or embedding models, please jump to the Adding a Custom LLM, Embedder or Document Store to Wren AI section to add your preferred LLM or embedding models to Wren AI.

Copy and Rename the Configuration Files

First, you need to copy the example configuration file and rename it. This file will be used to configure your custom provider.

• Replace <WRENAI_VERSION_NUMBER> with the version number of the Wren AI you are using.

• For MacOS or Linux Users: Open your terminal and run the following command:

  wget -O config.example.yaml https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml && \
  mkdir -p ~/.wrenai && cp config.example.yaml ~/.wrenai/config.yaml

  wget -O .env.example https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/.env.example && \
  mkdir -p ~/.wrenai && cp .env.example ~/.wrenai/.env

• For Windows Users: Open PowerShell and execute these commands:

  wget -O config.example.yaml https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
  mkdir -p ~/.wrenai
  cp config.example.yaml ~/.wrenai/config.yaml
  notepad ~/.wrenai/config.yaml  # Fill in required configurations

  wget -O .env.example https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/.env.example
  mkdir -p ~/.wrenai
  cp .env.example ~/.wrenai/.env.example.txt
  notepad ~/.wrenai/.env.example.txt  # Fill in required configurations
  mv ~/.wrenai/.env.example.txt ~/.wrenai/.env  # Rename the file

Update Your Configurations

Open the ~/.wrenai/config.yaml file and update it to match your custom LLM, Embedder, or Document Store settings. You can refer to the section for guidance on how to configure these settings. Also, you may need to update the .env file to fill in required API keys.

• 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 Ollama

    1. add the following configuration to your config.yaml under the litellm_llm section:

    type: llm
    provider: litellm_llm
    models:
    - api_base: http://host.docker.internal:11434/v1  # if you are using mac/windows, don't change this; if you are using linux, please search "Run Ollama in docker container" in this page: https://docs.getwren.ai/oss/installation/custom_llm#running-wren-ai-with-your-custom-llm-embedder
      model: ollama_chat/llama3.1:8b  # ollama_chat/<ollama_model_name>
      timeout: 600
      kwargs:
        n: 1
        temperature: 0

  • 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

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

  • For example, if you want to use nomic-embed-text from Ollama, add the following configuration to your config.yaml under the litellm_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: litellm_embedder
    models:
    - model: ollama/nomic-embed-text  # put your ollama embedder model name here, openai/<ollama_model_name>
      api_base: http://host.docker.internal:11434/v1  # if you are using mac/windows, don't change this; if you are using linux, please search "Run Ollama in docker container" in this page: https://docs.getwren.ai/oss/installation/custom_llm#running-wren-ai-with-your-custom-llm-embedder
      timeout: 600
    ---
    type: document_store
    provider: qdrant
    location: http://qdrant:6333
    embedding_model_dim: 768  # put your embedding model dimension here
    timeout: 120
    recreate_index: true

• After updating llm and embedder configurations, you need to replace corresponding llm and embedder values in the pipeline section. Be careful that the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.ollama_chat/llama3.1:8b and litellm_llm.ollama/nomic-embed-text.

Launch Wren AI

Launch Wren AI by running the launcher app and select the "Custom" option from the dropdown menu.

note

For Ollama Integration:

• Run Ollama in deskstop application:
  • Only for Windows/MacOS users.
  • Install Ollama from ollama.com.
  • Start the Ollama desktop application or run ollama serve in your terminal.
  • Pull your desired model using the command: ollama pull <model_name>.
  • Set the url in the ollama_embedder/ollama_llm section of config.yaml to point to your Ollama server (default: http://docker.host.internal:11434).
• Run Ollama in docker container:
  • For Windows/MacOS/Linux users.
  • Run Ollama in docker container using the following command: docker run -d --network wrenai_wren -v ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama.
  • Set the url in the ollama_embedder/ollama_llm section of config.yaml to point to your Ollama server (default: http://ollama:11434).