scmcp-shared 0.4.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

scmcp_shared/__init__.py

@@ -1,3 +1 @@
-
-__version__ = "0.4.0"
-
+__version__ = "0.6.0"

scmcp_shared/agent.py

@@ -1,30 +1,47 @@
-import instructor
-from openai import OpenAI
-from scmcp_shared.schema.tool import ToolList
+from .schema.tool import ToolList
 import os
 
 
+from agno.agent import Agent
+from agno.models.openai import OpenAILike
+from scmcp_shared.kb import load_kb
+
+model = OpenAILike(
+    id=os.getenv("MODEL"),
+    base_url=os.getenv("BASE_URL"),
+    api_key=os.getenv("API_KEY"),
+)
 
-def select_tool(query):
 
-    API_KEY = os.environ.get("API_KEY",  None)
-    BASE_URL = os.environ.get("BASE_URL", None)
-    MODEL = os.environ.get("MODEL", None)
+def rag_agent(task, software=None):
+    knowledge_base = load_kb(software=software)
+    agent = Agent(
+        model=model,
+        knowledge=knowledge_base,
+        show_tool_calls=True,
+        search_knowledge=True,
+    )
+    query = f"""
+    <task>
+    {task}
+    </task>
+    查询知识库，给出一个用于解决任务的代码示例。返回结果格式为：
+    <code_example>
+        [code_example]
+    </code_example>
+    """
+    rep = agent.run(query)
+    return rep.content
 
-    client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
-    client = instructor.from_openai(client)
 
-    response = client.chat.completions.create(
-        model=MODEL,
-        messages=[
-            {
-                "role": "system", 
-                "content": f"you are a bioinformatician, you are given a task and a list of tools, you need to select the most directly relevant tools to use to solve the task"},
-            {
-                "role": "user",
-                "content": query
-            },  
-        ],
+def select_tool(query):
+    agent = Agent(
+        model=model,
         response_model=ToolList,
+        use_json_mode=True,
+        instructions="""
+        you are a bioinformatician, you are given a task and a list of tools, you need to select the most directly relevant tools to use to solve the task
+        """,
     )
-    return response.tools
+    rep = agent.run(query)
+    return rep.content

scmcp_shared/backend.py

@@ -0,0 +1,44 @@
+from collections.abc import Iterable
+from abcoder.backend import NotebookManager
+
+__all__ = ["AdataManager", "NotebookManager"]
+
+
+class AdataManager:
+    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

scmcp_shared/cli.py

@@ -1,51 +1,69 @@
 import argparse
-from typing import Optional, Union, Type, Dict, Callable
-from enum import Enum
+from typing import Optional, Dict, Callable
 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.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')
+        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)
-        
+        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('--tool-mode', choices=["auto", "normal"], default="normal", help='tool selection mode')
-        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:
+        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(
+            "--run-mode", choices=["tool", "code"], default="code", help="run mode"
+        )
+        parser.add_argument(
+            "--tool-mode",
+            choices=["auto", "normal"],
+            default="normal",
+            help="tool selection mode",
+        )
+        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
         """
@@ -53,24 +71,24 @@ class MCPCLI:
         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 hasattr(args, "module"):
             if "all" in args.module:
                 modules = None
             elif isinstance(args.module, list) and bool(args.module):
@@ -78,16 +96,31 @@ class MCPCLI:
         else:
             modules = None
         if self.manager is not None:
-            self.mcp = self.manager(self.name, include_modules=modules).mcp
-            all_tools = self.mcp._tool_manager._tools
-            auto_tools = {tool: all_tools[tool] for tool in all_tools if all_tools[tool].name in ["search_tool", "run_tool"]}
-            if args.tool_mode == "auto":
-                all_tools = self.mcp._tool_manager._tools
-                self.mcp._tool_manager._all_tools = all_tools
-                self.mcp._tool_manager._tools = auto_tools
+            from .backend import NotebookManager, AdataManager
+
+            if args.run_mode == "code":
+                backend = NotebookManager
+                self.mcp = self.manager(
+                    self.name, include_tags=["nb", "rag"], backend=backend
+                ).mcp
             else:
-                for name in auto_tools:
-                    self.mcp._tool_manager.remove_tool(name)
+                backend = AdataManager
+                self.mcp = self.manager(
+                    self.name, include_modules=modules, backend=backend
+                ).mcp
+                all_tools = self.mcp._tool_manager._tools
+                auto_tools = {
+                    tool: all_tools[tool]
+                    for tool in all_tools
+                    if "auto" in all_tools[tool].tags
+                }
+                if args.tool_mode == "auto":
+                    all_tools = self.mcp._tool_manager._tools
+                    self.mcp._tool_manager._all_tools = all_tools
+                    self.mcp._tool_manager._tools = auto_tools
+                else:
+                    for name in auto_tools:
+                        self.mcp._tool_manager.remove_tool(name)
         elif self.mcp is not None:
             pass
         else:
@@ -98,19 +131,15 @@ class MCPCLI:
     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"
-            )
- 
+            self.mcp.run(transport=transport, host=host, port=port, log_level="info")
+
     def run(self):
         """Run the CLI application."""
         args = self.parser.parse_args()

scmcp_shared/kb.py

@@ -0,0 +1,139 @@
+from agno.document.chunking.agentic import AgenticChunking
+from agno.embedder.openai import OpenAIEmbedder
+from agno.models.deepseek import DeepSeek
+from agno.vectordb.lancedb import LanceDb
+from agno.knowledge.agent import AgentKnowledge
+import importlib.resources
+import os
+import requests
+import zipfile
+import tempfile
+import shutil
+from pathlib import Path
+import logging
+
+logger = logging.getLogger(__name__)
+
+embedder_id = os.getenv("EMBEDDER_MODEL")
+embedder_api_key = os.getenv("EMBEDDER_API_KEY")
+embedder_base_url = os.getenv("EMBEDDER_BASE_URL")
+model_id = os.getenv("MODEL")
+model_api_key = os.getenv("API_KEY")
+model_base_url = os.getenv("BASE_URL")
+
+# 配置信息
+config = {
+    "local_dir": "vector_db",
+    "huggingface_url": "https://huggingface.co/datasets/huangshing/scmcp_vector_db/resolve/main/vector_db.zip",
+}
+
+
+def download_vector_db(source="huggingface"):
+    """
+    下载向量数据库文件
+
+    Args:
+        source: 下载源 ("huggingface" 或 "github")
+    """
+
+    # 获取本地存储路径
+    package_path = importlib.resources.path("scmcp_shared", "")
+    local_dir = Path(package_path) / config["local_dir"]
+    local_dir.mkdir(exist_ok=True)
+
+    # 检查是否已存在
+    if (local_dir / "scmcp.lance").exists():
+        logger.info("Vector database already exists locally")
+        return str(local_dir)
+
+    logger.info(f"Downloading vector database from {source}...")
+
+    # 创建临时目录用于下载和解压
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_path = Path(temp_dir)
+        zip_path = temp_path / "vector_db.zip"
+
+        try:
+            # 下载文件
+            if source == "huggingface":
+                url = config["huggingface_url"]
+            else:
+                raise ValueError(f"Unsupported source: {source}")
+
+            logger.info(f"Downloading from: {url}")
+            response = requests.get(url, stream=True)
+            response.raise_for_status()
+
+            with open(zip_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+
+            # 解压文件
+            logger.info("Extracting downloaded archive...")
+            _extract_archive(zip_path, local_dir)
+
+            logger.info(f"Vector database downloaded and extracted to: {local_dir}")
+            return str(local_dir)
+
+        except requests.RequestException as e:
+            raise RuntimeError(f"Failed to download vector database: {e}")
+        except Exception as e:
+            raise RuntimeError(f"Failed to process vector database: {e}")
+
+
+def _extract_archive(archive_path, extract_dir):
+    """解压归档文件"""
+    with zipfile.ZipFile(archive_path, "r") as zip_ref:
+        zip_ref.extractall(extract_dir)
+
+    # 如果解压后只有一个子目录，移动内容到目标目录
+    extracted_items = list(Path(extract_dir).iterdir())
+    if len(extracted_items) == 1 and extracted_items[0].is_dir():
+        subdir = extracted_items[0]
+        for item in subdir.iterdir():
+            shutil.move(str(item), str(extract_dir / item.name))
+        subdir.rmdir()
+
+
+def load_kb(software=None, auto_download=True, download_source="huggingface"):
+    """
+    加载知识库
+
+    Args:
+        software: 软件名称
+        auto_download: 是否自动下载向量数据库
+        download_source: 下载源 ("huggingface" 或 "github")
+    """
+    # 获取向量数据库路径
+    try:
+        vector_db_path = importlib.resources.path("scmcp_shared", "vector_db")
+    except FileNotFoundError:
+        if auto_download:
+            logger.info("Vector database not found in package, attempting download...")
+            vector_db_path = download_vector_db(download_source)
+        else:
+            raise FileNotFoundError(
+                "Vector database not found. Set auto_download=True to download automatically, "
+                "or manually place the vector database in the scmcp_shared package."
+            )
+
+    vector_db = LanceDb(
+        table_name=software,
+        uri=vector_db_path,
+        embedder=OpenAIEmbedder(
+            id=embedder_id,
+            base_url=embedder_base_url,
+            api_key=embedder_api_key,
+        ),
+    )
+    model = DeepSeek(
+        id=model_id,
+        base_url=model_base_url,
+        api_key=model_api_key,
+    )
+    knowledge_base = AgentKnowledge(
+        chunking_strategy=AgenticChunking(model=model),
+        vector_db=vector_db,
+    )
+
+    return knowledge_base

scmcp_shared/logging_config.py

@@ -1,18 +1,17 @@
 import logging
 import sys
-import os
+
 from .util import get_env
 
-def setup_logger(name="sc-mcp-server", log_file=None):
 
+def setup_logger(name="sc-mcp-server", log_file=None):
     logger = logging.getLogger(name)
     logger.setLevel(logging.INFO)
     if logger.handlers:
         return logger
-    
+
     formatter = logging.Formatter(
-        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-        '%Y-%m-%d %H:%M:%S'
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s", "%Y-%m-%d %H:%M:%S"
     )
     if log_file is None:
         log_file = get_env("LOG_FILE")
@@ -20,12 +19,11 @@ def setup_logger(name="sc-mcp-server", log_file=None):
         log_handler = logging.FileHandler(log_file)
         log_handler.setFormatter(formatter)
         logger.addHandler(log_handler)
-        
+
         logger.info(f"logging output: {log_file}")
     else:
         log_handler = logging.StreamHandler(sys.stdout)
         log_handler.setFormatter(formatter)
         logger.addHandler(log_handler)
-        logger.info(f"loggin file output: stdout")
+        logger.info("loggin file output: stdout")
     return logger
-    

scmcp_shared/mcp_base.py

@@ -0,0 +1,184 @@
+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
+from abcoder.backend import NotebookManager
+
+
+# class BaseMCP:
+#     """Base class for all Scanpy MCP classes."""
+
+#     def __init__(self, name: str, instructions=None, lifespan=None, 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, instructions=instructions, lifespan=lifespan)
+#         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 = [tl_method() for name, tl_method in methods if name.startswith('_tool_')]
+
+#         # Filter tools based on include/exclude lists
+#         if self.include_tools is not None:
+#             tool_methods = [tl for tl in tool_methods if tl.name in self.include_tools]
+
+#         if self.exclude_tools is not None:
+#             tool_methods = [tl for tl in tool_methods if tl.name not in self.exclude_tools]
+
+#         # Register filtered tools
+#         for tool in tool_methods:
+#             # Get the function returned by the tool method
+#             if tool is not None:
+#                 self.mcp.add_tool(tool)
+
+
+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 = [
+            tl_method() for name, tl_method in methods if name.startswith("_tool_")
+        ]
+
+        # Filter tools based on include/exclude lists
+        if self.include_tools is not None:
+            tool_methods = [tl for tl in tool_methods if tl.name in self.include_tools]
+
+        if self.exclude_tools is not None:
+            tool_methods = [
+                tl for tl in tool_methods if tl.name not in self.exclude_tools
+            ]
+
+        # Register filtered tools
+        for tool in tool_methods:
+            # Get the function returned by the tool method
+            if tool is not None:
+                self.mcp.add_tool(tool)
+
+
+class BaseMCPManager:
+    """Base class for MCP module management."""
+
+    def __init__(
+        self,
+        name: str,
+        instructions: Optional[str] = None,
+        include_modules: Optional[List[str]] = None,
+        exclude_modules: Optional[List[str]] = None,
+        include_tools: Optional[List[str]] = None,
+        exclude_tools: Optional[List[str]] = None,
+        include_tags: Optional[List[str]] = None,
+        exclude_tags: Optional[List[str]] = None,
+        backend: Optional[NotebookManager] = 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.exec_backend = backend()
+        self.mcp = FastMCP(
+            name,
+            instructions=instructions,
+            lifespan=self.exec_lifespan,
+            include_tags=include_tags,
+            exclude_tags=exclude_tags,
+        )
+        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_mcp()
+        self.register_mcp()
+
+    def init_mcp(self):
+        """Initialize available modules. To be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement init_mcp")
+
+    def register_mcp(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 exec_lifespan(self, server: FastMCP) -> AsyncIterator[Any]:
+        yield self.exec_backend