Skip to main content
ClaudeChatGPT

How to Set Up 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 whether Wren AI already supports your preferred LLM or embedding models in the LiteLLM documentation. You can also check whether Wren AI already provides a configuration example in the configuration examples here.

  • If Wren AI already supports your preferred LLM or embedding models and provides a configuration example, you can set up the configuration by referring to the example and then go to step 4 to launch Wren AI.
  • If Wren AI already supports your preferred LLM or embedding models but does not provide a configuration example, please follow the steps below to set up the configuration. We need your help to improve the documentation and add configuration examples — contributions to Wren AI are welcome!
  • 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 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

Configure Your Provider​

Open the ~/.wrenai/config.yaml file and update it to match your custom LLM, Embedder, or Document Store settings. Use the examples below as a starting point. You may also need to update the .env file with required API keys.

  • For custom LLM

    • We are now using LiteLLM to support LLMs, so you can use any LLM 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.

      • You need to fill in a new model configuration under the litellm_llm section — for example, the model name, api_base, api_key_name, keyword arguments, etc. — according to the LiteLLM documentation. (Please refer to the example configuration above.)

  • For custom Embedder

    • We are now using LiteLLM to support embedding models, so you can use any embedding model 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 embedding_model_dim under the document_store section is set to the dimension of the embedding model:

      ---
      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 the LLM and embedder configurations, you need to replace the corresponding llm and embedder values in the pipeline section. Note that the format of llm and embedder must 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 selecting the "Custom" option from the dropdown menu.

note

For Ollama Integration:

  • Run Ollama as a desktop 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 a Docker container:
    • For Windows/MacOS/Linux users.
    • Run Ollama in a 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).