mcp-stata 1.18.0__cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

mcp_stata/graph_detector.py

@@ -0,0 +1,401 @@
+"""
+Graph creation detection for streaming Stata output.
+
+This module provides functionality to detect when graphs are created
+during Stata command execution and automatically cache them.
+"""
+
+import asyncio
+import contextlib
+import inspect
+import re
+import threading
+import time
+from typing import List, Set, Callable, Dict, Any
+import logging
+
+
+# SFI is always available
+SFI_AVAILABLE = True
+
+logger = logging.getLogger(__name__)
+
+
+class GraphCreationDetector:
+    """Detects graph creation using SFI-only detection with pystata integration."""
+    
+    def __init__(self, stata_client=None):
+        self._lock = threading.Lock()
+        self._detected_graphs: Set[str] = set()
+        self._removed_graphs: Set[str] = set()
+        self._unnamed_graph_counter = 0  # Track unnamed graphs for identification
+        self._stata_client = stata_client
+        self._last_graph_state: Dict[str, Any] = {}  # Track graph state changes
+
+    def _describe_graph_signature(self, graph_name: str) -> str:
+        """Return a stable signature for a graph.
+
+        We avoid using Stata calls like 'graph describe' here because they are slow 
+        (each call takes ~35ms) and would be called for every graph on every poll,
+        bottlenecking the streaming output.
+        
+        Instead, we use name-based tracking tied to the Stata command execution 
+        context. The signature is stable within a single command execution but 
+        changes when a new command starts, allowing us to detect modifications 
+        between commands without any Stata overhead.
+        """
+        if not self._stata_client:
+            return ""
+        
+        # Access command_idx from stata_client if available
+        # NOTE: We only use command_idx for the default 'Graph' name to detect
+        # modifications. For named graphs, we only detect creation (name change)
+        # to avoid triggering redundant notifications for all existing graphs 
+        # on every command (since command_idx changes globally).
+        cmd_idx = getattr(self._stata_client, "_command_idx", 0)
+        if graph_name.lower() == "graph":
+            return f"{graph_name}_{cmd_idx}"
+        return graph_name
+    
+    def _detect_graphs_via_pystata(self) -> List[str]:
+        """Detect newly created graphs using direct pystata state access."""
+        if not self._stata_client:
+            return []
+        
+        try:
+            # Get current graph state using pystata's sfi interface
+            current_graphs = self._get_current_graphs_from_pystata()
+            current_state = self._get_graph_state_from_pystata()
+            
+            # Compare with last known state to detect new graphs
+            new_graphs = []
+            
+            # Check for new graph names
+            for graph_name in current_graphs:
+                if graph_name not in self._last_graph_state and graph_name not in self._removed_graphs:
+                    new_graphs.append(graph_name)
+            
+            # Check for state changes in existing graphs (modifications)
+            for graph_name, state in current_state.items():
+                if graph_name in self._last_graph_state:
+                    last_state = self._last_graph_state[graph_name]
+                    # Compare stable signature only.
+                    if state.get("signature") != last_state.get("signature"):
+                        if graph_name not in self._removed_graphs:
+                            new_graphs.append(graph_name)
+            
+            # Update cached state
+            self._last_graph_state = current_state.copy()
+            
+            return new_graphs
+            
+        except (ImportError, RuntimeError, ValueError, AttributeError) as e:
+            # These are expected exceptions when SFI is not available or Stata state is inaccessible
+            logger.debug(f"Failed to detect graphs via pystata (expected): {e}")
+            return []
+        except Exception as e:
+            # Unexpected errors should be logged as errors
+            logger.error(f"Unexpected error in pystata graph detection: {e}")
+            return []
+    
+    def _get_current_graphs_from_pystata(self) -> List[str]:
+        """Get current list of graphs using pystata's sfi interface."""
+        try:
+            # Use pystata to get graph list directly
+            if self._stata_client and hasattr(self._stata_client, 'list_graphs'):
+                return self._stata_client.list_graphs(force_refresh=True)
+            else:
+                # Fallback to sfi Macro interface - only if stata is available
+                if self._stata_client and hasattr(self._stata_client, 'stata'):
+                    # Access the lock from client to prevent concurrency issues with pystata
+                    exec_lock = getattr(self._stata_client, "_exec_lock", None)
+                    ctx = exec_lock if exec_lock else contextlib.nullcontext()
+                    
+                    with ctx:
+                        try:
+                            from sfi import Macro
+                            hold_name = f"_mcp_detector_hold_{int(time.time() * 1000 % 1000000)}"
+                            self._stata_client.stata.run(f"capture _return hold {hold_name}", echo=False)
+                            try:
+                                self._stata_client.stata.run("macro define mcp_graph_list \"\"", echo=False)
+                                self._stata_client.stata.run("quietly graph dir, memory", echo=False)
+                                self._stata_client.stata.run("macro define mcp_graph_list `r(list)'", echo=False)
+                                graph_list_str = Macro.getGlobal("mcp_graph_list")
+                            finally:
+                                self._stata_client.stata.run(f"capture _return restore {hold_name}", echo=False)
+                            return graph_list_str.split() if graph_list_str else []
+                        except ImportError:
+                            logger.warning("sfi.Macro not available for fallback graph detection")
+                            return []
+                else:
+                    return []
+        except Exception as e:
+            logger.warning(f"Failed to get current graphs from pystata: {e}")
+            return []
+    
+    def _get_graph_state_from_pystata(self) -> Dict[str, Any]:
+        """Get detailed graph state information using pystata's sfi interface."""
+        graph_state = {}
+        
+        try:
+            current_graphs = self._get_current_graphs_from_pystata()
+            
+            for graph_name in current_graphs:
+                try:
+                    signature = self._describe_graph_signature(graph_name)
+                    state_info = {
+                        "name": graph_name,
+                        "exists": True,
+                        "valid": bool(signature),
+                        "signature": signature,
+                    }
+
+                    # Only update timestamps when the signature changes.
+                    prev = self._last_graph_state.get(graph_name)
+                    if prev is None or prev.get("signature") != signature:
+                        state_info["timestamp"] = time.time()
+                    else:
+                        state_info["timestamp"] = prev.get("timestamp", time.time())
+                    
+                    graph_state[graph_name] = state_info
+                    
+                except Exception as e:
+                    logger.warning(f"Failed to get state for graph {graph_name}: {e}")
+                    graph_state[graph_name] = {"name": graph_name, "timestamp": time.time(), "exists": False, "signature": ""}
+            
+        except Exception as e:
+            logger.warning(f"Failed to get graph state from pystata: {e}")
+        
+        return graph_state
+    
+        
+        
+    def detect_graph_modifications(self, text: str = None) -> dict:
+        """Detect graph modification/removal using SFI state comparison."""
+        modifications = {"dropped": [], "renamed": [], "cleared": False}
+        
+        if not self._stata_client:
+            return modifications
+        
+        try:
+            # Get current graph state via SFI
+            current_graphs = set(self._get_current_graphs_from_pystata())
+            
+            # Compare with last known state to detect modifications
+            if self._last_graph_state:
+                last_graphs = set(self._last_graph_state.keys())
+                
+                # Detect dropped graphs (in last state but not current)
+                dropped_graphs = last_graphs - current_graphs
+                modifications["dropped"].extend(dropped_graphs)
+                
+                # Detect clear all (no graphs remain when there were some before)
+                if last_graphs and not current_graphs:
+                    modifications["cleared"] = True
+            
+            # Update last known state for next comparison (stable signatures)
+            new_state: Dict[str, Any] = {}
+            for graph in current_graphs:
+                sig = self._describe_graph_signature(graph)
+                new_state[graph] = {
+                    "name": graph,
+                    "exists": True,
+                    "valid": bool(sig),
+                    "signature": sig,
+                    "timestamp": time.time(),
+                }
+            self._last_graph_state = new_state
+            
+        except Exception as e:
+            logger.debug(f"SFI modification detection failed: {e}")
+        
+        return modifications
+    
+        
+    def should_cache_graph(self, graph_name: str) -> bool:
+        """Determine if a graph should be cached."""
+        with self._lock:
+            # Don't cache if already detected or removed
+            if graph_name in self._detected_graphs or graph_name in self._removed_graphs:
+                return False
+            
+            # Mark as detected
+            self._detected_graphs.add(graph_name)
+            return True
+    
+    def mark_graph_removed(self, graph_name: str) -> None:
+        """Mark a graph as removed."""
+        with self._lock:
+            self._removed_graphs.add(graph_name)
+            self._detected_graphs.discard(graph_name)
+    
+    def mark_all_cleared(self) -> None:
+        """Mark all graphs as cleared."""
+        with self._lock:
+            self._detected_graphs.clear()
+            self._removed_graphs.clear()
+    
+    def clear_detection_state(self) -> None:
+        """Clear all detection state."""
+        with self._lock:
+            self._detected_graphs.clear()
+            self._removed_graphs.clear()
+            self._unnamed_graph_counter = 0
+    
+    def process_modifications(self, modifications: dict) -> None:
+        """Process detected modifications."""
+        with self._lock:
+            # Handle dropped graphs
+            for graph_name in modifications.get("dropped", []):
+                self.mark_graph_removed(graph_name)
+            
+            # Handle renamed graphs
+            for old_name, new_name in modifications.get("renamed", []):
+                self.mark_graph_removed(old_name)
+                self._detected_graphs.discard(new_name)  # Allow re-detection with new name
+            
+            # Handle clear all
+            if modifications.get("cleared", False):
+                self.mark_all_cleared()
+
+
+class StreamingGraphCache:
+    """Integrates graph detection with caching during streaming."""
+    
+    def __init__(self, stata_client, auto_cache: bool = False):
+        self.stata_client = stata_client
+        self.auto_cache = auto_cache
+        # Use persistent detector from client if available, else create local one
+        if hasattr(stata_client, "_graph_detector"):
+            self.detector = stata_client._graph_detector
+        else:
+            self.detector = GraphCreationDetector(stata_client)
+        self._lock = threading.Lock()
+        self._cache_callbacks: List[Callable[[str, bool], None]] = []
+        self._graphs_to_cache: List[str] = []
+        self._cached_graphs: Set[str] = set()
+        self._removed_graphs = set()  # Track removed graphs directly
+        self._initial_graphs: Set[str] = set()  # Captured before execution starts
+    
+    def add_cache_callback(self, callback: Callable[[str, bool], None]) -> None:
+        """Add callback for graph cache events."""
+        with self._lock:
+            self._cache_callbacks.append(callback)
+
+    async def _notify_cache_callbacks(self, graph_name: str, success: bool) -> None:
+        for callback in self._cache_callbacks:
+            try:
+                result = callback(graph_name, success)
+                if inspect.isawaitable(result):
+                    await result
+            except Exception as e:
+                logger.warning(f"Cache callback failed for {graph_name}: {e}")
+
+    
+    async def cache_detected_graphs_with_pystata(self) -> List[str]:
+        """Enhanced caching method that uses pystata for real-time graph detection."""
+        if not self.auto_cache:
+            return []
+        
+        cached_names = []
+        
+        # First, try to get any newly detected graphs via pystata state
+        if self.stata_client:
+            try:
+                # Get current state and check for new graphs
+
+                pystata_detected = self.detector._detect_graphs_via_pystata()
+                
+                # Add any newly detected graphs to cache queue
+                for graph_name in pystata_detected:
+                    if graph_name not in self._cached_graphs and graph_name not in self._removed_graphs:
+                        self._graphs_to_cache.append(graph_name)
+                        
+            except Exception as e:
+                logger.warning(f"Failed to get pystata graph updates: {e}")
+        
+        # Process the cache queue
+        with self._lock:
+            graphs_to_process = self._graphs_to_cache.copy()
+            self._graphs_to_cache.clear()
+        
+        # Get current graph list for verification
+        try:
+            current_graphs = self.stata_client.list_graphs()
+        except Exception as e:
+            logger.warning(f"Failed to get current graph list: {e}")
+            return cached_names
+        
+        for graph_name in graphs_to_process:
+            if graph_name in current_graphs and graph_name not in self._cached_graphs:
+                try:
+                    success = await asyncio.to_thread(self.stata_client.cache_graph_on_creation, graph_name)
+                    if success:
+                        cached_names.append(graph_name)
+                        with self._lock:
+                            self._cached_graphs.add(graph_name)
+                    
+                    # Notify callbacks
+                    await self._notify_cache_callbacks(graph_name, success)
+                
+                except Exception as e:
+                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
+                    # Still notify callbacks of failure
+                    await self._notify_cache_callbacks(graph_name, False)
+        
+        return cached_names
+    
+    async def cache_detected_graphs(self) -> List[str]:
+        """Cache all detected graphs."""
+        if not self.auto_cache:
+            return []
+        
+        cached_names = []
+        
+        with self._lock:
+            graphs_to_process = self._graphs_to_cache.copy()
+            self._graphs_to_cache.clear()
+        
+        # Get current graph list for verification
+        try:
+            current_graphs = self.stata_client.list_graphs()
+        except Exception as e:
+            logger.warning(f"Failed to get current graph list: {e}")
+            return cached_names
+        
+        for graph_name in graphs_to_process:
+            if graph_name in current_graphs and graph_name not in self._cached_graphs:
+                try:
+                    success = await asyncio.to_thread(self.stata_client.cache_graph_on_creation, graph_name)
+                    if success:
+                        cached_names.append(graph_name)
+                        with self._lock:
+                            self._cached_graphs.add(graph_name)
+                    
+                    # Notify callbacks
+                    await self._notify_cache_callbacks(graph_name, success)
+                
+                except Exception as e:
+                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
+                    # Still notify callbacks of failure
+                    await self._notify_cache_callbacks(graph_name, False)
+        
+        return cached_names
+    
+    def get_cache_stats(self) -> dict:
+        """Get caching statistics."""
+        with self._lock:
+            return {
+                "auto_cache_enabled": self.auto_cache,
+                "pending_cache_count": len(self._graphs_to_cache),
+                "cached_graphs_count": len(self._cached_graphs),
+                "detected_graphs_count": len(self.detector._detected_graphs),
+                "removed_graphs_count": len(self.detector._removed_graphs),
+            }
+    
+    def reset(self) -> None:
+        """Reset the cache state."""
+        with self._lock:
+            self._graphs_to_cache.clear()
+            self._cached_graphs.clear()
+        self.detector.clear_detection_state()

mcp_stata/models.py

@@ -0,0 +1,62 @@
+from typing import List, Optional, Dict, Any
+from pydantic import BaseModel
+
+
+class ErrorEnvelope(BaseModel):
+    message: str
+    rc: Optional[int] = None
+    line: Optional[int] = None
+    command: Optional[str] = None
+    log_path: Optional[str] = None
+    context: Optional[str] = None
+    stdout: Optional[str] = None
+    stderr: Optional[str] = None
+    snippet: Optional[str] = None
+    trace: Optional[bool] = None
+    smcl_output: Optional[str] = None
+
+
+class CommandResponse(BaseModel):
+    command: str
+    rc: int
+    stdout: str
+    stderr: Optional[str] = None
+    log_path: Optional[str] = None
+    success: bool
+    error: Optional[ErrorEnvelope] = None
+    smcl_output: Optional[str] = None
+
+
+class DataResponse(BaseModel):
+    start: int
+    count: int
+    data: List[Dict[str, Any]]
+
+
+class VariableInfo(BaseModel):
+    name: str
+    label: Optional[str] = None
+    type: Optional[str] = None
+
+
+class VariablesResponse(BaseModel):
+    variables: List[VariableInfo]
+
+
+class GraphInfo(BaseModel):
+    name: str
+    active: bool = False
+
+
+class GraphListResponse(BaseModel):
+    graphs: List[GraphInfo]
+
+
+class GraphExport(BaseModel):
+    name: str
+    file_path: Optional[str] = None
+
+
+class GraphExportResponse(BaseModel):
+    graphs: List[GraphExport]
+

mcp_stata/native_ops.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import logging
+from typing import Iterable, Any, Tuple
+
+logger = logging.getLogger(__name__)
+
+try:
+    from mcp_stata import _native_ops as _native
+except Exception:  # pragma: no cover - optional module
+    _native = None
+
+
+def argsort_numeric(
+    columns: Iterable["numpy.ndarray"],
+    descending: list[bool],
+    nulls_last: list[bool],
+) -> list[int] | None:
+    if _native is None:
+        return None
+    cols = list(columns)
+    if not cols:
+        return []
+    try:
+        return _native.argsort_numeric(cols, descending, nulls_last)
+    except Exception as e:
+        logger.warning(f"Native numeric sort failed: {e}")
+        return None
+
+
+def argsort_mixed(
+    columns: Iterable[object],
+    is_string: list[bool],
+    descending: list[bool],
+    nulls_last: list[bool],
+) -> list[int] | None:
+    if _native is None:
+        return None
+    cols = list(columns)
+    if not cols:
+        return []
+    try:
+        return _native.argsort_mixed(cols, is_string, descending, nulls_last)
+    except Exception as e:
+        logger.warning(f"Native mixed sort failed: {e}")
+        return None
+
+
+def smcl_to_markdown(smcl_text: str) -> str | None:
+    if _native is None:
+        return None
+    try:
+        return _native.smcl_to_markdown(smcl_text)
+    except Exception as e:
+        logger.warning(f"Native SMCL conversion failed: {e}")
+        return None
+
+
+def fast_scan_log(smcl_content: str, rc_default: int) -> Tuple[str, str, int | None] | None:
+    if _native is None:
+        return None
+    try:
+        return _native.fast_scan_log(smcl_content, rc_default)
+    except Exception as e:
+        logger.warning(f"Native log scanning failed: {e}")
+        return None
+
+
+def compute_filter_indices(
+    data_numeric: dict[str, "numpy.ndarray"],
+    data_string: dict[str, list[str]],
+    filter_json: str,
+    row_count: int,
+) -> list[int] | None:
+    if _native is None:
+        return None
+    try:
+        return _native.compute_filter_indices(
+            data_numeric, 
+            data_string, 
+            filter_json, 
+            row_count
+        )
+    except Exception as e:
+        logger.warning(f"Native filtering failed: {e}")
+        return None
+