flowbook-python 0.0.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

flowbook/__init__.py

@@ -1,31 +1,82 @@
-"""
-flowbook-python: Placeholder package for PyPI name reservation.
+try:
+    from ._version import __version__
+except ImportError:
+    # Fallback when using the package in dev mode without installing
+    # in editable mode with pip. It is highly recommended to install
+    # the package from a stable release or in editable mode: https://pip.pypa.io/en/stable/topics/local-project-installs/#editable-installs
+    import warnings
 
-This is a placeholder package to reserve the 'flowbook-python' name on PyPI.
-The actual FlowBook project is under active development.
+    warnings.warn("Importing 'flowbook' outside a proper installation.")
+    __version__ = "dev"
 
-For more information, visit: https://github.com/stephenfreund/FlowBook
-"""
+from jupyter_server.extension.application import ExtensionApp
+from traitlets import Unicode
+from flowbook.util.output import timer
 
-__version__ = "0.0.1"
 
-__all__ = ["__version__"]
+class FlowBookExtension(ExtensionApp):
+    """FlowBook server extension."""
 
+    name = "flowbook"
+    load_other_extensions = True
 
-def main():
-    """Placeholder function that provides information about the package."""
-    print("=" * 60)
-    print("flowbook-python - Placeholder Package")
-    print("=" * 60)
-    print()
-    print("This is a placeholder package to reserve the PyPI name.")
-    print("The actual FlowBook project is under development.")
-    print()
-    print("For more information:")
-    print("  GitHub: https://github.com/stephenfreund/FlowBook")
-    print()
-    print("=" * 60)
+    model = Unicode(
+        default_value="gpt-4o",
+        help="The model to use for the extension",
+    ).tag(config=True)
 
+    fast_model = Unicode(
+        default_value="gpt-4o-mini",
+        help="The fast model to use for the extension",
+    ).tag(config=True)
 
-if __name__ == "__main__":
-    main()
+    aliases = {
+        "model": "FlowBookExtension.model",
+        "fast-model": "FlowBookExtension.fast_model",
+    }
+
+    def initialize_settings(self):
+        """Initialize settings for the extension."""
+        with timer(message="Initializing FlowBook settings..."):
+            self.log.info(f"Initializing {self.name} extension")
+            self.serverapp.web_app.settings["flowbook"] = {
+                "ext": self,
+                "model": self.model,
+                "fast-model": self.fast_model,
+            }
+            km = self.serverapp.kernel_manager
+            if km.default_kernel_name == "python3":
+                km.default_kernel_name = "flowbook_kernel"
+            ksm = self.serverapp.kernel_spec_manager
+            if not ksm.allowed_kernelspecs:
+                ksm.allowed_kernelspecs = {"flowbook_kernel", "python3"}
+        make_kernels()
+
+    def initialize_handlers(self):
+        """Register HTTP handlers for the extension."""
+        with timer(message="Initializing FlowBook handlers..."):
+            from flowbook.server.handlers import setup_handlers
+            setup_handlers(self.serverapp.web_app)
+            self.log.info(f"Registered {self.name} server extension handlers")
+
+
+def _jupyter_labextension_paths():
+    """Provide the location of the labextension."""
+    with timer(message="JupyterLab extension paths..."):
+        return [{"src": "labextension", "dest": "flowbook"}]
+
+
+def _jupyter_server_extension_points():
+    """Define the server extension entry point."""
+    with timer(message="Jupyter server extension points..."):
+        return [{"module": "flowbook", "app": FlowBookExtension}]
+
+
+def make_kernels():
+    from flowbook.util.kernel_installer import install_kernel_spec
+    from pathlib import Path
+
+    base = Path(__file__).parent
+    install_kernel_spec("flowbook_kernel", base / "kernel" / "kernelspec")
+    install_kernel_spec("checkpoint_kernel", base / "checkpoint_kernel" / "kernelspec")
+    install_kernel_spec("baseline_kernel", base / "baseline_kernel" / "kernelspec")

flowbook/__main__.py

@@ -0,0 +1,44 @@
+import sys
+from typing import List
+from flowbook.util.output import timer
+import argparse
+import asyncio
+from pathlib import Path
+from flowbook.kernel_support.experimental_client import FlowbookClient
+
+
+async def handle_time(args):
+    from nbformat import read
+
+    with timer(message=f"Timing {args.path}"):
+        nb = read(args.path, as_version=4)
+        client = FlowbookClient(
+            nb,
+            kernel_name="flowbook_kernel",
+            allow_errors=False,
+            timeout=60,
+        )
+        await client.async_execute()
+
+
+def make_parser():
+    parser = argparse.ArgumentParser(description="FlowBook - A data analysis tool")
+    parser.add_argument("--model", default="gpt-4.1-mini", help="Base Model to use")
+    parser.add_argument("path", help="Path to the notebook or directory of notebooks")
+
+    return parser
+
+
+async def async_main():
+    parser = make_parser()
+    args = parser.parse_args(sys.argv[1:])
+
+    await handle_time(args)
+
+
+def main():
+    asyncio.run(async_main())
+
+
+if __name__ == "__main__":
+    main()

flowbook/_version.py

@@ -0,0 +1,4 @@
+# This file is auto-generated by Hatchling. As such, do not:
+#   - modify
+#   - track in version control e.g. be sure to add to .gitignore
+__version__ = VERSION = '0.1.2'

flowbook/baseline_kernel/__init__.py

@@ -0,0 +1,37 @@
+"""
+Baseline Kernel - For comparison testing with FlowBook.
+
+This package provides a baseline IPython kernel for fair A/B comparison
+with FlowBook. It has standard IPython behavior without reproducibility
+tracking or checkpointing.
+"""
+
+import os
+
+from flowbook.kernel_support.install import install_kernel
+
+from flowbook.baseline_kernel.baseline_kernel import (
+    BaselineKernel,
+)
+
+__all__ = [
+    "BaselineKernel",
+    "install_baseline_kernel",
+]
+
+
+def install_baseline_kernel() -> str:
+    """
+    Install the Baseline kernel spec.
+
+    Returns:
+        Path to installed kernel spec directory
+    """
+    return install_kernel(os.path.dirname(__file__), "baseline_kernel")
+
+
+# Install kernel on import
+try:
+    install_baseline_kernel()
+except Exception:
+    pass  # Don't fail import if kernel install fails

flowbook/baseline_kernel/__main__.py

@@ -0,0 +1,13 @@
+"""Entry point for running BaselineKernel as a module."""
+
+from ipykernel.kernelapp import IPKernelApp
+from traitlets.config import Config
+
+from flowbook.baseline_kernel.baseline_kernel import (
+    BaselineKernel,
+)
+
+if __name__ == "__main__":
+    config = Config()
+    config.IPKernelApp.capture_fd_output = False
+    IPKernelApp.launch_instance(kernel_class=BaselineKernel, config=config)

flowbook/baseline_kernel/baseline_kernel.py

@@ -0,0 +1,238 @@
+"""Baseline kernel for fair comparison with FlowBook.
+
+This kernel provides the same memory measurement capabilities as FlowBook
+but without any reproducibility tracking, checkpointing, or other FlowBook features.
+It's used for fair A/B comparison of memory overhead.
+
+Features:
+- Standard IPython kernel behavior
+- %memory magic command (same as FlowBook)
+- Memory reporting in same format as FlowBook for comparison
+- Kernel-side timing reported via metadata (for fair timing comparison)
+
+Does NOT include:
+- Reproducibility tracking
+- Checkpointing
+- Variable access tracking
+- Any FlowBook-specific features
+"""
+
+import os
+import time
+import types
+from typing import Any, Dict, Optional
+
+from ipykernel.ipkernel import IPythonKernel
+from flowbook.kernel_support.virtual_fs import VirtualFileSystem
+from IPython.core.magic import line_magic, magics_class, Magics
+
+
+@magics_class
+class MemoryMagics(Magics):
+    """Memory magic commands for baseline comparison using HeapSizer."""
+
+    @line_magic
+    def memory(self, line: str) -> None:
+        """Show memory usage using HeapSizer.
+
+        Usage:
+            %memory         - Show namespace memory summary
+            %memory vars    - Show per-variable breakdown
+            %memory vars 10 - Show top 10 variables
+        """
+        from flowbook.kernel_support.heap_size import HeapSizer
+
+        args = line.strip().lower().split()
+        if not args:
+            args = [""]
+
+        cmd = args[0]
+
+        if cmd == "" or cmd == "?" or cmd == "status":
+            # Show namespace summary
+            sizer = HeapSizer()
+            user_ns = self.shell.user_ns
+            # Filter to user variables
+            user_vars = {
+                k: v for k, v in user_ns.items()
+                if not k.startswith('_') and not isinstance(v, (type, types.FunctionType, types.BuiltinFunctionType, types.ModuleType))
+            }
+            ns_size = sizer.sizeof_namespace(user_vars)
+            print(f"Namespace memory: {ns_size.total_bytes / (1024*1024):.1f} MB")
+            print(f"Variables: {len(ns_size.by_variable)}")
+
+        elif cmd == "vars":
+            # Show per-variable breakdown
+            limit = 20
+            if len(args) > 1:
+                try:
+                    limit = int(args[1])
+                except ValueError:
+                    pass
+
+            sizer = HeapSizer()
+            user_ns = self.shell.user_ns
+            user_vars = {
+                k: v for k, v in user_ns.items()
+                if not k.startswith('_') and not isinstance(v, (type, types.FunctionType, types.BuiltinFunctionType, types.ModuleType))
+            }
+            ns_size = sizer.sizeof_namespace(user_vars)
+
+            # Build list of (name, type, size)
+            var_sizes = []
+            for name, size in ns_size.by_variable.items():
+                type_name = type(user_vars[name]).__name__
+                var_sizes.append((name, type_name, size))
+
+            # Sort by size descending
+            var_sizes.sort(key=lambda x: x[2], reverse=True)
+            var_sizes = var_sizes[:limit]
+
+            # Format output
+            print("Variable         Type            Size")
+            print("─" * 50)
+            for name, type_name, size in var_sizes:
+                size_str = self._format_bytes(size)
+                print(f"{name:<16} {type_name:<15} {size_str:>10}")
+
+        else:
+            print(f"Unknown command: {cmd}")
+            print("Usage: %memory [vars [limit]]")
+
+    def _format_bytes(self, size: int) -> str:
+        """Format bytes as human-readable string."""
+        if size < 1024:
+            return f"{size} B"
+        elif size < 1024 * 1024:
+            return f"{size / 1024:.1f} KB"
+        elif size < 1024 * 1024 * 1024:
+            return f"{size / (1024 * 1024):.1f} MB"
+        else:
+            return f"{size / (1024 * 1024 * 1024):.1f} GB"
+
+    @line_magic
+    def diagnostic(self, line: str) -> None:
+        """No-op magic for compatibility with FlowBook notebooks.
+
+        In FlowBook kernel, %diagnostic marks a cell as diagnostic-only
+        (no reproducibility tracking). In baseline kernel, this is a no-op
+        since there's no reproducibility tracking anyway.
+        """
+        pass
+
+
+
+class BaselineKernel(IPythonKernel):
+    """IPython kernel for baseline comparison with FlowBook.
+
+    This kernel provides standard IPython behavior with memory introspection
+    via HeapSizer, but without any FlowBook features like reproducibility
+    tracking or checkpointing.
+    """
+
+    implementation = "baseline_kernel"
+    implementation_version = "2.0"
+    banner = "Baseline Kernel - For comparison testing"
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._memory_metadata: Optional[Dict[str, Any]] = None
+
+        # Virtual filesystem: redirect writes to overlay to preserve real FS
+        self._vfs = VirtualFileSystem()
+        self._vfs.set_notebook_dir(os.getcwd())
+        if os.environ.get("FLOWBOOK_NO_VIRTUAL_FS", "0") == "1" or \
+                os.environ.get("FLOWBOOK_VIRTUAL_FS") == "0":
+            pass  # VFS disabled — writes go to real FS
+        else:
+            self._vfs.enable()
+        self._vfs_namespace_patched = False
+
+    def start(self) -> None:
+        """Start the kernel and register magic commands."""
+        super().start()
+
+        # Register memory magics
+        if self.shell:
+            self.shell.register_magics(MemoryMagics)
+            self._ensure_vfs_namespace_patched()
+
+            # Pre-import pandas and enable same options as FlowBook kernel
+            # These are set in memory_checkpoint.py for FlowBook - we need them
+            # here for fair comparison
+            try:
+                import pandas as pd
+
+                # Enable copy-on-write mode for better performance with DataFrame copies
+                # (always enabled in pandas >= 3.0, but needs to be set for pandas 2.x)
+                if hasattr(pd.options.mode, 'copy_on_write'):
+                    pd.options.mode.copy_on_write = True
+
+                # Enable string inference so read_csv() returns StringDtype instead of object dtype
+                # (always enabled in pandas >= 3.0, but needs to be set for pandas 2.x)
+                if hasattr(pd.options, 'future') and hasattr(pd.options.future, 'infer_string'):
+                    pd.options.future.infer_string = True
+            except ImportError:
+                pass
+
+    def _ensure_vfs_namespace_patched(self) -> None:
+        """Patch VFS open into the user namespace if VFS is active."""
+        if self._vfs_namespace_patched:
+            return
+        if self.shell is None:
+            return
+        if self._vfs.enabled:
+            self._vfs.patch_namespace(self.shell.user_global_ns)
+        self._vfs_namespace_patched = True
+
+    def do_shutdown(self, restart: bool) -> dict:
+        """Clean up VFS overlay on shutdown."""
+        if self._vfs.enabled:
+            self._vfs.disable()
+        return super().do_shutdown(restart)
+
+    async def do_execute(
+        self,
+        code: str,
+        silent: bool,
+        store_history: bool = True,
+        user_expressions: Optional[Dict[str, Any]] = None,
+        allow_stdin: bool = False,
+        *,
+        cell_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Execute code with standard IPython behavior and timing.
+
+        Reports kernel-side execution time via display_data metadata
+        for fair comparison with FlowBook kernel.
+        """
+        start_time = time.perf_counter()
+
+        # Execute code normally (IPythonKernel.do_execute is async)
+        result = await super().do_execute(
+            code,
+            silent,
+            store_history,
+            user_expressions,
+            allow_stdin,
+            cell_id=cell_id,
+        )
+
+        code_duration_ms = (time.perf_counter() - start_time) * 1000
+
+        # Emit timing metadata (same pattern as FlowBook kernel)
+        if not silent:
+            self.send_response(
+                self.iopub_socket,
+                "display_data",
+                {
+                    "data": {"text/plain": f"✓ Code: {code_duration_ms:.0f} ms"},
+                    "metadata": {
+                        "baseline": {
+                            "code_duration_ms": code_duration_ms,
+                        }
+                    },
+                },
+            )
+
+        return result

flowbook/baseline_kernel/kernelspec/kernel.json

@@ -0,0 +1,14 @@
+{
+  "argv": [
+    "python",
+    "-m",
+    "flowbook.baseline_kernel",
+    "-f",
+    "{connection_file}"
+  ],
+  "display_name": "Baseline Kernel",
+  "language": "python",
+  "metadata": {
+    "debugger": false
+  }
+}

flowbook/checkpoint_kernel/__init__.py

@@ -0,0 +1,37 @@
+"""
+Checkpoint Kernel - Execution and checkpoint timing measurement.
+
+This package provides:
+- CheckpointKernel: IPython kernel that measures execution and checkpoint time
+- CheckpointKernelClient: Client that sends cell_id with executions
+"""
+
+import os
+
+from flowbook.kernel_support.install import install_kernel
+
+from flowbook.checkpoint_kernel.checkpoint_client import CheckpointKernelClient
+from flowbook.checkpoint_kernel.checkpoint_kernel import CheckpointKernel
+
+__all__ = [
+    "CheckpointKernel",
+    "CheckpointKernelClient",
+    "install_checkpoint_kernel",
+]
+
+
+def install_checkpoint_kernel() -> str:
+    """
+    Install the Checkpoint kernel spec.
+
+    Returns:
+        Path to installed kernel spec directory
+    """
+    return install_kernel(os.path.dirname(__file__), "checkpoint_kernel")
+
+
+# Install kernel on import
+try:
+    install_checkpoint_kernel()
+except Exception:
+    pass  # Don't fail import if kernel install fails

flowbook/checkpoint_kernel/__main__.py

@@ -0,0 +1,11 @@
+"""Entry point for running CheckpointKernel as a module."""
+
+from ipykernel.kernelapp import IPKernelApp
+from traitlets.config import Config
+
+from flowbook.checkpoint_kernel.checkpoint_kernel import CheckpointKernel
+
+if __name__ == "__main__":
+    config = Config()
+    config.IPKernelApp.capture_fd_output = False
+    IPKernelApp.launch_instance(kernel_class=CheckpointKernel, config=config)

flowbook/checkpoint_kernel/checkpoint_client.py

@@ -0,0 +1,19 @@
+"""
+CheckpointKernelClient - Client that sends cell_id with executions.
+
+Extends BaseFlowbookClient to inject cell_id into execution requests.
+"""
+
+from flowbook.kernel_support.base_client import BaseFlowbookClient
+
+
+class CheckpointKernelClient(BaseFlowbookClient):
+    """
+    Kernel client that sends cell_id with executions.
+
+    Usage:
+        client = CheckpointKernelClient()
+        client.execute(code, cell_id='cell_abc')
+    """
+
+    pass

flowbook/checkpoint_kernel/checkpoint_kernel.py

@@ -0,0 +1,123 @@
+"""
+CheckpointKernel - IPython kernel that measures cell execution and checkpoint time.
+
+This is a minimal kernel that:
+1. Measures cell execution time
+2. Takes a checkpoint after each execution
+3. Measures checkpoint (commit) time
+4. Reports timing via flowbook_update IOPub messages
+
+No reproducibility tracking, no variable tracking - just execution and checkpoint timing.
+"""
+
+import time
+from typing import Optional
+
+from ipykernel.kernelapp import IPKernelApp
+
+from flowbook.kernel.protocol import IOPUB_MSG_TYPE
+from flowbook.kernel_support.base_kernel import BaseFlowbookKernel
+
+
+class CheckpointKernel(BaseFlowbookKernel):
+    """
+    IPython kernel with checkpoint timing measurement.
+
+    Features:
+    - Measures cell execution time
+    - Measures checkpoint (state save) time
+    - Reports timing via flowbook_update IOPub messages
+    """
+
+    implementation = "checkpoint_kernel"
+    implementation_version = "0.1"
+    banner = "FlowBook Checkpoint Kernel - Execution and checkpoint timing"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        # Expose checkpoint object to user code for rerun trials
+        self.shell.user_ns["_flowbook_checkpoint"] = self._checkpoint
+
+    async def _do_execute_impl(
+        self,
+        code: str,
+        silent: bool,
+        store_history: bool,
+        user_expressions: Optional[dict],
+        allow_stdin: bool,
+        cell_meta: Optional[dict],
+    ) -> dict:
+        """
+        Execute code and measure timing.
+        """
+        self._ensure_fs_magics()
+        self._ensure_vfs_namespace_patched()
+
+        # For empty code or pure magic, still report timing (with 0 values)
+        is_trivial = not code.strip() or self._is_pure_magic(code)
+
+        # Measure cell execution time
+        exec_start = time.perf_counter()
+        result = await self._ipython_do_execute(
+            code,
+            silent,
+            store_history,
+            user_expressions,
+            allow_stdin,
+            cell_meta=cell_meta,
+            cell_id=self._cell_id,
+        )
+        exec_end = time.perf_counter()
+
+        # For trivial cells, report 0 timings and skip checkpoint
+        if is_trivial:
+            cell_runtime_s = 0.0
+            commit_time_s = 0.0
+        else:
+            cell_runtime_s = exec_end - exec_start
+
+            # Measure checkpoint time
+            commit_start = time.perf_counter()
+            _, uncopyable = self._take_checkpoint(f"post_{self._cell_id}")
+            # Checkpoint kernel uses old behavior: remove uncopyable vars
+            for k in uncopyable:
+                if k in self.shell.user_ns:
+                    del self.shell.user_ns[k]
+            commit_end = time.perf_counter()
+            commit_time_s = commit_end - commit_start
+
+        # Check for execution errors
+        error_msg = None
+        if result.get("status") == "error":
+            error_msg = result.get("evalue", "Unknown error")
+
+        # Send timing metadata
+        if not silent:
+            msg = {
+                "type": "checkpoint_timing",
+                "cell_id": self._cell_id or "",
+                "execution_count": self.execution_count,
+                "cell_runtime_s": cell_runtime_s,
+                "commit_time_s": commit_time_s,
+            }
+            if error_msg:
+                msg["error"] = error_msg
+            self.send_response(
+                self.iopub_socket,
+                IOPUB_MSG_TYPE,
+                {"flowbook": msg},
+            )
+
+            # Display timing info
+            self._display.display_icon_and_text(
+                "ok" if error_msg is None else "error",
+                f"Run: {cell_runtime_s*1000:.1f}ms | Commit: {commit_time_s*1000:.1f}ms",
+            )
+
+        return result
+
+
+# Entry point
+if __name__ == "__main__":
+    IPKernelApp.launch_instance(kernel_class=CheckpointKernel)

flowbook/checkpoint_kernel/kernelspec/kernel.json

@@ -0,0 +1,14 @@
+{
+  "argv": [
+    "python",
+    "-m",
+    "flowbook.checkpoint_kernel",
+    "-f",
+    "{connection_file}"
+  ],
+  "display_name": "Checkpoint Kernel",
+  "language": "python",
+  "metadata": {
+    "debugger": false
+  }
+}

flowbook/cli/__init__.py

@@ -0,0 +1 @@
+"""CLI tools for FlowBook."""