Capabilities

SQLsaber exposes its database behavior as pydantic-ai capabilities. Attach SQLsaber to an agent you own, or add your capabilities to SQLsaber’s managed agent.

Bundled handoff

Handoff is installed with SQLsaber, but SDK sessions opt in through the same capabilities option as external plugins. SDK sessions do not read CLI plugin settings or SQLSABER_HANDOFF_MODEL.

from sqlsaber import SQLSaber, SQLSaberOptions
from sqlsaber.bundled.handoff import capability as handoff

async with SQLSaber(options=SQLSaberOptions(
    database="analytics",
    capabilities=(handoff,),
)) as saber:
    await saber.query("Explain the revenue trend")
    draft = await saber.draft_handoff("Investigate the largest decline")

The main model can also invoke draft_handoff(goal). Its history comes from the current run, and child-model usage counts toward the query’s usage and request limits. Neither invocation starts a thread or clears history; the caller decides whether to use the draft.

To customize the child model, pass partial(handoff, model_name="openai:gpt-5-mini") from functools.partial in capabilities. An explicit child model uses credentials from the session’s Config.auth; optionally pass api_key to the factory. Without an explicit model, handoff uses the session model and inherits its explicit API key.

Migration: saber.draft_handoff() now requires the handoff capability. It raises ValueError with configuration instructions when absent, rather than constructing a separate globally configured agent. Other SDK capabilities stay explicit and unchanged.

Add SQL tools to your agent

import asyncio

from pydantic_ai import Agent
from sqlsaber import SqlTools


async def main() -> None:
    sql = SqlTools(database="postgresql://user:password@localhost/analytics")
    agent = Agent(
        "anthropic:claude-sonnet-4-6",
        instructions="You are our analytics copilot.",
        capabilities=[sql],
    )

    async with agent:
        result = await agent.run("Top 5 customers by revenue")
        print(result.output)


asyncio.run(main())

SqlTools provides list_tables, introspect_schema, and execute_sql. Multi-database agents also receive list_dbs, and SQL tool schemas require a valid db_name.

Close owned connections

When SqlTools receives database=, it owns the resulting connections. Enter the agent with async with agent: as shown above, or call await sql.close() explicitly. A registry passed with registry= is borrowed and is never closed by the capability.

Multiple databases

Pass saved connection names or DSNs as a sequence:

sql = SqlTools(database=["production", "warehouse"])
agent = Agent("openai:gpt-5.2", capabilities=[sql])

SQLsaber includes a database catalog in the capability instructions. Cross-database joins are not attempted. The model queries each database separately.

Experimental CSV tool results

JSON is the default. Opt in to CSV tables for model-facing SQL tool results:

sql = SqlTools(database="analytics", csv_tool_results=True)

The option applies to all SQL tools in this capability, including list_dbs. Metadata and errors remain JSON. Retained query data and terminal rendering do not change. Token savings do not establish equivalent answer quality. See CSV tool result behavior for details. Managed sessions use SQLSaberOptions(csv_tool_results=True) instead.

Allow writes

SQL execution is read-only by default. Enable guarded write operations explicitly:

sql = SqlTools(database="production", allow_dangerous=True)

Dangerous mode still blocks destructive and administrative operations such as DROP, TRUNCATE, and role management. Only enable it for agents and users you trust.

Keep complete query results

execute_sql writes the full payload to a QueryResultStore and returns a bounded preview to the model. Pass the store into SqlTools. Load rows after agent.run().

from pydantic_ai import Agent
from sqlsaber import InMemoryQueryResultStore, SqlTools
from sqlsaber.query_result_resolution import query_result_references_from_messages
from sqlsaber.query_results import QueryResultContext

store = InMemoryQueryResultStore()
sql = SqlTools(database="analytics", query_result_store=store)
agent = Agent("anthropic:claude-sonnet-4-6", capabilities=[sql])

async with agent:
    result = await agent.run("Top 5 customers by revenue")

references = query_result_references_from_messages(result.new_messages())
loaded = await store.get(
    references[0].descriptor.id,
    context=QueryResultContext(),
)
print(loaded.rows())

Pass the same store into PluginContext.query_result_store when you construct notebook, viz, or sandbox capabilities. These capabilities use the store to load complete SQL results.

See Build an agent that keeps complete SQL results for a step-by-step agent that prints a truncated model preview and then loads all 250 stored rows.

Add knowledge search

Knowledge exposes saved SQL patterns and domain definitions. Share a manager when your application already owns one:

from pydantic_ai import Agent
from sqlsaber import Knowledge, SqlTools
from sqlsaber.knowledge.manager import KnowledgeManager

manager = KnowledgeManager()
sql = SqlTools(database="postgresql://localhost/analytics")
knowledge = Knowledge(
    knowledge_manager=manager,
    registry=sql.registry,
)
agent = Agent(
    "anthropic:claude-sonnet-4-6",
    capabilities=[sql, knowledge],
)

async with agent:
    result = await agent.run("How do we define active customers?")

await manager.close()

A supplied KnowledgeManager is borrowed. If omitted, Knowledge creates and closes its own manager when the agent context exits (or when await knowledge.close() is called).

Use your capabilities inside SQLsaber

Pass pydantic-ai capabilities and plugin factories through SQLSaberOptions.capabilities. SQLSaber always includes SQL and knowledge tools. It does not load installed plugins unless you list them.

from pydantic_ai.capabilities import WebSearch
from sqlsaber import SQLSaber, SQLSaberOptions
from sqlsaber_notebook.capability import capability as notebook
from sqlsaber_viz import capability as viz

options = SQLSaberOptions(
    database="analytics",
    capabilities=[notebook, viz, WebSearch()],
)

async with SQLSaber(options=options) as saber:
    result = await saber.query("Compare our latest revenue with public guidance")

List only the plugins you want. SQLSaber does not scan installed entry points. Factories run later with the session PluginContext, so they share the session query_result_store. For sandbox analysis, see Persistent sandbox analysis.

Configure notebook resources and the base image with a plugin-owned NotebookConfig:

from functools import partial
from sqlsaber_notebook import NotebookConfig, WorkspaceLimits

notebook_config = NotebookConfig(
    backend="docker",
    image="registry.example.com/analysis-with-ml:2026-09",
    workspace=WorkspaceLimits(max_total_bytes=2 * 1024**3),
    cell_seconds=3600,
)
options = SQLSaberOptions(
    database="analytics",
    capabilities=[partial(notebook, config=notebook_config)],
)

The direct sqlsaber_notebook.analyze(..., config=notebook_config) API accepts the same configuration. Explicit image and backend selectors take precedence over environment settings. Workspace budgets exclude the generated manifest; resource and artifact budgets do not increase model-preview budgets. Only the host can change this configuration, not the model calling analyze_data.

Constructed capabilities stay free to use their own dependency type when attached directly to your agent. SQLsaber’s built-in capabilities do not claim ctx.deps.

Persistent sandbox analysis

SandboxSession gives an analysis agent a persistent Python kernel. Follow-up goals reuse variables, installed packages, files, and the analysis agent’s history.

Start a session

Select a provider and pass input files through Workspace:

from pathlib import Path

from pydantic_ai.usage import UsageLimits
from sqlsaber_sandbox import SandboxConfig, SandboxSession, Workspace

config = SandboxConfig(provider="e2b", idle_seconds=300)
workspace = Workspace.from_files([
    ("observations.csv", Path("observations.csv").read_bytes()),
])

async with SandboxSession(
    model="anthropic:claude-sonnet-4-6",
    config=config,
) as session:
    first = await session.analyze(
        "Fit a model and save a diagnostic plot",
        workspace=workspace,
        usage_limits=UsageLimits(request_limit=12),
    )
    follow_up = await session.analyze(
        "Explain the largest residual using the same fitted model",
        usage_limits=UsageLimits(request_limit=6),
    )

The session closes when the async with block exits. To add files on a later call, pass another Workspace. Existing filenames must have the same contents and provenance.

To execute a cell without a model request, call await session.execute(code). To export files and the executed notebook, call await session.snapshot(). Snapshots do not rerun cells.

Both analyze and snapshot return an AnalysisResult with notebook bytes and generated files. Results remain usable after the session closes. publish_analysis(result, store=artifact_store, context=artifact_context) stores the notebook and files in your application’s ArtifactStore.

For a custom backend, pass a SandboxBackend implementation through SandboxSession(..., backend=my_backend). The session closes that backend on exit.

Add the managed capability

Register the factory with partial to supply SandboxConfig. SQLsaber supplies the stores and other PluginContext values:

from functools import partial

from sqlsaber import SQLSaberOptions
from sqlsaber_sandbox import SandboxConfig, capability as sandbox

options = SQLSaberOptions(
    database="analytics",
    query_result_store=store,
    artifact_store=artifact_store,
    capabilities=[
        partial(sandbox, config=SandboxConfig(provider="e2b")),
    ],
)

The main agent receives these tools:

• analyze_in_sandbox accepts a goal and automatically publishes completed artifacts to the configured ArtifactStore. It returns an answer, session ID, and artifact references. Passing session_id continues that analysis.
• close_sandbox retries pending publication before releasing the session’s environment. It returns the artifact references without duplicating an existing publication.
• publish_sandbox_artifacts retries publication without rerunning code or closing the session. Normal analysis does not need this separate call.

With the default artifact_failure_mode="required", close_sandbox returns an error without closing the session if publication fails. In best_effort mode, the session closes, but its downloaded snapshot remains available for a publication retry during the same runtime.

Capability shutdown also retries pending publication. It attempts to close every environment even if storage fails, and raises a cleanup error for required publication failures. Unpublished snapshots remain in memory, not durable storage.

analyze_in_sandbox replaces run_python and accepts a goal rather than Python code. The analysis agent keeps its code iterations in a separate conversation.

To accept non-SQL inputs, set SQLSaberOptions.workspace_input_resolver to your application’s attachment resolver. The tool then accepts attachment_refs.

For a Pydantic AI agent you own, use async with Sandbox(context, config=...) around the agent’s lifetime to close its sandbox sessions on exit.

Sandbox providers

SandboxConfig.provider takes precedence over SQLSABER_SANDBOX_PROVIDER. Without either setting, selection requires exactly one complete cloud credential set. Docker and microsandbox require explicit selection.

Provider	Package	Credentials or prerequisites
docker	sqlsaber-sandbox	Docker CLI and a running daemon
microsandbox	sqlsaber-sandbox[microsandbox]	Host virtualization. On Linux, readable and writable /dev/kvm
modal	sqlsaber-sandbox[modal]	MODAL_TOKEN_ID and MODAL_TOKEN_SECRET
e2b	sqlsaber-sandbox[e2b]	E2B_API_KEY
daytona	sqlsaber-sandbox[daytona]	DAYTONA_API_KEY
sprites	sqlsaber-sandbox[sprites]	SPRITES_TOKEN

Resources and usage limits

One immutable SandboxConfig applies to each session or capability. The main agent cannot change its provider, resources, or file limits through tool arguments.

image accepts a registry image for Docker, microsandbox, Modal, and Daytona. These providers default to DEFAULT_SANDBOX_IMAGE, a digest-pinned quay.io/jupyter/scipy-notebook image with NumPy, pandas, SciPy, scikit-learn, Matplotlib, and the Jupyter kernel packages. For E2B, image is a template ID. Sprites does not accept custom images.

Startup uses /opt/conda/bin/python when present, or resolves python from PATH. Compatible kernel packages are reused without running pip. If packages are missing or incompatible, startup installs them once in that session. Follow-up calls reuse the same environment.

cpu_cores and memory_mb set CPU and memory resources. Microsandbox and Sprites require whole CPU counts. Daytona requires whole CPU counts and memory in multiples of 1024 MiB. E2B takes both settings from its template.

gpu is a provider-specific string:

• Modal accepts a GPU selector such as "L4".
• Docker accepts a --gpus selector such as "all". The host needs a GPU and the NVIDIA container runtime.
• Daytona accepts a positive GPU-count string such as "1".

E2B, microsandbox, and Sprites reject GPU requests. Python code must use a GPU-enabled library to run computations on an allocated GPU.

Workspace, artifact, and model-preview limits are separate. Increasing max_artifact_bytes does not increase max_output_chars, max_image_bytes, or max_history_image_bytes.

Each direct analyze(..., usage_limits=...) call starts a fresh usage count unless you pass parent_usage. Managed analysis shares the current saber.query() budget. Independent sessions can run concurrently. Calls within one session run sequentially, as do managed calls that share a finite request or tool-call budget.

Session lifetime

Sessions start their sandbox on first use. idle_seconds measures time between operations, excluding active model calls and file transfers. max_lifetime_seconds limits the total session lifetime.

With no configured lifetime, E2B uses one hour and Modal uses 24 hours. Other providers retain their default lifetime. A host crash can prevent cleanup, so an environment can remain billable until its provider stops it.

Sessions do not recover after a host process restart. Errors can leave partial changes in Python state. SQLsaber does not replay code when an execution outcome is unknown.

Managed session IDs are scoped to a conversation. Sandbox code has outbound network access and can modify its input files.

Build a plugin capability

Third-party plugins expose a factory through the sqlsaber.capabilities entry-point group:

[project.entry-points."sqlsaber.capabilities"]
my_plugin = "my_plugin:capability"

from pydantic_ai.capabilities import Capability
from sqlsaber.capabilities.plugins import PluginContext


def capability(context: PluginContext):
    return Capability(
        id="my-plugin",
        description="Use my application's specialist tool.",
        tools=[my_tool],
    )

PluginContext supplies the database registry, knowledge manager, dangerous-mode flag, normalized tool model overrides, the required shared QueryResultStore, and the optional application-owned artifact store. A factory may return one capability or a sequence of capabilities. Managed SQLSaber invokes the factories you list. It does not scan installed entry points. On a host agent that already has a PluginContext, construct the plugin classes you want, such as Notebook(context) and Visualization(context).