use local vega-lite schema

Canner · Dec 16, 2024 · f6ad7d3 · f6ad7d3
1 parent de2782f
commit f6ad7d3
Show file tree

Hide file tree

Showing 3 changed files with 32,526 additions and 9 deletions.
diff --git a/wren-ai-service/src/pipelines/generation/chart_adjustment.py b/wren-ai-service/src/pipelines/generation/chart_adjustment.py
@@ -4,7 +4,6 @@
 from typing import Any, Dict
 
 import orjson
-import requests
 from hamilton import base
 from hamilton.async_driver import AsyncDriver
 from haystack import component
@@ -219,12 +218,13 @@ def __init__(
             "chart_data_preprocessor": ChartDataPreprocessor(),
             "post_processor": ChartAdjustmentPostProcessor(),
         }
+
+        with open("src/pipelines/generation/utils/vega-lite-schema-v5.json", "r") as f:
+            _vega_schema = orjson.loads(f.read())
+
         self._configs = {
-            "vega_schema": requests.get(
-                "https://vega.github.io/schema/vega-lite/v5.json"
-            ).json(),
+            "vega_schema": _vega_schema,
         }
-
         super().__init__(
             AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
         )

diff --git a/wren-ai-service/src/pipelines/generation/chart_generation.py b/wren-ai-service/src/pipelines/generation/chart_generation.py
@@ -4,7 +4,6 @@
 from typing import Any, Dict
 
 import orjson
-import requests
 from hamilton import base
 from hamilton.async_driver import AsyncDriver
 from haystack import component
@@ -194,10 +193,12 @@ def __init__(
             "chart_data_preprocessor": ChartDataPreprocessor(),
             "post_processor": ChartGenerationPostProcessor(),
         }
+
+        with open("src/pipelines/generation/utils/vega-lite-schema-v5.json", "r") as f:
+            _vega_schema = orjson.loads(f.read())
+
         self._configs = {
-            "vega_schema": requests.get(
-                "https://vega.github.io/schema/vega-lite/v5.json"
-            ).json(),
+            "vega_schema": _vega_schema,
         }
 
         super().__init__(