How to use your preferred LLM, Embedder or Document Store in Wren AI

Adding a Custom LLM, Embedder or Document Store to Wren AI

Wren AI team is working hard to make text-to-sql user experience and performance better. However, we hope to leverage the power of the community to make faster progress. We warmly welcome contributors to provide feedback, create issues, or open pull requests.

Currently we receive feedbacks from our users that they want to use their preferred LLM or Document Store. We are happy to announce that we have made it possible for you to add your preferred LLM or Document Store to Wren AI.

Here is how you can add your preferred LLM or Document Store and contribute on this topic:

Decide on the LLM, Embedder or Document Store you would like to add

Underneath Wren AI, we are using Haystack to provide the LLM and Document Store functionalities. You can find the list of supported LLMs and Document Stores.

Haystack provides a wide range of LLMs and Document Stores, and it has simple APIs and great developer experience that we can easily add custom components to Wren AI.

For embedders, please make sure it is supported by the Document Store you choose. For example, these are embedding models supported by Qdrant. Also, you can refer to the supported LLMs from Haystack to check if the corresponding embedder is supported.

Create a provider definition file under the llm, embedder or document_store package.

The file structure should look like this:

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

For example, if you would like to add Mistral as a new LLM provider, you might add a new file called mistral.py under the llm package.

Create a class that inherits from `LLMProvider`, `EmbedderProvider` or `DocumentStoreProvider`.

Below is an example of OpenAILLMProvider implementation, and there are several things you need to consider:

The class should be inherited from LLMProvider and implemented the necessary methods
We use the async version of the generator class; otherwise there will be a performance issue.
Please make sure the provider name is the same as the file name with _llm as suffix.
Please make sure you define several default variables for the provider such as GENERATION_MODEL, GENERATION_MODEL_KWARGS, etc., and these variables should also be defined in the env files.

OPENAI_API_BASE = "https://api.openai.com/v1"
GENERATION_MODEL_NAME = "gpt-3.5-turbo"
GENERATION_MODEL_KWARGS = {
    "temperature": 0,
    "n": 1,
    "max_tokens": 4096,
    "response_format": {"type": "json_object"},
}

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

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

Other providers such as EmbedderProvider and DocumentStoreProvider should follow the similar pattern. You can check out the official implementations for reference here.

Setup the environment variables for the provider

In the last step, update the env files to set the LLM_PROVIDER, EMBEDDING_PROVIDER or DOCUMENT_STORE_PROVIDER variable to your specified provider. Additionally, add any custom variables for your provider into the env files, namely wren-ai-service/.env.dev.example, docker/.env.ai.example.

Running Wren AI with your Custom LLM or Document Store

Copy the content of .env.ai.example, rename the file name to .env.ai and put it to ~/.wrenai folder.
- Using MacOS or Linux with terminal
```
wget -O .env.ai.example https://raw.githubusercontent.com/canner/WrenAI/main/docker/.env.ai.example && \
mkdir -p ~/.wrenai && cp .env.ai.example ~/.wrenai/.env.ai
```
- Using Windows with PowerShell
  - wget -O .env.ai.example https://raw.githubusercontent.com/canner/WrenAI/main/docker/.env.ai.example
  - mkdir -p ~/.wrenai
  - cp .env.ai.example ~/.wrenai/.env.ai.txt
  - Fill in required configurations: notepad ~/.wrenai/.env.ai.txt
  - Rename the file: mv ~/.wrenai/.env.ai.txt ~/.wrenai/.env.ai
Update any variables in the .env.ai file to match your custom LLM or Document Store.
- To use your preferred LLM
  - set LLM_PROVIDER to the name of your custom LLM provider, e.g. openai_llm, ollama_llm, azure_openai_llm, etc.
  - Fill in necessary environment variables based on the provider you choose
  - Fill in GENERATION_MODEL, GENERATION_MODEL_KWARGS
- To use your preferred Embedder
  - set EMBEDDER_PROVIDER to the name of your custom Embedder provider, e.g. openai_embedder, ollama_embedder, azure_openai_embedder, etc.
  - Fill in necessary environment variables based on the provider you choose
  - Fill in EMBEDDING_MODEL, EMBEDDING_MODEL_DIMENSION
- NOTES: If you would like to use Ollama, you can refer to the Ollama documentation for more information if needed.
  - Install Ollama
  - Start Ollama desktop application or run ollama serve in the terminal
  - Pull a model by running ollama pull <model_name>
Launch Wren AI by clicking the launcher application.

warning

If you use a custom LLM, you may need an LLM that is at least as powerful as the OpenAI GPT-3.5-turbo model; for example, meta-llama/Meta-Llama-3-70B. Otherwise, the performance may be affected.

Currently the prompts are designed based on OpenAI's LLMs, so you may need to adjust the prompts to fit your custom LLM.

How to use your preferred LLM, Embedder or Document Store in Wren AI

Adding a Custom LLM, Embedder or Document Store to Wren AI​

Decide on the LLM, Embedder or Document Store you would like to add​

Create a provider definition file under the llm, embedder or document_store package.​

Create a class that inherits from LLMProvider, EmbedderProvider or DocumentStoreProvider.​

Setup the environment variables for the provider​

Running Wren AI with your Custom LLM or Document Store​