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

scmcp_shared/__init__.py

@@ -1,3 +1,3 @@
 
-__version__ = "0.2.5"
+__version__ = "0.3.5"
 

scmcp_shared/cli.py

@@ -0,0 +1,111 @@
+import argparse
+from typing import Optional, Union, Type, Dict, Callable
+from enum import Enum
+from .util import add_figure_route, set_env
+import os
+
+
+class MCPCLI:
+    """Base class for CLI applications with support for dynamic modules and parameters."""
+    
+    def __init__(self, name: str, help_text: str, mcp=None, manager=None):
+        self.name = name
+        self.mcp = mcp
+        self.manager = manager
+        self.parser = argparse.ArgumentParser(
+            description=help_text,
+            prog=name
+        )
+        self.subcommands: Dict[str, tuple[argparse.ArgumentParser, Callable]] = {}
+        self._setup_commands()
+    
+    def _setup_commands(self):
+        """Setup the main commands for the CLI."""
+        subparsers = self.parser.add_subparsers(dest='command', help='Available commands')        
+        run_parser = subparsers.add_parser('run', help='Start the server with the specified configuration')
+        self._setup_run_command(run_parser)
+        self.subcommands['run'] = (run_parser, self._run_command)
+        
+    def _setup_run_command(self, parser: argparse.ArgumentParser):
+        """Setup run command arguments."""
+        parser.add_argument('-t', '--transport', default="stdio",
+                          choices=["stdio", "shttp", "sse"],
+                          help='specify transport type')
+        parser.add_argument('-p', '--port', type=int, default=8000, help='transport port')
+        parser.add_argument('--host', default='127.0.0.1', help='transport host')
+        parser.add_argument('-f', '--forward', help='forward request to another server')
+        parser.add_argument('-wd', '--working-dir', default=".", help='working directory')
+        parser.add_argument('--log-file', help='log file path, use stdout if None')
+    
+    def add_command(self, name: str, help_text: str, handler: Callable) -> argparse.ArgumentParser:
+        """add new subcommand
+        
+        Args:
+            name: subcommand name
+            help_text: help text
+            handler: handler function
+            
+        Returns:
+            ArgumentParser: parser for the subcommand
+        """
+        subparsers = self.parser._subparsers._group_actions[0]
+        parser = subparsers.add_parser(name, help=help_text)
+        self.subcommands[name] = (parser, handler)
+        return parser
+    
+    def get_command_parser(self, name: str) -> Optional[argparse.ArgumentParser]:
+        """get the parser for the subcommand
+        
+        Args:
+            name: subcommand name
+            
+        Returns:
+            ArgumentParser: parser for the subcommand, return None if the subcommand does not exist
+        """
+        if name in self.subcommands:
+            return self.subcommands[name][0]
+        return None
+    
+    def _run_command(self, args):
+        """Start the server with the specified configuration."""
+        os.chdir(args.working_dir)
+        if hasattr(args, 'module'):
+            if "all" in args.module:
+                modules = None
+            elif isinstance(args.module, list) and bool(args.module):
+                modules = args.module
+        else:
+            modules = None
+        if self.manager is not None:
+            self.mcp = self.manager(self.name, include_modules=modules).mcp
+        elif self.mcp is not None:
+            pass
+        else:
+            raise ValueError("No manager or mcp provided")
+        transport = args.transport
+        self.run_mcp(args.log_file, args.forward, transport, args.host, args.port)
+
+    def run_mcp(self, log_file, forward, transport, host, port):
+        set_env(log_file, forward, transport, host, port)
+        from .logging_config import setup_logger
+        setup_logger(log_file)
+        if transport == "stdio":
+            self.mcp.run()
+        elif transport in ["sse", "shttp"]:
+            transport = "streamable-http" if transport == "shttp" else transport
+            add_figure_route(self.mcp)
+            self.mcp.run(
+                transport=transport,
+                host=host, 
+                port=port, 
+                log_level="info"
+            )
+ 
+    def run(self):
+        """Run the CLI application."""
+        args = self.parser.parse_args()
+        if args.command in self.subcommands:
+            handler = self.subcommands[args.command][1]
+            handler(args)
+        else:
+            self.parser.print_help()

scmcp_shared/schema/__init__.py

@@ -11,3 +11,12 @@ class AdataModel(BaseModel):
     model_config = ConfigDict(
         extra="ignore"
     )
+
+class AdataInfo(BaseModel):
+    """Input schema for the adata tool."""
+    sampleid: str | None = Field(default=None, description="adata sampleid")
+    adtype: str = Field(default="exp", description="The input adata.X data type for preprocess/analysis/plotting")
+
+    model_config = ConfigDict(
+        extra="ignore"
+    )

scmcp_shared/schema/io.py

@@ -22,10 +22,6 @@ class ReadModel(BaseModel):
         default=None,
         description="Name of sheet/table in hdf5 or Excel file."
     )
-    ext: str = Field(
-        default=None,
-        description="Extension that indicates the file type. If None, uses extension of filename."
-    )
     delimiter: str = Field(
         default=None,
         description="Delimiter that separates data within text file. If None, will split at arbitrary number of white spaces, which is different from enforcing splitting at any single white space."

scmcp_shared/schema/tl.py

@@ -300,11 +300,6 @@ class LeidenModel(BaseModel):
         description="Which package's implementation to use."
     )
     
-    clustering_args: Optional[Dict[str, Any]] = Field(
-        default=None,
-        description="Any further arguments to pass to the clustering algorithm."
-    )
-    
     @field_validator('resolution')
     def validate_resolution(cls, v: float) -> float:
         """Validate resolution is positive"""
@@ -943,6 +938,11 @@ class PCAModel(BaseModel):
         gt=0
     )
     
+    key_added: str = Field(
+        default="X_pca",
+        description="PCA embedding stored key in adata.obsm."
+    )
+    
     @field_validator('n_comps', 'chunk_size')
     def validate_positive_integers(cls, v: Optional[int]) -> Optional[int]:
         """Validate positive integers"""

scmcp_shared/server/__init__.py

@@ -1,51 +1,13 @@
 from __future__ import annotations
 
-from collections.abc import Iterable
-from typing import Any
-
-from .io import io_mcp
-from .util import ul_mcp
-from .pl import pl_mcp
-from .pp import pp_mcp
-from .tl import tl_mcp
-
-
-
-class AdataState:
-    def __init__(self, add_adtypes=None):
-        self.adata_dic = {"exp": {}, "activity": {}, "cnv": {}, "splicing": {}}
-        if isinstance(add_adtypes, str):
-            self.adata_dic[add_adtypes] = {}
-        elif isinstance(add_adtypes, Iterable):
-            self.adata_dic.update({adtype: {} for adtype in add_adtypes})
-        self.active_id = None
-        self.metadatWa = {}
-        self.cr_kernel = {}
-        self.cr_estimator = {}
-        
-    def get_adata(self, sampleid=None, adtype="exp", adinfo=None):
-        if adinfo is not None:
-            kwargs = adinfo.model_dump()
-            sampleid = kwargs.get("sampleid", None)
-            adtype = kwargs.get("adtype", "exp")
-        try:
-            if self.active_id is None:
-                return None
-            sampleid = sampleid or self.active_id
-            return self.adata_dic[adtype][sampleid]
-        except KeyError as e:
-            raise KeyError(f"Key {e} not found in adata_dic[{adtype}].Please check the sampleid or adtype.")
-        except Exception as e:
-            raise Exception(f"fuck {e} {type(e)}")
-    
-    def set_adata(self, adata, sampleid=None, sdtype="exp", adinfo=None):
-        if adinfo is not None:
-            kwargs = adinfo.model_dump()
-            sampleid = kwargs.get("sampleid", None)
-            sdtype = kwargs.get("adtype", "exp")
-        sampleid = sampleid or self.active_id
-        if sdtype not in self.adata_dic:
-            self.adata_dic[sdtype] = {}
-        self.adata_dic[sdtype][sampleid] = adata
-                          
-
+from collections.abc import Iterable, AsyncIterator
+from typing import Any, Dict, List, Optional
+from contextlib import asynccontextmanager
+import asyncio
+
+from .base import BaseMCP,AdataState,BaseMCPManager
+from .io import ScanpyIOMCP, io_mcp
+from .util import ScanpyUtilMCP
+from .pl import ScanpyPlottingMCP
+from .pp import ScanpyPreprocessingMCP
+from .tl import ScanpyToolsMCP

scmcp_shared/server/base.py

@@ -0,0 +1,153 @@
+import inspect
+from fastmcp import FastMCP
+from ..schema import AdataInfo
+from ..util import filter_tools
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+import asyncio
+from typing import Optional, List, Any, Iterable
+
+
+class BaseMCP:
+    """Base class for all Scanpy MCP classes."""
+    
+    def __init__(self, name: str, include_tools: list = None, exclude_tools: list = None, AdataInfo = AdataInfo):
+        """
+        Initialize BaseMCP with optional tool filtering.
+        
+        Args:
+            name (str): Name of the MCP server
+            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.
+        """
+        self.mcp = FastMCP(name)
+        self.include_tools = include_tools
+        self.exclude_tools = exclude_tools
+        self.AdataInfo = AdataInfo
+        self._register_tools()
+
+    def _register_tools(self):
+        """Register all tool methods with the FastMCP instance based on include/exclude filters"""
+        # Get all methods of the class
+        methods = inspect.getmembers(self, predicate=inspect.ismethod)
+        
+        # Filter methods that start with _tool_
+        tool_methods = {
+            name[6:]: method  # Remove '_tool_' prefix
+            for name, method in methods
+            if name.startswith('_tool_')
+        }
+        
+        # Filter tools based on include/exclude lists
+        if self.include_tools is not None:
+            tool_methods = {k: v for k, v in tool_methods.items() if k in self.include_tools}
+        
+        if self.exclude_tools is not None:
+            tool_methods = {k: v for k, v in tool_methods.items() if k not in self.exclude_tools}
+
+        # Register filtered tools
+        for tool_name, tool_method in tool_methods.items():
+            # Get the function returned by the tool method
+            tool_func = tool_method()
+            if tool_func is not None:
+                self.mcp.add_tool(tool_func, name=tool_name)
+
+
+class AdataState:
+    def __init__(self, add_adtypes=None):
+        self.adata_dic = {"exp": {}, "activity": {}, "cnv": {}, "splicing": {}}
+        if isinstance(add_adtypes, str):
+            self.adata_dic[add_adtypes] = {}
+        elif isinstance(add_adtypes, Iterable):
+            self.adata_dic.update({adtype: {} for adtype in add_adtypes})
+        self.active_id = None
+        self.metadatWa = {}
+        self.cr_kernel = {}
+        self.cr_estimator = {}
+        
+    def get_adata(self, sampleid=None, adtype="exp", adinfo=None):
+        if adinfo is not None:
+            kwargs = adinfo.model_dump()
+            sampleid = kwargs.get("sampleid", None)
+            adtype = kwargs.get("adtype", "exp")
+        try:
+            if self.active_id is None:
+                return None
+            sampleid = sampleid or self.active_id
+            return self.adata_dic[adtype][sampleid]
+        except KeyError as e:
+            raise KeyError(f"Key {e} not found in adata_dic[{adtype}].Please check the sampleid or adtype.")
+        except Exception as e:
+            raise Exception(f"fuck {e} {type(e)}")
+    
+    def set_adata(self, adata, sampleid=None, sdtype="exp", adinfo=None):
+        if adinfo is not None:
+            kwargs = adinfo.model_dump()
+            sampleid = kwargs.get("sampleid", None)
+            sdtype = kwargs.get("adtype", "exp")
+        sampleid = sampleid or self.active_id
+        if sdtype not in self.adata_dic:
+            self.adata_dic[sdtype] = {}
+        self.adata_dic[sdtype][sampleid] = adata
+
+
+class BaseMCPManager:
+    """Base class for MCP module management."""
+    
+    def __init__(self, 
+        name: str, 
+        include_modules: Optional[List[str]] = None, 
+        exclude_modules: Optional[List[str]] = None,
+        include_tools: Optional[List[str]] = None,
+        exclude_tools: Optional[List[str]] = None,
+    ):
+        """
+        Initialize BaseMCPManager with optional module filtering.
+        
+        Args:
+            name (str): Name of the MCP server
+            include_modules (List[str], optional): List of module names to include. If None, all modules are included.
+            exclude_modules (List[str], optional): List of module names to exclude. If None, no modules are excluded.
+            include_tools (List[str], optional): List of tool names to include. If None, all tools are included.
+            exclude_tools (List[str], optional): List of tool names to exclude. If None, no tools are excluded.
+        """
+        self.ads = AdataState()
+        self.mcp = FastMCP(name, lifespan=self.adata_lifespan)
+        self.include_modules = include_modules
+        self.exclude_modules = exclude_modules
+        self.include_tools = include_tools
+        self.exclude_tools = exclude_tools
+        self.available_modules = {}
+        self._init_modules()
+        self._register_modules()
+
+    def _init_modules(self):
+        """Initialize available modules. To be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement _init_modules")
+
+    def _register_modules(self):
+        """Register modules based on include/exclude filters."""
+        # Filter modules based on include/exclude lists
+        if self.include_modules is not None:
+            self.available_modules = {k: v for k, v in self.available_modules.items() if k in self.include_modules}
+        
+        if self.exclude_modules is not None:
+            self.available_modules = {k: v for k, v in self.available_modules.items() if k not in self.exclude_modules}
+
+        # Register each module
+        for module_name, mcpi in self.available_modules.items():
+            if isinstance(mcpi, FastMCP):
+                if self.include_tools is not None and module_name in self.include_tools:
+                    mcpi = filter_tools(mcpi, include_tools= self.include_tools[module_name])
+                if self.exclude_tools is not None and module_name in self.exclude_tools:
+                    mcpi = filter_tools(mcpi, exclude_tools=self.exclude_tools[module_name])
+
+                asyncio.run(self.mcp.import_server(module_name, mcpi))
+            else:
+                asyncio.run(self.mcp.import_server(module_name, mcpi().mcp))
+
+    @asynccontextmanager
+    async def adata_lifespan(self, server: FastMCP) -> AsyncIterator[Any]:
+        """Context manager for AdataState lifecycle."""
+        yield self.ads

scmcp_shared/server/io.py

@@ -2,83 +2,83 @@ import os
 import inspect
 from pathlib import Path
 import scanpy as sc
-from fastmcp import FastMCP , Context
+from fastmcp import FastMCP, Context
 from fastmcp.exceptions import ToolError
-from ..schema import AdataModel
+from ..schema import AdataInfo
 from ..schema.io import *
 from ..util import filter_args, forward_request, get_ads, generate_msg
+from .base import BaseMCP
 
 
-io_mcp = FastMCP("SCMCP-IO-Server")
+class ScanpyIOMCP(BaseMCP):
+    def __init__(self, include_tools: list = None, exclude_tools: list = None, AdataInfo = AdataInfo):
+        """Initialize ScanpyIOMCP with optional tool filtering."""
+        super().__init__("SCMCP-IO-Server", include_tools, exclude_tools, AdataInfo)
 
+    def _tool_read(self):
+        def _read(request: ReadModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """
+            Read data from 10X directory or various file formats (h5ad, 10x, text files, etc.).
+            """
+            try:
+                res = forward_request("io_read", request, adinfo)
+                if res is not None:
+                    return res
+                kwargs = request.model_dump()
+                file = Path(kwargs.get("filename", None))
+                if file.is_dir():
+                    kwargs["path"] = kwargs["filename"]
+                    func_kwargs = filter_args(request, sc.read_10x_mtx)
+                    adata = sc.read_10x_mtx(kwargs["path"], **func_kwargs)
+                elif file.is_file():
+                    func_kwargs = filter_args(request, sc.read)
+                    adata = sc.read(**func_kwargs)
+                    if not kwargs.get("first_column_obs", True):
+                        adata = adata.T
+                else:
+                    raise FileNotFoundError(f"{kwargs['filename']} does not exist")
 
-@io_mcp.tool()
-async def read(
-    request: ReadModel, 
-    adinfo: AdataModel = AdataModel()
-):
-    """
-    Read data from 10X directory or various file formats (h5ad, 10x, text files, etc.).
-    """
-    try:
-        result = await forward_request("io_read", request, adinfo)
-        if result is not None:
-            return result        
-        kwargs = request.model_dump()
+                ads = get_ads()
+                if adinfo.sampleid is not None:
+                    ads.active_id = adinfo.sampleid
+                else:
+                    ads.active_id = f"adata{len(ads.adata_dic[adinfo.adtype])}"
+                    
+                adata.layers["counts"] = adata.X
+                adata.var_names_make_unique()
+                adata.obs_names_make_unique()
+                ads.set_adata(adata, adinfo=adinfo)
+                return generate_msg(adinfo, adata, ads)
+            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 _read
 
-        file = Path(kwargs.get("filename", None))
-        if file.is_dir():
-            kwargs["path"] = kwargs["filename"]
-            func_kwargs = filter_args(request, sc.read_10x_mtx)
-            adata = sc.read_10x_mtx(kwargs["path"], **func_kwargs)
-        elif file.is_file():
-            func_kwargs = filter_args(request, sc.read)
-            adata = sc.read(**func_kwargs)
-            if not kwargs.get("first_column_obs", True):
-                adata = adata.T
-        else:
-            raise FileNotFoundError(f"{kwargs['filename']} does not exist")
+    def _tool_write(self):
+        def _write(request: WriteModel, adinfo: self.AdataInfo=self.AdataInfo()):
+            """save adata into a file."""
+            try:
+                res = forward_request("io_write", request, adinfo)
+                if res is not None:
+                    return res   
+                ads = get_ads()
+                adata = ads.get_adata(adinfo=adinfo)
+                kwargs = request.model_dump()
+                sc.write(kwargs["filename"], adata)
+                return {"filename": kwargs["filename"], "msg": "success to save file"}
+            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 _write
 
-        ads = get_ads()
-        if adinfo.sampleid is not None:
-            ads.active_id = adinfo.sampleid
-        else:
-            ads.active_id = f"adata{len(ads.adata_dic[adinfo.adtype])}"
-            
-        adata.layers["counts"] = adata.X
-        adata.var_names_make_unique()
-        adata.obs_names_make_unique()
-        ads.set_adata(adata, adinfo=adinfo)
-        return generate_msg(adinfo, adata, ads)
-    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)
 
-
-@io_mcp.tool()
-async def write(
-    request: WriteModel, 
-    adinfo: AdataModel = AdataModel()
-):
-    """save adata into a file.
-    """
-    try:
-        result = await forward_request("io_write", request, adinfo)
-        if result is not None:
-            return result
-        ads = get_ads()
-        adata = ads.get_adata(adinfo=adinfo)
-        kwargs = request.model_dump()
-        sc.write(kwargs["filename"], adata)
-        return {"filename": kwargs["filename"], "msg": "success to save file"}
-    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)
+# Create an instance of the class
+io_mcp = ScanpyIOMCP().mcp