sirchmunk 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

sirchmunk_mcp/service.py

@@ -0,0 +1,327 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""
+Sirchmunk Service Wrapper for MCP Server.
+
+Provides a high-level interface to Sirchmunk's AgenticSearch functionality,
+managing initialization, configuration, and session state.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import logging
+import os
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from .config import Config
+
+if TYPE_CHECKING:
+    from sirchmunk.schema.knowledge import KnowledgeCluster
+
+
+logger = logging.getLogger(__name__)
+
+
+@contextlib.contextmanager
+def suppress_stdout():
+    """Context manager to suppress stdout output.
+    
+    Used during initialization to prevent third-party libraries
+    (ModelScope, transformers, etc.) from printing to stdout,
+    which would break MCP stdio protocol.
+    """
+    # Check if we're in stdio MCP mode (stdout should be protected)
+    if os.environ.get("MCP_TRANSPORT") == "stdio":
+        old_stdout = sys.stdout
+        sys.stdout = io.StringIO()
+        try:
+            yield
+        finally:
+            sys.stdout = old_stdout
+    else:
+        yield
+
+
+class SirchmunkService:
+    """Service wrapper for AgenticSearch with lifecycle management.
+    
+    This class manages the AgenticSearch instance and provides a clean interface
+    for MCP tool implementations.
+    
+    Attributes:
+        config: Configuration object
+        search: AgenticSearch instance
+        initialized: Whether the service is initialized
+    """
+    
+    def __init__(self, config: Config):
+        """Initialize Sirchmunk service.
+        
+        Args:
+            config: Configuration object
+        
+        Raises:
+            RuntimeError: If initialization fails
+        """
+        self.config = config
+        self.searcher: Optional[AgenticSearch] = None
+        self.initialized = False
+        
+        logger.info(f"Initializing Sirchmunk service with config: {config.sirchmunk.work_path}")
+        
+        try:
+            self._initialize_search()
+            self.initialized = True
+            logger.info("Sirchmunk service initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize Sirchmunk service: {e}")
+            raise RuntimeError(f"Sirchmunk service initialization failed: {e}") from e
+    
+    def _initialize_search(self) -> None:
+        """Initialize AgenticSearch instance with configuration.
+        
+        Raises:
+            Exception: If AgenticSearch initialization fails
+        """
+        # Import sirchmunk modules inside function to allow stdout suppression
+        # These imports may trigger model downloads that print to stdout
+        with suppress_stdout():
+            from sirchmunk.search import AgenticSearch
+            from sirchmunk.llm.openai_chat import OpenAIChat
+        
+        # Create LLM client
+        llm = OpenAIChat(
+            base_url=self.config.llm.base_url,
+            api_key=self.config.llm.api_key,
+            model=self.config.llm.model_name,
+        )
+        
+        # Create AgenticSearch instance with stdout suppression
+        # AgenticSearch may load embedding models which print progress
+        with suppress_stdout():
+            self.searcher = AgenticSearch(
+                llm=llm,
+                work_path=self.config.sirchmunk.work_path,
+                verbose=False,  # Disable verbose in stdio mode to prevent stdout pollution
+                reuse_knowledge=self.config.sirchmunk.enable_cluster_reuse,
+                cluster_sim_threshold=self.config.sirchmunk.cluster_similarity.threshold,
+                cluster_sim_top_k=self.config.sirchmunk.cluster_similarity.top_k,
+            )
+        
+        logger.info("AgenticSearch instance created")
+    
+    async def search(
+        self,
+        query: str,
+        search_paths: Union[str, List[str]],
+        mode: str = "DEEP",
+        max_depth: Optional[int] = None,
+        top_k_files: Optional[int] = None,
+        keyword_levels: Optional[int] = None,
+        include: Optional[List[str]] = None,
+        exclude: Optional[List[str]] = None,
+        return_cluster: bool = False,
+    ) -> Union[str, List[Dict[str, Any]], KnowledgeCluster, None]:
+        """Search and retrieve various types of raw documents using AgenticSearch.
+        
+        This method performs intelligent search across code, documentation, and 
+        other text-based documents. It directly retrieves and analyzes raw content
+        from multiple file types including source code, markdown files, PDFs, 
+        text files, and other document formats supported by ripgrep-all.
+        
+        Args:
+            query: Search query or question to find relevant documents
+            search_paths: Paths to search in (files or directories)
+            mode: Search mode (DEEP, FAST, FILENAME_ONLY)
+            max_depth: Maximum directory depth to search
+            top_k_files: Number of top files to return
+            keyword_levels: Number of keyword granularity levels (DEEP mode)
+            include: File patterns to include (glob)
+            exclude: File patterns to exclude (glob)
+            return_cluster: Whether to return full KnowledgeCluster object
+        
+        Returns:
+            Search results: str (summary), List[Dict] (FILENAME_ONLY), 
+            KnowledgeCluster (if return_cluster=True), or None (if no results)
+        
+        Raises:
+            RuntimeError: If service is not initialized
+            ValueError: If parameters are invalid
+        """
+        if not self.initialized or self.searcher is None:
+            raise RuntimeError("Sirchmunk service is not initialized")
+        
+        # Validate mode
+        if mode not in ("DEEP", "FAST", "FILENAME_ONLY"):
+            raise ValueError(f"Invalid mode: {mode}. Must be DEEP, FAST, or FILENAME_ONLY")
+        
+        # Normalize search_paths
+        if isinstance(search_paths, str):
+            search_paths = [search_paths]
+        
+        # Validate search paths
+        for path in search_paths:
+            path_obj = Path(path)
+            if not path_obj.exists():
+                logger.warning(f"Search path does not exist: {path}")
+        
+        # Apply defaults from configuration
+        max_depth = max_depth or self.config.sirchmunk.search_defaults.max_depth
+        top_k_files = top_k_files or self.config.sirchmunk.search_defaults.top_k_files
+        keyword_levels = keyword_levels or self.config.sirchmunk.search_defaults.keyword_levels
+        
+        logger.info(
+            f"Starting search: mode={mode}, query='{query[:50]}...', "
+            f"paths={len(search_paths)}, max_depth={max_depth}"
+        )
+        
+        try:
+            # Perform search
+            result = await self.searcher.search(
+                query=query,
+                search_paths=search_paths,
+                mode=mode,
+                max_depth=max_depth,
+                top_k_files=top_k_files,
+                keyword_levels=keyword_levels,
+                include=include,
+                exclude=exclude,
+                verbose=self.config.sirchmunk.verbose,
+                grep_timeout=self.config.sirchmunk.search_defaults.grep_timeout,
+                return_cluster=return_cluster,
+            )
+            
+            logger.info(f"Search completed: mode={mode}, result_type={type(result).__name__}")
+            return result
+        
+        except Exception as e:
+            logger.error(f"Search failed: {e}", exc_info=True)
+            raise
+    
+    async def get_cluster(self, cluster_id: str) -> Optional[KnowledgeCluster]:
+        """Retrieve a knowledge cluster by ID.
+        
+        Args:
+            cluster_id: Cluster ID (e.g., 'C1007')
+        
+        Returns:
+            KnowledgeCluster if found, None otherwise
+        
+        Raises:
+            RuntimeError: If service is not initialized
+        """
+        if not self.initialized or self.searcher is None:
+            raise RuntimeError("Sirchmunk service is not initialized")
+        
+        try:
+            cluster = await self.searcher.knowledge_manager.get(cluster_id)
+            if cluster:
+                logger.info(f"Retrieved cluster: {cluster_id}")
+            else:
+                logger.warning(f"Cluster not found: {cluster_id}")
+            return cluster
+        except Exception as e:
+            logger.error(f"Failed to get cluster {cluster_id}: {e}")
+            raise
+    
+    async def list_clusters(
+        self,
+        limit: int = 10,
+        sort_by: str = "last_modified",
+    ) -> List[Dict[str, Any]]:
+        """List saved knowledge clusters with optional filtering.
+        
+        Args:
+            limit: Maximum number of clusters to return
+            sort_by: Sort field (hotness, confidence, last_modified)
+        
+        Returns:
+            List of cluster metadata dictionaries
+        
+        Raises:
+            RuntimeError: If service is not initialized
+        """
+        if not self.initialized or self.searcher is None:
+            raise RuntimeError("Sirchmunk service is not initialized")
+        
+        try:
+            # Get all cluster IDs
+            all_clusters = await self.searcher.knowledge_manager.list_all()
+            
+            # Sort clusters
+            if sort_by == "hotness":
+                all_clusters.sort(key=lambda c: c.hotness or 0.0, reverse=True)
+            elif sort_by == "confidence":
+                all_clusters.sort(key=lambda c: c.confidence or 0.0, reverse=True)
+            else:  # last_modified
+                all_clusters.sort(key=lambda c: c.last_modified, reverse=True)
+            
+            # Limit results
+            result_clusters = all_clusters[:limit]
+            
+            # Convert to dictionaries
+            results = []
+            for cluster in result_clusters:
+                results.append({
+                    "id": cluster.id,
+                    "name": cluster.name,
+                    "confidence": cluster.confidence,
+                    "hotness": cluster.hotness,
+                    "lifecycle": cluster.lifecycle.value,
+                    "version": cluster.version,
+                    "last_modified": cluster.last_modified.isoformat() if cluster.last_modified else None,
+                    "queries": cluster.queries,
+                    "evidences_count": len(cluster.evidences),
+                })
+            
+            logger.info(f"Listed {len(results)} clusters (limit={limit}, sort_by={sort_by})")
+            return results
+        
+        except Exception as e:
+            logger.error(f"Failed to list clusters: {e}")
+            raise
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get service statistics.
+        
+        Returns:
+            Dictionary with service statistics
+        
+        Raises:
+            RuntimeError: If service is not initialized
+        """
+        if not self.initialized or self.searcher is None:
+            raise RuntimeError("Sirchmunk service is not initialized")
+        
+        try:
+            # Get knowledge manager stats
+            stats = self.searcher.knowledge_manager.get_stats()
+            
+            # Add service-level stats
+            stats["service"] = {
+                "initialized": self.initialized,
+                "work_path": str(self.config.sirchmunk.work_path),
+                "cluster_reuse_enabled": self.config.sirchmunk.enable_cluster_reuse,
+            }
+            
+            return stats
+        except Exception as e:
+            logger.error(f"Failed to get stats: {e}")
+            return {"error": str(e)}
+    
+    async def shutdown(self) -> None:
+        """Gracefully shutdown the service.
+        
+        Performs cleanup operations like closing connections and saving state.
+        """
+        logger.info("Shutting down Sirchmunk service")
+        
+        try:
+            # Currently no cleanup needed, but this provides extension point
+            self.initialized = False
+            logger.info("Sirchmunk service shutdown complete")
+        except Exception as e:
+            logger.error(f"Error during shutdown: {e}")

sirchmunk_mcp/setup.py

@@ -0,0 +1,15 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""
+Setup script for Sirchmunk MCP Server.
+
+For backwards compatibility with older pip versions.
+Modern installations should use pyproject.toml.
+"""
+
+from setuptools import setup, find_packages
+
+if __name__ == "__main__":
+    setup(
+        packages=find_packages(),
+        include_package_data=True,
+    )