Add timeout to requests calls

src/vanna/__init__.py

@@ -62,8 +62,8 @@ def __unauthenticated_rpc_call(method, params):
     data = {"method": method, "params": [__dataclass_to_dict(obj) for obj in params]}
 
     response = requests.post(
-        _unauthenticated_endpoint, headers=headers, data=json.dumps(data)
-    )
+        _unauthenticated_endpoint, headers=headers, data=json.dumps(data), 
+    timeout=60)
     return response.json()
 
 
@@ -397,4 +397,4 @@ def connect_to_bigquery(cred_file_path: str = None, project_id: str = None):
     error_deprecation()
 
 def connect_to_duckdb(url: str="memory", init_sql: str = None):
-    error_deprecation()
+    error_deprecation()

src/vanna/base/base.py

@@ -808,7 +808,7 @@ def connect_to_sqlite(self, url: str):
 
         # Download the database if it doesn't exist
         if not os.path.exists(url):
-            response = requests.get(url)
+            response = requests.get(url, timeout=60)
             response.raise_for_status()  # Check that the request was successful
             with open(path, "wb") as f:
                 f.write(response.content)
@@ -1297,7 +1297,7 @@ def connect_to_duckdb(self, url: str, init_sql: str = None):
                 path = os.path.basename(urlparse(url).path)
                 # Download the database if it doesn't exist
                 if not os.path.exists(path):
-                    response = requests.get(url)
+                    response = requests.get(url, timeout=60)
                     response.raise_for_status()  # Check that the request was successful
                     with open(path, "wb") as f:
                         f.write(response.content)

src/vanna/flask/__init__.py

@@ -736,7 +736,7 @@ def proxy_assets(filename):
         @self.flask_app.route("/vanna.svg")
         def proxy_vanna_svg():
             remote_url = "https://vanna.ai/img/vanna.svg"
-            response = requests.get(remote_url, stream=True)
+            response = requests.get(remote_url, stream=True, timeout=60)
 
             # Check if the request to the remote URL was successful
             if response.status_code == 200:

src/vanna/vannadb/vannadb_vector.py

@@ -60,7 +60,7 @@ def _rpc_call(self, method, params):
             "params": [self._dataclass_to_dict(obj) for obj in params],
         }
 
-        response = requests.post(self._endpoint, headers=headers, data=json.dumps(data))
+        response = requests.post(self._endpoint, headers=headers, data=json.dumps(data), timeout=60)
         return response.json()
 
     def _dataclass_to_dict(self, obj):
@@ -85,7 +85,7 @@ def get_all_functions(self) -> list:
             }
         """
 
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query})
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query}, timeout=60)
         response_json = response.json()
         if response.status_code == 200 and 'data' in response_json and 'get_all_sql_functions' in response_json['data']:
             self.log(response_json['data']['get_all_sql_functions'])
@@ -123,7 +123,7 @@ def get_function(self, question: str, additional_data: dict = {}) -> dict:
         """
         static_function_arguments = [{"name": key, "value": str(value)} for key, value in additional_data.items()]
         variables = {"question": question, "staticFunctionArguments": static_function_arguments}
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query, 'variables': variables})
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query, 'variables': variables}, timeout=60)
         response_json = response.json()
         if response.status_code == 200 and 'data' in response_json and 'get_and_instantiate_function' in response_json['data']:
             self.log(response_json['data']['get_and_instantiate_function'])
@@ -153,7 +153,7 @@ def create_function(self, question: str, sql: str, plotly_code: str, **kwargs) -
         }
         """
         variables = {"question": question, "sql": sql, "plotly_code": plotly_code}
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query, 'variables': variables})
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': query, 'variables': variables}, timeout=60)
         response_json = response.json()
         if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'generate_and_create_sql_function' in response_json['data']:
             resp = response_json['data']['generate_and_create_sql_function']
@@ -216,7 +216,7 @@ def validate_arguments(args):
 
         print("variables", variables)
 
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': mutation, 'variables': variables})
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': mutation, 'variables': variables}, timeout=60)
         response_json = response.json()
         if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'update_sql_function' in response_json['data']:
             return response_json['data']['update_sql_function']
@@ -230,7 +230,7 @@ def delete_function(self, function_name: str) -> bool:
         }
         """
         variables = {"function_name": function_name}
-        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': mutation, 'variables': variables})
+        response = requests.post(self._graphql_endpoint, headers=self._graphql_headers, json={'query': mutation, 'variables': variables}, timeout=60)
         response_json = response.json()
         if response.status_code == 200 and 'data' in response_json and response_json['data'] is not None and 'delete_sql_function' in response_json['data']:
             return response_json['data']['delete_sql_function']

src/vanna/vllm/vllm.py

@@ -78,11 +78,11 @@ def submit_prompt(self, prompt, **kwargs) -> str:
             'Authorization': f'Bearer {self.auth_key}' 
             }
 
-            response = requests.post(url, headers=headers,json=data)
+            response = requests.post(url, headers=headers,json=data, timeout=60)
 
 
         else:
-            response = requests.post(url, json=data)
+            response = requests.post(url, json=data, timeout=60)
 
         response_dict = response.json()