feat(wren-ai-service): add litellm embedder

.github/workflows/ai-service-test.yaml

@@ -42,7 +42,7 @@ jobs:
           python-version-file: ./wren-ai-service/pyproject.toml
           cache: "poetry"
       - name: Install the project dependencies
-        run: poetry install
+        run: poetry install --without eval
       - name: Install Just
         uses: extractions/setup-just@v2
         with:

wren-ai-service/docs/config_examples/config.azure.yaml

@@ -1,12 +1,11 @@
 type: llm
-provider: litellm_llm  # litellm supports Azure through its provider
+provider: litellm_llm
 timeout: 120
 models:
+# put AZURE_API_KEY=<your_api_key> in ~/.wrenai/.env
 - model: azure/gpt-4     # Your Azure deployment name, put 'azure/' before deployment name
-  api_base: https://endpoint.openai.azure.com/      #Replace with your custom Azure endpoint
-  api_key_name: LLM_AZURE_OPENAI_API_KEY
+  api_base: https://endpoint.openai.azure.com  # Replace with your custom Azure endpoint
   api_version: 2024-02-15-preview
-
   kwargs:
     temperature: 0
     n: 1
@@ -16,14 +15,13 @@ models:
 
 ---
 type: embedder
-provider: azure_openai_embedder  
+provider: litellm_embedder  
 models:
-  - model: text-embedding-ada-002  # Your Azure deployment name
-      # Must match model output check for your model
-
-api_base: https://endpoint.openai.azure.com/     # Replace with your custom Azure endpoint
-api_version: 2023-05-15  # Your Azure deployment name
-timeout: 300
+# put AZURE_API_KEY=<your_api_key> in ~/.wrenai/.env
+- model: azure/text-embedding-ada-002  # Your Azure deployment name, put 'azure/' before deployment name
+  api_base: https://endpoint.openai.azure.com  # Replace with your custom Azure endpoint
+  api_version: 2023-05-15
+  timeout: 300
 
 ---
 type: engine
@@ -32,36 +30,34 @@ endpoint: http://wren-ui:3000
 
 ---
 type: document_store
-#name: qdrant
 provider: qdrant
 location: http://qdrant:6333  # Donot set the QDRANT_API_KEY if you are using the qdrant from docker
-
 embedding_model_dim: 1536  # Must match model dimension from embedder
 timeout: 120
 recreate_index: true
 
-# For each pipe line component 
-# Replace llm with Azure deployed LLM model
-# Replace Embeddings with Azure deployed Embedding model
 ---
+# please change the llm and embedder names to the ones you want to use
+# the format of llm and embedder should be <provider>.<model_name> such as litellm_llm.gpt-4o-2024-08-06
+# the pipes may be not the latest version, please refer to the latest version: https://raw.githubusercontent.com/canner/WrenAI/<WRENAI_VERSION_NUMBER>/docker/config.example.yaml
 type: pipeline
 pipes:
   - name: db_schema_indexing
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     document_store: qdrant  # Match document_store name
     llm: litellm_llm.azure/gpt-4
   - name: historical_question_indexing
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     document_store: qdrant
   - name: table_description_indexing
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     document_store: qdrant
   - name: db_schema_retrieval
     llm: litellm_llm.azure/gpt-4
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     document_store: qdrant
   - name: historical_question_retrieval
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     document_store: qdrant
   - name: sql_generation
     llm: litellm_llm.azure/gpt-4
@@ -97,20 +93,20 @@ pipes:
     llm: litellm_llm.azure/gpt-4
   - name: intent_classification
     llm: litellm_llm.azure/gpt-4
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     document_store: qdrant
   - name: data_assistance
     llm: litellm_llm.azure/gpt-4
   - name: sql_pairs_preparation
     document_store: qdrant
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     llm: litellm_llm.azure/gpt-4
   - name: sql_pairs_deletion
     document_store: qdrant
-    embedder: azure_openai_embedder.text-embedding-ada-002 
+    embedder: litellm_embedder.azure/text-embedding-ada-002 
   - name: sql_pairs_retrieval
     document_store: qdrant
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     llm: litellm_llm.azure/gpt-4
   - name: preprocess_sql_data
     llm: litellm_llm.azure/gpt-4
@@ -122,12 +118,12 @@ pipes:
     llm: litellm_llm.azure/gpt-4
   - name: sql_pairs_indexing
     document_store: qdrant
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
   - name: sql_generation_reasoning
     llm: litellm_llm.azure/gpt-4
   - name: question_recommendation_db_schema_retrieval
     llm: litellm_llm.azure/gpt-4
-    embedder: azure_openai_embedder.text-embedding-ada-002
+    embedder: litellm_embedder.azure/text-embedding-ada-002
     document_store: qdrant
   - name: question_recommendation_sql_generation
     llm: litellm_llm.azure/gpt-4

wren-ai-service/docs/config_examples/config.deepseek.yaml

@@ -30,12 +30,13 @@ models:
 
 ---
 type: embedder
-provider: openai_embedder
+provider: litellm_embedder
 models:
-  # find EMBEDDER_OPENAI_API_KEY and fill in value of api key in ~/.wrenai/.env
-  - model: text-embedding-3-large  # put your openai compatible embedder model name here
-url: https://api.openai.com/v1  # change this according to your openai compatible embedder model
-timeout: 120
+  # define OPENAI_API_KEY=<api_key> in ~/.wrenai/.env if you are using openai embedding model
+  # please refer to LiteLLM documentation for more details: https://docs.litellm.ai/docs/providers
+  - model: text-embedding-3-large  # put your embedding model name here, if it is not openai embedding model, should be <provider>/<model_name>
+    api_base: https://api.openai.com/v1  # change this according to your embedding model
+    timeout: 120
 
 ---
 type: engine
@@ -57,20 +58,20 @@ recreate_index: false
 type: pipeline
 pipes:
   - name: db_schema_indexing
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: historical_question_indexing
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: table_description_indexing
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: db_schema_retrieval
     llm: litellm_llm.deepseek/deepseek-coder
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: historical_question_retrieval
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: sql_generation
     llm: litellm_llm.deepseek/deepseek-coder
@@ -106,7 +107,7 @@ pipes:
     llm: litellm_llm.deepseek/deepseek-coder
   - name: question_recommendation_db_schema_retrieval
     llm: litellm_llm.deepseek/deepseek-coder
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: question_recommendation_sql_generation
     llm: litellm_llm.deepseek/deepseek-coder
@@ -117,19 +118,19 @@ pipes:
     llm: litellm_llm.deepseek/deepseek-coder
   - name: intent_classification
     llm: litellm_llm.deepseek/deepseek-coder
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: data_assistance
     llm: litellm_llm.deepseek/deepseek-chat
   - name: sql_pairs_indexing
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
   - name: sql_pairs_deletion
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-3-large 
+    embedder: litellm_embedder.text-embedding-3-large 
   - name: sql_pairs_retrieval
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     llm: litellm_llm.deepseek/deepseek-coder
   - name: preprocess_sql_data
     llm: litellm_llm.deepseek/deepseek-coder

wren-ai-service/docs/config_examples/config.google_ai_studio.yaml

@@ -13,12 +13,12 @@ models:
 
 ---
 type: embedder
-provider: openai_embedder
+provider: litellm_embedder
 models:
-  # find EMBEDDER_OPENAI_API_KEY and fill in value of api key in ~/.wrenai/.env
-  - model: text-embedding-004  # put your openai compatible embedder model name here
-url: https://generativelanguage.googleapis.com/v1beta/openai  # change this according to your openai compatible embedder model
-timeout: 120
+  # put GEMINI_API_KEY=<your_api_key> in ~/.wrenai/.env
+  - model: gemini/text-embedding-004  # gemini/<gemini_model_name>
+    api_base: https://generativelanguage.googleapis.com/v1beta/openai  # change this according to your embedding model
+    timeout: 120
 
 ---
 type: engine
@@ -40,20 +40,20 @@ recreate_index: false
 type: pipeline
 pipes:
   - name: db_schema_indexing
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
     document_store: qdrant
   - name: historical_question_indexing
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
     document_store: qdrant
   - name: table_description_indexing
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
     document_store: qdrant
   - name: db_schema_retrieval
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
     document_store: qdrant
   - name: historical_question_retrieval
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
     document_store: qdrant
   - name: sql_generation
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
@@ -89,7 +89,7 @@ pipes:
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
   - name: question_recommendation_db_schema_retrieval
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
     document_store: qdrant
   - name: question_recommendation_sql_generation
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
@@ -100,19 +100,19 @@ pipes:
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
   - name: intent_classification
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
     document_store: qdrant
   - name: data_assistance
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
   - name: sql_pairs_indexing
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
   - name: sql_pairs_deletion
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-004 
+    embedder: litellm_embedder.text-embedding-004 
   - name: sql_pairs_retrieval
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-004
+    embedder: litellm_embedder.text-embedding-004
     llm: litellm_llm.gemini/gemini-2.0-flash-exp
   - name: preprocess_sql_data
     llm: litellm_llm.gemini/gemini-2.0-flash-exp

wren-ai-service/docs/config_examples/config.groq.yaml

@@ -14,12 +14,13 @@ models:
 
 ---
 type: embedder
-provider: openai_embedder
+provider: litellm_embedder
 models:
-  # find EMBEDDER_OPENAI_API_KEY and fill in value of api key in ~/.wrenai/.env
-  - model: text-embedding-3-large  # put your openai compatible embedder model name here
-url: https://api.openai.com/v1  # change this according to your openai compatible embedder model
-timeout: 120
+  # define OPENAI_API_KEY=<api_key> in ~/.wrenai/.env if you are using openai embedding model
+  # please refer to LiteLLM documentation for more details: https://docs.litellm.ai/docs/providers
+  - model: text-embedding-3-large  # put your embedding model name here, if it is not openai embedding model, should be <provider>/<model_name>
+    api_base: https://api.openai.com/v1  # change this according to your embedding model
+    timeout: 120
 
 ---
 type: engine
@@ -41,20 +42,20 @@ recreate_index: false
 type: pipeline
 pipes:
   - name: db_schema_indexing
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: historical_question_indexing
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: table_description_indexing
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: db_schema_retrieval
     llm: litellm_llm.groq/llama-3.3-70b-specdec
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: historical_question_retrieval
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: sql_generation
     llm: litellm_llm.groq/llama-3.3-70b-specdec
@@ -90,7 +91,7 @@ pipes:
     llm: litellm_llm.groq/llama-3.3-70b-specdec
   - name: question_recommendation_db_schema_retrieval
     llm: litellm_llm.groq/llama-3.3-70b-specdec
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: question_recommendation_sql_generation
     llm: litellm_llm.groq/llama-3.3-70b-specdec
@@ -101,19 +102,19 @@ pipes:
     llm: litellm_llm.groq/llama-3.3-70b-specdec
   - name: intent_classification
     llm: litellm_llm.groq/llama-3.3-70b-specdec
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     document_store: qdrant
   - name: data_assistance
     llm: litellm_llm.groq/llama-3.3-70b-specdec
   - name: sql_pairs_indexing
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
   - name: sql_pairs_deletion
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-3-large 
+    embedder: litellm_embedder.text-embedding-3-large 
   - name: sql_pairs_retrieval
     document_store: qdrant
-    embedder: openai_embedder.text-embedding-3-large
+    embedder: litellm_embedder.text-embedding-3-large
     llm: litellm_llm.groq/llama-3.3-70b-specdec
   - name: preprocess_sql_data
     llm: litellm_llm.groq/llama-3.3-70b-specdec