Merge pull request #487 from vanna-ai/colab-function-fix

fix for colab create_function
vanna-ai · Jun 7, 2024 · 7aa353c · 7aa353c
2 parents 202d730 + eb07108
commit 7aa353c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 8 deletions.
diff --git a/src/vanna/flask/__init__.py b/src/vanna/flask/__init__.py
@@ -585,13 +585,10 @@ def add_training_data(user: any):
                 print("TRAINING ERROR", e)
                 return jsonify({"type": "error", "error": str(e)})
 
-        @self.flask_app.route("/api/v0/create_function", methods=["POST"])
+        @self.flask_app.route("/api/v0/create_function", methods=["GET"])
         @self.requires_auth
-        def create_function(user: any):
-            question = flask.request.json.get("question")
-            sql = flask.request.json.get("sql")
-            id = flask.request.json.get("id")
-
+        @self.requires_cache(["question", "sql"])
+        def create_function(user: any, id: str, question: str, sql: str):
             plotly_code = self.cache.get(id=id, field="plotly_code")
 
             if plotly_code is None: