chore(wren-ai-service): minor update (#592)

* adjust

* update

* add env

* fix bug

* add qdrant_host

* update kustomizations configs

* change AI_SERVICE_ENV to SHOULD_FORCE_DEPLOY

* remove

* update

* update.gitignore

* remove file

* adjust entrypoint

deployment/kustomizations/.gitignore

@@ -1,2 +1,2 @@
-*.kustimized.yaml
+*.kustomized.yaml
 charts/*

deployment/kustomizations/base/cm.yaml

@@ -26,6 +26,9 @@ data:
   EMBEDDER_PROVIDER: "openai_embedder"
   EMBEDDER_OPENAI_API_BASE: "https://api.openai.com/v1"
 
+  # Document store related
+  QDRANT_HOST: "wren-qdrant"
+  DOCUMENT_STORE_PROVIDER: "qdrant"
 
   # Telemetry
   POSTHOG_HOST: "https://app.posthog.com"

deployment/kustomizations/base/deploy-wren-ai-service.yaml

@@ -58,13 +58,13 @@ spec:
                 key: GENERATION_MODEL
           - name: QDRANT_HOST
             valueFrom:
-              secretKeyRef:
-                name: wrenai-secrets
+              configMapKeyRef:
+                name: wren-config
                 key: QDRANT_HOST
           - name: DOCUMENT_STORE_PROVIDER
             valueFrom:
-              secretKeyRef:
-                name: wrenai-secrets
+              configMapKeyRef:
+                name: wren-config
                 key: DOCUMENT_STORE_PROVIDER
           - name: WREN_ENGINE_ENDPOINT
             valueFrom:

deployment/kustomizations/examples/secret-wren_example.yaml

@@ -30,9 +30,6 @@ data:
 
   POSTHOG_API_KEY: cGhjX2tleS1wbGFjZWhvbGRlcg==
   POSTHOG_HOST: aHR0cHM6Ly9hcHAucG9zdGhvZy5jb20=
-  QDRANT_HOST: d3Jlbi1xZHJhbnQ=
-  QDRANT_PORT: NjMzMw==
-  DOCUMENT_STORE_PROVIDER: cWRyYW50
   USER_UUID: MDAwMDAwMDAtMDAwMC0wMDAwLTAwMDAtMDAwMDAwMDAwMDAw
 ---
 apiVersion: v1

deployment/kustomizations/kustomization.yaml

@@ -39,13 +39,13 @@ images:
   - name: ghcr.io/canner/wren-bootstrap
     newTag: 0.1.5 # WREN_BOOTSTRAP_VERSION
   - name: ghcr.io/canner/wren-engine
-    newTag: 0.7.0 # WREN_ENGINE_VERSION
+    newTag: 0.9.0 # WREN_ENGINE_VERSION
   - name: ghcr.io/canner/wren-ui
-    newTag: 0.10.0 # WREN_UI_VERSION
+    newTag: 0.9.2 # WREN_UI_VERSION
   - name: ghcr.io/canner/wren-ai-service
-    newTag: 0.7.2 # WREN_AI_SERVICE_VERSION
+    newTag: 0.8.2 # WREN_AI_SERVICE_VERSION
   - name: ghcr.io/canner/wren-engine-ibis
-    newTag: 0.7.0 # IBIS_SERVER_VERSION
+    newTag: 0.9.0 # IBIS_SERVER_VERSION
 
 resources:
   - base/cm.yaml
@@ -55,8 +55,8 @@ resources:
   - base/pvc.yaml
   - base/svc.yaml
 ### Modify these examples first and uncomment them:
-  - examples/ingress-wren_example.yaml
-  - examples/certificate-wren_example.yaml
+  # - examples/ingress-wren_example.yaml
+  # - examples/certificate-wren_example.yaml
 ### Usually you do not need to generate a certificate for Qdrant
 #  - examples/certificate-qdrant_example.yaml
 ### Best practice is to create and deploy Secrets manually, not as part of kustomization or GitOps!

deployment/kustomizations/patches/cm.yaml

@@ -10,7 +10,7 @@
     WREN_PRODUCT_VERSION: "0.7.5"
     #fix:
     WREN_ENGINE_VERSION: "0.9.0"
-    WREN_AI_SERVICE_VERSION: "0.8.0"
+    WREN_AI_SERVICE_VERSION: "0.8.2"
     #fix:
     WREN_UI_VERSION: "0.9.2"
 

docker/.env.example

@@ -30,6 +30,8 @@ WREN_BOOTSTRAP_VERSION=0.1.5
 # AI service related env variables
 AI_SERVICE_ENABLE_TIMER=
 AI_SERVICE_LOGGING_LEVEL=INFO
+SHOULD_FORCE_DEPLOY=true
+QDRANT_HOST=qdrant
 
 # user id (uuid v4)
 USER_UUID=

docker/docker-compose-dev.yaml

@@ -43,6 +43,7 @@ services:
       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}
@@ -50,6 +51,7 @@ services:
       GENERATION_MODEL: ${GENERATION_MODEL}
       ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER}
       LOGGING_LEVEL: ${AI_SERVICE_LOGGING_LEVEL}
+      SHOULD_FORCE_DEPLOY: ${SHOULD_FORCE_DEPLOY}
       # sometimes the console won't show print messages,
       # using PYTHONUNBUFFERED: 1 can fix this
       PYTHONUNBUFFERED: 1
@@ -57,7 +59,7 @@ services:
       - wren
     depends_on:
       - qdrant
-      - wren-engine
+      - wren-ui
 
   ibis-server:
     image: ghcr.io/canner/wren-engine-ibis:${IBIS_SERVER_VERSION}

docker/docker-compose.yaml

@@ -55,6 +55,7 @@ services:
       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}
@@ -62,13 +63,13 @@ services:
       GENERATION_MODEL: ${GENERATION_MODEL}
       ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER}
       LOGGING_LEVEL: ${AI_SERVICE_LOGGING_LEVEL}
+      SHOULD_FORCE_DEPLOY: ${SHOULD_FORCE_DEPLOY}
       # sometimes the console won't show print messages,
       # using PYTHONUNBUFFERED: 1 can fix this
       PYTHONUNBUFFERED: 1
     networks:
       - wren
     depends_on:
-      - wren-engine
       - qdrant
 
   qdrant:

wren-ai-service/.env.dev.example

@@ -1,6 +1,7 @@
 # app related
 WREN_AI_SERVICE_HOST=127.0.0.1
 WREN_AI_SERVICE_PORT=5556
+SHOULD_FORCE_DEPLOY=
 
 ## LLM
 LLM_PROVIDER=openai_llm # openai_llm, azure_openai_llm, ollama_llm

wren-ai-service/entrypoint.sh

@@ -1,24 +1,58 @@
 #!/bin/bash
 set -e
 
+INTERVAL=1
+TIMEOUT=10
+
+# Wait for qdrant to be responsive
+echo "Waiting for qdrant to start..."
+current=0
+
+while ! nc -z $QDRANT_HOST 6333; do
+    sleep $INTERVAL
+    current=$((current + INTERVAL))
+    if [ $current -eq $TIMEOUT ]; then
+        echo "Timeout: qdrant did not start within $TIMEOUT seconds"
+        exit 1
+    fi
+done
+echo "qdrant has started."
+
 # Start wren-ai-service in the background
 uvicorn src.__main__:app --host 0.0.0.0 --port $WREN_AI_SERVICE_PORT --loop uvloop --http httptools &
 
-# Wait for wren-ui to be responsive
-echo "Waiting for wren-ui to start..."
-while ! nc -z -w 5 wren-ui $WREN_UI_PORT && ! nc -z -w 5 host.docker.internal $WREN_UI_PORT; do   
-    sleep 1  # wait for 1 second before check again
-done
-
 # Wait for the server to be responsive
 echo "Waiting for wren-ai-service to start..."
-while ! nc -z localhost $WREN_AI_SERVICE_PORT; do   
-    sleep 1  # wait for 1 second before check again
-done
+current=0
 
+while ! nc -z localhost $WREN_AI_SERVICE_PORT; do   
+    sleep $INTERVAL
+    current=$((current + INTERVAL))
+    if [ $current -eq $TIMEOUT ]; then
+        echo "Timeout: wren-ai-service did not start within $TIMEOUT seconds"
+        exit 1
+    fi
+done
 echo "wren-ai-service has started."
 
-python -m src.force_deploy
+if [[ -n "$SHOULD_FORCE_DEPLOY" ]]; then
+    # Wait for wren-ui to be responsive
+    echo "Waiting for wren-ui to start..."
+    current=0
+
+    while ! nc -z -w 5 wren-ui $WREN_UI_PORT && ! nc -z -w 5 host.docker.internal $WREN_UI_PORT; do   
+        sleep $INTERVAL
+        current=$((current + INTERVAL))
+        if [ $current -eq $TIMEOUT ]; then
+            echo "Timeout: wren-ui did not start within $TIMEOUT seconds"
+            exit 1
+        fi
+    done
+    echo "wren-ui has started."
+
+    echo "Forcing deployment..."
+    python -m src.force_deploy
+fi
 
 # Bring wren-ai-service to the foreground
 wait

wren-ai-service/src/__main__.py

@@ -36,7 +36,9 @@ async def lifespan(app: FastAPI):
     providers = init_providers(
         engine_config=EngineConfig(provider=os.getenv("ENGINE", "wren_ui"))
     )
-    container.init_globals(*providers)
+    container.init_globals(
+        *providers, should_force_deploy=bool(os.getenv("SHOULD_FORCE_DEPLOY", ""))
+    )
     service_metadata(*providers)
     init_langfuse()
 

wren-ai-service/src/globals.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from src.core.engine import Engine
 from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider
 from src.pipelines.ask import (
@@ -45,6 +47,7 @@ def init_globals(
     embedder_provider: EmbedderProvider,
     document_store_provider: DocumentStoreProvider,
     engine: Engine,
+    should_force_deploy: Optional[str] = None,
 ):
     global \
         INDEXING_SERVICE, \
@@ -53,12 +56,11 @@ def init_globals(
         SQL_EXPLANATION_SERVICE, \
         SQL_REGENERATION_SERVICE
 
-    # Recreate the document store to ensure a clean slate
-    # TODO: for SaaS, we need to use a flag to prevent this collection_recreation
-    document_store_provider.get_store(recreate_index=True)
-    document_store_provider.get_store(
-        dataset_name="view_questions", recreate_index=True
-    )
+    if should_force_deploy:
+        document_store_provider.get_store(recreate_index=True)
+        document_store_provider.get_store(
+            dataset_name="view_questions", recreate_index=True
+        )
 
     INDEXING_SERVICE = IndexingService(
         pipelines={