Skip to main content

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

Adding a Custom LLM 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 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.

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

The file structure should look like this:

src
|__ providers
|   |__ llm
|   |__ 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 or DocumentStoreProvider.​

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

  1. The class should be inherited from LLMProvider and implemented the necessary methods
  2. We use the async version of the generator, text embedder, and document embedder classes; otherwise there will be a performance issue.
  3. Please make sure the provider name is the same as the file name.
  4. Please make sure you define several default variables for the provider such as EMBEDDING_MODEL_NAME, EMBEDDING_MODEL_DIMENSION, 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"},
}
EMBEDDING_MODEL_NAME = "text-embedding-3-large"
EMBEDDING_MODEL_DIMENSION = 3072

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

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

    def get_text_embedder(self):
        return AsyncTextEmbedder(
            ...
        )

    def get_document_embedder(self):
        return AsyncDocumentEmbedder(
            ...
        )

Setup the environment variables for the provider​

In the last step, update the env files to set the LLM_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, wren-ai-service/.env.prod.example, docker/.env.ai.example.

Running Wren AI with your Custom LLM or Document Store​

  1. Copy the content of .env.ai.example, rename the file name to .env.ai and put it to ~/.wrenai folder.
  2. 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, ollama, azure_openai, etc.
      • If you would like to use Ollama, please start Ollama and pull models first. You can refer to the Ollama documentation for more information.
      • Fill in necessary environment variables based on the provider you choose
      • Fill in GENERATION_MODEL, EMBEDDING_MODEL and EMBEDDING_MODEL_DIMENSION
  3. 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.