sqlsaber-notebook 0.2.0__tar.gz

sqlsaber_notebook-0.2.0/.gitignore

@@ -0,0 +1,12 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+*.local.*

sqlsaber_notebook-0.2.0/CHANGELOG.md

@@ -0,0 +1,10 @@
+# Changelog
+
+## [0.2.0](https://github.com/SarthakJariwala/sqlsaber/compare/sqlsaber-notebook-v0.1.0...sqlsaber-notebook-v0.2.0) (2026-07-22)
+
+
+### Features
+
+* add sandboxed notebook analysis ([#199](https://github.com/SarthakJariwala/sqlsaber/issues/199)) ([c82590f](https://github.com/SarthakJariwala/sqlsaber/commit/c82590ffa8cb93649f021d548ff4d4c12171e54e))
+* trigger release ([145d163](https://github.com/SarthakJariwala/sqlsaber/commit/145d163c91989e62b722a266842b13e9145be638))
+* visualization model override support via api and prompt improvements ([35a506b](https://github.com/SarthakJariwala/sqlsaber/commit/35a506b23337dd64f8f20606e7e547755eae6f21))

sqlsaber_notebook-0.2.0/PKG-INFO

@@ -0,0 +1,124 @@
+Metadata-Version: 2.4
+Name: sqlsaber-notebook
+Version: 0.2.0
+Summary: Notebook analysis subagent for SQLsaber
+Requires-Python: >=3.12
+Requires-Dist: nbformat<6,>=5.10
+Requires-Dist: pillow<13,>=11
+Requires-Dist: sqlsaber>=0.69.0
+Provides-Extra: modal
+Requires-Dist: modal<2,>=1.4; extra == 'modal'
+Description-Content-Type: text/markdown
+
+# SQLsaber Notebook
+
+Notebook-specific data-analysis subagent for SQLsaber.
+
+Implemented components:
+
+- provider-neutral notebook execution contract,
+- hardened local Docker execution (default),
+- explicit remote Modal Sandbox execution through the optional `modal` extra,
+- fresh-kernel transactional notebook sessions,
+- bounded notebook/image rendering and history collapse,
+- `list_workspace` and `edit_cell` analyst tools,
+- a Pydantic AI notebook analyst,
+- a managed SQLsaber `analyze_data` capability, and
+- the standalone `sqlsaber-notebook` CLI.
+
+When installed with SQLsaber, the main agent can hand prior successful SQL results to
+`analyze_data` for multi-step calculations, statistics, transformations, and plots.
+The terminal displays the bounded executed notebook and plot previews before the main
+agent's text response. Notebook bytes and images are not sent to the parent model.
+Managed SDK applications can persist the notebook, plots, and generated files through
+`SQLSaberOptions.artifact_publisher`; only the publisher's durable references are
+stored in tool metadata and exposed through `SQLSaberResult.artifacts`.
+
+The default balanced runtime targets larger EDA and classical ML: 4 CPUs, 8 GiB
+memory, and up to 100 MiB per input/250 MiB total. SQLsaber does not cap model
+requests, notebook cell count, the analyst loop, or the whole operation. Individual
+cells retain a 10-minute timeout so a stuck computation can be diagnosed without
+ending the overall analysis. These are fixed product defaults rather than CLI tuning
+flags. Use an immutable custom image through `SQLSABER_NOTEBOOK_IMAGE` when
+additional ML libraries are required.
+
+## Managed SQLsaber usage
+
+```bash
+uv tool install --with sqlsaber-notebook sqlsaber
+saber
+```
+
+Docker is the default local backend. Select Modal explicitly because query results
+will be uploaded to a third party:
+
+```bash
+SQLSABER_NOTEBOOK_BACKEND=modal saber
+```
+
+Configure a dedicated analyst model with:
+
+```bash
+saber models set --agent notebook
+```
+
+For a web backend, inject an application-owned artifact publisher and pass tenant
+scope as run metadata:
+
+```python
+from sqlsaber import FilesystemArtifactPublisher, SQLSaber, SQLSaberOptions
+
+options = SQLSaberOptions(
+    database="sqlite:///analytics.db",
+    artifact_publisher=FilesystemArtifactPublisher("/private/artifacts"),
+)
+
+async with SQLSaber(options=options) as saber:
+    result = await saber.query(
+        "Analyze and plot revenue anomalies",
+        conversation_id="conversation-123",
+        metadata={"tenant_id": "acme"},
+    )
+    print(result.artifacts)
+```
+
+Implement the cloud-neutral `ArtifactPublisher` protocol to use S3, GCS, Azure
+Blob Storage, or another bucket. Return stable object references rather than
+expiring signed URLs.
+
+## Standalone usage
+
+```bash
+uv run sqlsaber-notebook \
+  --model anthropic:claude-sonnet-4-6 \
+  --backend docker \
+  --output analysis.ipynb \
+  "Compare revenue by region and explain material anomalies" data.csv
+```
+
+Modal is never selected as an automatic fallback. Select it explicitly because local
+files will be uploaded to Modal:
+
+```bash
+modal setup
+SQLSABER_NOTEBOOK_BACKEND=modal uv run sqlsaber-notebook \
+  --model anthropic:claude-sonnet-4-6 \
+  "Analyze this dataset" data.csv
+```
+
+## Development
+
+```bash
+uv sync
+uv run pytest plugins/notebook/tests -q
+```
+
+Run live backend integration tests explicitly:
+
+```bash
+SQLSABER_RUN_DOCKER_INTEGRATION=1 \
+  uv run pytest plugins/notebook/tests/test_notebook_docker_integration.py -q
+
+SQLSABER_RUN_MODAL_INTEGRATION=1 \
+  uv run pytest plugins/notebook/tests/test_notebook_modal_integration.py -q
+```

sqlsaber_notebook-0.2.0/README.md

@@ -0,0 +1,112 @@
+# SQLsaber Notebook
+
+Notebook-specific data-analysis subagent for SQLsaber.
+
+Implemented components:
+
+- provider-neutral notebook execution contract,
+- hardened local Docker execution (default),
+- explicit remote Modal Sandbox execution through the optional `modal` extra,
+- fresh-kernel transactional notebook sessions,
+- bounded notebook/image rendering and history collapse,
+- `list_workspace` and `edit_cell` analyst tools,
+- a Pydantic AI notebook analyst,
+- a managed SQLsaber `analyze_data` capability, and
+- the standalone `sqlsaber-notebook` CLI.
+
+When installed with SQLsaber, the main agent can hand prior successful SQL results to
+`analyze_data` for multi-step calculations, statistics, transformations, and plots.
+The terminal displays the bounded executed notebook and plot previews before the main
+agent's text response. Notebook bytes and images are not sent to the parent model.
+Managed SDK applications can persist the notebook, plots, and generated files through
+`SQLSaberOptions.artifact_publisher`; only the publisher's durable references are
+stored in tool metadata and exposed through `SQLSaberResult.artifacts`.
+
+The default balanced runtime targets larger EDA and classical ML: 4 CPUs, 8 GiB
+memory, and up to 100 MiB per input/250 MiB total. SQLsaber does not cap model
+requests, notebook cell count, the analyst loop, or the whole operation. Individual
+cells retain a 10-minute timeout so a stuck computation can be diagnosed without
+ending the overall analysis. These are fixed product defaults rather than CLI tuning
+flags. Use an immutable custom image through `SQLSABER_NOTEBOOK_IMAGE` when
+additional ML libraries are required.
+
+## Managed SQLsaber usage
+
+```bash
+uv tool install --with sqlsaber-notebook sqlsaber
+saber
+```
+
+Docker is the default local backend. Select Modal explicitly because query results
+will be uploaded to a third party:
+
+```bash
+SQLSABER_NOTEBOOK_BACKEND=modal saber
+```
+
+Configure a dedicated analyst model with:
+
+```bash
+saber models set --agent notebook
+```
+
+For a web backend, inject an application-owned artifact publisher and pass tenant
+scope as run metadata:
+
+```python
+from sqlsaber import FilesystemArtifactPublisher, SQLSaber, SQLSaberOptions
+
+options = SQLSaberOptions(
+    database="sqlite:///analytics.db",
+    artifact_publisher=FilesystemArtifactPublisher("/private/artifacts"),
+)
+
+async with SQLSaber(options=options) as saber:
+    result = await saber.query(
+        "Analyze and plot revenue anomalies",
+        conversation_id="conversation-123",
+        metadata={"tenant_id": "acme"},
+    )
+    print(result.artifacts)
+```
+
+Implement the cloud-neutral `ArtifactPublisher` protocol to use S3, GCS, Azure
+Blob Storage, or another bucket. Return stable object references rather than
+expiring signed URLs.
+
+## Standalone usage
+
+```bash
+uv run sqlsaber-notebook \
+  --model anthropic:claude-sonnet-4-6 \
+  --backend docker \
+  --output analysis.ipynb \
+  "Compare revenue by region and explain material anomalies" data.csv
+```
+
+Modal is never selected as an automatic fallback. Select it explicitly because local
+files will be uploaded to Modal:
+
+```bash
+modal setup
+SQLSABER_NOTEBOOK_BACKEND=modal uv run sqlsaber-notebook \
+  --model anthropic:claude-sonnet-4-6 \
+  "Analyze this dataset" data.csv
+```
+
+## Development
+
+```bash
+uv sync
+uv run pytest plugins/notebook/tests -q
+```
+
+Run live backend integration tests explicitly:
+
+```bash
+SQLSABER_RUN_DOCKER_INTEGRATION=1 \
+  uv run pytest plugins/notebook/tests/test_notebook_docker_integration.py -q
+
+SQLSABER_RUN_MODAL_INTEGRATION=1 \
+  uv run pytest plugins/notebook/tests/test_notebook_modal_integration.py -q
+```

sqlsaber_notebook-0.2.0/pyproject.toml

@@ -0,0 +1,37 @@
+[project]
+name = "sqlsaber-notebook"
+version = "0.2.0"
+description = "Notebook analysis subagent for SQLsaber"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "sqlsaber>=0.69.0",
+    "nbformat>=5.10,<6",
+    "pillow>=11,<13",
+]
+
+[project.optional-dependencies]
+modal = ["modal>=1.4,<2"]
+
+[tool.uv.sources]
+sqlsaber = { path = "../..", editable = true }
+
+[project.scripts]
+sqlsaber-notebook = "sqlsaber_notebook.cli:main"
+
+[project.entry-points."sqlsaber.capabilities"]
+notebook = "sqlsaber_notebook.capability:capability"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/sqlsaber_notebook"]
+
+[dependency-groups]
+dev = [
+    "pytest>=9.0.2",
+    "pytest-asyncio>=1.0.0",
+    "modal>=1.4,<2",
+]

sqlsaber_notebook-0.2.0/pytest.ini

@@ -0,0 +1,5 @@
+[pytest]
+asyncio_mode = auto
+asyncio_default_fixture_loop_scope = function
+markers =
+    integration: exercises a real external execution backend

sqlsaber_notebook-0.2.0/src/sqlsaber_notebook/__init__.py

@@ -0,0 +1,29 @@
+"""SQLsaber notebook analysis plugin."""
+
+from .execution import (
+    DEFAULT_NOTEBOOK_BACKEND,
+    DEFAULT_NOTEBOOK_IMAGE,
+    ArtifactInfo,
+    DockerNotebookBackend,
+    ExecutionLimits,
+    NotebookBackend,
+    NotebookEnvironment,
+    NotebookExecutionResult,
+    NotebookInput,
+    resolve_notebook_backend,
+    resolve_notebook_image,
+)
+
+__all__ = [
+    "ArtifactInfo",
+    "DEFAULT_NOTEBOOK_BACKEND",
+    "DEFAULT_NOTEBOOK_IMAGE",
+    "DockerNotebookBackend",
+    "ExecutionLimits",
+    "NotebookBackend",
+    "NotebookEnvironment",
+    "NotebookExecutionResult",
+    "NotebookInput",
+    "resolve_notebook_backend",
+    "resolve_notebook_image",
+]

sqlsaber_notebook-0.2.0/src/sqlsaber_notebook/_shared.py

@@ -0,0 +1,41 @@
+"""Central limits and defaults for notebook analysis."""
+
+from .execution import DEFAULT_NOTEBOOK_IMAGE, ExecutionLimits
+
+# Balanced defaults target larger EDA and classical ML while retaining finite
+# model-context, host-memory, runtime, and cost boundaries.
+MAX_DEFAULT_RESULTS = 20
+MAX_WORKSPACE_FILES = 50
+MAX_WORKSPACE_FILE_BYTES = 100 * 1024 * 1024
+MAX_WORKSPACE_TOTAL_BYTES = 250 * 1024 * 1024
+MAX_CELL_SOURCE_CHARS = 30_000
+MAX_TOTAL_SOURCE_CHARS = 200_000
+MAX_OUTPUT_CHARS = 6_000
+MAX_SNAPSHOT_CHARS = 80_000
+MAX_SNAPSHOT_IMAGES = 12
+MAX_IMAGE_BYTES = 4 * 1024 * 1024
+MAX_SNAPSHOT_IMAGE_BYTES = 24 * 1024 * 1024
+MAX_HISTORY_IMAGE_BYTES = 48 * 1024 * 1024
+MAX_IMAGE_WIDTH = 2048
+MAX_IMAGE_HEIGHT = 2048
+MAX_GOAL_CHARS = 20_000
+DEFAULT_EXECUTION_LIMITS = ExecutionLimits()
+
+__all__ = [
+    "DEFAULT_EXECUTION_LIMITS",
+    "DEFAULT_NOTEBOOK_IMAGE",
+    "MAX_CELL_SOURCE_CHARS",
+    "MAX_GOAL_CHARS",
+    "MAX_HISTORY_IMAGE_BYTES",
+    "MAX_IMAGE_BYTES",
+    "MAX_IMAGE_HEIGHT",
+    "MAX_IMAGE_WIDTH",
+    "MAX_OUTPUT_CHARS",
+    "MAX_SNAPSHOT_CHARS",
+    "MAX_SNAPSHOT_IMAGE_BYTES",
+    "MAX_SNAPSHOT_IMAGES",
+    "MAX_TOTAL_SOURCE_CHARS",
+    "MAX_WORKSPACE_FILES",
+    "MAX_WORKSPACE_FILE_BYTES",
+    "MAX_WORKSPACE_TOTAL_BYTES",
+]

sqlsaber_notebook-0.2.0/src/sqlsaber_notebook/analyst.py

@@ -0,0 +1,202 @@
+"""SQLsaber-independent notebook analyst orchestration."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import hashlib
+from functools import partial
+
+from pydantic_ai import Agent
+from pydantic_ai.capabilities import ProcessHistory
+from pydantic_ai.models import Model
+from pydantic_ai.models.anthropic import AnthropicModelSettings
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.usage import RunUsage, UsageLimits
+
+from ._shared import (
+    DEFAULT_EXECUTION_LIMITS,
+    MAX_GOAL_CHARS,
+    MAX_SNAPSHOT_IMAGE_BYTES,
+    MAX_SNAPSHOT_IMAGES,
+)
+from .execution import (
+    ExecutionLimits,
+    NotebookBackend,
+    resolve_notebook_backend,
+    resolve_notebook_image,
+)
+from .history import collapse_old_snapshots
+from .prompts import ANALYST_SYSTEM_PROMPT, goal_prompt
+from .rendering import extract_notebook_images, normalize_png_bytes
+from .result import AnalysisResult, ArtifactRef, Workspace
+from .session import NotebookSession
+from .tools import analyst_toolset
+
+
+def build_analyst_agent(
+    model: Model | str,
+    *,
+    model_provider: str,
+) -> Agent[NotebookSession, str]:
+    provider = model_provider.strip().lower()
+    return Agent(
+        model,
+        deps_type=NotebookSession,
+        output_type=str,
+        instructions=ANALYST_SYSTEM_PROMPT,
+        model_settings=_cache_settings(provider),
+        capabilities=[
+            ProcessHistory(
+                partial(
+                    collapse_old_snapshots,
+                    cache=provider == "anthropic",
+                )
+            )
+        ],
+        toolsets=[analyst_toolset()],
+    )
+
+
+async def analyze(
+    goal: str,
+    workspace: Workspace,
+    *,
+    model: Model | str,
+    model_provider: str,
+    backend: NotebookBackend | str | None = None,
+    image: str | None = None,
+    include_snapshot_images: bool = False,
+    collect_files: bool = True,
+    execution_limits: ExecutionLimits = DEFAULT_EXECUTION_LIMITS,
+    usage_limits: UsageLimits | None = None,
+    parent_usage: RunUsage | None = None,
+) -> AnalysisResult:
+    """Run one bounded notebook-analysis environment and always clean it up."""
+
+    if not goal.strip():
+        raise ValueError("Analysis goal cannot be empty")
+    if len(goal) > MAX_GOAL_CHARS:
+        raise ValueError(f"Analysis goal exceeds {MAX_GOAL_CHARS} characters")
+    selected_backend = (
+        backend
+        if isinstance(backend, NotebookBackend)
+        else resolve_notebook_backend(backend)
+    )
+    session = NotebookSession(
+        workspace=workspace,
+        backend=selected_backend,
+        image=resolve_notebook_image(image),
+        execution_limits=execution_limits,
+        include_snapshot_images=include_snapshot_images,
+    )
+    agent = build_analyst_agent(model, model_provider=model_provider)
+    child_usage = RunUsage()
+    try:
+        await session.ensure_environment()
+        result = await agent.run(
+            goal_prompt(goal),
+            deps=session,
+            usage=child_usage,
+            usage_limits=usage_limits,
+        )
+        images, files = await _harvest_artifacts(
+            session,
+            collect_files=collect_files,
+        )
+        return AnalysisResult(
+            answer=result.output,
+            notebook=session.notebook_bytes(),
+            images=images,
+            files=files,
+            provenance=_infer_provenance(session),
+        )
+    finally:
+        if parent_usage is not None:
+            parent_usage.incr(child_usage)
+        await _bounded_shielded_close(session)
+
+
+def supports_notebook_images(model_name: str, model_provider: str) -> bool:
+    """Conservative allowlist for analyst snapshot image input."""
+
+    provider = model_provider.strip().lower()
+    name = model_name.lower()
+    if provider == "anthropic":
+        return "claude" in name
+    if provider == "google":
+        return "gemini" in name
+    if provider == "openai":
+        return any(token in name for token in ("gpt-4", "gpt-5", "o3", "o4"))
+    return False
+
+
+def _cache_settings(provider: str) -> ModelSettings:
+    if provider == "anthropic":
+        return AnthropicModelSettings(
+            parallel_tool_calls=False,
+            anthropic_cache_instructions=True,
+            anthropic_cache_tool_definitions=True,
+        )
+    return ModelSettings(parallel_tool_calls=False)
+
+
+async def _harvest_artifacts(
+    session: NotebookSession,
+    *,
+    collect_files: bool,
+) -> tuple[list[bytes], list[ArtifactRef]]:
+    images = extract_notebook_images(session.outputs)
+    image_hashes = {hashlib.sha256(image).hexdigest() for image in images}
+    image_bytes = sum(map(len, images))
+    files: list[ArtifactRef] = []
+    if session.environment is None:
+        return images, files
+
+    for artifact in session.artifacts:
+        is_png = artifact.media_type == "image/png" or artifact.path.lower().endswith(
+            ".png"
+        )
+        if not is_png and not collect_files:
+            continue
+        data = await session.environment.read_artifact(artifact)
+        if is_png:
+            normalized = normalize_png_bytes(data)
+            if normalized is None:
+                continue
+            digest = hashlib.sha256(normalized).hexdigest()
+            if digest in image_hashes or len(images) >= MAX_SNAPSHOT_IMAGES:
+                continue
+            if image_bytes + len(normalized) > MAX_SNAPSHOT_IMAGE_BYTES:
+                continue
+            image_hashes.add(digest)
+            image_bytes += len(normalized)
+            images.append(normalized)
+        elif collect_files:
+            files.append(
+                ArtifactRef(
+                    name=artifact.path,
+                    data=data,
+                    media_type=artifact.media_type or "application/octet-stream",
+                )
+            )
+    return images, files
+
+
+def _infer_provenance(session: NotebookSession) -> list[str]:
+    provenance: list[str] = []
+    for item in session.workspace.files:
+        token = f"../inputs/{item.name}"
+        if any(token in source or item.name in source for source in session.cells):
+            provenance.append(f"input:{item.name}")
+    provenance.extend(
+        f"cell:{index}" for index, source in enumerate(session.cells) if source.strip()
+    )
+    return provenance
+
+
+async def _bounded_shielded_close(session: NotebookSession) -> None:
+    task = asyncio.create_task(session.close())
+    with contextlib.suppress(Exception):
+        async with asyncio.timeout(30):
+            await asyncio.shield(task)