scmcp-shared 0.2.5__py3-none-any.whl → 0.3.0__py3-none-any.whl

scmcp_shared/server/util.py

@@ -5,257 +5,247 @@ import scanpy as sc
 from fastmcp import FastMCP , Context
 from fastmcp.exceptions import ToolError
 from ..schema.util import *
-from ..schema import AdataModel
+from ..schema import AdataModel, AdataInfo
 from ..util import filter_args, forward_request, get_ads, generate_msg,add_op_log
+from .base import BaseMCP
 
 
-ul_mcp = FastMCP("SCMCP-Util-Server")
-
-
-@ul_mcp.tool()
-async def query_op_log(
-    request: QueryOpLogModel = QueryOpLogModel(),
-    adinfo: AdataModel = AdataModel()
-):
-    """Query the adata operation log"""
-    adata = get_ads().get_adata(adinfo=adinfo)
-    op_dic = adata.uns["operation"]["op"]
-    opids = adata.uns["operation"]["opid"][-n:]
-    op_list = []
-    for opid in opids:
-        op_list.append(op_dic[opid])
-    return op_list
-
-
-@ul_mcp.tool()
-async def mark_var(
-    request: MarkVarModel = MarkVarModel(),
-    adinfo: AdataModel = AdataModel()
-):
-    """
-    Determine if each gene meets specific conditions and store results in adata.var as boolean values.
-    For example: mitochondrion genes startswith MT-.
-    The tool should be called first when calculate quality control metrics for mitochondrion, ribosomal, harhemoglobin genes, or other qc_vars.
-    """
-    try:
-        result = await forward_request("ul_mark_var", request, adinfo)
-        if result is not None:
-            return result
-        adata = get_ads().get_adata(adinfo=adinfo)
-        var_name = request.var_name
-        gene_class = request.gene_class
-        pattern_type = request.pattern_type
-        patterns = request.patterns
-        if gene_class is not None:
-            if gene_class == "mitochondrion":
-                adata.var["mt"] = adata.var_names.str.startswith(('MT-', 'Mt','mt-'))
-                var_name = "mt"
-            elif gene_class == "ribosomal":
-                adata.var["ribo"] = adata.var_names.str.startswith(("RPS", "RPL", "Rps", "Rpl"))
-                var_name = "ribo"
-            elif gene_class == "hemoglobin":
-                adata.var["hb"] = adata.var_names.str.contains("^HB[^(P)]", case=False)
-                var_name = "hb"    
-        elif pattern_type is not None and patterns is not None:
-            if pattern_type == "startswith":
-                adata.var[var_name] = adata.var_names.str.startswith(patterns)
-            elif pattern_type == "endswith":
-                adata.var[var_name] = adata.var_names.str.endswith(patterns)
-            elif pattern_type == "contains":
-                adata.var[var_name] = adata.var_names.str.contains(patterns)
-            else:
-                raise ValueError(f"Did not support pattern_type: {pattern_type}")
-        else:
-            raise ValueError(f"Please provide validated parameter")
-    
-        res = {var_name: adata.var[var_name].value_counts().to_dict(), "msg": f"add '{var_name}' column in adata.var"}
-        func_kwargs = {"var_name": var_name, "gene_class": gene_class, "pattern_type": pattern_type, "patterns": patterns}
-        add_op_log(adata, "mark_var", func_kwargs, adinfo)
-        return res
-    except ToolError as e:
-        raise ToolError(e)
-    except Exception as e:
-        if hasattr(e, '__context__') and e.__context__:
-            raise ToolError(e.__context__)
-        else:
-            raise ToolError(e)
-
-
-@ul_mcp.tool()
-async def list_var(
-    request: ListVarModel = ListVarModel(),
-    adinfo: AdataModel = AdataModel()
-):
-    """List key columns in adata.var. It should be called for checking when other tools need var key column names as input."""
-    try:
-        result = await forward_request("ul_list_var", request, adinfo)
-        if result is not None:
-            return result
-        adata = get_ads().get_adata(adinfo=adinfo)
-        columns = list(adata.var.columns)
-        add_op_log(adata, list_var, {}, adinfo)
-        return columns
-    except ToolError as e:
-        raise ToolError(e)
-    except Exception as e:
-        if hasattr(e, '__context__') and e.__context__:
-            raise ToolError(e.__context__)
-        else:
-            raise ToolError(e)
-
-@ul_mcp.tool()
-async def list_obs(
-    request: ListObsModel = ListObsModel(),
-    adinfo: AdataModel = AdataModel()
-):
-    """List key columns in adata.obs. It should be called before other tools need obs key column names input."""
-    try:
-        result = await forward_request("ul_list_obs", request, adinfo)
-        if result is not None:
-            return result    
-        adata = get_ads().get_adata(adinfo=adinfo)
-        columns = list(adata.obs.columns)
-        add_op_log(adata, list_obs, {}, adinfo)
-        return columns
-    except ToolError as e:
-        raise ToolError(e)
-    except Exception as e:
-        if hasattr(e, '__context__') and e.__context__:
-            raise ToolError(e.__context__)
-        else:
-            raise ToolError(e)
-
-@ul_mcp.tool()
-async def check_var(
-    request: VarNamesModel = VarNamesModel(),
-    adinfo: AdataModel = AdataModel()
-):
-    """Check if genes/variables exist in adata.var_names. This tool should be called before gene expression visualizations or color by genes."""
-    try:
-        result = await forward_request("ul_check_var", request, adinfo)
-        if result is not None:
-            return result     
-        adata = get_ads().get_adata(adinfo=adinfo)
-        var_names = request.var_names
-        result = {v: v in adata.var_names for v in var_names}
-        add_op_log(adata, check_var, {"var_names": var_names}, adinfo)
-        return result
-    except ToolError as e:
-        raise ToolError(e)
-    except Exception as e:
-        if hasattr(e, '__context__') and e.__context__:
-            raise ToolError(e.__context__)
-        else:
-            raise ToolError(e)
-
-@ul_mcp.tool()
-async def merge_adata(
-    request: ConcatBaseModel = ConcatBaseModel(),
-    adinfo: AdataModel = AdataModel()
-):
-    """Merge multiple adata objects."""
-     
-    try:
-        result = await forward_request("ul_merge_adata", request, adinfo)
-        if result is not None:
-            return result
-        ads = get_ads()
-        adata = ads.get_adata(adinfo=adinfo)
-        kwargs = {k: v for k, v in request.model_dump().items() if v is not None}
-        merged_adata = adata.concat(list(ads.adata_dic[dtype].values()), **kwargs)
-        ads.adata_dic[dtype] = {}
-        ads.active_id = "merged_adata"
-        add_op_log(merged_adata, ad.concat, kwargs, adinfo)
-        ads.adata_dic[ads.active_id] = merged_adata
-        return {"status": "success", "message": "Successfully merged all AnnData objects"}
-    except ToolError as e:
-        raise ToolError(e)
-    except Exception as e:
-        if hasattr(e, '__context__') and e.__context__:
-            raise ToolError(e.__context__)
-        else:
-            raise ToolError(e)
-
-
-@ul_mcp.tool()
-async def set_dpt_iroot(
-    request: DPTIROOTModel,
-    adinfo: AdataModel = AdataModel()
-):
-    """Set the iroot cell"""
-    try:
-        result = await forward_request("ul_set_dpt_iroot", request, adinfo)
-        if result is not None:
-            return result     
-        adata = get_ads().get_adata(adinfo=adinfo)
-        diffmap_key = request.diffmap_key
-        dimension = request.dimension
-        direction = request.direction
-        if diffmap_key not in adata.obsm:
-            raise ValueError(f"Diffusion map key '{diffmap_key}' not found in adata.obsm")
-        if direction == "min":
-            adata.uns["iroot"] = adata.obsm[diffmap_key][:, dimension].argmin()
-        else:  
-            adata.uns["iroot"] = adata.obsm[diffmap_key][:, dimension].argmax()
-        
-        func_kwargs = {"diffmap_key": diffmap_key, "dimension": dimension, "direction": direction}
-        add_op_log(adata, "set_dpt_iroot", func_kwargs, adinfo)
-        
-        return {"status": "success", "message": f"Successfully set root cell for DPT using {direction} of dimension {dimension}"}
-    except ToolError as e:
-        raise ToolError(e)
-    except Exception as e:
-        if hasattr(e, '__context__') and e.__context__:
-            raise ToolError(e.__context__)
-        else:
-            raise ToolError(e)
-
-@ul_mcp.tool()
-async def add_layer(
-    request: AddLayerModel,
-    adinfo: AdataModel = AdataModel()
-):
-    """Add a layer to the AnnData object.
-    """
-    try:
-        result = await forward_request("ul_add_layer", request, adinfo)
-        if result is not None:
-            return result         
-        adata = get_ads().get_adata(adinfo=adinfo)
-        layer_name = request.layer_name
-        
-        # Check if layer already exists
-        if layer_name in adata.layers:
-            raise ValueError(f"Layer '{layer_name}' already exists in adata.layers")
-        # Add the data as a new layer
-        adata.layers[layer_name] = adata.X.copy()
-
-        func_kwargs = {"layer_name": layer_name}
-        add_op_log(adata, "add_layer", func_kwargs, adinfo)
+class ScanpyUtilMCP(BaseMCP):
+    def __init__(self, include_tools: list = None, exclude_tools: list = None, AdataInfo = AdataInfo):
+        """
+        Initialize ScanpyUtilMCP with optional tool filtering.
         
-        return {
-            "status": "success", 
-            "message": f"Successfully added layer '{layer_name}' to adata.layers"
-        }
-    except ToolError as e:
-        raise ToolError(e)
-    except Exception as e:
-        if hasattr(e, '__context__') and e.__context__:
-            raise ToolError(e.__context__)
-        else:
-            raise ToolError(e)
-
-@ul_mcp.tool()
-async def check_samples():
-    """check the stored samples    
-    """
-    try:
-        ads = get_ads()
-        return {"sampleid": [list(ads.adata_dic[dk].keys()) for dk in ads.adata_dic.keys()]}
-    except ToolError as e:
-        raise ToolError(e)
-    except Exception as e:
-        if hasattr(e, '__context__') and e.__context__:
-            raise ToolError(e.__context__)
-        else:
-            raise ToolError(e)
+        Args:
+            include_tools (list, optional): List of tool names to include. If None, all tools are included.
+            exclude_tools (list, optional): List of tool names to exclude. If None, no tools are excluded.
+            AdataInfo: The AdataInfo class to use for type annotations.
+        """
+        super().__init__("SCMCP-Util-Server", include_tools, exclude_tools, AdataInfo)
+
+
+    def _tool_query_op_log(self):
+        def _query_op_log(request: QueryOpLogModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """Query the adata operation log"""
+            adata = get_ads().get_adata(adinfo=adinfo)
+            op_dic = adata.uns["operation"]["op"]
+            opids = adata.uns["operation"]["opid"][-request.n:]
+            op_list = []
+            for opid in opids:
+                op_list.append(op_dic[opid])
+            return op_list
+        return _query_op_log
+
+    def _tool_mark_var(self):
+        def _mark_var(request: MarkVarModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """
+            Determine if each gene meets specific conditions and store results in adata.var as boolean values.
+            For example: mitochondrion genes startswith MT-.
+            The tool should be called first when calculate quality control metrics for mitochondrion, ribosomal, harhemoglobin genes, or other qc_vars.
+            """
+            try:
+                result = forward_request("ul_mark_var", request, adinfo)
+                if result is not None:
+                    return result
+                adata = get_ads().get_adata(adinfo=adinfo)
+                var_name = request.var_name
+                gene_class = request.gene_class
+                pattern_type = request.pattern_type
+                patterns = request.patterns
+                if gene_class is not None:
+                    if gene_class == "mitochondrion":
+                        adata.var["mt"] = adata.var_names.str.startswith(('MT-', 'Mt','mt-'))
+                        var_name = "mt"
+                    elif gene_class == "ribosomal":
+                        adata.var["ribo"] = adata.var_names.str.startswith(("RPS", "RPL", "Rps", "Rpl"))
+                        var_name = "ribo"
+                    elif gene_class == "hemoglobin":
+                        adata.var["hb"] = adata.var_names.str.contains("^HB[^(P)]", case=False)
+                        var_name = "hb"    
+                elif pattern_type is not None and patterns is not None:
+                    if pattern_type == "startswith":
+                        adata.var[var_name] = adata.var_names.str.startswith(patterns)
+                    elif pattern_type == "endswith":
+                        adata.var[var_name] = adata.var_names.str.endswith(patterns)
+                    elif pattern_type == "contains":
+                        adata.var[var_name] = adata.var_names.str.contains(patterns)
+                    else:
+                        raise ValueError(f"Did not support pattern_type: {pattern_type}")
+                else:
+                    raise ValueError(f"Please provide validated parameter")
+            
+                res = {var_name: adata.var[var_name].value_counts().to_dict(), "msg": f"add '{var_name}' column in adata.var"}
+                func_kwargs = {"var_name": var_name, "gene_class": gene_class, "pattern_type": pattern_type, "patterns": patterns}
+                add_op_log(adata, "mark_var", func_kwargs, adinfo)
+                return res
+            except ToolError as e:
+                raise ToolError(e)
+            except Exception as e:
+                if hasattr(e, '__context__') and e.__context__:
+                    raise ToolError(e.__context__)
+                else:
+                    raise ToolError(e)
+        return _mark_var
+
+    def _tool_list_var(self):
+        def _list_var(request: ListVarModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """List key columns in adata.var. It should be called for checking when other tools need var key column names as input."""
+            try:
+                result = forward_request("ul_list_var", request, adinfo)
+                if result is not None:
+                    return result
+                adata = get_ads().get_adata(adinfo=adinfo)
+                columns = list(adata.var.columns)
+                add_op_log(adata, list_var, {}, adinfo)
+                return columns
+            except ToolError as e:
+                raise ToolError(e)
+            except Exception as e:
+                if hasattr(e, '__context__') and e.__context__:
+                    raise ToolError(e.__context__)
+                else:
+                    raise ToolError(e)
+        return _list_var
+
+    def _tool_list_obs(self):
+        def _list_obs(request: ListObsModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """List key columns in adata.obs. It should be called before other tools need obs key column names input."""
+            try:
+                result = forward_request("ul_list_obs", request, adinfo)
+                if result is not None:
+                    return result    
+                adata = get_ads().get_adata(adinfo=adinfo)
+                columns = list(adata.obs.columns)
+                add_op_log(adata, list_obs, {}, adinfo)
+                return columns
+            except ToolError as e:
+                raise ToolError(e)
+            except Exception as e:
+                if hasattr(e, '__context__') and e.__context__:
+                    raise ToolError(e.__context__)
+                else:
+                    raise ToolError(e)
+        return _list_obs
+
+    def _tool_check_var(self):
+        def _check_var(request: VarNamesModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """Check if genes/variables exist in adata.var_names. This tool should be called before gene expression visualizations or color by genes."""
+            try:
+                result = forward_request("ul_check_var", request, adinfo)
+                if result is not None:
+                    return result     
+                adata = get_ads().get_adata(adinfo=adinfo)
+                var_names = request.var_names
+                result = {v: v in adata.var_names for v in var_names}
+                add_op_log(adata, check_var, {"var_names": var_names}, adinfo)
+                return result
+            except ToolError as e:
+                raise ToolError(e)
+            except Exception as e:
+                if hasattr(e, '__context__') and e.__context__:
+                    raise ToolError(e.__context__)
+                else:
+                    raise ToolError(e)
+        return _check_var
+
+    def _tool_merge_adata(self):
+        def _merge_adata(request: ConcatBaseModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """Merge multiple adata objects."""
+            try:
+                result = forward_request("ul_merge_adata", request, adinfo)
+                if result is not None:
+                    return result
+                ads = get_ads()
+                adata = ads.get_adata(adinfo=adinfo)
+                kwargs = {k: v for k, v in request.model_dump().items() if v is not None}
+                merged_adata = adata.concat(list(ads.adata_dic[dtype].values()), **kwargs)
+                ads.adata_dic[dtype] = {}
+                ads.active_id = "merged_adata"
+                add_op_log(merged_adata, ad.concat, kwargs, adinfo)
+                ads.adata_dic[ads.active_id] = merged_adata
+                return {"status": "success", "message": "Successfully merged all AnnData objects"}
+            except ToolError as e:
+                raise ToolError(e)
+            except Exception as e:
+                if hasattr(e, '__context__') and e.__context__:
+                    raise ToolError(e.__context__)
+                else:
+                    raise ToolError(e)
+        return _merge_adata
+
+    def _tool_set_dpt_iroot(self):
+        def _set_dpt_iroot(request: DPTIROOTModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """Set the iroot cell"""
+            try:
+                result = forward_request("ul_set_dpt_iroot", request, adinfo)
+                if result is not None:
+                    return result     
+                adata = get_ads().get_adata(adinfo=adinfo)
+                diffmap_key = request.diffmap_key
+                dimension = request.dimension
+                direction = request.direction
+                if diffmap_key not in adata.obsm:
+                    raise ValueError(f"Diffusion map key '{diffmap_key}' not found in adata.obsm")
+                if direction == "min":
+                    adata.uns["iroot"] = adata.obsm[diffmap_key][:, dimension].argmin()
+                else:  
+                    adata.uns["iroot"] = adata.obsm[diffmap_key][:, dimension].argmax()
+                
+                func_kwargs = {"diffmap_key": diffmap_key, "dimension": dimension, "direction": direction}
+                add_op_log(adata, "set_dpt_iroot", func_kwargs, adinfo)
+                
+                return {"status": "success", "message": f"Successfully set root cell for DPT using {direction} of dimension {dimension}"}
+            except ToolError as e:
+                raise ToolError(e)
+            except Exception as e:
+                if hasattr(e, '__context__') and e.__context__:
+                    raise ToolError(e.__context__)
+                else:
+                    raise ToolError(e)
+        return _set_dpt_iroot
+
+    def _tool_add_layer(self):
+        def _add_layer(request: AddLayerModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """Add a layer to the AnnData object."""
+            try:
+                result = forward_request("ul_add_layer", request, adinfo)
+                if result is not None:
+                    return result         
+                adata = get_ads().get_adata(adinfo=adinfo)
+                layer_name = request.layer_name
+                
+                # Check if layer already exists
+                if layer_name in adata.layers:
+                    raise ValueError(f"Layer '{layer_name}' already exists in adata.layers")
+                # Add the data as a new layer
+                adata.layers[layer_name] = adata.X.copy()
+
+                func_kwargs = {"layer_name": layer_name}
+                add_op_log(adata, "add_layer", func_kwargs, adinfo)
+                
+                return {
+                    "status": "success", 
+                    "message": f"Successfully added layer '{layer_name}' to adata.layers"
+                }
+            except ToolError as e:
+                raise ToolError(e)
+            except Exception as e:
+                if hasattr(e, '__context__') and e.__context__:
+                    raise ToolError(e.__context__)
+                else:
+                    raise ToolError(e)
+        return _add_layer
+
+    def _tool_check_samples(self):
+        def _check_samples(request: None, adinfo: self.AdataInfo=self.AdataInfo()):
+            """check the stored samples"""
+            try:
+                ads = get_ads()
+                return {"sampleid": [list(ads.adata_dic[dk].keys()) for dk in ads.adata_dic.keys()]}
+            except ToolError as e:
+                raise ToolError(e)
+            except Exception as e:
+                if hasattr(e, '__context__') and e.__context__:
+                    raise ToolError(e.__context__)
+                else:
+                    raise ToolError(e)
+        return _check_samples

scmcp_shared/util.py

@@ -1,8 +1,11 @@
 import inspect
 import os
+from enum import Enum
 from pathlib import Path
 from fastmcp.server.dependencies import get_context
 from fastmcp.exceptions import ToolError
+import asyncio
+import nest_asyncio
 
 
 def get_env(key):
@@ -58,7 +61,7 @@ def add_op_log(adata, func, kwargs, adinfo):
 
 
 
-def savefig(axes, file):
+def save_fig_path(axes, file):
     from matplotlib.axes import Axes
 
     try:
@@ -84,7 +87,7 @@ def savefig(axes, file):
         raise e
 
 
-def set_fig_path(axes, func=None, **kwargs):
+def savefig(axes, func=None, **kwargs):
     if hasattr(func, "func") and hasattr(func.func, "__name__"):
         # For partial functions, use the original function name
         func_name = func.func.__name__
@@ -107,9 +110,9 @@ def set_fig_path(axes, func=None, **kwargs):
     args_str = "_".join(args).replace(" ", "")
     fig_path = fig_dir / f"{func_name}_{args_str}.png"
     try:
-        savefig(axes, fig_path)
+        save_fig_path(axes, fig_path)
     except PermissionError:
-        raise PermissionError("You don't have permission to rename this file")
+        raise PermissionError("You don't have permission to save figure")
     except Exception as e:
         raise e
     transport = get_env("TRANSPORT") 
@@ -140,7 +143,7 @@ def add_figure_route(server):
     server._additional_http_routes = [Route("/figures/{figure_name}", endpoint=get_figure)]
 
 
-async def forward_request(func, request, adinfo, **kwargs):
+async def async_forward_request(func, request, adinfo, **kwargs):
     from fastmcp import Client
     forward_url = get_env("FORWARD")
     request_kwargs = request.model_dump()
@@ -149,8 +152,6 @@ async def forward_request(func, request, adinfo, **kwargs):
         "request": {k: request_kwargs.get(k) for k in request_args},
         "adinfo": adinfo.model_dump()
     }
-    print(func_kwargs)
-    # func_kwargs.update({k:v for k,v in kwargs.items() if v is not None})
     if not forward_url:
         return None
         
@@ -170,13 +171,35 @@ async def forward_request(func, request, adinfo, **kwargs):
                 raise e
 
 
+
+def forward_request(func, request, adinfo, **kwargs):
+    """Synchronous wrapper for forward_request"""
+    try:
+        # Apply nest_asyncio to allow nested event loops
+        nest_asyncio.apply()
+        loop = asyncio.get_event_loop()
+        if loop.is_running():
+            # If we're in a running event loop, use create_task
+            async def _run():
+                return await async_forward_request(func, request, adinfo, **kwargs)
+            return loop.run_until_complete(_run())
+        else:
+            # If no event loop is running, use asyncio.run()
+            return asyncio.run(async_forward_request(func, request, adinfo, **kwargs))
+    except Exception as e:
+        if hasattr(e, '__context__') and e.__context__:
+            raise Exception(f"{str(e.__context__)}")
+        else:
+            raise e
+
+
 def obsm2adata(adata, obsm_key):
     from anndata import AnnData
 
     if obsm_key not in adata.obsm_keys():
         raise ValueError(f"key {obsm_key} not found in adata.obsm")
     else:
-        return AnnData(adata.obsm[obsm_key], obs=adata.obs, obsm=adata.obsm)
+        return AnnData(adata.obsm[obsm_key], obs=adata.obs, obsm=adata.obsm, uns=adata.uns)
 
 
 def get_ads():
@@ -196,16 +219,58 @@ def sc_like_plot(plot_func, adata, request, adinfo, **kwargs):
     axes = plot_func(adata, **func_kwargs)
     if axes is None:
         axes = plt.gca()
-    fig_path = set_fig_path(axes, plot_func, **func_kwargs)
+    fig_path = savefig(axes, plot_func, **func_kwargs)
     add_op_log(adata, plot_func, func_kwargs, adinfo)
     return fig_path
 
 
-async def filter_tools(mcp, include_tools=None, exclude_tools=None):
-    tools = await mcp.get_tools()
-    for tool in tools:
-        if exclude_tools and tool in exclude_tools:
-            mcp.remove_tool(tool)
-        if include_tools and tool not in include_tools:
-            mcp.remove_tool(tool)
-    return mcp
+def filter_tools(mcp, include_tools=None, exclude_tools=None):
+    import asyncio
+    import copy
+    mcp = copy.deepcopy(mcp)
+    async def _filter_tools(mcp, include_tools=None, exclude_tools=None):
+        tools = await mcp.get_tools()
+        for tool in tools:
+            if exclude_tools and tool in exclude_tools:
+                mcp.remove_tool(tool)
+            if include_tools and tool not in include_tools:
+                mcp.remove_tool(tool)
+        return mcp
+    return asyncio.run(_filter_tools(mcp, include_tools, exclude_tools))
+
+
+ 
+def set_env(log_file, forward, transport, host, port):
+    if log_file is not None:
+        os.environ['SCMCP_LOG_FILE'] = log_file
+    if forward is not None:
+        os.environ['SCMCP_FORWARD'] = forward           
+    os.environ['SCMCP_TRANSPORT'] = transport
+    os.environ['SCMCP_HOST'] = host
+    os.environ['SCMCP_PORT'] = str(port)
+
+
+
+def setup_mcp(mcp, sub_mcp_dic, modules=None):
+    import asyncio
+    if modules is None or modules == "all":
+        modules = sub_mcp_dic.keys()
+    for module in modules:
+        asyncio.run(mcp.import_server(module, sub_mcp_dic[module]))
+    return mcp
+
+def _update_args(mcp, func, args_dic : dict):
+    defs = mcp._tool_manager._tools[func].parameters['$defs']
+    model_names = list(defs.keys())
+    args_model = model_names[0] if  model_names[0] != "AdataModel" else model_names[1]
+    for args, property_dic in args_dic.items():
+        for pk, v in property_dic.items():
+            for model in model_names:
+                if args in mcp._tool_manager._tools[func].parameters['$defs'][model]["properties"]:
+                    mcp._tool_manager._tools[func].parameters['$defs'][model]["properties"][args][pk] = v
+
+
+def update_mcp_args(mcp, tool_args : dict):
+    tools = mcp._tool_manager._tools.keys()
+    for tool in tool_args: 
+        _update_args(mcp, tool, tool_args[tool])

{scmcp_shared-0.2.5.dist-info → scmcp_shared-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: scmcp_shared
-Version: 0.2.5
+Version: 0.3.0
 Summary: A shared function libray for scmcphub
 Author-email: shuang <hsh-me@outlook.com>
 License: BSD 3-Clause License