opensource/WrenAI
3
0
Fork 0
mirror of https://github.com/Canner/WrenAI.git synced 2026-09-01 15:34:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(wren-ai-service): multi-llm deployment and setting loading mechanism (#792)

Browse Source 
Co-authored-by: ChihYu Yeh <chihyu.jimmy.yeh@gmail.com>
Co-authored-by: Ted Yan <tedyyan@hotmail.com>
Co-authored-by: Teddy Yan <xuan.yan@nokia.com>
This commit is contained in:
Pao Sheng
2024-11-15 11:02:21 +08:00
parent 15db25265f
commit 0f1d8d857c
44 changed files with 1729 additions and 615 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ai-service-test.yaml
+1 -1
View File
@@ -46,7 +46,7 @@ jobs:
- name: Install Just
uses: extractions/setup-just@v2
with:
just-version: "1.31.0"
just-version: "1.36.0"
- name: Prepare testing environment and Run tests
run: |
just test
.gitignore
+2 -1
View File
@@ -26,7 +26,8 @@ wren-ai-service/demo/poetry.lock
wren-ai-service/demo/custom_dataset
wren-ai-service/demo/.env
wren-ai-service/tools/dev/etc/**
wren-ai-service/.deepeval-cache.json
.deepeval-cache.json
docker/config.yaml
# python
.python-version
deployment/kustomizations/base/cm.yaml
+117 -15
View File
@@ -3,12 +3,10 @@ kind: ConfigMap
metadata:
name: wren-config
data:
# Wren Engine Service Port
WREN_ENGINE_PORT: "8080"
# Wren AI Service Port
WREN_AI_SERVICE_PORT: "5555"
WREN_AI_SERVICE_ENABLE_TIMER: "1"
WREN_UI_ENDPOINT: http://wren-ui-svc:3000
@@ -18,25 +16,14 @@ data:
WREN_AI_SERVICE_VERSION: "0.10.4"
WREN_UI_VERSION: "0.16.0"
# LLM and Embedder Configurations
LLM_PROVIDER: "openai_llm"
LLM_OPENAI_API_BASE: "https://api.openai.com/v1"
GENERATION_MODEL: "gpt-4o-mini"
EMBEDDER_PROVIDER: "openai_embedder"
EMBEDDER_OPENAI_API_BASE: "https://api.openai.com/v1"
# Document store related
QDRANT_HOST: "wren-qdrant"
DOCUMENT_STORE_PROVIDER: "qdrant"
# Langfuse: for LLM tracing
LANGFUSE_HOST: "https://cloud.langfuse.com"
LANGFUSE_ENABLE: ""
# Telemetry
POSTHOG_HOST: "https://app.posthog.com"
TELEMETRY_ENABLED: "false"
# this is for telemetry to know the model, i think ai-service might be able to provide a endpoint to get the information
GENERATION_MODEL: "gpt-4o-mini"
# service endpoints of AI service & engine service
WREN_ENGINE_ENDPOINT: "http://wren-engine-svc:8080"
@@ -55,3 +42,118 @@ data:
LOGGING_LEVEL: INFO
IBIS_SERVER_ENDPOINT: http://wren-ibis-server-svc:8000
---
apiVersion: v1
kind: ConfigMap
metadata:
name: wren-ai-service-config
data:
config.yaml: |
type: llm
provider: openai_llm
models:
- model: gpt-4o-mini
kwargs:
{
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
- model: gpt-4o
kwargs:
{
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
api_base: https://api.openai.com/v1
timeout: 120
---
type: embedder
provider: openai_embedder
models:
- model: text-embedding-3-large
dimension: 3072
api_base: https://api.openai.com/v1
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui-svc:3000
---
type: document_store
provider: qdrant
location: http://wren-qdrant:6333
embedding_model_dim: 3072
timeout: 120
---
type: pipeline
pipes:
- name: indexing
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: retrieval
llm: openai_llm.gpt-4o-mini
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: historical_question
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_correction
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: followup_sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_summary
llm: openai_llm.gpt-4o-mini
- name: sql_answer
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_breakdown
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_expansion
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_explanation
llm: openai_llm.gpt-4o-mini
- name: sql_regeneration
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: semantics_description
llm: openai_llm.gpt-4o-mini
- name: relationship_recommendation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: question_recommendation
llm: openai_llm.gpt-4o-mini
- name: intent_classification
llm: openai_llm.gpt-4o-mini
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: data_assistance
llm: openai_llm.gpt-4o-mini
---
settings:
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_using_db_schemas_without_pruning: false
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
enable_timer: false
logging_level: DEBUG
development: false
deployment/kustomizations/base/deploy-wren-ai-service.yaml
+50 -93
View File
@@ -14,99 +14,56 @@ spec:
spec:
containers:
- name: wren-ai-service
image: ghcr.io/canner/wren-ai-service:0.3.7
image: ghcr.io/canner/wren-ai-service:latest
volumeMounts:
- name: config-volume
mountPath: /app/data
env:
- name: WREN_AI_SERVICE_PORT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_AI_SERVICE_PORT
- name: LLM_PROVIDER
valueFrom:
configMapKeyRef:
name: wren-config
key: LLM_PROVIDER
- name: LLM_OPENAI_API_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LLM_OPENAI_API_KEY
- name: LLM_OPENAI_API_BASE
valueFrom:
configMapKeyRef:
name: wren-config
key: LLM_OPENAI_API_BASE
- name: EMBEDDER_PROVIDER
valueFrom:
configMapKeyRef:
name: wren-config
key: EMBEDDER_PROVIDER
- name: EMBEDDER_OPENAI_API_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: EMBEDDER_OPENAI_API_KEY
- name: EMBEDDER_OPENAI_API_BASE
valueFrom:
configMapKeyRef:
name: wren-config
key: EMBEDDER_OPENAI_API_BASE
- name: GENERATION_MODEL
valueFrom:
configMapKeyRef:
name: wren-config
key: GENERATION_MODEL
- name: QDRANT_HOST
valueFrom:
configMapKeyRef:
name: wren-config
key: QDRANT_HOST
- name: DOCUMENT_STORE_PROVIDER
valueFrom:
configMapKeyRef:
name: wren-config
key: DOCUMENT_STORE_PROVIDER
- name: WREN_ENGINE_ENDPOINT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_ENGINE_ENDPOINT
- name: LOGGING_LEVEL
valueFrom:
configMapKeyRef:
name: wren-config
key: LOGGING_LEVEL
- name: WREN_UI_ENDPOINT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_UI_ENDPOINT
- name: ENABLE_TIMER
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_AI_SERVICE_ENABLE_TIMER
- name: PYTHONUNBUFFERED
value: "1"
- name: LANGFUSE_ENABLE
valueFrom:
configMapKeyRef:
name: wren-config
key: LANGFUSE_ENABLE
- name: LANGFUSE_HOST
valueFrom:
configMapKeyRef:
name: wren-config
key: LANGFUSE_HOST
- name: LANGFUSE_PUBLIC_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LANGFUSE_PUBLIC_KEY
- name: LANGFUSE_SECRET_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LANGFUSE_SECRET_KEY
- name: WREN_AI_SERVICE_PORT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_AI_SERVICE_PORT
- name: OPENAI_API_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LLM_OPENAI_API_KEY
- name: QDRANT_HOST
valueFrom:
configMapKeyRef:
name: wren-config
key: QDRANT_HOST
- name: LOGGING_LEVEL
valueFrom:
configMapKeyRef:
name: wren-config
key: LOGGING_LEVEL
- name: WREN_UI_ENDPOINT
valueFrom:
configMapKeyRef:
name: wren-config
key: WREN_UI_ENDPOINT
- name: PYTHONUNBUFFERED
value: "1"
- name: LANGFUSE_PUBLIC_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LANGFUSE_PUBLIC_KEY
- name: LANGFUSE_SECRET_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LANGFUSE_SECRET_KEY
- name: CONFIG_PATH
value: /app/data/config.yaml
ports:
- containerPort: 5555
volumes:
- name: config-volume
configMap:
name: wren-ai-service-config
items:
- key: config.yaml
path: config.yaml
deployment/kustomizations/examples/secret-wren_example.yaml
-1
View File
@@ -7,7 +7,6 @@ type: Opaque
data:
# LLM_OPENAI_API_KEY and EMBEDDER_OPENAI_API_KEY are REQUIRED: without a valid key the wren-ai-service-deployment pod will not start
LLM_OPENAI_API_KEY: UkVRVUlSRUQ6IHNrLXByb2otYWxsLWFjY2Vzcy1wbGFjZWhvbGRlci00LXdyZW4tYWktc2VydmljZS1kZXBsb3ltZW50
EMBEDDER_OPENAI_API_KEY: UkVRVUlSRUQ6IHNrLXByb2otYWxsLWFjY2Vzcy1wbGFjZWhvbGRlci00LXdyZW4tYWktc2VydmljZS1kZXBsb3ltZW50
# Azure openai env
AZURE_CHAT_BASE: bi9h
deployment/kustomizations/kustomization.yaml
+1 -1
View File
@@ -43,7 +43,7 @@ images:
- name: ghcr.io/canner/wren-ui
newTag: 0.9.2 # WREN_UI_VERSION
- name: ghcr.io/canner/wren-ai-service
newTag: 0.8.2 # WREN_AI_SERVICE_VERSION
newTag: commit-391ea49 # WREN_AI_SERVICE_VERSION
- name: ghcr.io/canner/wren-engine-ibis
newTag: 0.9.0 # IBIS_SERVER_VERSION
docker/.env.ai.example
-57
View File
@@ -1,57 +0,0 @@
## LLM
# openai_llm, azure_openai_llm, ollama_llm
LLM_PROVIDER=openai_llm
LLM_TIMEOUT=120
GENERATION_MODEL=gpt-4o-mini
GENERATION_MODEL_KWARGS={"temperature": 0, "n": 1, "max_tokens": 4096, "response_format": {"type": "json_object"}}
COLUMN_INDEXING_BATCH_SIZE=50
TABLE_RETRIEVAL_SIZE=10
TABLE_COLUMN_RETRIEVAL_SIZE=1000
QUERY_CACHE_TTL=3600
# openai or openai-api-compatible
LLM_OPENAI_API_KEY=sk-xxxx
LLM_OPENAI_API_BASE=https://api.openai.com/v1
# azure_openai
LLM_AZURE_OPENAI_API_KEY=
LLM_AZURE_OPENAI_API_BASE=
LLM_AZURE_OPENAI_VERSION=
# ollama
LLM_OLLAMA_URL=http://host.docker.internal:11434
## EMBEDDER
# openai_embedder, azure_openai_embedder, ollama_embedder
EMBEDDER_PROVIDER=openai_embedder
EMBEDDER_TIMEOUT=120
# supported embedding models providers by qdrant: https://qdrant.tech/documentation/embeddings/
EMBEDDING_MODEL=text-embedding-3-large
EMBEDDING_MODEL_DIMENSION=3072
# openai or openai-api-compatible
EMBEDDER_OPENAI_API_KEY=sk-xxxx
EMBEDDER_OPENAI_API_BASE=https://api.openai.com/v1
# azure_openai
EMBEDDER_AZURE_OPENAI_API_KEY=
EMBEDDER_AZURE_OPENAI_API_BASE=
EMBEDDER_AZURE_OPENAI_VERSION=
# ollama
EMBEDDER_OLLAMA_URL=http://host.docker.internal:11434
## DOCUMENT_STORE
DOCUMENT_STORE_PROVIDER=qdrant
QDRANT_HOST=qdrant
QDRANT_TIMEOUT=120
## Langfuse: https://langfuse.com/
# empty means disabled
LANGFUSE_ENABLE=
LANGFUSE_SECRET_KEY=
LANGFUSE_PUBLIC_KEY=
LANGFUSE_HOST=https://cloud.langfuse.com
docker/.env.example
+13 -12
View File
@@ -1,4 +1,4 @@
COMPOSE_PROJECT_NAME=wren
COMPOSE_PROJECT_NAME=wrenai
PLATFORM=linux/amd64
PROJECT_DIR=.
@@ -13,11 +13,14 @@ IBIS_SERVER_PORT=8000
# service endpoint (for docker-compose-dev.yaml file)
WREN_UI_ENDPOINT=http://docker.for.mac.localhost:3000
# LLM
LLM_OPENAI_API_KEY=
EMBEDDER_OPENAI_API_KEY=
# gpt-4o-mini, gpt-4o
GENERATION_MODEL=gpt-4o-mini
# ai service settings
QDRANT_HOST=qdrant
SHOULD_FORCE_DEPLOY=1
# vendor keys
OPENAI_API_KEY=
AZURE_OPENAI_API_KEY=
QDRANT_API_KEY=
# version
# CHANGE THIS TO THE LATEST VERSION
@@ -28,12 +31,6 @@ IBIS_SERVER_VERSION=0.11.2
WREN_UI_VERSION=0.16.0
WREN_BOOTSTRAP_VERSION=0.1.5
# AI service related env variables
AI_SERVICE_ENABLE_TIMER=
AI_SERVICE_LOGGING_LEVEL=INFO
SHOULD_FORCE_DEPLOY=1
QDRANT_HOST=qdrant
# user id (uuid v4)
USER_UUID=
@@ -41,6 +38,10 @@ USER_UUID=
POSTHOG_API_KEY=phc_nhF32aj4xHXOZb0oqr2cn4Oy9uiWzz6CCP4KZmRq9aE
POSTHOG_HOST=https://app.posthog.com
TELEMETRY_ENABLED=true
# this is for telemetry to know the model, i think ai-service might be able to provide a endpoint to get the information
GENERATION_MODEL=gpt-4o-mini
LANGFUSE_SECRET_KEY=
LANGFUSE_PUBLIC_KEY=
# the port exposes to the host
# OPTIONAL: change the port if you have a conflict
docker/README.md
+27 -18
View File
@@ -1,33 +1,42 @@
## Service
* `wren-engine`: the engine service. check out example here: [wren-engine
/example](https://github.com/Canner/wren-engine/tree/main/example)
* `wren-ai-service`: the AI service.
* `qdrant`: the vector store ai service is using.
* `wren-ui`: the UI service.
* `bootstrap`: put required files to volume for engine service.
- `wren-engine`: the engine service. check out example here: [wren-engine
/example](https://github.com/Canner/wren-engine/tree/main/example)
- `wren-ai-service`: the AI service.
- `qdrant`: the vector store ai service is using.
- `wren-ui`: the UI service.
- `bootstrap`: put required files to volume for engine service.
## Volume
Shared data using `data` volume.
Path structure as following:
* `/mdl`
* `*.json` (will put `sample.json` during bootstrap)
* `accounts`
* `config.properties`
- `/mdl`
- `*.json` (will put `sample.json` during bootstrap)
- `accounts`
- `config.properties`
## Network
* Check out [Network drivers overview](https://docs.docker.com/engine/network/drivers/) to learn more about `bridge` network driver.
- Check out [Network drivers overview](https://docs.docker.com/engine/network/drivers/) to learn more about `bridge` network driver.
## How to start with OpenAI
1. copy `.env.example` to `.env.local` and modify the OpenAI API key.
2. (optional) if your port 3000 is occupied, you can modify the `HOST_PORT` in `.env.local`.
2. copy `config.example.yaml` to `config.yaml` for AI service configuration.
3. start all services: `docker-compose --env-file .env.local up -d`.
4. stop all services: `docker-compose --env-file .env.local down`.
## How to start with custom LLM
1. copy `.env.example` to `.env.local` and modify the OpenAI API key.
2. copy `.env.ai.example` to `.env.ai` and fill in necessary information if you would like to use custom LLM.
3. start all services(with custom LLM): `docker-compose -f docker-compose.yaml -f docker-compose.llm.yaml --env-file .env.local --env-file .env.ai up -d`.
4. stop all services(with custom LLM): `docker-compose -f docker-compose.yaml -f docker-compose.llm.yaml --env-file .env.local --env-file .env.ai down`.
### Optional
>Note: If your port 3000 is occupied, you can modify the `HOST_PORT` in `.env.local`.
- If your port 3000 is occupied, you can modify the `HOST_PORT` in `.env.local`.
## How to start with custom LLM
To start with a custom LLM, the process is similar to starting with OpenAI. The main difference is that you need to modify the `config.yaml` file
that we created on the previous step. After modifying the file, you can restart the services by running `docker-compose --env-file .env.local up -d --force-recreate wren-ai-service`.
For detailed information on how to modify the configuration for different LLM providers and models, please refer to the [AI Service Configuration](../wren-ai-service/docs/configuration.md).
This guide provides comprehensive instructions on setting up various LLM providers, embedders, and other components of the AI service.
docker/bootstrap/init.sh
+3 -3
View File
@@ -6,7 +6,7 @@ data_path=${DATA_PATH:-"./"}
# put a content into config.properties if not exists
if [ ! -f ${data_path}/config.properties ]; then
echo "init config.properties"
echo "node.environment=production" > ${data_path}/config.properties
echo "node.environment=production" >${data_path}/config.properties
fi
# after the config.properties is created, check if config properties properly set
@@ -14,7 +14,7 @@ fi
# check if wren.experimental-enable-dynamic-fields is set, otherwise append it with true
if ! grep -q "wren.experimental-enable-dynamic-fields" ${data_path}/config.properties; then
echo "wren.experimental-enable-dynamic-fields is not set, set it to true"
echo "wren.experimental-enable-dynamic-fields=true" >> ${data_path}/config.properties
echo "wren.experimental-enable-dynamic-fields=true" >>${data_path}/config.properties
fi
# create a folder mdl if not exists
@@ -26,5 +26,5 @@ fi
# put a emtpy sample.json if not exists
if [ ! -f ${data_path}/mdl/sample.json ]; then
echo "init mdl/sample.json"
echo "{\"catalog\": \"test_catalog\", \"schema\": \"test_schema\", \"models\": []}" > ${data_path}/mdl/sample.json
echo "{\"catalog\": \"test_catalog\", \"schema\": \"test_schema\", \"models\": []}" >${data_path}/mdl/sample.json
fi
docker/config.example.yaml
+127
View File
@@ -0,0 +1,127 @@
type: llm
provider: openai_llm
models:
- model: gpt-4o-mini
kwargs:
{
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
- model: gpt-4o
kwargs:
{
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
api_base: https://api.openai.com/v1
timeout: 120
---
type: llm
provider: ollama_llm
models:
- model: gemma2:9b
kwargs: { "temperature": 0 }
url: http://host.docker.internal:11434
timeout: 120
---
type: embedder
provider: openai_embedder
models:
- model: text-embedding-3-large
dimension: 3072
api_base: https://api.openai.com/v1
timeout: 120
---
type: embedder
provider: ollama_embedder
models:
- model: nomic-embed-text
dimension: 786
url: http://host.docker.internal:11434
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://wren-ui:3000
---
type: document_store
provider: qdrant
location: http://qdrant:6333
embedding_model_dim: 3072
timeout: 120
recreate_index: true
---
type: pipeline
pipes:
- name: indexing
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: retrieval
llm: openai_llm.gpt-4o-mini
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: historical_question
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_correction
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: followup_sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_summary
llm: openai_llm.gpt-4o-mini
- name: sql_answer
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_breakdown
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_expansion
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_explanation
llm: openai_llm.gpt-4o-mini
- name: sql_regeneration
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: semantics_description
llm: openai_llm.gpt-4o-mini
- name: relationship_recommendation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: question_recommendation
llm: openai_llm.gpt-4o-mini
- name: intent_classification
llm: openai_llm.gpt-4o-mini
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: data_assistance
llm: openai_llm.gpt-4o-mini
---
settings:
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_using_db_schemas_without_pruning: false
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
enable_timer: false
logging_level: DEBUG
development: false
docker/docker-compose-dev.yaml
+7 -8
View File
@@ -42,19 +42,18 @@ services:
environment:
WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT}
WREN_UI_PORT: ${WREN_UI_PORT}
WREN_UI_ENDPOINT: ${WREN_UI_ENDPOINT}
QDRANT_HOST: ${QDRANT_HOST}
LLM_OPENAI_API_KEY: ${LLM_OPENAI_API_KEY}
EMBEDDER_OPENAI_API_KEY: ${EMBEDDER_OPENAI_API_KEY}
LLM_AZURE_OPENAI_API_KEY: ${LLM_AZURE_OPENAI_API_KEY}
EMBEDDER_AZURE_OPENAI_API_KEY: ${EMBEDDER_AZURE_OPENAI_API_KEY}
GENERATION_MODEL: ${GENERATION_MODEL}
ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER}
LOGGING_LEVEL: ${AI_SERVICE_LOGGING_LEVEL}
WREN_UI_ENDPOINT: ${WREN_UI_ENDPOINT}
OPENAI_API_KEY: ${OPENAI_API_KEY}
AZURE_OPENAI_API_KEY: ${AZURE_OPENAI_API_KEY}
QDRANT_API_KEY: ${QDRANT_API_KEY}
SHOULD_FORCE_DEPLOY: ${SHOULD_FORCE_DEPLOY}
# sometimes the console won't show print messages,
# using PYTHONUNBUFFERED: 1 can fix this
PYTHONUNBUFFERED: 1
CONFIG_PATH: /app/data/config.yaml
volumes:
- ${PROJECT_DIR}/config.yaml:/app/data/config.yaml
networks:
- wren
depends_on:
docker/docker-compose.llm.yaml
-3
View File
@@ -1,3 +0,0 @@
services:
wren-ai-service:
env_file: ${PROJECT_DIR}/.env.ai
docker/docker-compose.yaml
+8 -10
View File
@@ -54,19 +54,18 @@ services:
environment:
WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT}
WREN_UI_PORT: ${WREN_UI_PORT}
WREN_UI_ENDPOINT: http://wren-ui:${WREN_UI_PORT}
QDRANT_HOST: ${QDRANT_HOST}
LLM_OPENAI_API_KEY: ${LLM_OPENAI_API_KEY}
EMBEDDER_OPENAI_API_KEY: ${EMBEDDER_OPENAI_API_KEY}
LLM_AZURE_OPENAI_API_KEY: ${LLM_AZURE_OPENAI_API_KEY}
EMBEDDER_AZURE_OPENAI_API_KEY: ${EMBEDDER_AZURE_OPENAI_API_KEY}
GENERATION_MODEL: ${GENERATION_MODEL}
ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER}
LOGGING_LEVEL: ${AI_SERVICE_LOGGING_LEVEL}
WREN_UI_ENDPOINT: http://wren-ui:${WREN_UI_PORT}
OPENAI_API_KEY: ${OPENAI_API_KEY}
AZURE_OPENAI_API_KEY: ${AZURE_OPENAI_API_KEY}
QDRANT_API_KEY: ${QDRANT_API_KEY}
SHOULD_FORCE_DEPLOY: ${SHOULD_FORCE_DEPLOY}
# sometimes the console won't show print messages,
# using PYTHONUNBUFFERED: 1 can fix this
PYTHONUNBUFFERED: 1
CONFIG_PATH: /app/data/config.yaml
volumes:
- ${PROJECT_DIR}/config.yaml:/app/data/config.yaml
networks:
- wren
depends_on:
@@ -94,8 +93,7 @@ services:
WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT}
WREN_AI_ENDPOINT: http://wren-ai-service:${WREN_AI_SERVICE_PORT}
IBIS_SERVER_ENDPOINT: http://ibis-server:${IBIS_SERVER_PORT}
EMBEDDING_MODEL: ${EMBEDDING_MODEL}
EMBEDDING_MODEL_DIMENSION: ${EMBEDDING_MODEL_DIMENSION}
# this is for telemetry to know the model, i think ai-service might be able to provide a endpoint to get the information
GENERATION_MODEL: ${GENERATION_MODEL}
# telemetry
WREN_ENGINE_PORT: ${WREN_ENGINE_PORT}
wren-ai-service/.env.dev.example
-86
View File
@@ -1,86 +0,0 @@
# app related
WREN_AI_SERVICE_HOST=127.0.0.1
WREN_AI_SERVICE_PORT=5556
SHOULD_FORCE_DEPLOY=
COLUMN_INDEXING_BATCH_SIZE=50
TABLE_RETRIEVAL_SIZE=10
TABLE_COLUMN_RETRIEVAL_SIZE=100
QUERY_CACHE_TTL=3600
## LLM
# openai_llm, azure_openai_llm, ollama_llm
LLM_PROVIDER=openai_llm
LLM_TIMEOUT=120
# gpt-4o-mini, gpt-4o
GENERATION_MODEL=gpt-4o-mini
# openai or openai-api-compatible
LLM_OPENAI_API_KEY=sk-1234567890
LLM_OPENAI_API_BASE=https://api.openai.com/v1
# azure_openai
LLM_AZURE_OPENAI_API_KEY=
LLM_AZURE_OPENAI_API_BASE=
LLM_AZURE_OPENAI_VERSION=
# ollama
LLM_OLLAMA_URL=http://localhost:11434
## EMBEDDER
# openai_embedder, azure_openai_embedder, ollama_embedder
EMBEDDER_PROVIDER=openai_embedder
EMBEDDER_TIMEOUT=120
EMBEDDING_MODEL=text-embedding-3-large
EMBEDDING_MODEL_DIMENSION=3072
# openai or openai-api-compatible
EMBEDDER_OPENAI_API_KEY=sk-1234567890
EMBEDDER_OPENAI_API_BASE=https://api.openai.com/v1
# azure_openai
EMBEDDER_AZURE_OPENAI_API_KEY=
EMBEDDER_AZURE_OPENAI_API_BASE=
EMBEDDER_AZURE_OPENAI_VERSION=
# ollama
EMBEDDER_OLLAMA_URL=http://localhost:11434
## DOCUMENT_STORE
DOCUMENT_STORE_PROVIDER=qdrant
QDRANT_HOST=http://localhost:6333
QDRANT_TIMEOUT=120
QDRANT_API_KEY=
# wren_ui, wren_ibis, wren_engine
ENGINE=wren_ui
## when using wren_ui as the engine
WREN_UI_ENDPOINT=http://localhost:3000
## when using wren_ibis as the engine
WREN_IBIS_ENDPOINT=http://localhost:8000
WREN_IBIS_SOURCE=bigquery
### this is a base64 encoded string of the MDL
WREN_IBIS_MANIFEST=
### this is a base64 encode string of the connection info
WREN_IBIS_CONNECTION_INFO=
## when using wren_engine as the engine
WREN_ENGINE_ENDPOINT=http://localhost:8080
WREN_ENGINE_MANIFEST=
# Evaluation
DATASET_NAME=book_2
# empty means disabled
LANGFUSE_ENABLE=
LANGFUSE_SECRET_KEY=
LANGFUSE_PUBLIC_KEY=
LANGFUSE_HOST=https://cloud.langfuse.com
# Debugging
ENABLE_TIMER=
LOGGING_LEVEL=DEBUG
wren-ai-service/Justfile
+25 -3
View File
@@ -1,8 +1,30 @@
GREEN := "\u{001b}[32m"
YELLOW := "\u{001b}[33m"
RESET := "\u{001b}[0m"
## todo: consider to support --override flag to override existing files
init dev='--dev':
@if [ ! -f config.yaml ]; then \
echo "{{GREEN}}config.yaml does not exist. Creating from example...{{RESET}}"; \
cp tools/config/config.example.yaml config.yaml; \
else \
echo "{{YELLOW}}config.yaml already exists. Skipping creation.{{RESET}}"; \
fi
@if [ {{dev}} = "--dev" ] || [ {{dev}} != "--non-dev" ]; then \
if [ ! -f .env.dev ]; then \
echo "{{GREEN}}env.dev does not exist. Creating from example...{{RESET}}"; \
cp tools/config/.env.dev.example .env.dev; \
else \
echo "{{YELLOW}}env.dev already exists. Skipping creation.{{RESET}}"; \
fi \
fi
up: prepare-wren-engine
docker compose -f ./tools/dev/docker-compose-dev.yaml --env-file ./tools/dev/.env.example up -d
docker compose -f ./tools/dev/docker-compose-dev.yaml --env-file ./tools/dev/.env up -d
down:
docker compose -f ./tools/dev/docker-compose-dev.yaml --env-file ./tools/dev/.env.example down
docker compose -f ./tools/dev/docker-compose-dev.yaml --env-file ./tools/dev/.env down
start:
poetry run python -m src.__main__
@@ -40,4 +62,4 @@ prepare-wren-engine:
cp tools/dev/config.properties.example tools/dev/etc/config.properties
mkdir -p tools/dev/etc/mdl
echo "{\"catalog\": \"test_catalog\", \"schema\": \"test_schema\", \"models\": []}" \\
> tools/dev/etc/mdl/sample.json
> tools/dev/etc/mdl/sample.json
wren-ai-service/README.md
+80 -17
View File
@@ -6,26 +6,89 @@ Please read the [documentation](https://docs.getwren.ai/oss/concept/wren_ai_serv
## Setup for Local Development
### Environment Setup
### Prerequisites
- Python 3.12.\*, recommended to use [`pyenv`](https://github.com/pyenv/pyenv?tab=readme-ov-file#installation) to manage the Python versions
- install `poetry` with version 1.8.3: `curl -sSL https://install.python-poetry.org | python3 - --version 1.8.3`
- execute `poetry install` to install the dependencies
- copy `.env.dev.example` file to `.env.dev` and fill in the environment variables
- [for development] execute `poetry run pre-commit install` to install the pre-commit hooks and `poetry run pre-commit run --all-files` to run the pre-commit checks at the first time to check if everything is set up correctly
- [for development] install [Just](https://github.com/casey/just?tab=readme-ov-file#packages)
- [for development] to run the tests, execute `just test`
1. **Python**: Install Python 3.12.\*
### Start the service for development
- Recommended: Use [`pyenv`](https://github.com/pyenv/pyenv?tab=readme-ov-file#installation) to manage Python versions
The following commands can quickly start the service for development:
2. **Poetry**: Install Poetry 1.8.3
- `just up` to start needed containers
- `just start` to start the service
- go to `http://WREN_AI_SERVICE_HOST:WREN_AI_SERVICE_PORT`(default is http://localhost:5556) to see the API
documentation and try them.
- go to `http://WREN_UI_HOST:WREN_UI_PORT`(default is http://localhost:3000) to interact interact from the UI
- `just down` to stop the needed containers
```bash
curl -sSL https://install.python-poetry.org | python3 - --version 1.8.3
```
3. **Just**: Install [Just](https://github.com/casey/just?tab=readme-ov-file#packages) command runner (version 1.36 or higher)
### Step-by-Step Setup
1. **Install Dependencies**:
```bash
poetry install
```
2. **Generate Configuration Files**:
```bash
just init
```
This creates both `.env.dev` and `config.yaml`. Use `just init --non-dev` to generate only `config.yaml`.
3. **Configure Environment**:
- Edit `.env.dev` to set environment variables
- Modify `config.yaml` to configure components, pipelines, and other settings
- Refer to [AI Service Configuration](./docs/configuration.md) for detailed setup instructions
4. **Set Up Development Environment** (optional):
- Install pre-commit hooks:
```bash
poetry run pre-commit install
```
- Run initial pre-commit checks:
```bash
poetry run pre-commit run --all-files
```
5. **Run Tests** (optional):
```bash
just test
```
### Starting the Service
1. **Start Required Containers**:
```bash
just up
```
2. **Launch the AI Service**:
```bash
just start
```
3. **Access the Service**:
- API Documentation: `http://WREN_AI_SERVICE_HOST:WREN_AI_SERVICE_PORT` (default: <http://localhost:5556>)
- User Interface: `http://WREN_UI_HOST:WREN_UI_PORT` (default: <http://localhost:3000>)
4. **Stop the Service**:
When finished, stop the containers:
```bash
just down
```
This setup ensures a consistent development environment and helps maintain code quality through pre-commit hooks and tests. Follow these steps to get started with local development of the Wren AI Service.
## Others
@@ -46,7 +109,7 @@ For a comprehensive understanding of how to evaluate the pipelines, please refer
- in wren-ai-service folder, run `just up` to start the docker containers
- in wren-ai-service folder, run `just start` to start the ai service
- run `just load-test`
- check reports in /outputs/locust folder, there are 3 files with filename **locust_report_{test_timestamp}**:
- check reports in /outputs/locust folder, there are 3 files with filename **locust*report*{test_timestamp}**:
- .json: test report in json format, including info like llm provider, version
- .html: test report in html format, showing tables and charts
- .log: test log
wren-ai-service/demo/utils.py
+12 -39
View File
@@ -579,47 +579,20 @@ CREATE TABLE dept_manager AS FROM read_parquet('https://assets.getwren.ai/sample
def _replace_wren_engine_env_variables(engine_type: str, data: dict):
assert engine_type in ("wren_engine", "wren_ibis")
if not Path("config.yaml").exists():
if engine_type == "wren_engine":
with open(".env.dev", "r") as f:
lines = f.readlines()
for i, line in enumerate(lines):
if line.startswith("ENGINE"):
lines[i] = "ENGINE=wren_engine\n"
elif line.startswith("WREN_ENGINE_MANIFEST"):
lines[i] = f"WREN_ENGINE_MANIFEST={data['manifest']}\n"
else:
with open(".env.dev", "r") as f:
lines = f.readlines()
for i, line in enumerate(lines):
if line.startswith("ENGINE"):
lines[i] = "ENGINE=wren_ibis\n"
elif line.startswith("WREN_IBIS_SOURCE"):
lines[i] = f"WREN_IBIS_SOURCE={data['source']}\n"
elif line.startswith("WREN_IBIS_MANIFEST"):
lines[i] = f"WREN_IBIS_MANIFEST={data['manifest']}\n"
elif line.startswith("WREN_IBIS_CONNECTION_INFO"):
lines[
i
] = f"WREN_IBIS_CONNECTION_INFO={data['connection_info']}\n"
with open("config.yaml", "r") as f:
configs = list(yaml.safe_load_all(f))
with open(".env.dev", "w") as f:
f.writelines(lines)
else:
with open("config.yaml", "r") as f:
configs = list(yaml.safe_load_all(f))
for config in configs:
if config.get("type") == "engine" and config.get("provider") == engine_type:
for key, value in data.items():
config[key] = value
if "pipes" in config:
for i, pipe in enumerate(config["pipes"]):
if "engine" in pipe:
config["pipes"][i]["engine"] = engine_type
for config in configs:
if config["type"] == "engine" and config["provider"] == engine_type:
for key, value in data.items():
config[key] = value
if "pipes" in config:
for i, pipe in enumerate(config["pipes"]):
if "engine" in pipe:
config["pipes"][i]["engine"] = engine_type
with open("config.yaml", "w") as f:
yaml.safe_dump_all(configs, f, default_flow_style=False)
with open("config.yaml", "w") as f:
yaml.safe_dump_all(configs, f, default_flow_style=False)
def prepare_semantics(mdl_json: dict):
wren-ai-service/docs/configuration.md
+144
View File
@@ -0,0 +1,144 @@
# AI Service Configuration
The AI service configuration is managed through a combination of environment variables and a configuration file, providing a flexible and secure approach to setting up the service.
1. **Environment Variables**:
- Used for configuring sensitive information such as vendor API keys
- Specify the configuration file to use
- Allow for partial settings to be configured directly, see [Settings Loading Mechanism](#settings-loading-mechanism) for more details
- Provide a way to override settings in different environments
2. **Configuration File**:
- Used for detailed configuration of components, pipelines, and other service settings
- Allows for more complex and structured configuration options
This dual approach ensures that sensitive data can be kept secure (using environment variables) while allowing for more detailed and shareable configuration through the configuration file. It also provides flexibility in deployment across different environments.
## Settings Loading Mechanism
The AI service uses a hierarchical approach to load settings, ensuring flexibility across different environments and deployment scenarios. The settings are loaded in the following order of precedence:
1. **Default Values**: These are defined as class attributes in the `Settings` class within [`config.py`](../src/config.py). They serve as the base configuration.
2. **Environment Variables**: Using [pydantic-settings](https://fastapi.tiangolo.com/advanced/settings/#pydantic-settings), the service checks for environment variables that match the setting names. If found, these override the default values. For example, `WREN_AI_SERVICE_HOST` can override the default `host` value.
3. **.env.dev File**: The service loads additional settings or overrides existing ones from a `.env.dev` file if present. This is particularly useful for development environments.
4. **config.yaml File**: This file provides the highest priority configuration. It can override all previous settings and is used to configure components, pipelines, and other detailed settings. See [Configuration File](#configuration-file) for more details.
This mechanism allows for easy configuration management across different environments, from development to production, while maintaining security for sensitive information like API keys.
## Configuration File
The configuration file (`config.yaml`) is structured into several sections, each defining different aspects of the AI service. Here's a breakdown of its main components:
1. **LLM Configuration**:
```yaml
type: llm
provider: <provider_name>
models:
- model: <model_name>
kwargs: {}
api_base: <api_endpoint>
```
This component initializes the LLM provider at runtime. You can specify multiple models with different parameters. The `kwargs` field allows for model-specific configurations. For example:
```yaml
type: llm
provider: openai_llm
models:
- model: gpt-4
kwargs:
temperature: 0
n: 1
max_tokens: 4096
response_format:
type: "json_object"
- model: gpt-4o-mini
kwargs: {}
api_base: https://api.openai.com/v1
```
For detailed parameter options, refer to the implementation of the specific LLM provider.
2. **Embedder Configuration**:
```yaml
type: embedder
provider: <provider_name>
models:
- model: <model_name>
dimension: <embedding_size>
api_base: <api_endpoint>
timeout: <timeout_in_seconds>
```
This component configures the embedder, which converts text into numerical vectors. The `provider` specifies the embedder service (e.g., OpenAI, Ollama). You can define multiple `models` with their parameters. The `dimension` parameter indicates the size of the embedding vector.
3. **Engine Configuration**:
```yaml
type: engine
provider: <provider_name>
endpoint: <engine_endpoint>
```
This component configures the engine responsible for generating SQL queries. The `provider` specifies the engine service (e.g., Wren UI).
4. **Document Store Configuration**:
```yaml
type: document_store
provider: <provider_name>
```
This component configures the document store, which is responsible for storing and retrieving embeddings. The `provider` specifies the document store service (e.g., Qdrant).
5. **Pipeline Configuration**:
```yaml
type: pipeline
pipes:
- name: <pipe_name>
llm: <provider>.<model_name>
embedder: <provider>.<model_name>
engine: <provider_name>
document_store: <provider_name>
```
This component configures each pipeline, specifying different LLM, embedder, engine, and document store combinations. For LLM and embedder, use `<provider>.<model_name>`. For engine and document store, use `<provider_name>`.
Example:
```yaml
type: pipeline
pipes:
- name: sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
```
6. **Settings**:
```yaml
settings:
host: <host_address>
port: <port_number>
column_indexing_batch_size: <batch_size>
table_retrieval_size: <retrieval_size>
table_column_retrieval_size: <column_retrieval_size>
query_cache_maxsize: <cache_size>
query_cache_ttl: <cache_ttl_in_seconds>
langfuse_host: <langfuse_endpoint>
langfuse_enable: <true/false>
enable_timer: <true/false>
logging_level: <log_level>
development: <true/false>
```
This section defines various service settings including host, port, indexing and retrieval parameters, cache settings, Langfuse configuration, logging level, and development mode.
This configuration file allows for detailed customization of the AI service components, pipelines, and overall behavior. It provides a centralized place to manage complex configurations while keeping sensitive information separate (managed through environment variables). See [Full Configuration File](../tools/config/config.full.yaml) for a complete example.
wren-ai-service/src/__main__.py
+11 -53
View File
@@ -1,5 +1,3 @@
import logging
import os
from contextlib import asynccontextmanager
import uvicorn
@@ -9,6 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import ORJSONResponse, RedirectResponse
from langfuse.decorators import langfuse_context
from src.config import settings
from src.globals import (
create_service_container,
create_service_metadata,
@@ -16,52 +15,20 @@ from src.globals import (
from src.providers import generate_components
from src.utils import (
init_langfuse,
load_env_vars,
setup_custom_logger,
)
from src.web.v1 import routers
env = load_env_vars()
setup_custom_logger(
"wren-ai-service",
level=(
logging.DEBUG if os.getenv("LOGGING_LEVEL", "INFO") == "DEBUG" else logging.INFO
),
)
setup_custom_logger("wren-ai-service", level_str=settings.logging_level)
# https://fastapi.tiangolo.com/advanced/events/#lifespan
@asynccontextmanager
async def lifespan(app: FastAPI):
# startup events
pipe_components = generate_components(force_deploy=os.getenv("SHOULD_FORCE_DEPLOY"))
app.state.service_container = create_service_container(
pipe_components,
column_indexing_batch_size=(
int(os.getenv("COLUMN_INDEXING_BATCH_SIZE"))
if os.getenv("COLUMN_INDEXING_BATCH_SIZE")
else 50
),
table_retrieval_size=(
int(os.getenv("TABLE_RETRIEVAL_SIZE"))
if os.getenv("TABLE_RETRIEVAL_SIZE")
else 10
),
table_column_retrieval_size=(
int(os.getenv("TABLE_COLUMN_RETRIEVAL_SIZE"))
if os.getenv("TABLE_COLUMN_RETRIEVAL_SIZE")
else 1000
),
query_cache={
# the maxsize is a necessary parameter to init cache, but we don't want to expose it to the user
# so we set it to 1_000_000, which is a large number
"maxsize": 1_000_000,
"ttl": int(os.getenv("QUERY_CACHE_TTL") or 120),
},
allow_using_db_schemas_without_pruning=bool(
os.getenv("ALLOW_USING_DB_SCHEMAS_WITHOUT_PRUNING", False)
),
)
pipe_components = generate_components(settings.components)
app.state.service_container = create_service_container(pipe_components, settings)
app.state.service_metadata = create_service_metadata(pipe_components)
init_langfuse()
@@ -86,14 +53,14 @@ app.add_middleware(
allow_headers=["*"],
)
app.include_router(routers.router, prefix="/v1", tags=["v1"])
if env == "dev":
if settings.development:
from src.web import development
app.include_router(development.router, prefix="/dev", tags=["dev"])
@app.exception_handler(Exception)
async def exception_handler(request, exc: Exception):
async def exception_handler(_, exc: Exception):
return ORJSONResponse(
status_code=500,
content={"detail": str(exc)},
@@ -101,7 +68,7 @@ async def exception_handler(request, exc: Exception):
@app.exception_handler(RequestValidationError)
async def request_exception_handler(request, exc: Exception):
async def request_exception_handler(_, exc: Exception):
return ORJSONResponse(
status_code=400,
content={"detail": str(exc)},
@@ -119,20 +86,11 @@ def health():
if __name__ == "__main__":
server_host = os.getenv("WREN_AI_SERVICE_HOST") or "127.0.0.1"
server_port = (
int(os.getenv("WREN_AI_SERVICE_PORT"))
if os.getenv("WREN_AI_SERVICE_PORT") is not None
else 8000
)
should_reload = env == "dev"
uvicorn.run(
"src.__main__:app",
host=server_host,
port=server_port,
reload=should_reload,
host=settings.host,
port=settings.port,
reload=settings.development,
reload_includes=["src/**/*.py", ".env.dev", "config.yaml"],
workers=1,
loop="uvloop",
wren-ai-service/src/config.py
+98
View File
@@ -0,0 +1,98 @@
import logging
import yaml
from dotenv import load_dotenv
from pydantic import Field
from pydantic_settings import BaseSettings
logger = logging.getLogger("wren-ai-service")
class Settings(BaseSettings):
"""
Configuration settings for the Wren AI service.
The settings are loaded in the following order of precedence:
1. Default values: Defined in the class attributes.
2. Environment variables: Overrides default values if set.
3. .env.dev file: Loads additional settings or overrides previous ones.
4. config.yaml file: Provides the highest priority configuration.
This hierarchical loading allows for flexible configuration management
across different environments and deployment scenarios.
"""
host: str = Field(default="127.0.0.1", alias="WREN_AI_SERVICE_HOST")
port: int = Field(default=5556, alias="WREN_AI_SERVICE_PORT")
# indexing and retrieval config
column_indexing_batch_size: int = Field(default=50)
table_retrieval_size: int = Field(default=10)
table_column_retrieval_size: int = Field(default=1000)
allow_using_db_schemas_without_pruning: bool = Field(default=False)
# service config
query_cache_ttl: int = Field(default=3600)
query_cache_maxsize: int = Field(
default=1_000_000,
comment="""
the maxsize is a necessary parameter to init cache, but we don't want to expose it to the user
so we set it to 1_000_000, which is a large number
""",
)
# langfuse config
langfuse_host: str = Field(default="https://cloud.langfuse.com")
langfuse_enable: bool = Field(default=True)
# debug config
enable_timer: bool = Field(default=False)
logging_level: str = Field(default="INFO")
development: bool = Field(default=False)
# this is used to store the config like type: llm, embedder, etc. and we will process them later
config_path: str = Field(default="config.yaml")
_components: list[dict]
def __init__(self):
load_dotenv(".env.dev", override=True)
super().__init__()
raw = self.config_loader()
self.override(raw)
self._components = [
component for component in raw if "settings" not in component
]
def config_loader(self):
try:
with open(self.config_path, "r") as file:
return list(yaml.load_all(file, Loader=yaml.SafeLoader))
except FileNotFoundError:
message = f"Warning: Configuration file {self.config_path} not found. Using default settings."
logger.warning(message)
return []
except yaml.YAMLError as e:
logger.exception(f"Error parsing YAML file: {e}")
return []
def override(self, raw: list[dict]) -> None:
override_settings = {}
for doc in raw:
if "settings" in doc:
override_settings = doc["settings"]
break
for key, value in override_settings.items():
if hasattr(self, key):
setattr(self, key, value)
else:
message = f"Warning: Unknown configuration key '{key}' in YAML file."
logger.warning(message)
@property
def components(self) -> list[dict]:
return self._components
settings = Settings()
wren-ai-service/src/globals.py
+15 -14
View File
@@ -1,9 +1,9 @@
import logging
from dataclasses import asdict, dataclass
from typing import Optional
import toml
from src.config import Settings
from src.core.pipeline import PipelineComponent
from src.core.provider import EmbedderProvider, LLMProvider
from src.pipelines.generation import (
@@ -60,12 +60,12 @@ class ServiceMetadata:
def create_service_container(
pipe_components: dict[str, PipelineComponent],
column_indexing_batch_size: Optional[int] = 50,
table_retrieval_size: Optional[int] = 10,
table_column_retrieval_size: Optional[int] = 100,
query_cache: Optional[dict] = {},
allow_using_db_schemas_without_pruning: Optional[bool] = False,
settings: Settings,
) -> ServiceContainer:
query_cache = {
"maxsize": settings.query_cache_maxsize,
"ttl": settings.query_cache_ttl,
}
return ServiceContainer(
semantics_description=SemanticsDescription(
pipelines={
@@ -79,7 +79,7 @@ def create_service_container(
pipelines={
"indexing": indexing.Indexing(
**pipe_components["indexing"],
column_indexing_batch_size=column_indexing_batch_size,
column_indexing_batch_size=settings.column_indexing_batch_size,
),
},
**query_cache,
@@ -94,9 +94,9 @@ def create_service_container(
),
"retrieval": retrieval.Retrieval(
**pipe_components["retrieval"],
table_retrieval_size=table_retrieval_size,
table_column_retrieval_size=table_column_retrieval_size,
allow_using_db_schemas_without_pruning=allow_using_db_schemas_without_pruning,
table_retrieval_size=settings.table_retrieval_size,
table_column_retrieval_size=settings.table_column_retrieval_size,
allow_using_db_schemas_without_pruning=settings.allow_using_db_schemas_without_pruning,
),
"historical_question": historical_question.HistoricalQuestion(
**pipe_components["historical_question"],
@@ -139,8 +139,8 @@ def create_service_container(
pipelines={
"retrieval": retrieval.Retrieval(
**pipe_components["retrieval"],
table_retrieval_size=table_retrieval_size,
table_column_retrieval_size=table_column_retrieval_size,
table_retrieval_size=settings.table_retrieval_size,
table_column_retrieval_size=settings.table_column_retrieval_size,
),
"sql_expansion": sql_expansion.SQLExpansion(
**pipe_components["sql_expansion"],
@@ -185,8 +185,9 @@ def create_service_container(
),
"retrieval": retrieval.Retrieval(
**pipe_components["retrieval"],
table_retrieval_size=table_retrieval_size,
table_column_retrieval_size=table_column_retrieval_size,
table_retrieval_size=settings.table_retrieval_size,
table_column_retrieval_size=settings.table_column_retrieval_size,
allow_using_db_schemas_without_pruning=settings.allow_using_db_schemas_without_pruning,
),
"sql_generation": sql_generation.SQLGeneration(
**pipe_components["sql_generation"],
wren-ai-service/src/providers/__init__.py
+264 -79
View File
@@ -1,11 +1,9 @@
import logging
import os
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Tuple
import yaml
from yaml.loader import SafeLoader
from src.core.engine import Engine, EngineConfig
from src.core.pipeline import PipelineComponent
from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider
@@ -17,29 +15,65 @@ logger = logging.getLogger("wren-ai-service")
def provider_factory(
config: dict = {},
) -> LLMProvider | EmbedderProvider | DocumentStoreProvider | Engine:
logger.info(f"initializing provider: {config.get('provider')}")
return loader.get_provider(config.get("provider"))(**config)
def load_config(path: str = "config.yaml") -> list[dict]:
if not os.path.exists(path):
return []
def llm_processor(entry: dict) -> dict:
"""
Process the LLM configuration entry.
with open(path, "r") as f:
return list(yaml.load_all(f, Loader=SafeLoader))
This function takes a dictionary containing LLM configuration and processes it
into a standardized format. The input dictionary is expected to have the following structure:
def process_llm(entry: dict) -> dict:
others = {
k: v
for k, v in entry.items()
if k not in ["type", "provider", "api_key", "models"]
{
"type": "llm",
"provider": "openai_llm",
"models": [
{
"model": "gpt-4o-mini",
"kwargs": {
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": {"type": "json_object"}
}
}
],
"api_base": "https://api.openai.com/v1"
}
The function processes this input and returns a dictionary with the following structure:
{
"openai_llm.gpt-4o-mini": {
"provider": "openai_llm",
"model": "gpt-4o-mini",
"kwargs": {
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": {"type": "json_object"}
},
"api_base": "https://api.openai.com/v1"
}
}
Args:
entry (dict): The input LLM configuration dictionary.
Returns:
dict: A processed dictionary with standardized LLM configuration.
Note:
The function does not handle the `api_key` field. It is to be handled by the provider itself.
"""
others = {k: v for k, v in entry.items() if k not in ["type", "provider", "models"]}
returned = {}
for model in entry["models"]:
model_name = f"{entry['provider']}.{model['model']}"
returned[model_name] = {
"provider": entry["provider"],
"api_key": entry["api_key"],
"model": model["model"],
"kwargs": model["kwargs"],
**others,
@@ -47,18 +81,49 @@ def process_llm(entry: dict) -> dict:
return returned
def process_embedder(entry: dict) -> dict:
others = {
k: v
for k, v in entry.items()
if k not in ["type", "provider", "api_key", "models"]
def embedder_processor(entry: dict) -> dict:
"""
Process the embedder configuration entry.
This function takes a dictionary containing embedder configuration and processes it
into a standardized format. The input dictionary is expected to have the following structure:
{
"type": "embedder",
"provider": "openai_embedder",
"models": [
{
"model": "text-embedding-ada-002",
"dimension": 1536
}
]
}
The function processes this input and returns a dictionary with the following structure:
{
"openai_embedder.text-embedding-ada-002": {
"provider": "openai_embedder",
"model": "text-embedding-ada-002",
"dimension": 1536
}
}
Args:
entry (dict): The input embedder configuration dictionary.
Returns:
dict: A processed dictionary with standardized embedder configuration.
Note:
The function does not handle the `api_key` field. It is to be handled by the provider itself.
"""
others = {k: v for k, v in entry.items() if k not in ["type", "provider", "models"]}
returned = {}
for model in entry["models"]:
model_name = f"{entry['provider']}.{model['model']}"
returned[model_name] = {
identifier = f"{entry['provider']}.{model['model']}"
returned[identifier] = {
"provider": entry["provider"],
"api_key": entry["api_key"],
"model": model["model"],
"dimension": model["dimension"],
**others,
@@ -67,15 +132,120 @@ def process_embedder(entry: dict) -> dict:
return returned
def process_document_store(entry: dict) -> dict:
def document_store_processor(entry: dict) -> dict:
"""
Process the document store configuration entry.
This function takes a dictionary containing document store configuration and processes it
into a standardized format. The input dictionary is expected to have the following structure:
{
"type": "document_store",
"provider": "qdrant",
"location": "http://localhost:6333",
"embedding_model_dim": 3072,
"timeout": 120,
"recreate_index": False,
}
The function processes this input and returns a dictionary with the following structure:
{
"qdrant": {
"provider": "qdrant",
"location": "http://localhost:6333",
"embedding_model_dim": 3072,
"timeout": 120,
"recreate_index": False,
}
}
Args:
entry (dict): The input document store configuration dictionary.
Returns:
dict: A processed dictionary with standardized document store configuration.
Note:
The function does not handle the `api_key` field. It is to be handled by the provider itself.
"""
return {entry["provider"]: {k: v for k, v in entry.items() if k not in ["type"]}}
def process_engine(entry: dict) -> dict:
def engine_processor(entry: dict) -> dict:
"""
Process the engine configuration entry.
This function takes a dictionary containing engine configuration and processes it
into a standardized format. The input dictionary is expected to have the following structure:
{
"type": "engine",
"provider": "wren_ui",
"kwargs": {
"host": "localhost",
"port": 8000
}
}
The function processes this input and returns a dictionary with the following structure:
{
"wren_ui": {
"provider": "wren_ui",
"kwargs": {
"host": "localhost",
"port": 8000
}
}
}
Args:
entry (dict): The input engine configuration dictionary.
Returns:
dict: A processed dictionary with standardized engine configuration.
"""
return {entry["provider"]: {k: v for k, v in entry.items() if k not in ["type"]}}
def process_pipeline(entry: dict) -> dict:
def pipeline_processor(entry: dict) -> dict:
"""
Process the pipeline configuration entry.
This function takes a dictionary containing pipeline configuration and processes it
into a standardized format. The input dictionary is expected to have the following structure:
{
"type": "pipeline",
"pipes": [
{
"name": "indexing",
"llm": "openai_llm.gpt-4o-mini",
"embedder": "openai_embedder.text-embedding-3-large",
"document_store": "qdrant",
"engine": "wren_ui"
}
]
}
The function processes this input and returns a dictionary with the following structure:
{
"indexing": {
"llm": "openai_llm.gpt-4o-mini",
"embedder": "openai_embedder.text-embedding-3-large",
"document_store": "qdrant",
"engine": "wren_ui",
}
}
Args:
entry (dict): The input pipeline configuration dictionary.
Returns:
dict: A processed dictionary with standardized pipeline configuration.
"""
return {
pipe["name"]: {
"llm": pipe.get("llm"),
@@ -87,7 +257,22 @@ def process_pipeline(entry: dict) -> dict:
}
def convert_data(config: list[dict]) -> dict:
_TYPE_TO_PROCESSOR = {
"llm": llm_processor,
"embedder": embedder_processor,
"document_store": document_store_processor,
"engine": engine_processor,
"pipeline": pipeline_processor,
}
@dataclass
class Configuration:
providers: dict
pipelines: dict
def transform(config: list[dict]) -> Configuration:
returned = {
"embedder": {},
"llm": {},
@@ -96,17 +281,9 @@ def convert_data(config: list[dict]) -> dict:
"pipeline": {},
}
type_to_processor = {
"llm": process_llm,
"embedder": process_embedder,
"document_store": process_document_store,
"engine": process_engine,
"pipeline": process_pipeline,
}
for entry in config:
type = entry["type"]
processor = type_to_processor.get(type)
processor = _TYPE_TO_PROCESSOR.get(type)
if not processor:
logger.error(f"Unknown type: {type}")
raise ValueError(f"Unknown type: {type}")
@@ -114,12 +291,16 @@ def convert_data(config: list[dict]) -> dict:
converted = processor(entry)
returned[type].update(converted)
return returned
return Configuration(
providers={k: v for k, v in returned.items() if k != "pipeline"},
pipelines=returned["pipeline"],
)
def init_providers(
engine_config: EngineConfig,
) -> Tuple[LLMProvider, EmbedderProvider, DocumentStoreProvider, Engine]:
# DEPRECATED: use generate_components instead
logger.info("Initializing providers...")
loader.import_mods()
@@ -136,15 +317,15 @@ def init_providers(
class Wrapper(Mapping):
def __init__(
self,
llm_provider: LLMProvider,
embedder_provider: EmbedderProvider,
document_store_provider: DocumentStoreProvider,
engine: Engine,
):
def __init__(self):
from src.utils import load_env_vars
load_env_vars()
self.value = PipelineComponent(
llm_provider, embedder_provider, document_store_provider, engine
*init_providers(
engine_config=EngineConfig(provider=os.getenv("ENGINE", "wren_ui"))
)
)
def __getitem__(self, key):
@@ -160,55 +341,59 @@ class Wrapper(Mapping):
return len(self.value)
def reset_document_store(document_store_provider: DocumentStoreProvider):
document_store_provider.get_store(recreate_index=True)
document_store_provider.get_store(
dataset_name="table_descriptions", recreate_index=True
)
document_store_provider.get_store(
dataset_name="view_questions", recreate_index=True
)
def generate_components(configs: list[dict]) -> dict[str, PipelineComponent]:
"""
Generate pipeline components from configuration.
This function takes a list of configuration dictionaries and generates pipeline components
based on the provided configurations. The configurations are processed into a standardized
format and then instantiated into actual provider objects.
def generate_components(force_deploy: bool = False) -> dict[str, PipelineComponent]:
raw = load_config()
if not raw:
(
llm_provider,
embedder_provider,
document_store_provider,
engine,
) = init_providers(EngineConfig(provider=os.getenv("ENGINE", "wren_ui")))
Args:
configs (list[dict]): A list of configuration dictionaries.
if force_deploy:
reset_document_store(document_store_provider)
Returns:
dict: A dictionary of pipeline components.
# if no config, initialize the providers from the environment variables
return Wrapper(llm_provider, embedder_provider, document_store_provider, engine)
Note:
instantiated_providers example:
{
"embedder": {
"openai_embedder.text-embedding-3-large": <EmbedderProvider>
},
"llm": {
"openai_llm.gpt-4o-mini": <LLMProvider>
},
...
}
config = convert_data(raw)
"""
loader.import_mods()
providers = {
"embedder": config.get("embedder", {}),
"llm": config.get("llm", {}),
"document_store": config.get("document_store", {}),
"engine": config.get("engine", {}),
}
# DEPRECATED: remove this fallback in the future
if not configs:
message = """
Warning: No configuration provided. Falling back to environment variables for settings.
This is a legacy approach and will be deprecated soon. Please refer to the README for
instructions on migrating to the new configuration format. It is strongly recommended
to update your configuration to ensure future compatibility and take advantage of new features.
"""
logger.warning(message)
return Wrapper()
if force_deploy:
reset_document_store(provider_factory(providers["document_store"]))
config = transform(configs)
instantiated_providers = {
category: {
type: {
identifier: provider_factory(config)
for identifier, config in configs.items()
}
for category, configs in providers.items()
for type, configs in config.providers.items()
}
def get(type: str, components: dict):
return instantiated_providers[type].get(components.get(type))
identifier = components.get(type)
return instantiated_providers[type].get(identifier)
def componentize(components: dict):
return PipelineComponent(
@@ -220,5 +405,5 @@ def generate_components(force_deploy: bool = False) -> dict[str, PipelineCompone
return {
pipe_name: componentize(components)
for pipe_name, components in config.get("pipeline", {}).items()
for pipe_name, components in config.pipelines.items()
}
wren-ai-service/src/providers/document_store/qdrant.py
+11
View File
@@ -345,12 +345,23 @@ class QdrantProvider(DocumentStoreProvider):
or get_default_embedding_model_dim(
os.getenv("EMBEDDER_PROVIDER", "openai_embedder")
),
recreate_index: bool = (
bool(os.getenv("SHOULD_FORCE_DEPLOY"))
if os.getenv("SHOULD_FORCE_DEPLOY")
else False
),
**_,
):
self._location = location
self._api_key = Secret.from_token(api_key) if api_key else None
self._timeout = timeout
self._embedding_model_dim = embedding_model_dim
self._reset_document_store(recreate_index)
def _reset_document_store(self, recreate_index: bool):
self.get_store(recreate_index=recreate_index)
self.get_store(dataset_name="table_descriptions", recreate_index=recreate_index)
self.get_store(dataset_name="view_questions", recreate_index=recreate_index)
def get_store(
self,
wren-ai-service/src/providers/embedder/azure_openai.py
+2 -1
View File
@@ -195,7 +195,8 @@ class AsyncDocumentEmbedder(AzureOpenAIDocumentEmbedder):
class AzureOpenAIEmbedderProvider(EmbedderProvider):
def __init__(
self,
api_key: Secret = Secret.from_env_var("EMBEDDER_AZURE_OPENAI_API_KEY"),
api_key: Secret = Secret.from_env_var("AZURE_OPENAI_API_KEY")
or Secret.from_env_var("EMBEDDER_AZURE_OPENAI_API_KEY"),
api_base: str = os.getenv("EMBEDDER_AZURE_OPENAI_API_BASE"),
api_version: str = os.getenv("EMBEDDER_AZURE_OPENAI_VERSION"),
model: str = os.getenv("EMBEDDING_MODEL") or EMBEDDING_MODEL,
wren-ai-service/src/providers/embedder/openai.py
+2 -1
View File
@@ -183,7 +183,8 @@ class AsyncDocumentEmbedder(OpenAIDocumentEmbedder):
class OpenAIEmbedderProvider(EmbedderProvider):
def __init__(
self,
api_key: str = os.getenv("EMBEDDER_OPENAI_API_KEY"),
api_key: str = os.getenv("OPENAI_API_KEY")
or os.getenv("EMBEDDER_OPENAI_API_KEY"),
api_base: str = os.getenv("EMBEDDER_OPENAI_API_BASE")
or EMBEDDER_OPENAI_API_BASE,
model: str = os.getenv("EMBEDDING_MODEL") or EMBEDDING_MODEL,
wren-ai-service/src/providers/engine/wren.py
+4 -5
View File
@@ -153,18 +153,17 @@ class WrenEngine(Engine):
def __init__(
self,
endpoint: str = os.getenv("WREN_ENGINE_ENDPOINT"),
manifest: str = os.getenv("WREN_ENGINE_MANIFEST"),
**_,
):
self._endpoint = endpoint
self._manifest = manifest
logger.info("Using Engine: wren_engine")
async def execute_sql(
self,
sql: str,
session: aiohttp.ClientSession,
properties: Dict[str, Any] = {
"manifest": os.getenv("WREN_ENGINE_MANIFEST"),
},
dry_run: bool = True,
timeout: float = 30.0,
**kwargs,
@@ -180,9 +179,9 @@ class WrenEngine(Engine):
api_endpoint,
json={
"manifest": orjson.loads(
base64.b64decode(properties.get("manifest"))
base64.b64decode(self._manifest)
)
if properties.get("manifest")
if self._manifest
else {},
"sql": remove_limit_statement(sql),
"limit": 1 if dry_run else 500,
wren-ai-service/src/providers/llm/azure_openai.py
+3 -4
View File
@@ -125,7 +125,8 @@ class AsyncGenerator(AzureOpenAIGenerator):
class AzureOpenAILLMProvider(LLMProvider):
def __init__(
self,
api_key: Secret = Secret.from_env_var("LLM_AZURE_OPENAI_API_KEY"),
api_key: Secret = Secret.from_env_var("AZURE_OPENAI_API_KEY")
or Secret.from_env_var("LLM_AZURE_OPENAI_API_KEY"),
api_base: str = os.getenv("LLM_AZURE_OPENAI_API_BASE"),
api_version: str = os.getenv("LLM_AZURE_OPENAI_VERSION"),
model: str = os.getenv("GENERATION_MODEL") or GENERATION_MODEL,
@@ -151,6 +152,7 @@ class AzureOpenAILLMProvider(LLMProvider):
logger.info(
f"Using AzureOpenAI LLM with API version: {self._generation_api_version}"
)
logger.info(f"Using AzureOpenAI LLM model kwargs: {self._model_kwargs}")
def get_generator(
self,
@@ -159,9 +161,6 @@ class AzureOpenAILLMProvider(LLMProvider):
generation_kwargs: Optional[Dict[str, Any]] = None,
streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
):
logger.info(
f"Creating Azure OpenAI generator with model kwargs: {self._model_kwargs}"
)
return AsyncGenerator(
api_key=self._generation_api_key,
model=self._generation_model,
wren-ai-service/src/providers/llm/ollama.py
+1 -3
View File
@@ -145,6 +145,7 @@ class OllamaLLMProvider(LLMProvider):
logger.info(f"Using Ollama LLM: {self._generation_model}")
logger.info(f"Using Ollama URL: {self._url}")
logger.info(f"Using Ollama model kwargs: {self._model_kwargs}")
def get_generator(
self,
@@ -153,9 +154,6 @@ class OllamaLLMProvider(LLMProvider):
generation_kwargs: Optional[Dict[str, Any]] = None,
streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
):
logger.info(
f"Creating Ollama generator with model kwargs: {self._model_kwargs}"
)
return AsyncGenerator(
model=self._generation_model,
url=f"{self._url}/api/generate",
wren-ai-service/src/providers/llm/openai.py
+5 -10
View File
@@ -157,7 +157,7 @@ class AsyncGenerator(OpenAIGenerator):
class OpenAILLMProvider(LLMProvider):
def __init__(
self,
api_key: str = os.getenv("LLM_OPENAI_API_KEY"),
api_key: str = os.getenv("OPENAI_API_KEY") or os.getenv("LLM_OPENAI_API_KEY"),
api_base: str = os.getenv("LLM_OPENAI_API_BASE") or LLM_OPENAI_API_BASE,
model: str = os.getenv("GENERATION_MODEL") or GENERATION_MODEL,
kwargs: Dict[str, Any] = (
@@ -179,8 +179,12 @@ class OpenAILLMProvider(LLMProvider):
logger.info(f"Using OpenAILLM provider with API base: {self._api_base}")
if self._api_base == LLM_OPENAI_API_BASE:
logger.info(f"Using OpenAI LLM: {self._generation_model}")
logger.info(f"Using OpenAI LLM model kwargs: {self._model_kwargs}")
else:
logger.info(f"Using OpenAI API-compatible LLM: {self._generation_model}")
logger.info(
f"Using OpenAI API-compatible LLM model kwargs: {self._model_kwargs}"
)
def get_generator(
self,
@@ -189,15 +193,6 @@ class OpenAILLMProvider(LLMProvider):
generation_kwargs: Optional[Dict[str, Any]] = None,
streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
):
if self._api_base == LLM_OPENAI_API_BASE:
logger.info(
f"Creating OpenAI generator {self._generation_model} with model kwargs: {self._model_kwargs}"
)
else:
logger.info(
f"Creating OpenAI API-compatible generator {self._generation_model} with model kwargs: {self._model_kwargs}"
)
return AsyncGenerator(
api_key=self._api_key,
api_base_url=self._api_base,
wren-ai-service/src/utils.py
+22 -11
View File
@@ -35,7 +35,14 @@ class CustomFormatter(logging.Formatter):
return formatter.format(record)
def setup_custom_logger(name, level=logging.INFO):
def setup_custom_logger(name, level_str: str):
level_str = level_str.upper()
if level_str not in logging._nameToLevel:
raise ValueError(f"Invalid logging level: {level_str}")
level = logging._nameToLevel[level_str]
handler = logging.StreamHandler()
handler.setFormatter(CustomFormatter())
@@ -46,6 +53,7 @@ def setup_custom_logger(name, level=logging.INFO):
def load_env_vars() -> str:
# DEPRECATED: This method is deprecated and will be removed in the future
if Path(".env.dev").exists():
load_dotenv(".env.dev", override=True)
return "dev"
@@ -56,7 +64,9 @@ def load_env_vars() -> str:
def timer(func):
@functools.wraps(func)
def wrapper_timer(*args, **kwargs):
if os.getenv("ENABLE_TIMER", False):
from src.config import settings
if settings.enable_timer:
startTime = time.perf_counter()
result = func(*args, **kwargs)
endTime = time.perf_counter()
@@ -80,7 +90,9 @@ def async_timer(func):
@functools.wraps(func)
async def wrapper_timer(*args, **kwargs):
if os.getenv("ENABLE_TIMER", False):
from src.config import settings
if settings.enable_timer:
startTime = time.perf_counter()
result = await process(func, *args, **kwargs)
endTime = time.perf_counter()
@@ -102,18 +114,17 @@ def remove_trailing_slash(endpoint: str) -> str:
def init_langfuse():
enabled = bool(os.getenv("LANGFUSE_ENABLE", False))
host = os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com")
from src.config import settings
langfuse_context.configure(
enabled=enabled,
public_key=os.getenv("LANGFUSE_PUBLIC_KEY", ""),
secret_key=os.getenv("LANGFUSE_SECRET_KEY", ""),
host=host,
enabled=settings.langfuse_enable,
host=settings.langfuse_host,
public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
secret_key=os.getenv("LANGFUSE_SECRET_KEY"),
)
logger.info(f"LANGFUSE_ENABLE: {enabled}")
logger.info(f"LANGFUSE_HOST: {host}")
logger.info(f"LANGFUSE_ENABLE: {settings.langfuse_enable}")
logger.info(f"LANGFUSE_HOST: {settings.langfuse_host}")
def trace_metadata(func):
wren-ai-service/src/web/v1/routers/semantics_preparations.py
+2 -7
View File
@@ -1,4 +1,3 @@
import uuid
from dataclasses import asdict
from fastapi import APIRouter, BackgroundTasks, Depends
@@ -9,10 +8,6 @@ from src.globals import (
get_service_container,
get_service_metadata,
)
router = APIRouter()
from src.web.v1.services.semantics_preparation import (
SemanticsPreparationRequest,
SemanticsPreparationResponse,
@@ -20,6 +15,8 @@ from src.web.v1.services.semantics_preparation import (
SemanticsPreparationStatusResponse,
)
router = APIRouter()
"""
Semantics Preparation Router
@@ -94,5 +91,3 @@ async def get_prepare_semantics_status(
return service_container.semantics_preparation_service.get_prepare_semantics_status(
SemanticsPreparationStatusRequest(mdl_hash=mdl_hash)
)
wren-ai-service/src/web/v1/routers/sql_explanations.py
-1
View File
@@ -98,4 +98,3 @@ async def get_sql_explanation_result(
return service_container.sql_explanation_service.get_sql_explanation_result(
SQLExplanationResultRequest(query_id=query_id)
)
wren-ai-service/src/web/v1/routers/sql_regenerations.py
+1
View File
@@ -80,6 +80,7 @@ Usage:
Note: The actual SQL generation occurs in the background using FastAPI's BackgroundTasks.
"""
@router.post("/sql-regenerations")
async def sql_regeneration(
sql_regeneration_request: SQLRegenerationRequest,
wren-ai-service/tests/data/config.test.yaml
+92
View File
@@ -0,0 +1,92 @@
type: llm
provider: openai_llm
models:
- model: gpt-4o-mini
kwargs:
{
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
api_base: https://api.openai.com/v1
---
type: embedder
provider: openai_embedder
models:
- model: text-embedding-3-large
dimension: 3072
api_base: https://api.openai.com/v1
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://localhost:3000
---
type: document_store
provider: qdrant
location: http://localhost:6333
embedding_model_dim: 3072
timeout: 120
---
type: pipeline
pipes:
- name: indexing
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: retrieval
llm: openai_llm.gpt-4o-mini
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: historical_question
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_correction
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: followup_sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_summary
llm: openai_llm.gpt-4o-mini
- name: sql_answer
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_breakdown
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_expansion
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_explanation
llm: openai_llm.gpt-4o-mini
- name: sql_regeneration
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: semantics_description
llm: openai_llm.gpt-4o-mini
- name: relationship_recommendation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
---
settings:
host: 127.0.0.1
port: 5556
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 1000
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: false
enable_timer: false
logging_level: INFO
development: false
wren-ai-service/tests/pytest/providers/test_providers.py
+121
View File
@@ -0,0 +1,121 @@
from pytest_mock import MockerFixture
from src.core.engine import Engine
from src.core.pipeline import PipelineComponent
from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider
from src.providers import Configuration, generate_components, transform
def test_transform():
config = [
{
"type": "llm",
"provider": "openai_llm",
"models": [
{"model": "gpt-4", "kwargs": {"temperature": 0, "max_tokens": 4096}}
],
},
{
"type": "embedder",
"provider": "openai_embedder",
"models": [{"model": "text-embedding-ada-002", "dimension": 1536}],
},
{
"type": "document_store",
"provider": "qdrant",
"kwargs": {"host": "localhost", "port": 6333},
},
{
"type": "engine",
"provider": "wren_ui",
"kwargs": {"host": "localhost", "port": 8000},
},
{
"type": "pipeline",
"pipes": [
{
"name": "indexing",
"llm": "openai_llm.gpt-4",
"embedder": "openai_embedder.text-embedding-ada-002",
"document_store": "qdrant",
"engine": "wren_ui",
}
],
},
]
result = transform(config)
assert isinstance(result, Configuration)
assert "openai_llm.gpt-4" in result.providers["llm"]
assert "openai_embedder.text-embedding-ada-002" in result.providers["embedder"]
assert "qdrant" in result.providers["document_store"]
assert "wren_ui" in result.providers["engine"]
assert "indexing" in result.pipelines
def test_generate_components(mocker: MockerFixture):
# Mock the provider_factory to return mock objects
mocker.patch(
"src.providers.provider_factory",
side_effect=[
mocker.Mock(spec=EmbedderProvider),
mocker.Mock(spec=LLMProvider),
mocker.Mock(spec=DocumentStoreProvider),
mocker.Mock(spec=Engine),
],
)
config = [
{
"type": "llm",
"provider": "openai_llm",
"models": [{"model": "gpt-4", "kwargs": {}}],
},
{
"type": "embedder",
"provider": "openai_embedder",
"models": [{"model": "text-embedding-ada-002", "dimension": 1536}],
},
{"type": "document_store", "provider": "qdrant", "kwargs": {}},
{"type": "engine", "provider": "wren_ui", "kwargs": {}},
{
"type": "pipeline",
"pipes": [
{
"name": "indexing",
"llm": "openai_llm.gpt-4",
"embedder": "openai_embedder.text-embedding-ada-002",
"document_store": "qdrant",
"engine": "wren_ui",
}
],
},
]
result = generate_components(config)
assert "indexing" in result
assert isinstance(result["indexing"], PipelineComponent)
assert isinstance(result["indexing"].embedder_provider, EmbedderProvider)
assert isinstance(result["indexing"].llm_provider, LLMProvider)
assert isinstance(result["indexing"].document_store_provider, DocumentStoreProvider)
assert isinstance(result["indexing"].engine, Engine)
def test_generate_components_empty_config(mocker: MockerFixture):
# Mock the Wrapper class
mock_wrapper = mocker.patch("src.providers.Wrapper")
# Create a mock PipelineComponent
mock_pipeline_component = mocker.Mock(spec=PipelineComponent)
# Set the return value of Wrapper() to be the mock PipelineComponent
mock_wrapper.return_value = mock_pipeline_component
# Mock the load_env_vars function to avoid side effects
mocker.patch("src.utils.load_env_vars")
result = generate_components([])
assert isinstance(result, PipelineComponent)
assert result == mock_pipeline_component
wren-ai-service/tests/pytest/test_config.py
+127
View File
@@ -0,0 +1,127 @@
from unittest.mock import mock_open, patch
import yaml
from src.config import Settings
def test_settings_default_values():
settings = Settings()
assert settings.host == "127.0.0.1"
assert settings.port == 5556
assert settings.column_indexing_batch_size == 50
assert settings.table_retrieval_size == 10
assert settings.table_column_retrieval_size == 1000
assert settings.query_cache_ttl == 3600
assert settings.query_cache_maxsize == 1_000_000
assert settings.langfuse_host == "https://cloud.langfuse.com"
assert settings.langfuse_enable is True
assert settings.logging_level == "INFO"
assert settings.enable_timer is False
assert settings.development is False
assert settings.config_path == "config.yaml"
def test_settings_env_var_override():
with patch.dict(
"os.environ",
{
"WREN_AI_SERVICE_HOST": "0.0.0.0",
"WREN_AI_SERVICE_PORT": "8000",
"LOGGING_LEVEL": "DEBUG",
},
):
settings = Settings()
assert settings.host == "0.0.0.0"
assert settings.port == 8000
assert settings.logging_level == "DEBUG"
def test_settings_env_dev_override():
# Mock the content of .env.dev file
mock_env_dev_content = """
WREN_AI_SERVICE_HOST=localhost
WREN_AI_SERVICE_PORT=7000
LOGGING_LEVEL=WARNING
"""
# Mock the load_dotenv function
with patch("src.config.load_dotenv") as mock_load_dotenv:
# Set up the mock to load our custom environment variables
def side_effect(path, override):
import os
for line in mock_env_dev_content.strip().split("\n"):
key, value = line.strip().split("=")
os.environ[key] = value
mock_load_dotenv.side_effect = side_effect
settings = Settings()
assert settings.host == "localhost"
assert settings.port == 7000
assert settings.logging_level == "WARNING"
def test_settings_yaml_config_override():
# Mock YAML config content
mock_yaml_content = """
settings:
host: 192.168.1.100
port: 9000
column_indexing_batch_size: 75
table_retrieval_size: 15
logging_level: ERROR
development: true
"""
# Patch the open function to return our mock YAML content
with patch("builtins.open", mock_open(read_data=mock_yaml_content)):
# Patch os.path.exists to return True for our config file
with patch("os.path.exists", return_value=True):
settings = Settings()
assert settings.host == "192.168.1.100"
assert settings.port == 9000
assert settings.column_indexing_batch_size == 75
assert settings.table_retrieval_size == 15
assert settings.logging_level == "ERROR"
assert settings.development is True
# Check that a value not in the YAML config remains at its default
assert settings.query_cache_maxsize == 1_000_000
def test_settings_components():
mock_config_content = [
{
"settings": {
"host": "0.0.0.0",
"port": 8000,
"column_indexing_batch_size": 100,
"table_retrieval_size": 20,
"logging_level": "DEBUG",
}
},
{
"type": "llm",
"provider": "openai_llm",
"models": [{"model": "gpt-4", "kwargs": {}}],
},
]
with patch(
"builtins.open",
new_callable=mock_open,
read_data=yaml.dump_all(mock_config_content),
):
settings = Settings()
assert len(settings._components) == 1
assert settings._components[0]["type"] == "llm"
assert settings._components[0]["provider"] == "openai_llm"
wren-ai-service/tests/pytest/test_main.py
+18 -7
View File
@@ -1,10 +1,21 @@
import json
import os
import uuid
import orjson
import pytest
from fastapi.testclient import TestClient
from src.__main__ import app
@pytest.fixture(scope="module", autouse=True)
def app():
os.environ["CONFIG_PATH"] = "tests/data/config.test.yaml"
from src.__main__ import app
yield app
# Clean up (if necessary)
del os.environ["CONFIG_PATH"]
GLOBAL_DATA = {
"semantics_preperation_id": str(uuid.uuid4()),
@@ -12,7 +23,7 @@ GLOBAL_DATA = {
}
def test_semantics_preparations():
def test_semantics_preparations(app):
with TestClient(app) as client:
semantics_preperation_id = GLOBAL_DATA["semantics_preperation_id"]
@@ -44,7 +55,7 @@ def test_semantics_preparations():
assert status == "finished"
def test_asks_with_successful_query():
def test_asks_with_successful_query(app):
with TestClient(app) as client:
semantics_preparation_id = GLOBAL_DATA["semantics_preperation_id"]
@@ -77,7 +88,7 @@ def test_asks_with_successful_query():
# assert r["summary"] is not None and r["summary"] != ""
def test_stop_asks():
def test_stop_asks(app):
with TestClient(app) as client:
query_id = GLOBAL_DATA["query_id"]
@@ -99,7 +110,7 @@ def test_stop_asks():
assert response.json()["status"] == "stopped"
def test_ask_details():
def test_ask_details(app):
with TestClient(app) as client:
response = client.post(
url="/v1/ask-details",
@@ -132,7 +143,7 @@ def test_ask_details():
assert step["cte_name"] == ""
def test_sql_regenerations():
def test_sql_regenerations(app):
with TestClient(app) as client:
response = client.post(
url="/v1/sql-regenerations",
@@ -222,7 +233,7 @@ LIMIT 1
assert response.json()["status"] == "finished" or "failed"
def test_web_error_handler():
def test_web_error_handler(app):
with TestClient(app) as client:
response = client.post(
url="/v1/asks",
wren-ai-service/tools/config/.env.dev.example
+10
View File
@@ -0,0 +1,10 @@
CONFIG_PATH=config.yaml
# vendor keys
OPENAI_API_KEY=
AZURE_OPENAI_API_KEY=
QDRANT_API_KEY=
# langfuse key
LANGFUSE_SECRET_KEY=
LANGFUSE_PUBLIC_KEY=
wren-ai-service/config.example.yaml → wren-ai-service/tools/config/config.example.yaml
+18 -11
View File
@@ -9,13 +9,6 @@ models:
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
api_key: sk-xxx
api_base: https://api.openai.com/v1
---
type: llm
provider: azure_openai_llm
models:
- model: gpt-4o
kwargs:
{
@@ -24,9 +17,8 @@ models:
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
api_key: sk-xxx
api_base: https://api.openai.com/v1
api_version: "2024-05-13"
timeout: 120
---
type: embedder
@@ -34,7 +26,6 @@ provider: openai_embedder
models:
- model: text-embedding-3-large
dimension: 3072
api_key: sk-xxx
api_base: https://api.openai.com/v1
timeout: 120
@@ -61,9 +52,9 @@ manifest: ""
type: document_store
provider: qdrant
location: http://localhost:6333
api_key: ""
embedding_model_dim: 3072
timeout: 120
recreate_index: false
---
type: pipeline
@@ -116,3 +107,19 @@ pipes:
document_store: qdrant
- name: data_assistance
llm: openai_llm.gpt-4o-mini
---
settings:
host: 127.0.0.1
port: 5556
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
allow_using_db_schemas_without_pruning: false
query_cache_maxsize: 1000
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
enable_timer: false
logging_level: DEBUG
development: true
wren-ai-service/tools/config/config.full.yaml
+159
View File
@@ -0,0 +1,159 @@
type: llm
provider: openai_llm
models:
- model: gpt-4o-mini
kwargs:
{
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
api_base: https://api.openai.com/v1
timeout: 120
---
type: llm
provider: azure_openai_llm
models:
- model: gpt-4o
kwargs:
{
"temperature": 0,
"n": 1,
"max_tokens": 4096,
"response_format": { "type": "json_object" },
}
api_base: https://<your-endpoint>.openai.azure.com/
api_version: "2024-05-13"
timeout: 120
---
type: llm
provider: ollama_llm
models:
- model: gemma2:9b
kwargs: { "temperature": 0 }
url: http://localhost:11434
timeout: 120
---
type: embedder
provider: openai_embedder
models:
- model: text-embedding-3-large
dimension: 3072
api_base: https://api.openai.com/v1
timeout: 120
---
type: embedder
provider: azure_openai_embedder
models:
- model: text-embedding-3-small
dimension: 256
api_base: https://<your-endpoint>.openai.azure.com/
api_version: "2024-05-13"
timeout: 120
---
type: embedder
provider: ollama_embedder
models:
- model: nomic-embed-text
dimension: 786
url: http://localhost:11434
timeout: 120
---
type: engine
provider: wren_ui
endpoint: http://localhost:3000
---
type: engine
provider: wren_ibis
endpoint: http://localhost:8000
source: bigquery
manifest: "" # base64 encoded string of the MDL
connection_info: "" # base64 encoded string of the connection info
---
type: engine
provider: wren_engine
endpoint: http://localhost:8080
manifest: ""
---
type: document_store
provider: qdrant
location: http://localhost:6333
embedding_model_dim: 3072
timeout: 120
recreate_index: false
---
type: pipeline
pipes:
- name: indexing
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: retrieval
llm: openai_llm.gpt-4o-mini
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: historical_question
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_correction
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: followup_sql_generation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_summary
llm: openai_llm.gpt-4o-mini
- name: sql_answer
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_breakdown
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_expansion
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: sql_explanation
llm: openai_llm.gpt-4o-mini
- name: sql_regeneration
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: semantics_description
llm: openai_llm.gpt-4o-mini
- name: relationship_recommendation
llm: openai_llm.gpt-4o-mini
engine: wren_ui
- name: intent_classification
llm: openai_llm.gpt-4o-mini
embedder: openai_embedder.text-embedding-3-large
document_store: qdrant
- name: data_assistance
llm: openai_llm.gpt-4o-mini
---
settings:
host: 127.0.0.1
port: 5556
column_indexing_batch_size: 50
table_retrieval_size: 10
table_column_retrieval_size: 100
query_cache_maxsize: 1000
allow_using_db_schemas_without_pruning: false
query_cache_ttl: 3600
langfuse_host: https://cloud.langfuse.com
langfuse_enable: true
enable_timer: false
logging_level: INFO
development: false
wren-ai-service/tools/dev/.env.example → wren-ai-service/tools/dev/.env
View File
wren-launcher/commands/launch.go
+15 -9
View File
@@ -114,12 +114,12 @@ func askForGenerationModel() (string, error) {
return result, nil
}
func isEnvFileValidForCustomLLM(projectDir string) error {
// validate if .env.ai file exists in ~/.wrenai
envFilePath := path.Join(projectDir, ".env.ai")
func isConfigFileValidForCustomLLM(projectDir string) error {
// validate if config.yaml file exists in ~/.wrenai
configFilePath := path.Join(projectDir, "config.yaml")
if _, err := os.Stat(envFilePath); os.IsNotExist(err) {
errMessage := fmt.Sprintf("Please create a .env.ai file in %s first, more details at https://docs.getwren.ai/oss/installation/custom_llm#running-wren-ai-with-your-custom-llm-or-document-store", projectDir)
if _, err := os.Stat(configFilePath); os.IsNotExist(err) {
errMessage := fmt.Sprintf("Please create a config.yaml file in %s first, more details at https://docs.getwren.ai/oss/installation/custom_llm#running-wren-ai-with-your-custom-llm-or-document-store", projectDir)
return errors.New(errMessage)
}
@@ -191,9 +191,15 @@ func Launch() {
if shouldReturn {
return
}
// prepare config.yaml file for OpenAI
err := utils.PrepareConfigFileForOpenAI(projectDir, openaiGenerationModel)
if err != nil {
panic(err)
}
} else {
// check if .env.ai file exists
err := isEnvFileValidForCustomLLM(projectDir)
// check if config.yaml file exists
err := isConfigFileValidForCustomLLM(projectDir)
if err != nil {
panic(err)
}
@@ -301,8 +307,8 @@ func getOpenaiGenerationModel() (string, bool) {
// validate if input args is a valid generation model
pterm.Info.Println("OpenAI generation model is provided")
validModels := map[string]bool{
"gpt-4o-mini": true,
"gpt-4o": true,
"gpt-4o-mini": true,
"gpt-4o": true,
}
if !validModels[openaiGenerationModel] {
pterm.Error.Println("Invalid generation model", openaiGenerationModel)
wren-launcher/utils/docker.go
+111 -31
View File
@@ -8,6 +8,7 @@ import (
"os"
"path"
"regexp"
"strings"
"github.com/Canner/WrenAI/wren-launcher/config"
"github.com/docker/cli/cli/command"
@@ -23,10 +24,10 @@ import (
const (
// please change the version when the version is updated
WREN_PRODUCT_VERSION string = "0.10.0"
DOCKER_COMPOSE_YAML_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/docker-compose.yaml"
DOCKER_COMPOSE_LLM_YAML_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/docker-compose.llm.yaml"
DOCKER_COMPOSE_ENV_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/.env.example"
WREN_PRODUCT_VERSION string = "0.10.0"
DOCKER_COMPOSE_YAML_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/docker-compose.yaml"
DOCKER_COMPOSE_ENV_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/.env.example"
AI_SERVICE_CONFIG_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/config.example.yaml"
)
func replaceEnvFileContent(content string, projectDir string, openaiApiKey string, openAIGenerationModel string, hostPort int, aiPort int, userUUID string, telemetryEnabled bool) string {
@@ -38,18 +39,12 @@ func replaceEnvFileContent(content string, projectDir string, openaiApiKey strin
reg = regexp.MustCompile(`SHOULD_FORCE_DEPLOY=(.*)`)
str = reg.ReplaceAllString(str, "SHOULD_FORCE_DEPLOY=1")
// replace LLM_OPENAI_API_KEY
// Might be overwritten by the .env.ai file
reg = regexp.MustCompile(`LLM_OPENAI_API_KEY=(.*)`)
str = reg.ReplaceAllString(str, "LLM_OPENAI_API_KEY="+openaiApiKey)
// replace EMBEDDER_OPENAI_API_KEY,
// Might be overwritten by the .env.ai file
reg = regexp.MustCompile(`EMBEDDER_OPENAI_API_KEY=(.*)`)
str = reg.ReplaceAllString(str, "EMBEDDER_OPENAI_API_KEY="+openaiApiKey)
// replace OPENAI_API_KEY
reg = regexp.MustCompile(`(?m)^OPENAI_API_KEY=(.*)`)
str = reg.ReplaceAllString(str, "OPENAI_API_KEY="+openaiApiKey)
// replace GENERATION_MODEL
// Might be overwritten by the .env.ai file
// it seems like using for telemetry to know the model, might be we can remove this in the future and provide a endpoint to get the information
reg = regexp.MustCompile(`GENERATION_MODEL=(.*)`)
str = reg.ReplaceAllString(str, "GENERATION_MODEL="+openAIGenerationModel)
@@ -136,6 +131,101 @@ func prepareUserUUID(projectDir string) (string, error) {
return userUUID, nil
}
func PrepareConfigFileForOpenAI(projectDir string, generationModel string) error {
// download config.yaml file
configPath := path.Join(projectDir, "config.yaml")
pterm.Info.Println("Downloading config.yaml file to", configPath)
err := downloadFile(configPath, AI_SERVICE_CONFIG_URL)
if err != nil {
return err
}
// read the config.yaml file
content, err := os.ReadFile(configPath)
if err != nil {
return err
}
// replace the generation model in config.yaml
config := string(content)
config = strings.ReplaceAll(config, "openai_llm.gpt-4o-mini", "openai_llm."+generationModel)
// disable the langfuse for starting wren-ai from the launcher
config = strings.ReplaceAll(config, "langfuse_enable: true", "langfuse_enable: false")
// write back to config.yaml
err = os.WriteFile(configPath, []byte(config), 0644)
if err != nil {
return err
}
return nil
}
func mergeEnvContent(newEnvFile string, envFileContent string) (string, error) {
// Check if .env file does not exist
if _, err := os.Stat(newEnvFile); err != nil {
return envFileContent, nil
}
// File exists, read existing content
existingContent, err := os.ReadFile(newEnvFile)
if err != nil {
return "", err
}
// Split both contents into lines
existingLines := strings.Split(string(existingContent), "\n")
newLines := strings.Split(envFileContent, "\n")
// Create map of existing env vars
existingEnvVars := make(map[string]string)
// Helper function to parse env var line
parseEnvVar := func(line string) (string, string, bool) {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
return "", "", false
}
parts := strings.SplitN(line, "=", 2)
if len(parts) != 2 {
return "", "", false
}
return parts[0], parts[1], true
}
// Parse existing env vars
for _, line := range existingLines {
if key, val, ok := parseEnvVar(line); ok {
existingEnvVars[key] = val
}
}
// Merge with new values
for _, line := range newLines {
if key, val, ok := parseEnvVar(line); ok && val != "" {
existingEnvVars[key] = val
}
}
// Build merged content
var mergedLines []string
for _, line := range newLines {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
mergedLines = append(mergedLines, line)
continue
}
if key, _, ok := parseEnvVar(line); ok {
if val, exists := existingEnvVars[key]; exists {
mergedLines = append(mergedLines, key+"="+val)
}
}
}
// Update envFileContent with merged content
envFileContent = strings.Join(mergedLines, "\n")
return envFileContent, nil
}
func PrepareDockerFiles(openaiApiKey string, openaiGenerationModel string, hostPort int, aiPort int, projectDir string, telemetryEnabled bool) error {
// download docker-compose file
composeFile := path.Join(projectDir, "docker-compose.yaml")
@@ -145,14 +235,6 @@ func PrepareDockerFiles(openaiApiKey string, openaiGenerationModel string, hostP
return err
}
// download docker-compose.llm.yaml file
composeLLMFile := path.Join(projectDir, "docker-compose.llm.yaml")
pterm.Info.Println("Downloading docker-compose.llm file to", composeLLMFile)
err = downloadFile(composeLLMFile, DOCKER_COMPOSE_LLM_YAML_URL)
if err != nil {
return err
}
userUUID, err := prepareUserUUID(projectDir)
if err != nil {
return err
@@ -184,6 +266,13 @@ func PrepareDockerFiles(openaiApiKey string, openaiGenerationModel string, hostP
telemetryEnabled,
)
newEnvFile := getEnvFilePath(projectDir)
// merge the env file content with the existing env file
envFileContent, err = mergeEnvContent(newEnvFile, envFileContent)
if err != nil {
return err
}
// write the file
err = os.WriteFile(newEnvFile, []byte(envFileContent), 0644)
if err != nil {
@@ -210,15 +299,6 @@ func RunDockerCompose(projectName string, projectDir string, llmProvider string)
envFiles := []string{envFile}
configPaths := []string{composeFilePath}
if llmProvider == "Custom" {
customEnvFile := path.Join(projectDir, ".env.ai")
llmComposeFile := path.Join(projectDir, "docker-compose.llm.yaml")
// Note: there are env variables with the same name in .env.ai and .env files
// Be aware of the order of the env files
envFiles = append(envFiles, customEnvFile)
configPaths = append(configPaths, llmComposeFile)
}
// docker-compose up
dockerCli, err := command.NewDockerCli()
if err != nil {