aigroup-stata-mcp 1.0.3__py3-none-any.whl

stata_mcp/server/prompts/core_prompts.py

@@ -0,0 +1,122 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import os
+from datetime import datetime
+from typing import Dict, Optional
+
+from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.session import ServerSession
+from pydantic import BaseModel, Field
+
+from ...utils.Prompt import pmp
+
+
+class PromptResult(BaseModel):
+   
+   
+    """Result model for prompt operations."""
+    
+    prompt_id: str = Field(description="The identifier of the prompt")
+    content: str = Field(description="The generated prompt content")
+    language: str = Field(description="The language of the prompt")
+    timestamp: str = Field(description="The timestamp when the prompt was generated")
+
+
+def register_core_prompts(server: FastMCP) -> None:
+    """Register core prompts with the MCP server."""
+    
+    @server.prompt()
+    def stata_assistant_role(ctx: Context[ServerSession, Dict], lang: Optional[str] = None) -> PromptResult:
+        """
+        Return the Stata assistant role prompt content.
+        
+        This function retrieves a predefined prompt that defines the role and capabilities
+        of a Stata analysis assistant. The prompt helps set expectations and context for
+        the assistant's behavior when handling Stata-related tasks.
+        
+        Args:
+            lang: Language code for localization of the prompt content.
+                If None, returns the default language version. Defaults to None.
+                Examples: "en" for English, "cn" for Chinese.
+                
+        Returns:
+            PromptResult: Structured result containing prompt content and metadata.
+        """
+        content = pmp.get_prompt(prompt_id="stata_assistant_role", lang=lang)
+        actual_lang = lang or "default"
+        
+        return PromptResult(
+            prompt_id="stata_assistant_role",
+            content=content,
+            language=actual_lang,
+            timestamp=datetime.now().isoformat()
+        )
+
+    @server.prompt()
+    def stata_analysis_strategy(ctx: Context[ServerSession, Dict], lang: Optional[str] = None) -> PromptResult:
+        """
+        Return the Stata analysis strategy prompt content.
+        
+        This function retrieves a predefined prompt that outlines the recommended
+        strategy for conducting data analysis using Stata. The prompt includes
+        guidelines for data preparation, code generation, results management,
+        reporting, and troubleshooting.
+        
+        Args:
+            lang: Language code for localization of the prompt content.
+                If None, returns the default language version. Defaults to None.
+                Examples: "en" for English, "cn" for Chinese.
+                
+        Returns:
+            PromptResult: Structured result containing prompt content and metadata.
+        """
+        content = pmp.get_prompt(prompt_id="stata_analysis_strategy", lang=lang)
+        actual_lang = lang or "default"
+        
+        return PromptResult(
+            prompt_id="stata_analysis_strategy",
+            content=content,
+            language=actual_lang,
+            timestamp=datetime.now().isoformat()
+        )
+
+    @server.prompt()
+    def results_doc_path(ctx: Context[ServerSession, Dict]) -> PromptResult:
+        """
+        Generate and return a result document storage path based on the current timestamp.
+        
+        This function performs the following operations:
+        1. Gets the current system time and formats it as a '%Y%m%d%H%M%S' timestamp string
+        2. Concatenates this timestamp string with the preset result_doc_path base path to form a complete path
+        3. Creates the directory corresponding to that path (no error if directory already exists)
+        4. Returns the complete path string of the newly created directory
+        
+        Returns:
+            PromptResult: Structured result containing the generated path and metadata.
+        """
+        stata_context = ctx.request_context.lifespan_context["stata_context"]
+        result_doc_path = stata_context.output_base_path / "stata-mcp-result"
+        
+        path = result_doc_path / datetime.now().strftime("%Y%m%d%H%M%S")
+        path.mkdir(exist_ok=True)
+        
+        content = f"""
+        The result document path has been created at:
+        {path}
+        
+        You can use this path for Stata commands that generate output files, such as:
+        - outreg2: outreg2 using "{path}/results", replace
+        - esttab: esttab using "{path}/regression_results.rtf", replace
+        - graph export: graph export "{path}/figure.png", replace
+        
+        Make sure to include the appropriate file extensions in your Stata commands.
+        """
+        
+        return PromptResult(
+            prompt_id="results_doc_path",
+            content=content,
+            language="en",
+            timestamp=datetime.now().isoformat()
+        )

stata_mcp/server/tools/__init__.py

@@ -0,0 +1,10 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+"""Stata MCP Tools package."""
+from .core_tools import register_core_tools
+from .file_tools import register_file_tools
+from .stata_tools import register_stata_tools
+
+__all__ = ["register_core_tools", "register_file_tools", "register_stata_tools"]

stata_mcp/server/tools/core_tools.py

@@ -0,0 +1,59 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.session import ServerSession
+from pydantic import BaseModel, Field
+
+from ...core.stata import StataController
+
+
+def register_core_tools(server: FastMCP) -> None:
+    """Register core Stata tools with the MCP server."""
+    
+    @server.tool()
+    def mk_dir(ctx: Context[ServerSession, Dict], path: str) -> bool:
+        """
+        Safely create a directory using pathvalidate for security validation.
+        
+        Args:
+            path: The path you want to create
+            
+        Returns:
+            bool: True if the path exists now, False if not successful
+            
+        Raises:
+            ValueError: if path is invalid or contains unsafe components
+            PermissionError: if insufficient permissions to create directory
+        """
+        from pathvalidate import ValidationError, sanitize_filepath
+
+        # Input validation
+        if not path or not isinstance(path, str):
+            raise ValueError("Path must be a non-empty string")
+
+        try:
+            # Use pathvalidate to sanitize and validate path
+            safe_path = sanitize_filepath(path, platform="auto")
+
+            # Get absolute path for further validation
+            absolute_path = os.path.abspath(safe_path)
+
+            # Create directory with reasonable permissions
+            os.makedirs(absolute_path, exist_ok=True, mode=0o755)
+
+            # Verify successful creation
+            return os.path.exists(absolute_path) and os.path.isdir(absolute_path)
+
+        except ValidationError as e:
+            raise ValueError(f"Invalid path detected: {e}")
+        except PermissionError:
+            raise PermissionError(f"Insufficient permissions to create directory: {path}")
+        except OSError as e:
+            raise OSError(f"Failed to create directory {path}: {str(e)}")

stata_mcp/server/tools/file_tools.py

@@ -0,0 +1,163 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from mcp.server.fastmcp import Context, FastMCP, Image
+from mcp.server.session import ServerSession
+from pydantic import BaseModel, Field
+
+
+class ReadFileResult(BaseModel):
+   
+   
+    """Result model for file reading operation."""
+    
+    file_path: str = Field(description="The path to the file that was read")
+    content: str = Field(description="The content of the file")
+    encoding: str = Field(description="The encoding used to read the file")
+
+
+class WriteDofileResult(BaseModel):
+   
+   
+    """Result model for dofile writing operation."""
+    
+    file_path: str = Field(description="The path to the created dofile")
+    content_length: int = Field(description="The length of the content written")
+    timestamp: str = Field(description="The timestamp when the file was created")
+
+
+class AppendDofileResult(BaseModel):
+   
+   
+    """Result model for dofile appending operation."""
+    
+    new_file_path: str = Field(description="The path to the new dofile")
+    original_exists: bool = Field(description="Whether the original file existed")
+    total_content_length: int = Field(description="Total length of content after appending")
+
+
+def register_file_tools(server: FastMCP) -> None:
+    """Register file-related tools with the MCP server."""
+    
+    @server.tool()
+    def read_file(ctx: Context[ServerSession, Dict], file_path: str, encoding: str = "utf-8") -> ReadFileResult:
+        """
+        Reads a file and returns its content as a string.
+        
+        Args:
+            file_path: The full path to the file to be read.
+            encoding: The encoding used to decode the file. Defaults to "utf-8".
+            
+        Returns:
+            ReadFileResult: Structured result containing file content and metadata.
+        """
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"The file at {file_path} does not exist.")
+
+        try:
+            with open(file_path, "r", encoding=encoding) as file:
+                log_content = file.read()
+            
+            return ReadFileResult(
+                file_path=file_path,
+                content=log_content,
+                encoding=encoding
+            )
+        except IOError as e:
+            raise IOError(f"An error occurred while reading the file: {e}")
+
+    @server.tool()
+    def write_dofile(ctx: Context[ServerSession, Dict], content: str, encoding: str = "utf-8") -> WriteDofileResult:
+        """
+        Write stata code to a dofile and return the do-file path.
+        
+        Args:
+            content: The stata code content which will be written to the designated do-file.
+            encoding: The encoding method for the dofile, default -> 'utf-8'
+            
+        Returns:
+            WriteDofileResult: Structured result containing file path and metadata.
+        """
+        stata_context = ctx.request_context.lifespan_context["stata_context"]
+        dofile_base_path = stata_context.output_base_path / "stata-mcp-dofile"
+        
+        file_path = dofile_base_path / f"{datetime.now().strftime('%Y%m%d%H%M%S')}.do"
+        
+        with open(file_path, "w", encoding=encoding) as f:
+            f.write(content)
+        
+        return WriteDofileResult(
+            file_path=str(file_path),
+            content_length=len(content),
+            timestamp=datetime.now().isoformat()
+        )
+
+    @server.tool()
+    def append_dofile(ctx: Context[ServerSession, Dict], original_dofile_path: str, content: str, encoding: str = "utf-8") -> AppendDofileResult:
+        """
+        Append stata code to an existing dofile or create a new one.
+        
+        Args:
+            original_dofile_path: Path to the original dofile to append to.
+                If empty or invalid, a new file will be created.
+            content: The stata code content which will be appended to the designated do-file.
+            encoding: The encoding method for the dofile, default -> 'utf-8'
+            
+        Returns:
+            AppendDofileResult: Structured result containing new file path and metadata.
+        """
+        stata_context = ctx.request_context.lifespan_context["stata_context"]
+        dofile_base_path = stata_context.output_base_path / "stata-mcp-dofile"
+        
+        # Create a new file path for the output
+        new_file_path = dofile_base_path / f"{datetime.now().strftime('%Y%m%d%H%M%S')}.do"
+        
+        # Check if original file exists and is valid
+        original_exists = False
+        original_content = ""
+        if original_dofile_path and os.path.exists(original_dofile_path):
+            try:
+                with open(original_dofile_path, "r", encoding=encoding) as f:
+                    original_content = f.read()
+                original_exists = True
+            except Exception:
+                # If there's any error reading the file, we'll create a new one
+                original_exists = False
+
+        # Write to the new file (either copying original content + new content, or just new content)
+        with open(new_file_path, "w", encoding=encoding) as f:
+            if original_exists:
+                f.write(original_content)
+                # Add a newline if the original file doesn't end with one
+                if original_content and not original_content.endswith("\n"):
+                    f.write("\n")
+            f.write(content)
+        
+        total_length = len(original_content) + len(content) if original_exists else len(content)
+        
+        return AppendDofileResult(
+            new_file_path=str(new_file_path),
+            original_exists=original_exists,
+            total_content_length=total_length
+        )
+
+    @server.tool()
+    def load_figure(ctx: Context[ServerSession, Dict], figure_path: str) -> Image:
+        """
+        Load figure from device.
+        
+        Args:
+            figure_path: the figure file path, only support png and jpg format
+            
+        Returns:
+            Image: the figure thumbnail
+        """
+        if not os.path.exists(figure_path):
+            raise FileNotFoundError(f"{figure_path} not found")
+        return Image(figure_path)

stata_mcp/server/tools/stata_tools.py

@@ -0,0 +1,221 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import os
+from datetime import datetime
+from pathlib import Path
+from typing import Callable, Dict, List, Optional
+
+from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.session import ServerSession
+from pydantic import BaseModel, Field
+
+from ...core.data_info import CsvDataInfo, DtaDataInfo
+from ...core.stata import StataDo
+
+
+class DataInfoResult(BaseModel):
+   
+   
+    """Result model for data information operation."""
+    
+    data_path: str = Field(description="The path to the data file")
+    file_extension: str = Field(description="The file extension of the data file")
+    summary: Dict[str, Dict] = Field(description="Summary statistics of the data")
+    info_file_path: Optional[str] = Field(description="Path to the saved info file, if applicable")
+
+
+class StataDoResult(BaseModel):
+   
+   
+    """Result model for Stata dofile execution."""
+    
+    dofile_path: str = Field(description="Path to the executed dofile")
+    log_file_path: str = Field(description="Path to the generated log file")
+    log_content: Optional[str] = Field(description="Content of the log file, if requested")
+    execution_time: str = Field(description="Timestamp of execution")
+
+
+class SscInstallResult(BaseModel):
+   
+   
+    """Result model for SSC package installation."""
+    
+    package: str = Field(description="The package that was installed")
+    success: bool = Field(description="Whether installation was successful")
+    log_content: str = Field(description="Installation log content")
+
+
+def register_stata_tools(server: FastMCP) -> None:
+    """Register Stata-specific tools with the MCP server."""
+    
+    @server.tool()
+    def get_data_info(
+        ctx: Context[ServerSession, Dict],
+        data_path: str,
+        vars_list: Optional[List[str]] = None,
+        encoding: str = "utf-8",
+        file_extension: Optional[str] = None,
+        is_save: bool = True,
+        save_path: Optional[str] = None,
+        info_file_encoding: str = "utf-8"
+    ) -> DataInfoResult:
+        """
+        Get descriptive statistics for the data file.
+        
+        Args:
+            data_path: the data file's absolute path.
+            vars_list: the vars you want to get info (default is None, means all vars).
+            encoding: data file encoding method (dta file is not supported this arg).
+            file_extension: the data file's extension, default is None, then would find it automatically.
+            is_save: whether save the result to a txt file.
+            save_path: the data-info saved file path.
+            info_file_encoding: the data-info saved file encoding.
+            
+        Returns:
+            DataInfoResult: Structured result containing data information.
+        """
+        stata_context = ctx.request_context.lifespan_context["stata_context"]
+        tmp_base_path = stata_context.output_base_path / "stata-mcp-tmp"
+        
+        EXTENSION_METHOD_MAPPING: Dict[str, Callable] = {
+            "dta": DtaDataInfo,
+            "csv": CsvDataInfo
+        }
+        
+        if file_extension is None:
+            file_extension = Path(data_path).suffix
+        file_extension = file_extension.split(".")[-1].lower()
+        
+        if file_extension not in EXTENSION_METHOD_MAPPING:
+            raise ValueError(f"Unsupported file extension: {file_extension}")
+        
+        cls = EXTENSION_METHOD_MAPPING.get(file_extension)
+        data_info = cls(
+            data_path=data_path,
+            vars_list=vars_list,
+            encoding=encoding,
+            is_save=is_save,
+            save_path=save_path,
+            info_file_encoding=info_file_encoding
+        ).info
+
+        info_file_path = None
+        if is_save:
+            if save_path is None:
+                data_name = Path(data_path).name.split(".")[0]
+                info_file_path = tmp_base_path / f"{data_name}.txt"
+            else:
+                info_file_path = Path(save_path)
+
+            info_file_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(info_file_path, "w", encoding=info_file_encoding) as f:
+                f.write(str(data_info))
+
+        return DataInfoResult(
+            data_path=data_path,
+            file_extension=file_extension,
+            summary=data_info,
+            info_file_path=str(info_file_path) if info_file_path else None
+        )
+
+    @server.tool()
+    def stata_do(
+        ctx: Context[ServerSession, Dict],
+        dofile_path: str,
+        log_file_name: Optional[str] = None,
+        is_read_log: bool = True
+    ) -> StataDoResult:
+        """
+        Execute a Stata do-file and return the log file path with optional log content.
+        
+        Args:
+            dofile_path: Absolute or relative path to the Stata do-file (.do) to execute.
+            log_file_name: Set log file name without a time-string. If None, using nowtime as filename.
+            is_read_log: Whether to read and return the log file content.
+            
+        Returns:
+            StataDoResult: Structured result containing execution details.
+        """
+        stata_context = ctx.request_context.lifespan_context["stata_context"]
+        stata_cli = stata_context.stata_finder.STATA_CLI
+        log_base_path = stata_context.output_base_path / "stata-mcp-log"
+        dofile_base_path = stata_context.output_base_path / "stata-mcp-dofile"
+        
+        # Initialize Stata executor with system configuration
+        stata_executor = StataDo(
+            stata_cli=stata_cli,
+            log_file_path=str(log_base_path),
+            dofile_base_path=str(dofile_base_path),
+            sys_os=platform.system()
+        )
+
+        # Execute the do-file and get log file path
+        log_file_path = stata_executor.execute_dofile(dofile_path, log_file_name)
+
+        # Return log content based on user preference
+        log_content = None
+        if is_read_log:
+            log_content = stata_executor.read_log(log_file_path)
+
+        return StataDoResult(
+            dofile_path=dofile_path,
+            log_file_path=log_file_path,
+            log_content=log_content,
+            execution_time=datetime.now().isoformat()
+        )
+
+    @server.tool()
+    def ssc_install(
+        ctx: Context[ServerSession, Dict],
+        command: str,
+        is_replace: bool = True
+    ) -> SscInstallResult:
+        """
+        Install a package from SSC.
+        
+        Args:
+            command: The name of the package to be installed from SSC.
+            is_replace: Whether to force replacement of an existing installation.
+            
+        Returns:
+            SscInstallResult: Structured result containing installation details.
+        """
+        stata_context = ctx.request_context.lifespan_context["stata_context"]
+        stata_cli = stata_context.stata_finder.STATA_CLI
+        log_base_path = stata_context.output_base_path / "stata-mcp-log"
+        dofile_base_path = stata_context.output_base_path / "stata-mcp-dofile"
+        
+        replace_clause = ", replace" if is_replace else ""
+        
+        # Create dofile content
+        content = f"ssc install {command}{replace_clause}"
+        file_path = dofile_base_path / f"{datetime.now().strftime('%Y%m%d%H%M%S')}.do"
+        
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(content)
+        
+        # Execute the dofile using StataDo directly
+        stata_executor = StataDo(
+            stata_cli=stata_cli,
+            log_file_path=str(log_base_path),
+            dofile_base_path=str(dofile_base_path),
+            sys_os=platform.system()
+        )
+        
+        log_file_path = stata_executor.execute_dofile(str(file_path))
+        log_content = stata_executor.read_log(log_file_path)
+        
+        # Check for success
+        success = "not found" not in log_content.lower() if log_content else False
+        
+        return SscInstallResult(
+            package=command,
+            success=success,
+            log_content=log_content or ""
+        )
+
+
+# Import platform for system detection
+import platform

stata_mcp/utils/Installer/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+from .installer import Installer
+
+__all__ = ["Installer"]

stata_mcp/utils/Installer/installer.py

@@ -0,0 +1,85 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import json
+import os
+
+from ...core.stata import StataFinder
+
+
+class Installer:
+   
+   
+   
+    def __init__(self, sys_os, is_env=False):
+        """简要描述函数功能"""
+        self.config_file_path: str = None
+        if sys_os == "Darwin":
+            self.config_file_path = os.path.expanduser(
+                "~/Library/Application Support/Claude/claude_desktop_config.json"
+            )
+        elif sys_os == "Linux":
+            print(
+                "There is not a Linux version of Claude yet, please use the Windows or macOS version."
+            )
+        elif sys_os == "Windows":
+            appdata = os.getenv(
+                "APPDATA", os.path.expanduser("~\\AppData\\Roaming"))
+            self.config_file_path = os.path.join(
+                appdata, "Claude", "claude_desktop_config.json"
+            )
+
+        os.makedirs(os.path.dirname(self.config_file_path), exist_ok=True)
+
+        # Create an empty file if it does not already exist
+        if not os.path.exists(self.config_file_path):
+            with open(self.config_file_path, "w", encoding="utf-8") as f:
+                # Or write the default configuration
+                f.write('{"mcpServers": {}}')
+
+        stata_cli = StataFinder().STATA_CLI
+        self.stata_mcp_config = {
+            "stata-mcp": {
+                "command": "uvx",
+                "args": ["aigroup-stata-mcp"],
+                "env": {"STATA_CLI": stata_cli},
+            }
+        }
+
+    def install(self):
+        """简要描述函数功能"""
+        server_cfg = self.stata_mcp_config["stata-mcp"]
+        stata_cli_path = server_cfg["env"]["STATA_CLI"]
+        print("About to install the following MCP server into your Claude config:\n")
+        print("  Server name:    stata-mcp")
+        print(f"  Command:        {server_cfg['command']}")
+        print(f"  Args:           {server_cfg['args']}")
+        print(f"  STATA_CLI path: {stata_cli_path}\n")
+        print(f"Configuration file to modify:\n  {self.config_file_path}\n")
+
+        # Ask the user for confirmation
+        choice = input(
+            "Do you want to proceed and add this configuration? [y/N]: ")
+        if choice.strip().lower() != "y":
+            print("Installation aborted.")
+            return
+
+        # Read the now config
+        try:
+            with open(self.config_file_path, "r", encoding="utf-8") as f:
+                config = json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            config = {"mcpServers": {}}
+
+        # Update MCP_Config
+        servers = config.setdefault("mcpServers", {})
+        servers.update(self.stata_mcp_config)
+
+        # Write it
+        with open(self.config_file_path, "w", encoding="utf-8") as f:
+            json.dump(config, f, ensure_ascii=False, indent=2)
+
+        print(
+            f"✅ Successfully wrote 'stata-mcp' configuration to: {self.config_file_path}"
+        )