flatmachines 1.0.0__py3-none-any.whl

flatmachines/persistence.py

@@ -0,0 +1,213 @@
+import json
+import fcntl
+import asyncio
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, List
+from dataclasses import dataclass, asdict, field
+from pathlib import Path
+from datetime import datetime, timezone
+import aiofiles
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class MachineSnapshot:
+    """Wire format for machine checkpoints."""
+    execution_id: str
+    machine_name: str
+    spec_version: str
+    current_state: str
+    context: Dict[str, Any]
+    step: int
+    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    event: Optional[str] = None  # The event that triggered this checkpoint (machine_start, etc)
+    output: Optional[Dict[str, Any]] = None  # Output if captured at state_exit/machine_end
+    total_api_calls: Optional[int] = None  # Cumulative API calls
+    total_cost: Optional[float] = None  # Cumulative cost
+    # Lineage (v0.4.0)
+    parent_execution_id: Optional[str] = None  # ID of launcher machine if this was launched
+    # Outbox pattern (v0.4.0)
+    pending_launches: Optional[List[Dict[str, Any]]] = None  # LaunchIntent dicts awaiting completion
+
+class PersistenceBackend(ABC):
+    """Abstract storage backend for checkpoints."""
+    
+    @abstractmethod
+    async def save(self, key: str, value: bytes) -> None:
+        pass
+    
+    @abstractmethod
+    async def load(self, key: str) -> Optional[bytes]:
+        pass
+    
+    @abstractmethod
+    async def delete(self, key: str) -> None:
+        pass
+
+class LocalFileBackend(PersistenceBackend):
+    """File-based persistence backend."""
+
+    def __init__(self, base_dir: str = ".checkpoints"):
+        self.base_dir = Path(base_dir)
+        self.base_dir.mkdir(parents=True, exist_ok=True)
+
+    def _validate_key(self, key: str) -> None:
+        """Validate key to prevent path traversal attacks."""
+        if '..' in key or key.startswith('/'):
+            raise ValueError(f"Invalid checkpoint key: {key}")
+
+    async def save(self, key: str, value: bytes) -> None:
+        self._validate_key(key)
+        path = self.base_dir / key
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write to temp file first for atomicity
+        temp_path = path.parent / f".{path.name}.tmp"
+        async with aiofiles.open(temp_path, 'wb') as f:
+            await f.write(value)
+
+        # Atomic rename (safe on POSIX and Windows)
+        temp_path.replace(path)
+
+    async def load(self, key: str) -> Optional[bytes]:
+        self._validate_key(key)
+        path = self.base_dir / key
+        if not path.exists():
+            return None
+        async with aiofiles.open(path, 'rb') as f:
+            return await f.read()
+
+    async def delete(self, key: str) -> None:
+        self._validate_key(key)
+        path = self.base_dir / key
+        path.unlink(missing_ok=True)
+
+class MemoryBackend(PersistenceBackend):
+    """In-memory backend for ephemeral executions."""
+    
+    def __init__(self):
+        self._store: Dict[str, bytes] = {}
+        
+    async def save(self, key: str, value: bytes) -> None:
+        self._store[key] = value
+        
+    async def load(self, key: str) -> Optional[bytes]:
+        return self._store.get(key)
+        
+    async def delete(self, key: str) -> None:
+        self._store.pop(key, None)
+
+class CheckpointManager:
+    """Manages saving and loading machine snapshots."""
+    
+    def __init__(self, backend: PersistenceBackend, execution_id: str):
+        self.backend = backend
+        self.execution_id = execution_id
+        
+    def _snapshot_key(self, event: str, step: int) -> str:
+        """Generate key for specific snapshot."""
+        return f"{self.execution_id}/step_{step:06d}_{event}.json"
+        
+    def _latest_pointer_key(self) -> str:
+        """Key that points to the latest snapshot."""
+        return f"{self.execution_id}/latest"
+
+    def _safe_serialize_value(self, value: Any, path: str, non_serializable: List[str]) -> Any:
+        """Recursively serialize a value, converting non-JSON types to strings."""
+        if isinstance(value, dict):
+            result = {}
+            for k, v in value.items():
+                try:
+                    json.dumps({k: v})
+                    result[k] = v
+                except (TypeError, OverflowError):
+                    result[k] = self._safe_serialize_value(v, f"{path}.{k}", non_serializable)
+            return result
+        elif isinstance(value, list):
+            result = []
+            for i, item in enumerate(value):
+                try:
+                    json.dumps(item)
+                    result.append(item)
+                except (TypeError, OverflowError):
+                    result.append(self._safe_serialize_value(item, f"{path}[{i}]", non_serializable))
+            return result
+        else:
+            try:
+                json.dumps(value)
+                return value
+            except (TypeError, OverflowError):
+                original_type = type(value).__name__
+                non_serializable.append(f"{path} ({original_type})")
+                return str(value)
+
+    def _safe_serialize(self, data: Dict[str, Any]) -> str:
+        """Safely serialize data to JSON, handling non-serializable objects."""
+        try:
+            return json.dumps(data)
+        except (TypeError, OverflowError):
+            # Identify and warn about specific non-serializable fields
+            safe_data = {}
+            non_serializable_fields: List[str] = []
+
+            for k, v in data.items():
+                if isinstance(v, dict):
+                    # Recursively check nested dicts
+                    try:
+                        json.dumps(v)
+                        safe_data[k] = v
+                    except (TypeError, OverflowError):
+                        safe_data[k] = self._safe_serialize_value(v, k, non_serializable_fields)
+                elif isinstance(v, list):
+                    # Recursively check lists
+                    try:
+                        json.dumps(v)
+                        safe_data[k] = v
+                    except (TypeError, OverflowError):
+                        safe_data[k] = self._safe_serialize_value(v, k, non_serializable_fields)
+                else:
+                    try:
+                        json.dumps({k: v})
+                        safe_data[k] = v
+                    except (TypeError, OverflowError):
+                        original_type = type(v).__name__
+                        safe_data[k] = str(v)
+                        non_serializable_fields.append(f"{k} ({original_type})")
+
+            if non_serializable_fields:
+                logger.warning(
+                    f"Context fields not JSON serializable, converted to strings: "
+                    f"{', '.join(non_serializable_fields)}. "
+                    f"These values will lose type information on restore."
+                )
+
+            return json.dumps(safe_data)
+
+    async def save_checkpoint(self, snapshot: MachineSnapshot) -> None:
+        """Save a snapshot and update latest pointer."""
+        data = asdict(snapshot)
+        json_bytes = self._safe_serialize(data).encode('utf-8')
+        
+        # Save the immutable snapshot
+        key = self._snapshot_key(snapshot.event or "unknown", snapshot.step)
+        await self.backend.save(key, json_bytes)
+        
+        # Update pointer to this key
+        await self.backend.save(self._latest_pointer_key(), key.encode('utf-8'))
+        
+    async def load_latest(self) -> Optional[MachineSnapshot]:
+        """Load the latest snapshot."""
+        # Get pointer
+        ptr_bytes = await self.backend.load(self._latest_pointer_key())
+        if not ptr_bytes:
+            return None
+            
+        # Get snapshot
+        key = ptr_bytes.decode('utf-8')
+        data_bytes = await self.backend.load(key)
+        if not data_bytes:
+            return None
+            
+        data = json.loads(data_bytes.decode('utf-8'))
+        return MachineSnapshot(**data)

flatmachines/run.py

@@ -0,0 +1,117 @@
+"""
+FlatMachines CLI Runner.
+
+Entry point for running machines via subprocess:
+    python -m flatmachines.run --config machine.yml --input '{"key": "value"}'
+
+Used by SubprocessInvoker for fire-and-forget machine execution.
+"""
+
+import argparse
+import asyncio
+import json
+import logging
+import sys
+from pathlib import Path
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+
+logger = logging.getLogger(__name__)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run a FlatMachine from the command line",
+        prog="python -m flatmachines.run"
+    )
+    parser.add_argument(
+        "--config", "-c",
+        required=True,
+        help="Path to machine config file (YAML or JSON)"
+    )
+    parser.add_argument(
+        "--input", "-i",
+        default="{}",
+        help="JSON string of input data"
+    )
+    parser.add_argument(
+        "--execution-id", "-e",
+        help="Predetermined execution ID"
+    )
+    parser.add_argument(
+        "--parent-id", "-p",
+        help="Parent execution ID for lineage tracking"
+    )
+    parser.add_argument(
+        "--max-steps",
+        type=int,
+        default=1000,
+        help="Maximum execution steps (default: 1000)"
+    )
+    parser.add_argument(
+        "--verbose", "-v",
+        action="store_true",
+        help="Enable verbose logging"
+    )
+    
+    args = parser.parse_args()
+    
+    if args.verbose:
+        logging.getLogger().setLevel(logging.DEBUG)
+    
+    # Parse input
+    try:
+        input_data = json.loads(args.input)
+    except json.JSONDecodeError as e:
+        logger.error(f"Invalid JSON input: {e}")
+        sys.exit(1)
+    
+    # Validate config path
+    config_path = Path(args.config)
+    if not config_path.exists():
+        logger.error(f"Config file not found: {config_path}")
+        sys.exit(1)
+    
+    # Import here to avoid circular imports
+    from .flatmachine import FlatMachine
+    
+    # Build machine with optional execution IDs
+    machine_kwargs = {
+        "config_file": str(config_path),
+    }
+    
+    if args.execution_id:
+        machine_kwargs["_execution_id"] = args.execution_id
+    if args.parent_id:
+        machine_kwargs["_parent_execution_id"] = args.parent_id
+    
+    try:
+        machine = FlatMachine(**machine_kwargs)
+        
+        # Run the machine
+        result = asyncio.run(
+            machine.execute(
+                input=input_data,
+                max_steps=args.max_steps
+            )
+        )
+        
+        # Output result as JSON
+        print(json.dumps(result, indent=2, default=str))
+        
+        # Log stats
+        logger.info(
+            f"Execution complete: {machine.total_api_calls} API calls, "
+            f"${machine.total_cost:.4f} cost"
+        )
+        
+    except Exception as e:
+        logger.exception(f"Execution failed: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

flatmachines/utils.py

@@ -0,0 +1,166 @@
+"""Utility functions for flatmachines."""
+
+import re
+from types import SimpleNamespace
+from typing import Any, Dict, List, Optional
+
+from .monitoring import get_logger
+
+logger = get_logger(__name__)
+
+
+def check_spec_version(config_version: Optional[str], sdk_version: str) -> str:
+    """
+    Check spec version compatibility and warn if mismatched.
+
+    Args:
+        config_version: Version from config file (may be None)
+        sdk_version: Current SDK version (__version__)
+
+    Returns:
+        The effective spec version (config_version or sdk_version as default)
+    """
+    effective_version = config_version or sdk_version
+    sdk_major_minor = '.'.join(sdk_version.split('.')[:2])
+    config_major_minor = '.'.join(effective_version.split('.')[:2])
+
+    if config_major_minor != sdk_major_minor:
+        logger.warning(
+            f"Config version {effective_version} may not be fully supported. "
+            f"Current SDK version is {sdk_version}."
+        )
+
+    return effective_version
+
+
+def strip_markdown_json(content: str) -> str:
+    """
+    Extract JSON from potentially wrapped response content.
+
+    LLMs sometimes wrap JSON responses in markdown code blocks like:
+    ```json
+    {"key": "value"}
+    ```
+
+    Or include explanatory text before/after the JSON:
+    "Here is the result:
+    ```json
+    {"key": "value"}
+    ```"
+
+    This function extracts the JSON so json.loads() can parse it.
+
+    Args:
+        content: Raw string that may contain markdown-wrapped JSON
+
+    Returns:
+        Extracted JSON string
+    """
+    if not content:
+        return content
+
+    text = content.strip()
+
+    # First, try to find JSON in a markdown code fence (anywhere in content)
+    fence_pattern = r'```(?:json|JSON)?\s*\n?([\s\S]*?)\n?```'
+    match = re.search(fence_pattern, text)
+    if match:
+        return match.group(1).strip()
+
+    # If no fence, try to find a raw JSON object or array
+    json_pattern = r'(\{[\s\S]*\}|\[[\s\S]*\])'
+    match = re.search(json_pattern, text)
+    if match:
+        return match.group(1)
+
+    return text
+
+
+def _get_attr(obj: Any, key: str, default: Any = None) -> Any:
+    if obj is None:
+        return default
+    if hasattr(obj, key):
+        return getattr(obj, key)
+    if isinstance(obj, dict):
+        return obj.get(key, default)
+    return default
+
+
+def _coerce_usage(usage: Any) -> Any:
+    if usage is None:
+        return None
+    if isinstance(usage, dict):
+        return SimpleNamespace(**usage)
+    return usage
+
+
+async def consume_litellm_stream(stream: Any) -> Any:
+    content_parts: List[str] = []
+    tool_calls: Dict[int, Dict[str, Any]] = {}
+    usage_data: Any = None
+    finish_reason: Optional[str] = None
+
+    async for chunk in stream:
+        if chunk is None:
+            continue
+        usage = _get_attr(chunk, "usage")
+        if usage:
+            usage_data = usage
+
+        choices = _get_attr(chunk, "choices")
+        if not choices:
+            continue
+        choice0 = choices[0]
+        finish = _get_attr(choice0, "finish_reason")
+        if finish:
+            finish_reason = finish
+
+        delta = _get_attr(choice0, "delta")
+        if not delta:
+            continue
+
+        content_piece = _get_attr(delta, "content")
+        if content_piece:
+            content_parts.append(content_piece)
+
+        delta_tool_calls = _get_attr(delta, "tool_calls")
+        if delta_tool_calls:
+            for tc in delta_tool_calls:
+                index = _get_attr(tc, "index", 0)
+                entry = tool_calls.setdefault(index, {"id": None, "name": None, "arguments": []})
+                tc_id = _get_attr(tc, "id")
+                if tc_id:
+                    entry["id"] = tc_id
+                function = _get_attr(tc, "function")
+                if function:
+                    name = _get_attr(function, "name")
+                    if name:
+                        entry["name"] = name
+                    arguments = _get_attr(function, "arguments")
+                    if arguments:
+                        entry["arguments"].append(arguments)
+
+    content = "".join(content_parts)
+    message_fields: Dict[str, Any] = {"content": content}
+    if tool_calls:
+        tool_call_objs = []
+        for index in sorted(tool_calls):
+            entry = tool_calls[index]
+            tool_call_objs.append(SimpleNamespace(
+                id=entry["id"],
+                function=SimpleNamespace(
+                    name=entry["name"],
+                    arguments="".join(entry["arguments"]) if entry["arguments"] else ""
+                )
+            ))
+        message_fields["tool_calls"] = tool_call_objs
+
+    message = SimpleNamespace(**message_fields)
+    choice = SimpleNamespace(message=message)
+    if finish_reason is not None:
+        choice.finish_reason = finish_reason
+
+    return SimpleNamespace(
+        choices=[choice],
+        usage=_coerce_usage(usage_data)
+    )

flatmachines/validation.py

@@ -0,0 +1,79 @@
+"""
+Schema validation for flatmachine configurations.
+
+Uses JSON Schema validation against the bundled schema.
+Validation errors are warnings by default to avoid breaking user configs.
+"""
+
+import json
+import warnings
+from importlib.resources import files
+from typing import Any, Dict, List, Optional
+
+_ASSETS = files("flatmachines.assets")
+
+
+class ValidationWarning(UserWarning):
+    """Warning for schema validation issues."""
+
+
+
+def _load_schema(filename: str) -> Optional[Dict[str, Any]]:
+    try:
+        content = (_ASSETS / filename).read_text()
+        return json.loads(content)
+    except FileNotFoundError:
+        return None
+
+
+def _validate_with_jsonschema(config: Dict[str, Any], schema: Dict[str, Any]) -> List[str]:
+    try:
+        import jsonschema
+    except ImportError:
+        return []
+
+    errors: List[str] = []
+    validator = jsonschema.Draft7Validator(schema)
+    for error in validator.iter_errors(config):
+        path = ".".join(str(p) for p in error.absolute_path) or "(root)"
+        errors.append(f"{path}: {error.message}")
+    return errors
+
+
+def validate_flatmachine_config(
+    config: Dict[str, Any],
+    warn: bool = True,
+    strict: bool = False,
+) -> List[str]:
+    """Validate a flatmachine configuration against the schema."""
+    schema = _load_schema("flatmachine.schema.json")
+    if schema is None:
+        return []
+
+    errors = _validate_with_jsonschema(config, schema)
+
+    if errors:
+        if strict:
+            raise ValueError(
+                "Flatmachine config validation failed:\n"
+                + "\n".join(f"  - {e}" for e in errors)
+            )
+        if warn:
+            warnings.warn(
+                "Flatmachine config has validation issues:\n"
+                + "\n".join(f"  - {e}" for e in errors),
+                ValidationWarning,
+                stacklevel=3,
+            )
+
+    return errors
+
+
+def get_flatmachine_schema() -> Optional[Dict[str, Any]]:
+    """Get the bundled flatmachine JSON schema."""
+    return _load_schema("flatmachine.schema.json")
+
+
+def get_asset(filename: str) -> str:
+    """Get the contents of a bundled asset file."""
+    return (_ASSETS / filename).read_text()