mcp-stata 1.2.2__py3-none-any.whl → 1.6.8__py3-none-any.whl

mcp_stata/graph_detector.py

@@ -0,0 +1,385 @@
+"""
+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 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 intentionally avoid using timestamps as the signature, since that makes
+        every poll look like a modification.
+        """
+        if not self._stata_client or not hasattr(self._stata_client, "stata"):
+            return ""
+        try:
+            # Capture output so we can hash it deterministically.
+            resp = self._stata_client.run_command_structured(f"graph describe {graph_name}", echo=False)
+            if resp.success and resp.stdout:
+                return resp.stdout
+            if resp.error and resp.error.snippet:
+                return resp.error.snippet
+        except Exception:
+            return ""
+        return ""
+    
+    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()
+            else:
+                # Fallback to sfi Macro interface - only if stata is available
+                if self._stata_client and hasattr(self._stata_client, 'stata'):
+                    try:
+                        from sfi import Macro
+                        self._stata_client.stata.run("quietly graph dir, memory")
+                        self._stata_client.stata.run("global mcp_graph_list `r(list)'")
+                        graph_list_str = Macro.getGlobal("mcp_graph_list")
+                        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
+        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 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
+                    for callback in self._cache_callbacks:
+                        try:
+                            callback(graph_name, success)
+                        except Exception as e:
+                            logger.warning(f"Cache callback failed for {graph_name}: {e}")
+                
+                except Exception as e:
+                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
+                    # Still notify callbacks of failure
+                    for callback in self._cache_callbacks:
+                        try:
+                            callback(graph_name, False)
+                        except Exception:
+                            pass
+        
+        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
+                    for callback in self._cache_callbacks:
+                        try:
+                            callback(graph_name, success)
+                        except Exception as e:
+                            logger.warning(f"Cache callback failed for {graph_name}: {e}")
+                
+                except Exception as e:
+                    logger.warning(f"Failed to cache graph {graph_name}: {e}")
+                    # Still notify callbacks of failure
+                    for callback in self._cache_callbacks:
+                        try:
+                            callback(graph_name, False)
+                        except Exception:
+                            pass
+        
+        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

@@ -7,6 +7,7 @@ class ErrorEnvelope(BaseModel):
     rc: Optional[int] = None
     line: Optional[int] = None
     command: Optional[str] = None
+    log_path: Optional[str] = None
     stdout: Optional[str] = None
     stderr: Optional[str] = None
     snippet: Optional[str] = None
@@ -18,6 +19,7 @@ class CommandResponse(BaseModel):
     rc: int
     stdout: str
     stderr: Optional[str] = None
+    log_path: Optional[str] = None
     success: bool
     error: Optional[ErrorEnvelope] = None
 
@@ -49,7 +51,8 @@ class GraphListResponse(BaseModel):
 
 class GraphExport(BaseModel):
     name: str
-    image_base64: str
+    file_path: Optional[str] = None
+    image_base64: Optional[str] = None
 
 
 class GraphExportResponse(BaseModel):