deep-db-agents 0.1.0__py3-none-any.whl

deep_db_agents/__init__.py

@@ -0,0 +1,36 @@
+"""deep-db-agents: factory for Deep Agents specialized on different databases."""
+
+from __future__ import annotations
+
+from .connection import ConnectionConfig
+from .exceptions import (
+    DeepDbAgentError,
+    GuardrailError,
+    InvalidDbUrlError,
+    InvalidMultiAgentConfigError,
+    QueryNotAllowedError,
+    UnsupportedSchemeError,
+)
+from .factory import create_db_agent, create_deep_db_agents, create_deep_db_multi_agent
+from .guardrails import GuardrailConfig
+from .observability import SessionMetrics, configure_logging
+from .registry import available_schemes
+
+__all__ = [
+    "create_deep_db_agents",
+    "create_deep_db_multi_agent",
+    "create_db_agent",
+    "GuardrailConfig",
+    "ConnectionConfig",
+    "SessionMetrics",
+    "configure_logging",
+    "available_schemes",
+    "DeepDbAgentError",
+    "InvalidDbUrlError",
+    "UnsupportedSchemeError",
+    "InvalidMultiAgentConfigError",
+    "QueryNotAllowedError",
+    "GuardrailError",
+]
+
+__version__ = "0.1.0"

deep_db_agents/backend_registry.py

@@ -0,0 +1,96 @@
+"""Singleton registry of ``BackendProtocol`` backends indexed by UUID.
+
+Allows registering a backend instance (derived from deepagents' ``BackendProtocol``),
+obtaining an opaque identifier that can later be used to retrieve or remove it,
+without letting the reference to the object circulate through the system.
+
+The registry is shared at the process level: whoever registers a backend (``add``)
+is responsible for removing it (``remove``) at the end of the agent's lifetime,
+otherwise the instance stays referenced for the entire process lifetime.
+"""
+
+from __future__ import annotations
+
+import threading
+import uuid
+
+from deepagents.backends.protocol import BackendProtocol
+
+
+class BERegistry:
+    """Singleton ``UUID -> BackendProtocol`` registry.
+
+    Every instantiation returns the same object, so the registry is shared at the
+    process level. Creation and mutations are protected by a lock: tools may run
+    in parallel threads (parallel tool calls in the tool node).
+
+    >>> reg = BERegistry()
+    >>> key = reg.add(my_backend)
+    >>> reg.get(key) is my_backend
+    True
+    >>> reg.remove(key)
+    >>> reg.get(key) is None
+    True
+    """
+
+    _instance: BERegistry | None = None
+    _instance_lock = threading.Lock()
+
+    def __new__(cls) -> BERegistry:
+        """Create or return the singleton instance.
+
+        Returns:
+            BERegistry: The single shared instance of the registry.
+        """
+        if cls._instance is None:
+            with cls._instance_lock:
+                if cls._instance is None:  # double-checked: avoids duplicate instances
+                    instance = super().__new__(cls)
+                    instance._backends = {}
+                    instance._lock = threading.Lock()
+                    cls._instance = instance
+        return cls._instance
+
+    def add(self, backend: BackendProtocol) -> str:
+        """Register a backend and return its identifier.
+
+        Args:
+            backend: The ``BackendProtocol`` instance to register.
+
+        Returns:
+            str: The UUID (as a string) associated with the registered backend.
+
+        Raises:
+            TypeError: If ``backend`` is not a ``BackendProtocol`` instance.
+        """
+        if not isinstance(backend, BackendProtocol):
+            raise TypeError(f"Expected a BackendProtocol, got {type(backend).__name__!r}.")
+        key = str(uuid.uuid4())
+        with self._lock:
+            self._backends[key] = backend
+        return key
+
+    def get(self, key: str) -> BackendProtocol | None:
+        """Retrieve a registered backend.
+
+        Args:
+            key: The UUID string previously returned by :meth:`add`.
+
+        Returns:
+            BackendProtocol | None: The registered backend, or ``None`` if
+            ``key`` is not registered.
+        """
+        with self._lock:
+            return self._backends.get(key)
+
+    def remove(self, key: str) -> None:
+        """Remove a registered backend.
+
+        Args:
+            key: The UUID string previously returned by :meth:`add`.
+
+        Returns:
+            None. This is a no-op if ``key`` does not exist.
+        """
+        with self._lock:
+            self._backends.pop(key, None)

deep_db_agents/base.py

@@ -0,0 +1,55 @@
+"""Abstract interface that every database dialect must implement."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+
+from langchain_core.tools import BaseTool
+
+from .connection import ConnectionConfig
+from .guardrails import GuardrailConfig
+from .observability import SessionMetrics
+
+
+class DbDialect(ABC):
+    """Specialization of the agent for a database type.
+
+    A dialect provides two things: the **generic system prompt** that instructs the
+    agent on how to operate efficiently on that database, and the set of **tools**
+    (with credentials injected) to interact with it.
+    """
+
+    #: URL schemes handled by this dialect (e.g. ``("postgres", "postgresql")``).
+    schemes: tuple[str, ...] = ()
+
+    #: Optional session counters, injected by the factory; tools update them if present.
+    _metrics: SessionMetrics | None = None
+
+    @abstractmethod
+    def system_prompt(self) -> str:
+        """Return the generic prompt for this database, concatenated with the user's own.
+
+        Returns:
+            str: The dialect-specific system prompt text.
+        """
+
+    @abstractmethod
+    def build_tools(
+        self,
+        conn: ConnectionConfig,
+        guardrails: GuardrailConfig,
+        materialize_enable: bool = False,
+    ) -> Sequence[BaseTool]:
+        """Build the LangChain tools bound to the connection and guardrails.
+
+        Args:
+            conn: Connection parameters (host/port/credentials or file path).
+            guardrails: Hard safety thresholds enforced by the tool wrappers.
+            materialize_enable: Whether to expose the tool(s) that materialize large
+                results to file instead of returning them inline.
+
+        Returns:
+            Sequence[BaseTool]: The tools built for this dialect, with credentials
+            captured in their closures.
+        """

deep_db_agents/connection.py

@@ -0,0 +1,57 @@
+"""Connection configuration passed to dialects."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ConnectionConfig:
+    """Connection parameters for a database.
+
+    ``credential`` is a free-form dictionary whose content depends on the specific DB
+    (e.g. ``{"user": ..., "password": ...}`` or ``{"secret_key": ...}``). Well-known keys
+    (``database``/``db``) are exposed as convenience properties.
+
+    Attributes:
+        scheme: The URL scheme identifying the dialect (e.g. ``"postgres"``).
+        host: Hostname for network databases, or ``None`` for file-based databases.
+        port: Port for network databases, or ``None`` for file-based databases.
+        credential: Free-form credential/connection dictionary.
+        path: File/directory path for file-based databases (``sqlite``/``duckdb``);
+            ``None`` for network databases. ``:memory:`` denotes an in-memory database.
+    """
+
+    scheme: str
+    host: str | None
+    port: int | None
+    credential: dict[str, Any] = field(default_factory=dict)
+    #: File/directory path for file-based databases (``sqlite``/``duckdb``); ``None`` for
+    #: network databases. ``:memory:`` denotes an in-memory database.
+    path: str | None = None
+
+    @property
+    def database(self) -> str | None:
+        """Return the logical database/schema name, if provided in the credentials.
+
+        Returns:
+            str | None: The value of the ``database`` or ``db`` credential key, or
+            ``None`` if neither is present.
+        """
+        return self.credential.get("database") or self.credential.get("db")
+
+    def __repr__(self) -> str:
+        """Return a repr with credentials masked.
+
+        Returns:
+            str: A string representation of this config where ``credential`` content
+            is masked so passwords never leak into logs, tracebacks, or error messages.
+        """
+        # Mask the contents of ``credential`` so passwords never end up in logs,
+        # tracebacks, or error messages when the object is printed.
+        cred = "{…}" if self.credential else "{}"
+        return (
+            f"ConnectionConfig(scheme={self.scheme!r}, host={self.host!r}, "
+            f"port={self.port!r}, credential={cred}, path={self.path!r})"
+        )

deep_db_agents/dialects/__init__.py

@@ -0,0 +1,31 @@
+"""Imports all dialects to populate the registry when the package is imported.
+
+The side-effect import is intentional: each module registers its own dialect via the
+``@register`` decorator.
+"""
+
+from __future__ import annotations
+
+from . import (  # noqa: F401
+    duckdb,
+    elasticsearch,
+    mariadb,
+    mongodb,
+    mysql,
+    neo4j,
+    opensearch,
+    postgres,
+    sqlite,
+)
+
+__all__ = [
+    "mysql",
+    "mariadb",
+    "postgres",
+    "mongodb",
+    "neo4j",
+    "sqlite",
+    "duckdb",
+    "elasticsearch",
+    "opensearch",
+]

deep_db_agents/dialects/duckdb/__init__.py

@@ -0,0 +1,7 @@
+"""DuckDB dialect."""
+
+from __future__ import annotations
+
+from .dialect import DuckDBDialect
+
+__all__ = ["DuckDBDialect"]

deep_db_agents/dialects/duckdb/dialect.py

@@ -0,0 +1,138 @@
+"""Full DuckDB dialect (``duckdb`` package).
+
+DuckDB is a file-based analytical SQL database: the URL ``duckdb://<path>`` indicates the
+file (SQLAlchemy style: ``duckdb:///rel.duckdb`` relative to the working dir,
+``duckdb:////abs.duckdb`` absolute) or a **folder** (trailing slash, e.g. ``duckdb:///lake/``)
+interpreted as a data lake: the parquet/csv/json files inside it are exposed as tables. It
+reuses ``SqlDialect``'s logic; the execution timeout is enforced by ``FileSqlDialect``'s
+watchdog.
+
+In data lake mode the views are created only once per dialect instance (a snapshot of the
+folder taken on the first call). Warning: with ``duckdb://:memory:`` every tool call opens a
+brand-new, empty database (the connection is per-call).
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any
+
+from ...connection import ConnectionConfig
+from ...guardrails import GuardrailConfig
+from ...registry import register
+from ..sql_base import FileSqlDialect
+from . import tools
+from .prompt import DUCKDB_SYSTEM_PROMPT
+
+
+@register("duckdb")
+class DuckDBDialect(FileSqlDialect):
+    """Agent specialized on DuckDB."""
+
+    schemes = ("duckdb",)
+
+    def __init__(self) -> None:
+        """Initializes the dialect with an empty data-lake connection cache."""
+        # Data-lake connection cache (per dialect instance, i.e. per agent): the folder glob
+        # and the view creation happen only once.
+        self._datalake_con: Any = None
+        self._datalake_folder: str | None = None
+        self._datalake_lock = threading.Lock()
+
+    def system_prompt(self) -> str:
+        """Returns the DuckDB-specific system prompt.
+
+        Returns:
+            The system prompt text for the DuckDB agent.
+        """
+        return DUCKDB_SYSTEM_PROMPT
+
+    def _connect(self, conn: ConnectionConfig, guardrails: GuardrailConfig):
+        """Opens a DuckDB connection, either to a file or to a cached data-lake view set.
+
+        Args:
+            conn: Connection configuration; ``conn.path`` is either a database file path or a
+                data-lake folder (trailing slash / existing directory).
+            guardrails: Guardrail configuration (unused here; the execution timeout is
+                enforced separately by the watchdog).
+
+        Returns:
+            A DuckDB connection (or cursor, in data-lake mode) ready for query execution.
+        """
+        if self.is_directory(conn):
+            folder = self.resolve_path(conn.path)
+            with self._datalake_lock:
+                if self._datalake_con is None or self._datalake_folder != folder:
+                    self._datalake_con = tools.connect_datalake(folder)
+                    self._datalake_folder = folder
+                # cursor() is an independent connection on the same in-memory DB: closing it
+                # at the end of the tool call does not close the parent connection that holds
+                # the views.
+                # Note: the schema is a snapshot of the folder taken on the first call; files
+                # added/removed during the session are not reflected.
+                return self._datalake_con.cursor()
+        path = self.resolve_path(conn.path)
+        # read_only allows concurrent connections (parallel tool calls); :memory: does not.
+        return tools.connect_file(path, read_only=path != ":memory:")
+
+    def _interrupt_target(self, connection: Any, cursor: Any) -> Any:
+        """Returns the object on which the timeout watchdog should call ``interrupt()``.
+
+        In DuckDB the cursor is an independent connection: the query runs there, not on the
+        parent connection.
+
+        Args:
+            connection: The parent DuckDB connection.
+            cursor: The DuckDB cursor (an independent connection) executing the query.
+
+        Returns:
+            The cursor, since that is where the running query lives.
+        """
+        return cursor
+
+    def _estimate_rows(self, cursor, sql: str) -> int:
+        """Estimates the number of rows a query would return.
+
+        Args:
+            cursor: An open DB-API cursor.
+            sql: The SELECT statement to estimate.
+
+        Returns:
+            The estimated row count from DuckDB's ``EXPLAIN`` plan, or 0 if it could not be
+            determined.
+        """
+        return tools.estimate_rows(cursor, sql)
+
+    def _list_tables_sql(self, conn: ConnectionConfig) -> str:
+        """Builds the SQL statement listing the tables of the current database.
+
+        Args:
+            conn: Connection configuration (unused; the query relies on
+                ``information_schema``, which is schema-agnostic).
+
+        Returns:
+            A SQL SELECT statement returning one table name per row.
+        """
+        return (
+            "SELECT table_name FROM information_schema.tables "
+            "WHERE table_schema NOT IN ('information_schema', 'pg_catalog') "
+            "ORDER BY table_name"
+        )
+
+    def _describe_table_sql(self, table: str) -> tuple[str, tuple[Any, ...]]:
+        """Builds the parameterized SQL statement describing a table's columns.
+
+        Args:
+            table: Name of the table to describe.
+
+        Returns:
+            A tuple ``(sql, params)`` where ``sql`` is a parameterized SELECT and ``params``
+            are the bind parameters (the table name).
+        """
+        return (
+            "SELECT column_name, data_type, is_nullable "
+            "FROM information_schema.columns "
+            "WHERE table_name = ? "
+            "ORDER BY ordinal_position",
+            (table,),
+        )

deep_db_agents/dialects/duckdb/prompt.py

@@ -0,0 +1,34 @@
+"""Generic system prompt for the DuckDB-specialized agent."""
+
+from __future__ import annotations
+
+from ...prompts.shared import CONTEXT_PRINCIPLES
+
+DUCKDB_SYSTEM_PROMPT = f"""\
+You are an expert DuckDB agent. You operate on a DuckDB database through a set of read-only \
+tools. The source can be a single ``.duckdb`` file or a "data lake" folder whose files \
+(parquet/csv/json) are exposed as tables. Your goal is to answer the user's questions \
+accurately while consuming as little context as possible.
+
+## Context management principles
+{CONTEXT_PRINCIPLES}
+
+## DuckDB specifics
+- Use DuckDB syntax (standard SQL, very close to PostgreSQL): double quotes for \
+  identifiers ("table"."column"), `LIMIT n OFFSET m`, analytical functions and `QUALIFY`.
+- DuckDB is columnar and optimized for analytics: ALWAYS push aggregations and filters into \
+  the engine (`GROUP BY`, `SUM`, `AVG`, window functions) instead of extracting and computing \
+  by hand.
+- In data lake mode the tables are views over the folder's files: use `list_tables` to \
+  discover which files are available and `describe_table` for the inferred columns.
+
+## Available tools
+- `list_tables`, `describe_table`: explore the schema (or the data lake files) before querying.
+- `count_rows`: count (with an optional filter) BEFORE extracting, to assess the volume.
+- `sample_rows`: small preview of a table.
+- `run_query`: runs a SELECT with forced LIMIT and pagination, plus an EXPLAIN estimate.
+- `materialize_query`: saves a large result to file (Parquet/CSV) and returns only metadata, \
+  preview and statistics — use it for analysis/charts on large volumes.
+
+Always proceed step by step: explore the schema, count, then extract or aggregate.
+"""

deep_db_agents/dialects/duckdb/tools.py

@@ -0,0 +1,135 @@
+"""Driver-specific helpers for DuckDB (``duckdb`` package)."""
+
+from __future__ import annotations
+
+import glob
+import os
+import re
+
+# File extension -> DuckDB reader function, for the data lake (path = folder).
+_READERS = {
+    ".parquet": "read_parquet",
+    ".csv": "read_csv_auto",
+    ".tsv": "read_csv_auto",
+    ".json": "read_json_auto",
+    ".ndjson": "read_json_auto",
+}
+
+# Captures cardinality estimates in the EXPLAIN plan text ("~N rows" or "EC: N").
+_ESTIMATE_RE = re.compile(r"(?:EC:\s*|~\s*)([\d,]+)")
+
+
+def _import_duckdb():
+    """Lazily imports the ``duckdb`` package.
+
+    Returns:
+        The imported ``duckdb`` module.
+
+    Raises:
+        ImportError: If the ``duckdb`` package is not installed.
+    """
+    try:
+        import duckdb
+    except ImportError as exc:  # pragma: no cover - depends on the installed extra
+        raise ImportError(
+            "The DuckDB dialect requires the 'duckdb' extra (pip install 'deep-db-agents[duckdb]')."
+        ) from exc
+    return duckdb
+
+
+def _sanitize_view_name(name: str) -> str:
+    """Turns a file name into a valid SQL identifier (for data-lake views).
+
+    Args:
+        name: The raw file name (without extension) to sanitize.
+
+    Returns:
+        A valid SQL identifier derived from ``name``, prefixed with ``_`` if it would
+        otherwise be empty or start with a digit.
+    """
+    safe = re.sub(r"\W", "_", name)
+    if not safe or safe[0].isdigit():
+        safe = "_" + safe
+    return safe
+
+
+def register_datalake_views(con, folder: str) -> list[str]:
+    """Creates a view for each data file (parquet/csv/json) in ``folder``.
+
+    Exposes the files as queryable tables with regular SQL, so the SQL dialect's tools work
+    unmodified. Each view is named after the (sanitized) file name without its extension.
+
+    Args:
+        con: An open DuckDB connection on which to create the views.
+        folder: Path to the data-lake folder to scan for data files.
+
+    Returns:
+        The list of created view names.
+    """
+    created: list[str] = []
+    for filepath in sorted(glob.glob(os.path.join(folder, "*"))):
+        ext = os.path.splitext(filepath)[1].lower()
+        reader = _READERS.get(ext)
+        if not reader:
+            continue
+        view = _sanitize_view_name(os.path.splitext(os.path.basename(filepath))[0])
+        escaped = filepath.replace("'", "''")
+        con.execute(f"CREATE OR REPLACE VIEW \"{view}\" AS SELECT * FROM {reader}('{escaped}')")
+        created.append(view)
+    return created
+
+
+def connect_file(path: str, *, read_only: bool = True):
+    """Opens a connection to a DuckDB file.
+
+    ``read_only=True`` allows concurrent connections to the same file (needed when tool calls
+    run in parallel): in write mode DuckDB holds an exclusive lock. Query duration is bounded
+    separately by the watchdog calling ``interrupt()`` (see ``FileSqlDialect``).
+
+    Args:
+        path: Path to the DuckDB database file.
+        read_only: If True, opens the connection in read-only mode.
+
+    Returns:
+        An open DuckDB connection.
+    """
+    duckdb = _import_duckdb()
+    return duckdb.connect(path, read_only=read_only)
+
+
+def connect_datalake(folder: str):
+    """Opens an in-memory DuckDB connection with a view for each data file in the folder.
+
+    Args:
+        folder: Path to the data-lake folder to expose as tables.
+
+    Returns:
+        An in-memory DuckDB connection with one view registered per recognized data file.
+    """
+    duckdb = _import_duckdb()
+    con = duckdb.connect(":memory:")
+    register_datalake_views(con, folder)
+    return con
+
+
+def estimate_rows(cursor, sql: str) -> int:
+    """Estimates the row count by reading the cardinality from the ``EXPLAIN`` plan.
+
+    This is best-effort: any failure to run or parse ``EXPLAIN`` results in an estimate of 0
+    rather than raising.
+
+    Args:
+        cursor: An open DuckDB cursor.
+        sql: The SELECT statement to estimate.
+
+    Returns:
+        The largest cardinality estimate found in the EXPLAIN plan text, or 0 if none could
+        be extracted.
+    """
+    try:
+        cursor.execute(f"EXPLAIN {sql}")
+        text = "\n".join(str(row[-1]) for row in cursor.fetchall())
+    except Exception:  # noqa: BLE001 - the estimate is best-effort, never blocking on its own
+        return 0
+    nums = [int(m.replace(",", "")) for m in _ESTIMATE_RE.findall(text)]
+    return max(nums) if nums else 0

deep_db_agents/dialects/elasticsearch/__init__.py

@@ -0,0 +1,7 @@
+"""Elasticsearch dialect."""
+
+from __future__ import annotations
+
+from .dialect import ElasticsearchDialect
+
+__all__ = ["ElasticsearchDialect"]

deep_db_agents/dialects/elasticsearch/dialect.py

@@ -0,0 +1,51 @@
+"""Complete Elasticsearch dialect (official ``elasticsearch`` driver).
+
+Reuses the entire tool logic from :mod:`..search_base` (shared by Elasticsearch and
+OpenSearch); this module only provides the driver-specific connection opening.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ...connection import ConnectionConfig
+from ...guardrails import GuardrailConfig
+from ...registry import register
+from ..search_base import SearchDialect
+from . import tools
+from .prompt import ELASTICSEARCH_SYSTEM_PROMPT
+
+
+@register("elasticsearch")
+class ElasticsearchDialect(SearchDialect):
+    """Agent specialized on Elasticsearch.
+
+    Access is restricted to the index(es) configured in ``credential["index"]``
+    (single name, CSV, or a ``*`` pattern); every tool inherited from
+    :class:`SearchDialect` validates the requested index against this scope before
+    querying the cluster.
+    """
+
+    schemes = ("elasticsearch",)
+
+    def system_prompt(self) -> str:
+        """Return the Elasticsearch-specific system prompt.
+
+        Returns:
+            str: The system prompt text for the Elasticsearch agent.
+        """
+        return ELASTICSEARCH_SYSTEM_PROMPT
+
+    def _connect(self, conn: ConnectionConfig, guardrails: GuardrailConfig | None = None) -> Any:
+        """Open the driver connection using the given configuration and guardrails.
+
+        Args:
+            conn: Connection configuration (host, port, credentials) for the cluster.
+            guardrails: Optional guardrail configuration; its ``query_timeout_s``, if
+                set, is forwarded as the request timeout for the driver client.
+
+        Returns:
+            Any: An initialized Elasticsearch client instance.
+        """
+        timeout = guardrails.query_timeout_s if guardrails else None
+        return tools.connect(conn, request_timeout=timeout)

deep_db_agents/dialects/elasticsearch/prompt.py

@@ -0,0 +1,28 @@
+"""Generic system prompt for the Elasticsearch-specialized agent."""
+
+from __future__ import annotations
+
+from ...prompts.shared import CONTEXT_PRINCIPLES
+
+ELASTICSEARCH_SYSTEM_PROMPT = f"""\
+You are an expert Elasticsearch agent. You operate on an Elasticsearch cluster through
+read-only tools (Query DSL). Your goal is to answer while consuming as little context as
+possible.
+
+## Context management principles
+{CONTEXT_PRINCIPLES}
+
+## Elasticsearch specifics
+- Push work into the cluster: use `count` to size a query before searching, and use
+  `aggregate` (terms/metric aggregations) to summarize data instead of downloading documents
+  with `run_query`/`search_query` and aggregating them by hand.
+- Always request only the fields you need and rely on the forced `size`/`from` pagination of
+  the search tools instead of asking for everything at once.
+- Explore the index scope first with `list_indices` and `describe_index` (field mapping)
+  before writing a Query DSL clause.
+- For simple text/filter searches prefer `search_query` (Lucene-like `query_string` syntax,
+  e.g. 'status:active AND price:[10 TO 50]'); use `run_query` with a full Query DSL clause
+  for anything more structured (bool/match/range/term...).
+- You can only operate on the index(es) configured for this agent; requests for indices
+  outside that scope are rejected by the tools.
+"""