Merge branch 'main' into feat/helm-chart

Canner · Dec 26, 2024 · ee8b360 · ee8b360
2 parents 59a520e + 81e4f2a
commit ee8b360
Show file tree

Hide file tree

Showing 54 changed files with 1,019 additions and 787 deletions.
diff --git a/deployment/kustomizations/base/cm.yaml b/deployment/kustomizations/base/cm.yaml
@@ -23,7 +23,7 @@ data:
   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"
+  GENERATION_MODEL: "gpt-4o-mini-2024-07-18"
 
   # service endpoints of AI service & engine service
   WREN_ENGINE_ENDPOINT: "http://wren-engine-svc:8080"
@@ -53,21 +53,23 @@ data:
     provider: litellm_llm
     timeout: 120
     models:
-    - model: gpt-4o-mini
+    - model: gpt-4o-mini-2024-07-18
       api_base: https://api.openai.com/v1
       api_key_name: LLM_OPENAI_API_KEY
       kwargs:
         temperature: 0
         n: 1
+        seed: 0
         max_tokens: 4096
         response_format:
           type: json_object
-    - model: gpt-4o
+    - model: gpt-4o-2024-08-06
       api_base: https://api.openai.com/v1
       api_key_name: LLM_OPENAI_API_KEY
       kwargs:
         temperature: 0
         n: 1
+        seed: 0
         max_tokens: 4096
         response_format:
           type: json_object
@@ -106,58 +108,58 @@ data:
         embedder: openai_embedder.text-embedding-3-large
         document_store: qdrant
       - name: db_schema_retrieval
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         embedder: openai_embedder.text-embedding-3-large
         document_store: qdrant
       - name: historical_question_retrieval
         embedder: openai_embedder.text-embedding-3-large
         document_store: qdrant
       - name: sql_generation
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
       - name: sql_correction
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
       - name: followup_sql_generation
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
       - name: sql_summary
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
       - name: sql_answer
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
       - name: preprocess_sql_data
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
       - name: sql_breakdown
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
       - name: sql_expansion
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
       - name: sql_explanation
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
       - name: sql_regeneration
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
       - name: semantics_description
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
       - name: relationship_recommendation
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         engine: wren_ui
       - name: question_recommendation
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
       - name: intent_classification
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
         embedder: openai_embedder.text-embedding-3-large
         document_store: qdrant
       - name: data_assistance
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
       - name: sql_executor
         engine: wren_ui
       - name: chart_generation
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
       - name: chart_adjustment
-        llm: litellm_llm.gpt-4o-mini
+        llm: litellm_llm.gpt-4o-mini-2024-07-18
     ---
     settings:
       column_indexing_batch_size: 50

diff --git a/docker/.env.example b/docker/.env.example
@@ -24,7 +24,7 @@ QDRANT_API_KEY=
 
 # version
 # CHANGE THIS TO THE LATEST VERSION
-WREN_PRODUCT_VERSION=0.13.1
+WREN_PRODUCT_VERSION=0.13.2
 WREN_ENGINE_VERSION=0.12.6
 WREN_AI_SERVICE_VERSION=0.13.4
 IBIS_SERVER_VERSION=0.12.6

diff --git a/docker/config.example.yaml b/docker/config.example.yaml
@@ -2,21 +2,25 @@ type: llm
 provider: litellm_llm
 timeout: 120
 models:
-- model: gpt-4o-mini
+- model: gpt-4o-mini-2024-07-18
   api_base: https://api.openai.com/v1
   api_key_name: LLM_OPENAI_API_KEY
   kwargs:
     temperature: 0
     n: 1
+    # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed
+    seed: 0
     max_tokens: 4096
     response_format:
       type: json_object
-- model: gpt-4o
+- model: gpt-4o-2024-08-06
   api_base: https://api.openai.com/v1
   api_key_name: LLM_OPENAI_API_KEY
   kwargs:
     temperature: 0
     n: 1
+    # for better consistency of llm response, refer: https://platform.openai.com/docs/api-reference/chat/create#chat-create-seed
+    seed: 0
     max_tokens: 4096
     response_format:
       type: json_object
@@ -56,58 +60,58 @@ pipes:
     embedder: openai_embedder.text-embedding-3-large
     document_store: qdrant
   - name: db_schema_retrieval
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     embedder: openai_embedder.text-embedding-3-large
     document_store: qdrant
   - name: historical_question_retrieval
     embedder: openai_embedder.text-embedding-3-large
     document_store: qdrant
   - name: sql_generation
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
   - name: sql_correction
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
   - name: followup_sql_generation
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
   - name: sql_summary
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
   - name: sql_answer
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
   - name: sql_breakdown
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
   - name: sql_expansion
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
   - name: sql_explanation
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
   - name: sql_regeneration
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
   - name: semantics_description
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
   - name: relationship_recommendation
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     engine: wren_ui
   - name: question_recommendation
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
   - name: intent_classification
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
     embedder: openai_embedder.text-embedding-3-large
     document_store: qdrant
   - name: data_assistance
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
   - name: preprocess_sql_data
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
   - name: sql_executor
     engine: wren_ui
   - name: chart_generation
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
   - name: chart_adjustment
-    llm: litellm_llm.gpt-4o-mini
+    llm: litellm_llm.gpt-4o-mini-2024-07-18
 ---
 settings:
   column_indexing_batch_size: 50

diff --git a/wren-ai-service/Justfile b/wren-ai-service/Justfile
@@ -26,7 +26,7 @@ up: prepare-wren-engine
 down:
 	docker compose -f ./tools/dev/docker-compose-dev.yaml --env-file ./tools/dev/.env down
 
-start:
+start: use-wren-ui-as-engine
 	poetry run python -m src.__main__
 
 curate_eval_data:
@@ -63,3 +63,6 @@ prepare-wren-engine:
 	mkdir -p tools/dev/etc/mdl
 	echo "{\"catalog\": \"test_catalog\", \"schema\": \"test_schema\", \"models\": []}" \\
 		> tools/dev/etc/mdl/sample.json
+
+use-wren-ui-as-engine:
+	poetry run python -m src.force_update_config
diff --git a/wren-ai-service/demo/utils.py b/wren-ai-service/demo/utils.py
@@ -659,7 +659,6 @@ def display_sql_answer(query_id: str):
         placeholder.markdown(markdown_content)
 
 
-@st.cache_data
 def get_sql_answer(
     query: str,
     sql: str,
@@ -708,7 +707,6 @@ def get_sql_answer(
         )
 
 
-@st.cache_data
 def ask_details():
     asks_details_response = requests.post(
         f"{WREN_AI_SERVICE_BASE_URL}/v1/ask-details",
@@ -842,7 +840,6 @@ def fill_vega_lite_values(vega_lite_schema: dict, df: pd.DataFrame) -> dict:
     return schema
 
 
-@st.cache_data
 def generate_chart(
     query: str,
     sql: str,
@@ -893,7 +890,6 @@ def generate_chart(
     return chart_response
 
 
-@st.cache_data
 def adjust_chart(
     query: str,
     sql: str,

diff --git a/wren-ai-service/pyproject.toml b/wren-ai-service/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "wren-ai-service"
-version = "0.13.4"
+version = "0.13.9"
 description = ""
 authors = ["Jimmy Yeh <jimmy.yeh@cannerdata.com>", "Pao Sheng Wang <david.wang@cannerdata.com>", "Aster Sun <aster.sun@cannerdata.com>"]
 license = "AGPL-3.0"

diff --git a/wren-ai-service/src/__main__.py b/wren-ai-service/src/__main__.py
@@ -92,6 +92,7 @@ def health():
         port=settings.port,
         reload=settings.development,
         reload_includes=["src/**/*.py", ".env.dev", "config.yaml"],
+        reload_excludes=["demo/*.py", "tests/**/*.py", "eval/**/*.py"],
         workers=1,
         loop="uvloop",
         http="httptools",

diff --git a/wren-ai-service/src/force_update_config.py b/wren-ai-service/src/force_update_config.py
@@ -0,0 +1,26 @@
+import yaml
+
+
+def update_config():
+    # Read the config file
+    with open("config.yaml", "r") as file:
+        # Load all documents from YAML file (since it has multiple documents separated by ---)
+        documents = list(yaml.safe_load_all(file))
+
+    # Find the pipeline configuration document
+    for doc in documents:
+        if doc.get("type") == "pipeline":
+            # Update engine name in all pipelines
+            for pipe in doc.get("pipes", []):
+                if "engine" in pipe:
+                    pipe["engine"] = "wren_ui"
+
+    # Write back to the file
+    with open("config.yaml", "w") as file:
+        yaml.safe_dump_all(documents, file, default_flow_style=False)
+
+    print("Successfully updated engine names to 'wren_ui' in all pipelines")
+
+
+if __name__ == "__main__":
+    update_config()