GML-2026 add mcp tool support

Author: chengbiao-jin

pyTigerGraph/__init__.py

@@ -2,6 +2,26 @@
 from pyTigerGraph.pytgasync.pyTigerGraph import AsyncTigerGraphConnection
 from pyTigerGraph.common.exception import TigerGraphException
 
-__version__ = "1.9.1"
+__version__ = "2.0.0"
 
 __license__ = "Apache 2"
+
+# Optional MCP support
+try:
+    from pyTigerGraph.mcp import serve, MCPServer, get_connection, ConnectionManager
+    __all__ = [
+        "TigerGraphConnection",
+        "AsyncTigerGraphConnection",
+        "TigerGraphException",
+        "serve",
+        "MCPServer",
+        "get_connection",
+        "ConnectionManager",
+    ]
+except ImportError:
+    # MCP dependencies not installed
+    __all__ = [
+        "TigerGraphConnection",
+        "AsyncTigerGraphConnection",
+        "TigerGraphException",
+    ]

pyTigerGraph/common/base.py

@@ -179,7 +179,7 @@ def __init__(self, host: str = "http://127.0.0.1", graphname: str = "",
         else:
             self.restppPort = restppPort
             self.restppUrl = self.host + ":" + self.restppPort
-            
+
         self.gsPort = gsPort
         if self.tgCloud and (gsPort == "14240" or gsPort == "443"):
             self.gsPort = sslPort
@@ -375,7 +375,7 @@ def customizeHeader(self, timeout: int = 16_000, responseSize: int = 3.2e+7):
         """
         self.responseConfigHeader = {
             "GSQL-TIMEOUT": str(timeout), "RESPONSE-LIMIT": str(responseSize)}
-        
+
     def _parse_get_ver(self, version, component, full):
         ret = ""
         for v in version:

pyTigerGraph/common/loading.py

@@ -49,7 +49,7 @@ def _prep_run_loading_job(gsUrl: str,
     '''url builder for runLoadingJob()'''
     url = gsUrl + "/gsql/v1/loading-jobs/run?graph=" + graphname
     data = {}
-    
+
     data["name"] = jobName
     data["dataSources"] = [data_source_config]
 
@@ -65,7 +65,7 @@ def _prep_run_loading_job(gsUrl: str,
         data["maxNumError"] = maxNumError
     if maxPercentError:
         data["maxPercentError"] = maxPercentError
-    
+
     return url, data
 
 def _prep_abort_loading_jobs(gsUrl: str, graphname: str, jobIds: list[str], pauseJob: bool):

pyTigerGraph/gds/dataloaders.py

@@ -676,7 +676,7 @@ def _read_data(
                 logger.error("Error parsing data: {}".format(raw))
                 logger.error("Parameters:\n  in_format={}\n  out_format={}\n  v_in_feats={}\n  v_out_labels={}\n  v_extra_feats={}\n  v_attr_types={}\n  e_in_feats={}\n  e_out_labels={}\n  e_extra_feats={}\n  e_attr_types={}\n  delimiter={}\n".format(
                     in_format, out_format, v_in_feats, v_out_labels, v_extra_feats, v_attr_types, e_in_feats, e_out_labels, e_extra_feats, e_attr_types, delimiter))
-                
+
             in_q.task_done()
 
     @staticmethod
@@ -766,7 +766,7 @@ def add_attributes(attr_names: list, attr_types: dict, attr_df: pd.DataFrame,
                         data = graph.ndata
 
             data[feat_name] = attr_to_tensor(attr_names, attr_types, attr_df)
-        
+
         def add_sep_attr(attr_names: list, attr_types: dict, attr_df: pd.DataFrame, 
                          graph, is_hetero: bool, mode: str,
                          target: 'Literal["edge", "vertex"]', vetype: str = None) -> None:
@@ -864,7 +864,7 @@ def add_sep_attr(attr_names: list, attr_types: dict, attr_df: pd.DataFrame,
                         )
                     elif mode == "spektral":
                         data[col] = attr_df[col].astype(dtype)
-        
+
         # Read in vertex and edge CSVs as dataframes              
         vertices, edges = None, None
         if in_format == "vertex":
@@ -1422,7 +1422,7 @@ def _generate_attribute_string(self, schema_type, attr_names, attr_types) -> str
                         for attr in attr_names
                     )
         return print_attr
-    
+
     def metadata(self, additional_v_types=None, additional_e_types=None) -> Tuple[list, list]:
         v_types = self._vtypes
         if additional_v_types:
@@ -1439,7 +1439,7 @@ def metadata(self, additional_v_types=None, additional_e_types=None) -> Tuple[li
             elif isinstance(additional_e_types, tuple):
                 edges.append(additional_e_types)
         return (v_types, edges)
-    
+
     def fetch(self, payload: dict) -> None:
         """Fetch the specific data instances for inference/prediction.
 
@@ -2478,7 +2478,7 @@ def _start(self) -> None:
         self._exit_event = Event()
 
         self._start_request(False, "vertex")
-            
+
         # Start reading thread.
         if not self.is_hetero:
             v_attr_types = next(iter(self._v_schema.values()))
@@ -3513,7 +3513,7 @@ def __init__(
             for tok in self.baseTokens + ancs:
                 self.idToIdx[str(tok)] = self.curIdx
                 self.curIdx += 1
-            
+
         self.num_tokens = len(self.idToIdx.keys())
 
     def saveTokens(self, filename) -> None:
@@ -3670,7 +3670,7 @@ def _start(self) -> None:
         self._exit_event = Event()
 
         self._start_request(False, "vertex")
-            
+
         # Start reading thread.
         if not self.is_hetero:
             v_attr_types = next(iter(self._v_schema.values()))
@@ -3782,7 +3782,7 @@ def precompute(self) -> None:
         resp = self._graph.runInstalledQuery(
             self.query_name, params=_payload, timeout=self.timeout, usePost=True
         )
-       
+
 
 class HGTLoader(BaseLoader):
     """HGTLoader

pyTigerGraph/gds/gds.py

@@ -46,7 +46,7 @@ def __init__(self, conn: "TigerGraphConnection") -> None:
             Args:
                 conn (TigerGraphConnection):
                     Accept a TigerGraphConnection to run queries with
-                    
+
             Returns:
                 None
         """
@@ -231,7 +231,7 @@ def neighborLoader(
         If specify both parameters, `batch_size` takes priority. 
         . It picks a specified number (`num_neighbors`) of neighbors of each seed at random.
         . It picks the same number of neighbors for each neighbor, and repeats this process until it finished performing a specified number of hops (`num_hops`).
-        
+
         This generates one subgraph. 
         As you loop through this data loader, every vertex will at some point be chosen as a seed and you will get the subgraph
         expanded from the seeds. 
@@ -949,7 +949,7 @@ def edgeNeighborLoader(
         If specify both parameters, `batch_size` takes priority. 
         . Starting from the vertices attached to the seed edges, it picks a specified number (`num_neighbors`) of neighbors of each vertex at random.
         . It picks the same number of neighbors for each neighbor, and repeats this process until it finished performing a specified number of hops (`num_hops`).
-        
+
         This generates one subgraph. 
         As you loop through this data loader, every edge will at some point be chosen as a seed and you will get the subgraph
         expanded from the seeds. 
@@ -1379,7 +1379,7 @@ def hgtLoader(
         If specify both parameters, `batch_size` takes priority. 
         . It picks a specified number of neighbors of each type (as specified by the dict `num_neighbors`) of each seed at random.
         . It picks the specified number of neighbors of every type for each neighbor, and repeats this process until it finished performing a specified number of hops (`num_hops`).
-        
+
         This generates one subgraph. 
         As you loop through this data loader, every vertex will at some point be chosen as a seed and you will get the subgraph
         expanded from the seeds. 
@@ -1565,7 +1565,7 @@ def hgtLoader(
             if reinstall_query:
                 loader.reinstall_query()
             return loader
-    
+
     def featurizer(
         self,
         repo: str = None, 
@@ -1650,7 +1650,7 @@ def edgeSplitter(self, e_types: List[str] = None, timeout: int = 600000, **split
         Make sure to create the appropriate attributes in the graph before using these functions.
 
         Usage:
-            
+
             * A random 60% of edges will have their attribute "attr_name" set to True, and 
             others False. `attr_name` can be any attribute that exists in the database (same below).
             Example:

pyTigerGraph/gds/splitters.py

@@ -142,7 +142,7 @@ class RandomVertexSplitter(BaseRandomSplitter):
         v_types (List[str], optional):
             List of vertex types to split. If not provided, all vertex types are used.
         timeout (int, optional):
-            Timeout value for the operation. Defaults to 600000.
+            Timeout value for the operation in milliseconds. Defaults to 600000 (10 minutes).
     """
 
     def __init__(
@@ -229,7 +229,7 @@ class RandomEdgeSplitter(BaseRandomSplitter):
         e_types (List[str], optional):
             List of edge types to split. If not provided, all edge types are used.
         timeout (int, optional):
-            Timeout value for the operation. Defaults to 600000.
+            Timeout value for the operation in milliseconds. Defaults to 600000 (10 minutes).
     """
 
     def __init__(