Skip to content

ARCHITECTURE.md

Latest commit

 

History

History
377 lines (304 loc) · 18.7 KB

ARCHITECTURE.md

File metadata and controls

377 lines (304 loc) · 18.7 KB
Raw

Beaker Kernel Architecture

Overview

Beaker Kernel is an extensible Jupyter kernel that provides enhanced computational capabilities through AI-powered contexts, multi-language subkernels, and intelligent code analysis. Built on top of Jupyter Server infrastructure, it supports standalone notebook mode and production server deployments.

High-Level Architecture

┌─────────────────────────────────────────────────────────────┐
│                     Client Layer                            │
│  (JupyterLab, Notebook, API clients)                        │
└─────────────────────┬───────────────────────────────────────┘
                      │ HTTP/WebSocket
┌─────────────────────▼───────────────────────────────────────┐
│                  Application Layer                          │
│  ┌─────────────────┬─────────────────┬─────────────────┐    │
│  │  Web Server     │   CLI Tools     │  API Handlers   │    │
│  │  (Tornado)      │                 │                 │    │
│  └─────────────────┼─────────────────┼─────────────────┘    │
└──────────────────┬─┼─────────────────┼─────────────────┬────┘
                   │ │                 │                 │
┌──────────────────▼─▼─────────────────▼─────────────────▼────┐
│                    Services Layer                           │
│  ┌──────────┬──────────┬──────────┬──────────┐             │
│  │ Session  │ Kernel   │ Context  │ Storage  │             │
│  │ Manager  │ Manager  │ Manager  │ & Auth   │             │
│  └──────────┴──────────┴──────────┴──────────┘             │
└─────────────────────┬───────────────────────────────────────┘
                      │
┌─────────────────────▼───────────────────────────────────────┐
│                  Core Kernel Layer                          │
│  ┌─────────────────┬─────────────────┬─────────────────┐    │
│  │  BeakerKernel   │   Contexts      │   SubKernels    │    │
│  │  (Proxy Mgr)    │   (AI Agents)   │   (Languages)   │    │
│  └─────────────────┼─────────────────┼─────────────────┘    │
└──────────────────┬─┼─────────────────┼─────────────────┬────┘
                   │ │                 │                 │
┌──────────────────▼─▼─────────────────▼─────────────────▼────┐
│                   Extension Layer                           │
│  (Autodiscovery system for contexts, subkernels, etc.)      │
└─────────────────────────────────────────────────────────────┘

Directory Structure

src/beaker_notebook/
├── kernel.py              # Main kernel implementation
├── app/                   # Application interface layer
│   ├── api/              # REST API handlers
│   ├── ui/               # Web UI assets
│   ├── base.py          # Base application class
│   ├── server_app.py    # Production server
│   ├── notebook_app.py  # Notebook server
│   ├── dev_app.py       # Development server with hot reload
│   └── handlers.py      # Web request handlers
├── services/             # Business logic services
│   ├── kernel/          # Kernel lifecycle management
│   ├── session/         # Session management
│   ├── context/         # Context discovery and management
│   ├── auth/            # Authentication providers
│   └── storage/         # Storage backends
├── lib/                 # Core library components
│   ├── code_analysis/   # Static code analysis
│   ├── integrations/    # Integration registry, skill providers, ad-hoc integrations
│   ├── exporters/       # Notebook exporters (chat, full-context, streamline)
│   ├── templates/       # Code generation templates (agent, context, subkernel, procedure, readme)
│   ├── context.py       # Context (AI agent) framework
│   ├── context_dump.py  # Context serialization for prompt assembly
│   ├── subkernel.py     # Language subkernel framework
│   ├── reflector.py     # Variable reflection abstraction
│   ├── kernel_state.py  # Kernel-side state representation
│   ├── notebook_state.py # Notebook-side state representation
│   ├── agent.py         # AI agent base classes
│   ├── agent_tasks.py   # Agent task definitions
│   ├── config.py        # Configuration management
│   ├── jupyter_kernel_proxy.py  # Kernel proxy infrastructure
│   ├── message_types.py # Message type definitions
│   ├── types.py         # Type definitions
│   ├── workflow.py      # Workflow management
│   ├── extension.py     # Extension system
│   ├── app.py           # Application utilities
│   ├── admin.py         # Administrative functions
│   ├── autodiscovery.py # Extension autodiscovery
│   └── utils.py         # Utility functions
├── contexts/            # Built-in AI contexts
├── subkernels/          # Built-in language subkernels (each is a package; see below)
├── cli/                 # Command-line interface
├── builder/             # Build system integration
└── util/                # Utility modules (includes import-redirect shim for legacy paths)
src/beaker_kernel/
└── __init__.py          # Shim to allow older packages to map to the new beaker_notebook namespace.

Core Components

1. Application Layer (app/)

The application layer provides different deployment modes and interfaces:

Base Application (base.py):

  • BaseBeakerApp: Foundation class for all Beaker applications
  • Configures Jupyter Server components with Beaker-specific services
  • Handles traitlet configuration and extension loading

Application Types:

  • Server App: Production-ready server with optimized defaults
  • Notebook App: Local development with notebook interface
  • Dev App: Development mode with file watching and auto-reload

API Layer (api/):

  • RESTful endpoints for kernel management, notebook operations
  • Integration with external services
  • Handler discovery and registration system

2. Services Layer (services/)

Provides business logic and manages system resources:

Session Management (session/):

  • BeakerSessionManager: Extends Jupyter's session management with in-memory session tracking
  • Handles session pruning and kernel environment setup
  • Manages session-specific configuration and state

Kernel Services (kernel/):

  • BeakerKernelManager: Manages individual kernel lifecycles
  • BeakerKernelMappingManager: Maps sessions to kernels
  • BeakerKernelSpecManager: Handles kernel specification discovery
  • BeakerProvisioner / BeakerLocalProvisioner / BeakerProvisionerFactory: Provisions and manages kernel resources, including support for distributed/remote kernels

Context Services (context/):

  • BeakerContextManager: Manages context lifecycle and state with in-memory tracking
  • ContextDiscoveryService: Pluggable discovery system for loading available contexts
  • Context request handlers for API integration

Authentication (auth/):

  • Pluggable authentication providers via BeakerIdentityProvider and BeakerAuthorizer
  • Notebook-based authentication for local development
  • Dummy authentication for testing

Storage (storage/):

  • Persistent storage backends for notebooks and state
  • BeakerLocalContentsManager provides user-aware file content management with per-user home isolation
  • Configurable storage providers

Service API Convention (services/__init__.py):

  • ServiceApi and ServiceApiHandler define a standardized REST handler pattern
  • Subclassing ServiceApi auto-registers handlers under /beaker/<prefix>/… via subclass introspection
  • Supports dataclass serialization with Jupyter JSON compatibility, simplifying the addition of new endpoints

3. Core Kernel Layer

The heart of Beaker's computational capabilities:

BeakerKernel (kernel.py at root):

  • Extends KernelProxyManager from lib/jupyter_kernel_proxy.py
  • Intercepts and enhances Jupyter messages with AI capabilities
  • Coordinates between contexts, subkernels, and the base Jupyter infrastructure
  • Manages execution context and message routing

Context System (lib/context.py):

  • BeakerContext: Base class for AI-powered execution contexts
  • Provides tools, workflows, and intelligent code assistance
  • Integrates with external AI services and models
  • Supports code analysis, generation, and interactive assistance

SubKernel System (lib/subkernel.py):

  • BeakerSubkernel: Base class for language-specific execution engines
  • Manages code execution, state persistence, and environment isolation
  • Includes a checkpointing framework (CheckpointableBeakerSubkernel) for snapshotting and rolling back kernel state across executions. Checkpointing is temporarily disabled in the current release because serializing kernel state can take several minutes in some cases; the implementation is retained and will be revisited in a future update.

Built-in subkernels live in subkernels/<language>/ as packages rather than single files. Each package contains:

  • subkernel.py: the BeakerSubkernel subclass for that language
  • procedures/: language-side helper scripts loaded into the running kernel (e.g. fetch_state.<ext>)
  • procedures/reflectors/: per-type variable reflectors invoked by the reflector system (see below)

Reflector System (lib/reflector.py + subkernels/*/procedures/reflectors/):

  • Provides a uniform way to introspect variables in a running subkernel and produce a structured summary suitable for inclusion in agent context
  • Each language ships per-type reflectors (e.g. Python: class, dataframe, function, mapping, ndarray, primitive, sequence, series, string; Julia: array, dict, function, module, type; R: data_frame, function, list, vector) plus a default fallback
  • Reflectors live as discrete files in each subkernel package so they are discoverable and extendable without modifying core code

State Representation (lib/kernel_state.py, lib/notebook_state.py, lib/context_dump.py):

  • kernel_state.py captures the current state of the subkernel (variables, imports, environment) for inclusion in the agent's prompt
  • notebook_state.py captures the current state of the notebook (cells, outputs, history) for the same purpose
  • context_dump.py assembles the serialized context payload that is fed to the LLM, with hooks for prompt-cache–friendly formatting

Integration System (lib/integrations/):

  • Provides a registry of named integrations (data sources, APIs, skills) that contexts can expose to their agents
  • registry.py defines the registry; base.py defines the integration base class
  • skill.py implements skill integrations — bundles of metadata, instructions, examples, and file resources that augment an agent's capabilities
  • adhoc.py supports user-defined integrations created at runtime
  • types.py defines integration resource types (metadata, instructions, files, examples)
  • Integrations are addressed by slug rather than by name, allowing stable references across renames

Agent System (lib/agent.py, lib/agent_tasks.py):

  • AI agent base classes and task definitions
  • Integration with agent frameworks
  • Task execution and management

Configuration (lib/config.py):

  • Centralized configuration management with environment support
  • Dynamic configuration updates and validation
  • Integration with Jupyter's traitlet system

Supporting Infrastructure (lib/):

  • jupyter_kernel_proxy.py: Kernel proxy pattern implementation
  • message_types.py: Message type definitions and handling
  • types.py: Shared type definitions
  • workflow.py: Workflow management system
  • extension.py: Extension system framework
  • autodiscovery.py: Dynamic extension discovery
  • utils.py: Common utility functions

4. Extension System

Autodiscovery (lib/autodiscovery.py):

  • Dynamic discovery and loading of contexts and subkernels
  • Entry point-based extension system
  • Graceful error handling for missing or broken extensions

Built-in Extensions:

  • Contexts (contexts/): Default AI-powered context with agents
  • SubKernels (subkernels/): Python, R, and Julia language support, each shipped as a package with its own procedures and reflectors

5. Backward Compatibility

The beaker_kernel/service/ (singular) module tree was reorganized in v2.0 into beaker_notebook/app/, beaker_notebook/services/, and beaker_notebook/lib/admin.py. To avoid breaking external code that imports from the old paths, beaker_kernel/__init__.py installs an import-redirect map (via beaker_noteboook.util.refactor.redirect_imports) that transparently forwards legacy beaker_notebook.service.* imports to their new locations. New code should import from the new paths; the redirect is provided for transition only.

Key Design Patterns

1. Proxy Pattern

The BeakerKernel acts as a proxy to standard Jupyter kernels, intercepting and enhancing messages with AI capabilities while maintaining compatibility.

2. Plugin Architecture

Extensible through autodiscovery of contexts and subkernels via Python entry points, allowing third-party extensions.

3. Service-Oriented Architecture

Clear separation between application interface, business services, and core kernel functionality.

4. Configuration-Driven

Extensive use of traitlets for runtime configuration, supporting both development and production deployments.

Message Flow

Client Request
    │
    ▼
Web Server (Tornado)
    │
    ▼
API Handler
    │
    ▼
Service Layer
    │
    ▼
BeakerKernel (Proxy)
    │
    ├─────────────────┬─────────────────┐
    ▼                 ▼                 ▼
Context System    SubKernel        Standard Jupyter
(AI Agent)       (Language)         Kernel
    │                 │                 │
    ▼                 ▼                 ▼
Response Processing & Enhancement
    │
    ▼
Client Response

Configuration System

Beaker uses a multi-layer configuration system:

  1. Default Configuration: Built-in defaults for all components
  2. Environment Variables: Runtime configuration via environment
  3. Configuration Files: Python-based configuration files
  4. CLI Arguments: Command-line overrides
  5. Runtime Updates: Dynamic configuration changes

Configuration is managed through:

  • Jupyter's traitlet system for type safety and validation
  • Environment variable integration
  • File-based configuration with Python syntax

Security Model

  • Authentication: Pluggable authentication providers
  • Code Analysis: Static analysis for security and trust evaluation
  • Sandboxing: Isolated execution environments for subkernels
  • Input Validation: Message and code validation before execution

Extension Development

Creating a Context

from beaker_notebook.lib.context import BeakerContext

class MyContext(BeakerContext):
    def __init__(self, beaker_notebook, config):
        super().__init__(beaker_notebook, config)
        # Initialize your AI agents, tools, etc.

    async def setup(self):
        # Async initialization
        pass

Creating a SubKernel

from beaker_notebook.lib.subkernel import BeakerSubkernel

class MySubkernel(BeakerSubkernel):
    language = "mylang"

    async def execute(self, code):
        # Execute code in your language
        return result

Registration

Extensions are registered via entry points in pyproject.toml:

[project.entry-points."beaker_notebook.contexts"]
my_context = "my_package.context:MyContext"

[project.entry-points."beaker_notebook.subkernels"]
my_subkernel = "my_package.subkernel:MySubkernel"

Deployment Modes

Development Mode

beaker dev watch  # File watching with auto-reload

Local Notebook Mode

beaker notebook  # Single-user notebook interface

Production Server Mode

beaker server start server --port 8888  # Multi-session server

Dependencies

Core Dependencies

  • Jupyter Infrastructure: jupyterlab, jupyterlab-server
  • Web Framework: tornado (async web server)
  • AI Integration: archytas (agent framework)
  • Configuration: Built on Jupyter's traitlet system
  • Development: watchdog (file watching), click (CLI)

Optional Dependencies

  • Language subkernels (Python, R, Julia specific packages)
  • AI model providers (OpenAI, Anthropic, etc.)
  • External integrations (databases, APIs)

Future Architecture Considerations

This architecture supports future enhancements including:

  • Enhanced AI model integration
  • Real-time collaboration features
  • Advanced security and sandboxing
  • Performance monitoring and optimization
  • Custom UI components and extensions

The modular design ensures that new capabilities can be added without disrupting existing functionality, maintaining backward compatibility while enabling innovation.