Chore/ai service/auto deploy at init (#457)

* allow force deploy during server start

* fix force_deploy and env file

docker/README.md

@@ -26,7 +26,7 @@ Path structure as following:
 ## 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 up -d`.
-4. start all services(with custom LLM): `docker-compose -f docker-compose.yaml -f docker-compose.llm.yaml --env-file .env.local up -d`.
+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`.
 
->Note: If your port 3000 is occupied, you can modify the `HOST_PORT` in `.env.example`.
+>Note: If your port 3000 is occupied, you can modify the `HOST_PORT` in `.env.local`.

docker/docker-compose-dev.yaml

@@ -44,16 +44,15 @@ services:
     environment:
       WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT}
       WREN_UI_ENDPOINT: ${WREN_UI_ENDPOINT}
-      OPENAI_API_KEY: ${OPENAI_API_KEY}
       GENERATION_MODEL: ${GENERATION_MODEL}
+      OPENAI_API_KEY: ${OPENAI_API_KEY}
+      AZURE_CHAT_KEY: ${AZURE_CHAT_KEY}
+      AZURE_EMBED_KEY: ${AZURE_EMBED_KEY}
       ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER}
       LOGGING_LEVEL: ${AI_SERVICE_LOGGING_LEVEL}
       # sometimes the console won't show print messages,
       # using PYTHONUNBUFFERED: 1 can fix this
       PYTHONUNBUFFERED: 1
-    env_file:
-      - path: ${PROJECT_DIR}/.env.ai
-        required: false
     networks:
       - wren
     depends_on:

docker/docker-compose.llm.yaml

@@ -1,5 +1,3 @@
 services:
   wren-ai-service:
-    env_file:
-      - path: ${PROJECT_DIR}/.env.ai
-        required: true
+    env_file: ${PROJECT_DIR}/.env.ai

docker/docker-compose.yaml

@@ -56,8 +56,10 @@ services:
     environment:
       WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT}
       WREN_UI_ENDPOINT: http://wren-ui:${WREN_UI_PORT}
-      OPENAI_API_KEY: ${OPENAI_API_KEY}
       GENERATION_MODEL: ${GENERATION_MODEL}
+      OPENAI_API_KEY: ${OPENAI_API_KEY}
+      AZURE_CHAT_KEY: ${AZURE_CHAT_KEY}
+      AZURE_EMBED_KEY: ${AZURE_EMBED_KEY}
       ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER}
       LOGGING_LEVEL: ${AI_SERVICE_LOGGING_LEVEL}
       # sometimes the console won't show print messages,
@@ -92,7 +94,7 @@ services:
       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_DIM: ${EMBEDDING_MODEL_DIM}
+      EMBEDDING_MODEL_DIMENSION: ${EMBEDDING_MODEL_DIMENSION}
       GENERATION_MODEL: ${GENERATION_MODEL}
       PG_USERNAME: ${PG_USERNAME}
       PG_PASSWORD: ${PG_PASSWORD}

wren-ai-service/.dockerignore

@@ -1,4 +1,5 @@
 *
 !src
+!entrypoint.sh
 !pyproject.toml
 src/eval

wren-ai-service/Makefile

@@ -12,7 +12,15 @@ dev-down:
 
 ## wren-ai-service related ##
 start:
-	poetry run python -m src.__main__
+	poetry run python -m src.__main__ & \
+	make force_deploy
+
+force_deploy:
+	while ! nc -z localhost 5556; do \
+		sleep 1; \
+	done; \
+	echo "wren-ai-service is up and running" && \
+	poetry run python src/force_deploy.py
 
 build:
 	docker compose -f docker/docker-compose.yaml --env-file .env.prod build

wren-ai-service/docker/Dockerfile

@@ -16,11 +16,15 @@ RUN poetry install --without dev --no-root && rm -rf $POETRY_CACHE_DIR
 
 FROM python:3.12.0-slim-bookworm as runtime
 
+RUN apt-get update && apt install -y netcat-traditional
+
 ENV VIRTUAL_ENV=/app/.venv \
     PATH="/app/.venv/bin:$PATH"
 
 COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
 
 COPY src src
+COPY entrypoint.sh /app/entrypoint.sh
+RUN chmod +x /app/entrypoint.sh
 
-ENTRYPOINT uvicorn src.__main__:app --host 0.0.0.0 --port $WREN_AI_SERVICE_PORT --loop uvloop --http httptools
+ENTRYPOINT [ "/app/entrypoint.sh" ]

wren-ai-service/entrypoint.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+set -e
+
+# 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 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
+
+echo "wren-ai-service has started."
+
+python src/force_deploy.py
+
+# Bring wren-ai-service to the foreground
+wait

wren-ai-service/src/force_deploy.py

@@ -0,0 +1,29 @@
+import asyncio
+import os
+
+import aiohttp
+import backoff
+from dotenv import load_dotenv
+
+load_dotenv(override=True)
+if is_dev_env := os.getenv("ENV") and os.getenv("ENV").lower() == "dev":
+    load_dotenv(".env.dev", override=True)
+
+
+@backoff.on_exception(backoff.expo, aiohttp.ClientError, max_time=60, max_tries=3)
+async def force_deploy():
+    async with aiohttp.ClientSession() as session:
+        async with session.post(
+            f"{os.getenv("WREN_UI_ENDPOINT", "http://wren-ui:3000")}/api/graphql",
+            json={
+                "query": "mutation Deploy($force: Boolean) { deploy(force: $force) }",
+                "variables": {"force": True},
+            },
+            timeout=aiohttp.ClientTimeout(total=60),  # 60 seconds
+        ) as response:
+            res = await response.json()
+            print(f"Forcing deployment: {res}")
+
+
+if os.getenv("ENGINE", "wren-ui") == "wren-ui":
+    asyncio.run(force_deploy())