feat(wren-ai-service): sql-expansion (#667)

docker/.env.ai.example

@@ -1,5 +1,6 @@
 ## LLM
-LLM_PROVIDER=openai_llm # openai_llm, azure_openai_llm, ollama_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"}}
@@ -22,7 +23,8 @@ LLM_OLLAMA_URL=http://host.docker.internal:11434
 
 
 ## EMBEDDER
-EMBEDDER_PROVIDER=openai_embedder # openai_embedder, azure_openai_embedder, ollama_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

docker/.env.example

@@ -16,7 +16,8 @@ WREN_UI_ENDPOINT=http://docker.for.mac.localhost:3000
 # LLM
 LLM_OPENAI_API_KEY=
 EMBEDDER_OPENAI_API_KEY=
-GENERATION_MODEL=gpt-4o-mini # gpt-4o-mini, gpt-4o, gpt-4-turbo, gpt-3.5-turbo
+# gpt-4o-mini, gpt-4o, gpt-4-turbo, gpt-3.5-turbo
+GENERATION_MODEL=gpt-4o-mini
 
 # version
 # CHANGE THIS TO THE LATEST VERSION
@@ -30,7 +31,7 @@ WREN_BOOTSTRAP_VERSION=0.1.5
 # AI service related env variables
 AI_SERVICE_ENABLE_TIMER=
 AI_SERVICE_LOGGING_LEVEL=INFO
-SHOULD_FORCE_DEPLOY=
+SHOULD_FORCE_DEPLOY=1
 QDRANT_HOST=qdrant
 
 # user id (uuid v4)

wren-ai-service/Justfile

@@ -19,8 +19,8 @@ eval prediction_result:
 demo:
 	poetry run streamlit run demo/app.py
 
-test: up && down
-	poetry run pytest -s $(args)
+test test_args='': up && down
+	poetry run pytest -s {{test_args}}
 
 load-test:
 	poetry run python -m tests.locust.locust_script

wren-ai-service/demo/utils.py

@@ -671,7 +671,7 @@ def get_sql_answer(
     sql_summary: str,
 ):
     sql_answer_response = requests.post(
-        f"{WREN_AI_SERVICE_BASE_URL}/v1/sql-answer",
+        f"{WREN_AI_SERVICE_BASE_URL}/v1/sql-answers",
         json={
             "query": query,
             "sql": sql,
@@ -687,7 +687,7 @@ def get_sql_answer(
         sql_answer_status != "finished" and sql_answer_status != "failed"
     ):
         sql_answer_status_response = requests.get(
-            f"{WREN_AI_SERVICE_BASE_URL}/v1/sql-answer/{query_id}/result"
+            f"{WREN_AI_SERVICE_BASE_URL}/v1/sql-answers/{query_id}/result"
         )
         assert sql_answer_status_response.status_code == 200
         sql_answer_status = sql_answer_status_response.json()["status"]

wren-ai-service/eval/pipelines.py

@@ -29,8 +29,9 @@
 )
 from src.core.engine import Engine
 from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider
-from src.pipelines.ask import generation, retrieval
+from src.pipelines.generation import sql_generation
 from src.pipelines.indexing import indexing
+from src.pipelines.retrieval import retrieval
 
 
 def deploy_model(mdl: str, pipe: indexing.Indexing) -> None:
@@ -228,7 +229,7 @@ def __init__(
     ):
         super().__init__(meta, 3)
         self._mdl = mdl
-        self._generation = generation.Generation(
+        self._generation = sql_generation.SQLGeneration(
             llm_provider=llm_provider,
             engine=engine,
         )
@@ -305,7 +306,7 @@ def __init__(
             embedder_provider=embedder_provider,
             document_store_provider=document_store_provider,
         )
-        self._generation = generation.Generation(
+        self._generation = sql_generation.SQLGeneration(
             llm_provider=llm_provider,
             engine=engine,
         )

wren-ai-service/src/__main__.py

@@ -9,13 +9,15 @@
 from fastapi.responses import ORJSONResponse, RedirectResponse
 from langfuse.decorators import langfuse_context
 
-import src.globals as container
 from src.core.engine import EngineConfig
+from src.globals import (
+    create_service_container,
+    create_service_metadata,
+)
 from src.utils import (
     init_langfuse,
     init_providers,
     load_env_vars,
-    service_metadata,
     setup_custom_logger,
 )
 from src.web.v1 import routers
@@ -36,7 +38,7 @@ async def lifespan(app: FastAPI):
     providers = init_providers(
         engine_config=EngineConfig(provider=os.getenv("ENGINE", "wren_ui"))
     )
-    container.init_globals(
+    app.state.service_container = create_service_container(
         *providers,
         should_force_deploy=bool(os.getenv("SHOULD_FORCE_DEPLOY", "")),
         column_indexing_batch_size=(
@@ -61,7 +63,7 @@ async def lifespan(app: FastAPI):
             "ttl": int(os.getenv("QUERY_CACHE_TTL") or 120),
         },
     )
-    service_metadata(*providers)
+    app.state.service_metadata = create_service_metadata(*providers)
     init_langfuse()
 
     yield
@@ -128,9 +130,6 @@ def health():
         port=server_port,
         reload=should_reload,
         reload_includes=["src/**/*.py", ".env.dev"],
-        reload_excludes=[
-            "./demo/*.py",
-        ],  # TODO: add eval folder when evaluation system is ready
         workers=1,
         loop="uvloop",
         http="httptools",