feature(ai-service): wren-ui dry-run SQL API instead of wren-engine A…

…PI (#400) * feat(ai-service): change dryrun api * feat: change the docker env
Canner · Jun 12, 2024 · b86cd6f · b86cd6f
1 parent a362757
commit b86cd6f
Show file tree

Hide file tree

Showing 9 changed files with 21 additions and 13 deletions.
diff --git a/docker/docker-compose-dev.yaml b/docker/docker-compose-dev.yaml
@@ -46,7 +46,7 @@ services:
       OPENAI_API_KEY: ${OPENAI_API_KEY}
       OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL}
       QDRANT_HOST: qdrant
-      WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT}
+      WREN_UI_ENDPOINT: http://wren-ui:${WREN_UI_PORT}
       REDIS_HOST: ${AI_SERVICE_REDIS_HOST}
       REDIS_PORT: ${AI_SERVICE_REDIS_PORT}
       ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER}

diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml
@@ -40,15 +40,15 @@ services:
     platform: ${PLATFORM}
     expose:
       - ${WREN_AI_SERVICE_PORT}
-    ports: 
+    ports:
       - ${AI_SERVICE_FORWARD_PORT}:${WREN_AI_SERVICE_PORT}
     environment:
       WREN_AI_SERVICE_PORT: ${WREN_AI_SERVICE_PORT}
       OPENAI_API_KEY: ${OPENAI_API_KEY}
       OPENAI_API_BASE: ${OPENAI_API_BASE}
       OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL}
       QDRANT_HOST: qdrant
-      WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT}
+      WREN_UI_ENDPOINT: http://wren-ui:${WREN_UI_PORT}
       REDIS_HOST: ${AI_SERVICE_REDIS_HOST}
       REDIS_PORT: ${AI_SERVICE_REDIS_PORT}
       ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER}

diff --git a/wren-ai-service/.env.dev.example b/wren-ai-service/.env.dev.example
@@ -28,6 +28,7 @@ AZURE_EMBED_VERSION=
 QDRANT_HOST=localhost
 
 WREN_ENGINE_ENDPOINT=http://localhost:8080
+WREN_UI_ENDPOINT=
 
 # evaluation related
 DATASET_NAME=book_2
@@ -36,4 +37,4 @@ REDIS_HOST=localhost
 REDIS_PORT=6379
 ENABLE_TIMER=
 LOGGING_LEVEL=INFO
-WORKERS=1
+WORKERS=1
diff --git a/wren-ai-service/.env.prod.example b/wren-ai-service/.env.prod.example
@@ -34,9 +34,10 @@ AZURE_EMBED_VERSION=
 QDRANT_HOST=qdrant
 
 WREN_ENGINE_ENDPOINT=http://wren-engine:8080
+WREN_UI_ENDPOINT=
 
 REDIS_HOST=redis
 REDIS_PORT=6379
 ENABLE_TIMER=
 LOGGING_LEVEL=INFO
-WORKERS=1
+WORKERS=1
diff --git a/wren-ai-service/docker/docker-compose.yml b/wren-ai-service/docker/docker-compose.yml
@@ -16,7 +16,7 @@ services:
       OPENAI_API_BASE: ${OPENAI_API_BASE}
       OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL}
       QDRANT_HOST: ${QDRANT_HOST}
-      WREN_ENGINE_ENDPOINT: ${WREN_ENGINE_ENDPOINT}
+      WREN_UI_ENDPOINT: ${WREN_UI_ENDPOINT}
       REDIS_HOST: ${REDIS_HOST}
       REDIS_PORT: ${REDIS_PORT}
       ENABLE_TIMER: ${ENABLE_TIMER}

diff --git a/wren-ai-service/src/engine.py b/wren-ai-service/src/engine.py
@@ -13,11 +13,17 @@ async def dry_run_sql(
     session: aiohttp.ClientSession,
     endpoint: str,
 ) -> Dict[str, str]:
-    async with session.get(
-        f"{endpoint}/v1/mdl/dry-run",
+    async with session.post(
+        f"{endpoint}/api/graphql",
         json={
-            "sql": _remove_limit_statement(add_quotes(sql)),
-            "limit": 1,
+            "query": "mutation PreviewSql($data: PreviewSQLDataInput) { previewSql(data: $data) }",
+            "variables": {
+                "data": {
+                    "dryRun": True,
+                    "limit": 1,
+                    "sql": _remove_limit_statement(add_quotes(sql)),
+                }
+            },
         },
     ) as response:
         return {"status": response.status, "body": await response.json()}

diff --git a/wren-ai-service/src/eval/ask/eval_sampledata.py b/wren-ai-service/src/eval/ask/eval_sampledata.py
@@ -25,7 +25,7 @@
 
 def get_mdl_from_wren_engine():
     response = requests.get(
-        f'{os.getenv("WREN_ENGINE_ENDPOINT")}/v1/mdl',
+        f'{os.getenv("WREN_UI_ENDPOINT")}/v1/mdl',
     )
     assert response.status_code == 200
 

diff --git a/wren-ai-service/src/pipelines/ask/components/post_processors.py b/wren-ai-service/src/pipelines/ask/components/post_processors.py
@@ -61,7 +61,7 @@ async def _task(result: Dict[str, str]):
             quoted_sql = add_quotes(result["sql"])
 
             response = await dry_run_sql(
-                quoted_sql, session, endpoint=os.getenv("WREN_ENGINE_ENDPOINT")
+                quoted_sql, session, endpoint=os.getenv("WREN_UI_ENDPOINT")
             )
 
             if response.get("status") == 200:

diff --git a/wren-ai-service/src/pipelines/ask_details/generation.py b/wren-ai-service/src/pipelines/ask_details/generation.py
@@ -109,7 +109,7 @@ async def _check_if_sql_executable(
     ):
         async with aiohttp.ClientSession() as session:
             response = await dry_run_sql(
-                sql, session, endpoint=os.getenv("WREN_ENGINE_ENDPOINT")
+                sql, session, endpoint=os.getenv("WREN_UI_ENDPOINT")
             )
 
         if response.get("status") != 200: