From f85ce2d893c4e002fd9358fd3686963e7c582023 Mon Sep 17 00:00:00 2001
From: Parker Erickson <parker.erickson@tigergraph.com>
Date: Thu, 25 Apr 2024 14:04:35 -0500
Subject: [PATCH] fix(updateCustomQuery): update json payload

---
 pyTigerGraph/ai/ai.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/pyTigerGraph/ai/ai.py b/pyTigerGraph/ai/ai.py
index cb68d105..a929d078 100644
--- a/pyTigerGraph/ai/ai.py
+++ b/pyTigerGraph/ai/ai.py
@@ -79,7 +79,8 @@ def registerCustomQuery(self, query_name: str, description: str = None, docstrin
                 "function_header": query_name,
                 "description": description,
                 "docstring": docstring,
-                "param_types": param_types
+                "param_types": param_types,
+                "graphname": self.conn.graphname
             }
             url = self.nlqs_host+"/"+self.conn.graphname+"/register_docs"
             return self.conn._req("POST", url, authMode="pwd", data = data, jsonData=True, resKey=None)
@@ -123,10 +124,11 @@ def updateCustomQuery(self, query_name: str, description: str = None, docstring:
                 "function_header": query_name,
                 "description": description,
                 "docstring": docstring,
-                "param_types": param_types
+                "param_types": param_types,
+                "graphname": self.conn.graphname
             }
 
-            json_payload = {"query_info": data}
+            json_payload = {"id": "", "query_info": data}
             url = self.nlqs_host+"/"+self.conn.graphname+"/upsert_docs"
             return self.conn._req("POST", url, authMode="pwd", data = json_payload, jsonData=True, resKey=None)