Skip to content

Dialects

Each dialect module wires up the tools and prompt fragment for a specific database backend.

deep_db_agents.dialects

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.

duckdb

DuckDB dialect.

DuckDBDialect

DuckDBDialect()

Bases: FileSqlDialect

Agent specialized on DuckDB.

Initializes the dialect with an empty data-lake connection cache.

Source code in src/deep_db_agents/dialects/duckdb/dialect.py
34
35
36
37
38
39
40
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()
system_prompt
system_prompt() -> str

Returns the DuckDB-specific system prompt.

Returns:

Type Description
str

The system prompt text for the DuckDB agent.

Source code in src/deep_db_agents/dialects/duckdb/dialect.py
42
43
44
45
46
47
48
def system_prompt(self) -> str:
    """Returns the DuckDB-specific system prompt.

    Returns:
        The system prompt text for the DuckDB agent.
    """
    return DUCKDB_SYSTEM_PROMPT

dialect

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).

DuckDBDialect
DuckDBDialect()

Bases: FileSqlDialect

Agent specialized on DuckDB.

Initializes the dialect with an empty data-lake connection cache.

Source code in src/deep_db_agents/dialects/duckdb/dialect.py
34
35
36
37
38
39
40
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()
system_prompt
system_prompt() -> str

Returns the DuckDB-specific system prompt.

Returns:

Type Description
str

The system prompt text for the DuckDB agent.

Source code in src/deep_db_agents/dialects/duckdb/dialect.py
42
43
44
45
46
47
48
def system_prompt(self) -> str:
    """Returns the DuckDB-specific system prompt.

    Returns:
        The system prompt text for the DuckDB agent.
    """
    return DUCKDB_SYSTEM_PROMPT

prompt

Generic system prompt for the DuckDB-specialized agent.

tools

Driver-specific helpers for DuckDB (duckdb package).

connect_datalake
connect_datalake(folder: str)

Opens an in-memory DuckDB connection with a view for each data file in the folder.

Parameters:

Name Type Description Default
folder str

Path to the data-lake folder to expose as tables.

required

Returns:

Type Description

An in-memory DuckDB connection with one view registered per recognized data file.

Source code in src/deep_db_agents/dialects/duckdb/tools.py
100
101
102
103
104
105
106
107
108
109
110
111
112
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
connect_file
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).

Parameters:

Name Type Description Default
path str

Path to the DuckDB database file.

required
read_only bool

If True, opens the connection in read-only mode.

True

Returns:

Type Description

An open DuckDB connection.

Source code in src/deep_db_agents/dialects/duckdb/tools.py
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
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)
estimate_rows
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.

Parameters:

Name Type Description Default
cursor

An open DuckDB cursor.

required
sql str

The SELECT statement to estimate.

required

Returns:

Type Description
int

The largest cardinality estimate found in the EXPLAIN plan text, or 0 if none could

int

be extracted.

Source code in src/deep_db_agents/dialects/duckdb/tools.py
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
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
register_datalake_views
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.

Parameters:

Name Type Description Default
con

An open DuckDB connection on which to create the views.

required
folder str

Path to the data-lake folder to scan for data files.

required

Returns:

Type Description
list[str]

The list of created view names.

Source code in src/deep_db_agents/dialects/duckdb/tools.py
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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

elasticsearch

Elasticsearch dialect.

ElasticsearchDialect

ElasticsearchDialect()

Bases: 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.

Source code in src/deep_db_agents/dialects/search_base.py
171
172
173
174
175
def __init__(self) -> None:
    """Initializes the dialect with an empty, lazily-populated client cache."""
    # The ES/OS client internally manages a thread-safe connection pool: it is reused
    # for the whole lifetime of the agent instead of being recreated on every tool call.
    self._client_cache = LazyClient()
system_prompt
system_prompt() -> str

Return the Elasticsearch-specific system prompt.

Returns:

Name Type Description
str str

The system prompt text for the Elasticsearch agent.

Source code in src/deep_db_agents/dialects/elasticsearch/dialect.py
31
32
33
34
35
36
37
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

dialect

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.

ElasticsearchDialect
ElasticsearchDialect()

Bases: 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.

Source code in src/deep_db_agents/dialects/search_base.py
171
172
173
174
175
def __init__(self) -> None:
    """Initializes the dialect with an empty, lazily-populated client cache."""
    # The ES/OS client internally manages a thread-safe connection pool: it is reused
    # for the whole lifetime of the agent instead of being recreated on every tool call.
    self._client_cache = LazyClient()
system_prompt
system_prompt() -> str

Return the Elasticsearch-specific system prompt.

Returns:

Name Type Description
str str

The system prompt text for the Elasticsearch agent.

Source code in src/deep_db_agents/dialects/elasticsearch/dialect.py
31
32
33
34
35
36
37
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

prompt

Generic system prompt for the Elasticsearch-specialized agent.

tools

Driver-specific helpers for Elasticsearch (official elasticsearch driver).

connect
connect(conn: ConnectionConfig, *, request_timeout: int | None = None) -> Any

Open an Elasticsearch client from the given credentials.

The driver is imported lazily so the package can be imported without the elasticsearch extra installed.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection configuration with host, port, and credentials. Recognized credential keys: use_ssl, verify_certs, ca_certs, api_key, user, password.

required
request_timeout int | None

Optional request timeout in seconds; forwarded to the client only if positive.

None

Returns:

Name Type Description
Any Any

An initialized Elasticsearch client instance.

Raises:

Type Description
ImportError

If the elasticsearch package is not installed.

Source code in src/deep_db_agents/dialects/elasticsearch/tools.py
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def connect(conn: ConnectionConfig, *, request_timeout: int | None = None) -> Any:
    """Open an Elasticsearch client from the given credentials.

    The driver is imported lazily so the package can be imported without the
    ``elasticsearch`` extra installed.

    Args:
        conn: Connection configuration with host, port, and credentials. Recognized
            ``credential`` keys: ``use_ssl``, ``verify_certs``, ``ca_certs``,
            ``api_key``, ``user``, ``password``.
        request_timeout: Optional request timeout in seconds; forwarded to the
            client only if positive.

    Returns:
        Any: An initialized ``Elasticsearch`` client instance.

    Raises:
        ImportError: If the ``elasticsearch`` package is not installed.
    """
    try:
        from elasticsearch import Elasticsearch
    except ImportError as exc:  # pragma: no cover - depends on the installed extra
        raise ImportError(
            "The Elasticsearch dialect requires the 'elasticsearch' extra "
            "(pip install 'deep-db-agents[elasticsearch]')."
        ) from exc

    cred = conn.credential
    host = conn.host or "localhost"
    port = conn.port or 9200
    scheme = "https" if cred.get("use_ssl") else "http"
    kwargs: dict[str, Any] = {
        "hosts": [f"{scheme}://{host}:{port}"],
        "verify_certs": cred.get("verify_certs", True),
    }
    if request_timeout and request_timeout > 0:
        kwargs["request_timeout"] = request_timeout
    if cred.get("ca_certs"):
        kwargs["ca_certs"] = cred["ca_certs"]
    if cred.get("api_key"):
        kwargs["api_key"] = cred["api_key"]
    elif cred.get("user"):
        kwargs["basic_auth"] = (cred["user"], cred.get("password", ""))
    return Elasticsearch(**kwargs)

mariadb

MariaDB dialect.

MariaDBDialect

Bases: MySQLDialect

Agent specialized on MariaDB.

system_prompt
system_prompt() -> str

Return the MariaDB-specific system prompt.

Returns:

Type Description
str

The full system prompt text for the MariaDB agent.

Source code in src/deep_db_agents/dialects/mariadb/dialect.py
22
23
24
25
26
27
28
def system_prompt(self) -> str:
    """Return the MariaDB-specific system prompt.

    Returns:
        The full system prompt text for the MariaDB agent.
    """
    return MARIADB_SYSTEM_PROMPT

dialect

Full MariaDB dialect.

MariaDB is compatible with the MySQL protocol, so it reuses the pymysql driver and most of MySQLDialect's logic. It differs in the query timeout mechanism: MariaDB uses the max_statement_time session variable (in seconds), while MySQL uses MAX_EXECUTION_TIME (in milliseconds).

MariaDBDialect

Bases: MySQLDialect

Agent specialized on MariaDB.

system_prompt
system_prompt() -> str

Return the MariaDB-specific system prompt.

Returns:

Type Description
str

The full system prompt text for the MariaDB agent.

Source code in src/deep_db_agents/dialects/mariadb/dialect.py
22
23
24
25
26
27
28
def system_prompt(self) -> str:
    """Return the MariaDB-specific system prompt.

    Returns:
        The full system prompt text for the MariaDB agent.
    """
    return MARIADB_SYSTEM_PROMPT

prompt

Generic system prompt for the MariaDB-specialized agent.

mongodb

MongoDB dialect (stub).

MongoDBDialect

MongoDBDialect()

Bases: DbDialect

Agent specialized on MongoDB.

Source code in src/deep_db_agents/dialects/mongodb/dialect.py
34
35
36
def __init__(self) -> None:
    # The MongoClient is a thread-safe pool: reuse it for the agent's whole lifetime.
    self._client_cache = LazyClient()
close
close() -> None

Closes the reused MongoClient. Call at the end of the agent's lifetime.

Source code in src/deep_db_agents/dialects/mongodb/dialect.py
41
42
43
def close(self) -> None:
    """Closes the reused MongoClient. Call at the end of the agent's lifetime."""
    self._client_cache.close()

dialect

Full MongoDB dialect (pymongo driver).

Carries the context-management principles over to the document model: exploration of collections and inferred schema, counting before extraction, read-only find/aggregate with forced limit and projection, and file materialization of large results.

MongoDBDialect
MongoDBDialect()

Bases: DbDialect

Agent specialized on MongoDB.

Source code in src/deep_db_agents/dialects/mongodb/dialect.py
34
35
36
def __init__(self) -> None:
    # The MongoClient is a thread-safe pool: reuse it for the agent's whole lifetime.
    self._client_cache = LazyClient()
close
close() -> None

Closes the reused MongoClient. Call at the end of the agent's lifetime.

Source code in src/deep_db_agents/dialects/mongodb/dialect.py
41
42
43
def close(self) -> None:
    """Closes the reused MongoClient. Call at the end of the agent's lifetime."""
    self._client_cache.close()

prompt

Generic system prompt for the agent specialized on MongoDB.

tools

Driver-specific helpers for MongoDB (pymongo).

connect
connect(conn: ConnectionConfig, *, socket_timeout_ms: int | None = None)

Opens a MongoClient from the credentials. Lazy import of the driver.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection configuration (host, port, credentials).

required
socket_timeout_ms int | None

Timeout enforced on network operations (socket read/write): without it, a find/aggregate whose socket blocks would hang forever even after the server has been selected. Should be derived from guardrails.query_timeout_s to align with the behavior of the SQL dialects.

None

Returns:

Type Description

A connected pymongo.MongoClient instance.

Raises:

Type Description
ImportError

If the mongodb extra is not installed.

Source code in src/deep_db_agents/dialects/mongodb/tools.py
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def connect(conn: ConnectionConfig, *, socket_timeout_ms: int | None = None):
    """Opens a MongoClient from the credentials. Lazy import of the driver.

    Args:
        conn: Connection configuration (host, port, credentials).
        socket_timeout_ms: Timeout enforced on network operations (socket read/write):
            without it, a ``find``/``aggregate`` whose socket blocks would hang forever even
            after the server has been selected. Should be derived from
            ``guardrails.query_timeout_s`` to align with the behavior of the SQL dialects.

    Returns:
        A connected ``pymongo.MongoClient`` instance.

    Raises:
        ImportError: If the ``mongodb`` extra is not installed.
    """
    try:
        import pymongo
    except ImportError as exc:  # pragma: no cover - depends on the installed extra
        raise ImportError(
            "The MongoDB dialect requires the 'mongodb' extra "
            "(pip install 'deep-db-agents[mongodb]')."
        ) from exc

    cred = conn.credential
    connect_timeout_ms = int(cred.get("connect_timeout", 10)) * 1000
    kwargs: dict[str, Any] = {
        "host": conn.host or "localhost",
        "port": conn.port or 27017,
        "serverSelectionTimeoutMS": connect_timeout_ms,
        "connectTimeoutMS": connect_timeout_ms,
    }
    if socket_timeout_ms and socket_timeout_ms > 0:
        kwargs["socketTimeoutMS"] = socket_timeout_ms
    if cred.get("user"):
        kwargs["username"] = cred["user"]
    if cred.get("password"):
        kwargs["password"] = cred["password"]
    auth_source = cred.get("authSource") or cred.get("auth_source")
    if auth_source:
        kwargs["authSource"] = auth_source
    return pymongo.MongoClient(**kwargs)
docs_to_table
docs_to_table(docs: list[dict]) -> tuple[list[str], list[list[Any]]]

Converts heterogeneous documents into (columns, rows) for file materialization.

Parameters:

Name Type Description Default
docs list[dict]

The list of documents to convert.

required

Returns:

Type Description
tuple[list[str], list[list[Any]]]

A tuple of (column names, row values).

Source code in src/deep_db_agents/dialects/mongodb/tools.py
175
176
177
178
179
180
181
182
183
184
def docs_to_table(docs: list[dict]) -> tuple[list[str], list[list[Any]]]:
    """Converts heterogeneous documents into (columns, rows) for file materialization.

    Args:
        docs: The list of documents to convert.

    Returns:
        A tuple of (column names, row values).
    """
    return tabular.docs_to_table(docs, scalar=_scalar)
ensure_read_only_filter
ensure_read_only_filter(filter_: Any) -> Any

Validates a filter/projection: no server-side JS operator ($where/$function).

Parameters:

Name Type Description Default
filter_ Any

The filter or projection document to validate.

required

Returns:

Type Description
Any

The same value, unchanged, if validation passes.

Raises:

Type Description
QueryNotAllowedError

If a forbidden operator is found.

Source code in src/deep_db_agents/dialects/mongodb/tools.py
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def ensure_read_only_filter(filter_: Any) -> Any:
    """Validates a filter/projection: no server-side JS operator (``$where``/``$function``).

    Args:
        filter_: The filter or projection document to validate.

    Returns:
        The same value, unchanged, if validation passes.

    Raises:
        QueryNotAllowedError: If a forbidden operator is found.
    """
    _reject_forbidden_keys(filter_, what="filter")
    return filter_
ensure_read_only_pipeline
ensure_read_only_pipeline(pipeline: Any) -> list

Validates the pipeline: a list with no write stages ($out/$merge) nor JS operators, at any depth (even nested in $facet/$lookup).

Parameters:

Name Type Description Default
pipeline Any

The aggregation pipeline to validate.

required

Returns:

Type Description
list

The same pipeline, unchanged, if validation passes.

Raises:

Type Description
QueryNotAllowedError

If pipeline is not a list, or contains a forbidden stage.

Source code in src/deep_db_agents/dialects/mongodb/tools.py
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def ensure_read_only_pipeline(pipeline: Any) -> list:
    """Validates the pipeline: a list with no write stages ($out/$merge) nor JS operators,
    at any depth (even nested in ``$facet``/``$lookup``).

    Args:
        pipeline: The aggregation pipeline to validate.

    Returns:
        The same pipeline, unchanged, if validation passes.

    Raises:
        QueryNotAllowedError: If ``pipeline`` is not a list, or contains a forbidden stage.
    """
    if not isinstance(pipeline, list):
        raise QueryNotAllowedError("The aggregation pipeline must be a list of stages.")
    _reject_forbidden_keys(pipeline, what="pipeline")
    return pipeline
parse_json
parse_json(text: str | None, *, what: str = 'argument') -> Any

Decodes JSON (extended, with ObjectId/date support) or returns an empty default.

Parameters:

Name Type Description Default
text str | None

The JSON string to decode, or None/empty string.

required
what str

Description of the value being parsed, used in the error message.

'argument'

Returns:

Type Description
Any

The decoded value, or None if text is empty.

Raises:

Type Description
QueryNotAllowedError

If text is not valid JSON.

Source code in src/deep_db_agents/dialects/mongodb/tools.py
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def parse_json(text: str | None, *, what: str = "argument") -> Any:
    """Decodes JSON (extended, with ObjectId/date support) or returns an empty default.

    Args:
        text: The JSON string to decode, or ``None``/empty string.
        what: Description of the value being parsed, used in the error message.

    Returns:
        The decoded value, or ``None`` if ``text`` is empty.

    Raises:
        QueryNotAllowedError: If ``text`` is not valid JSON.
    """
    if text is None or (isinstance(text, str) and not text.strip()):
        return None
    if not isinstance(text, str):
        return text
    try:
        from bson import json_util

        return json_util.loads(text)
    except ImportError:  # pragma: no cover - bson is bundled with pymongo
        return json.loads(text)
    except ValueError as exc:
        raise QueryNotAllowedError(f"Invalid {what} JSON: {exc}") from exc

mysql

MySQL dialect.

MySQLDialect

Bases: SqlDialect

Agent specialized on MySQL.

system_prompt
system_prompt() -> str

Return the MySQL-specific system prompt.

Returns:

Type Description
str

The full system prompt text for the MySQL agent.

Source code in src/deep_db_agents/dialects/mysql/dialect.py
22
23
24
25
26
27
28
def system_prompt(self) -> str:
    """Return the MySQL-specific system prompt.

    Returns:
        The full system prompt text for the MySQL agent.
    """
    return MYSQL_SYSTEM_PROMPT

dialect

Full MySQL dialect (pymysql driver).

MySQLDialect

Bases: SqlDialect

Agent specialized on MySQL.

system_prompt
system_prompt() -> str

Return the MySQL-specific system prompt.

Returns:

Type Description
str

The full system prompt text for the MySQL agent.

Source code in src/deep_db_agents/dialects/mysql/dialect.py
22
23
24
25
26
27
28
def system_prompt(self) -> str:
    """Return the MySQL-specific system prompt.

    Returns:
        The full system prompt text for the MySQL agent.
    """
    return MYSQL_SYSTEM_PROMPT

prompt

Generic system prompt for the MySQL-specialized agent.

tools

MySQL-specific driver helpers (pymysql).

connect
connect(conn: ConnectionConfig, *, network_timeout_s: int | None = None)

Open a pymysql connection from credentials, with lazy driver import.

network_timeout_s enforces a timeout on socket read/write: without it, a query whose socket stalls would hang forever (connect_timeout only covers connection setup). It should be derived from guardrails.query_timeout_s.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters (host, port, credentials, database).

required
network_timeout_s int | None

Read/write socket timeout in seconds, or None to fall back to the read_timeout credential (if any).

None

Returns:

Type Description

An open pymysql connection.

Raises:

Type Description
ImportError

If the pymysql driver is not installed.

Source code in src/deep_db_agents/dialects/mysql/tools.py
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def connect(conn: ConnectionConfig, *, network_timeout_s: int | None = None):
    """Open a pymysql connection from credentials, with lazy driver import.

    ``network_timeout_s`` enforces a timeout on socket read/write: without it, a query
    whose socket stalls would hang forever (``connect_timeout`` only covers connection
    setup). It should be derived from ``guardrails.query_timeout_s``.

    Args:
        conn: Connection parameters (host, port, credentials, database).
        network_timeout_s: Read/write socket timeout in seconds, or ``None`` to fall
            back to the ``read_timeout`` credential (if any).

    Returns:
        An open pymysql connection.

    Raises:
        ImportError: If the ``pymysql`` driver is not installed.
    """
    try:
        import pymysql
    except ImportError as exc:  # pragma: no cover - depends on the installed extra
        raise ImportError(
            "The MySQL dialect requires the 'mysql' extra (pip install 'deep-db-agents[mysql]')."
        ) from exc

    cred = conn.credential
    read_timeout = network_timeout_s if network_timeout_s is not None else cred.get("read_timeout")
    return pymysql.connect(
        host=conn.host or "localhost",
        port=conn.port or 3306,
        user=cred.get("user"),
        password=cred.get("password", ""),
        database=conn.database,
        connect_timeout=cred.get("connect_timeout", 10),
        read_timeout=read_timeout,
        write_timeout=read_timeout,
        cursorclass=pymysql.cursors.Cursor,
    )
estimate_rows
estimate_rows(cursor, sql: str) -> int

Estimate the rows returned by sql using EXPLAIN (rows column).

Parameters:

Name Type Description Default
cursor

DB-API cursor to run the EXPLAIN on.

required
sql str

SQL query to estimate.

required

Returns:

Type Description
int

Estimated row count, or 0 if it cannot be determined.

Source code in src/deep_db_agents/dialects/mysql/tools.py
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def estimate_rows(cursor, sql: str) -> int:
    """Estimate the rows returned by ``sql`` using ``EXPLAIN`` (``rows`` column).

    Args:
        cursor: DB-API cursor to run the ``EXPLAIN`` on.
        sql: SQL query to estimate.

    Returns:
        Estimated row count, or 0 if it cannot be determined.
    """
    cursor.execute(f"EXPLAIN {sql}")
    cols = [d[0] for d in (cursor.description or [])]
    try:
        idx = cols.index("rows")
    except ValueError:
        return 0
    total = 0
    for row in cursor.fetchall():
        value = row[idx]
        if value is not None:
            total += int(value)
    return total

neo4j

Neo4j dialect (stub).

Neo4jDialect

Neo4jDialect()

Bases: DbDialect

Agent specialized on Neo4j.

Source code in src/deep_db_agents/dialects/neo4j/dialect.py
78
79
80
81
def __init__(self) -> None:
    # The Neo4j Driver is a thread-safe pool: reuse it for the agent's whole lifetime,
    # opening a session (not thread-safe, but cheap) per individual tool call.
    self._driver_cache = LazyClient()
close
close() -> None

Closes the reused Neo4j Driver. Call at the end of the agent's lifetime.

Source code in src/deep_db_agents/dialects/neo4j/dialect.py
86
87
88
def close(self) -> None:
    """Closes the reused Neo4j Driver. Call at the end of the agent's lifetime."""
    self._driver_cache.close()

dialect

Full Neo4j dialect (official neo4j driver).

Carries the context-management principles over to the graph: exploration of labels, relationship types and property keys, counting before extraction, execution of read-only Cypher with a consumer-side forced limit, and file materialization of large results. Queries run in read transactions (execute_read), which prevent server-side writes on top of the clause whitelist.

Neo4jDialect
Neo4jDialect()

Bases: DbDialect

Agent specialized on Neo4j.

Source code in src/deep_db_agents/dialects/neo4j/dialect.py
78
79
80
81
def __init__(self) -> None:
    # The Neo4j Driver is a thread-safe pool: reuse it for the agent's whole lifetime,
    # opening a session (not thread-safe, but cheap) per individual tool call.
    self._driver_cache = LazyClient()
close
close() -> None

Closes the reused Neo4j Driver. Call at the end of the agent's lifetime.

Source code in src/deep_db_agents/dialects/neo4j/dialect.py
86
87
88
def close(self) -> None:
    """Closes the reused Neo4j Driver. Call at the end of the agent's lifetime."""
    self._driver_cache.close()

prompt

Generic system prompt for the agent specialized on Neo4j.

tools

Driver-specific helpers for Neo4j (official neo4j driver).

connect
connect(conn: ConnectionConfig, *, network_timeout_s: int | None = None)

Opens a Neo4j driver from the credentials. Lazy import of the driver.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection configuration (host, port, credentials).

required
network_timeout_s int | None

Client-side timeout (derived from guardrails.query_timeout_s) limiting the wait for a connection from the pool and the total retry time of managed transactions: without it, execute_read could retry/wait indefinitely, blocking the whole tool node when tool calls run in parallel.

None

Returns:

Type Description

A connected neo4j.Driver instance.

Raises:

Type Description
ImportError

If the neo4j extra is not installed.

Source code in src/deep_db_agents/dialects/neo4j/tools.py
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def connect(conn: ConnectionConfig, *, network_timeout_s: int | None = None):
    """Opens a Neo4j driver from the credentials. Lazy import of the driver.

    Args:
        conn: Connection configuration (host, port, credentials).
        network_timeout_s: Client-side timeout (derived from
            ``guardrails.query_timeout_s``) limiting the wait for a connection from the pool
            and the total retry time of managed transactions: without it, ``execute_read``
            could retry/wait indefinitely, blocking the whole tool node when tool calls run
            in parallel.

    Returns:
        A connected ``neo4j.Driver`` instance.

    Raises:
        ImportError: If the ``neo4j`` extra is not installed.
    """
    try:
        import neo4j
    except ImportError as exc:  # pragma: no cover - depends on the installed extra
        raise ImportError(
            "The Neo4j dialect requires the 'neo4j' extra (pip install 'deep-db-agents[neo4j]')."
        ) from exc

    cred = conn.credential
    host = conn.host or "localhost"
    port = conn.port or 7687
    uri = f"neo4j://{host}:{port}"
    auth = (cred["user"], cred.get("password", "")) if cred.get("user") else None
    config: dict[str, Any] = {"connection_timeout": int(cred.get("connect_timeout", 10))}
    if network_timeout_s and network_timeout_s > 0:
        config["connection_acquisition_timeout"] = network_timeout_s
        config["max_transaction_retry_time"] = network_timeout_s
    return neo4j.GraphDatabase.driver(uri, auth=auth, **config)
ensure_read_only
ensure_read_only(cypher: str) -> str

Validates that the Cypher is a single read-only statement.

Parameters:

Name Type Description Default
cypher str

The raw Cypher statement submitted by the agent.

required

Returns:

Type Description
str

The stripped statement, if validation passes.

Raises:

Type Description
QueryNotAllowedError

If the statement is empty, contains multiple statements, or uses a forbidden write clause.

Source code in src/deep_db_agents/dialects/neo4j/tools.py
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def ensure_read_only(cypher: str) -> str:
    """Validates that the Cypher is a single read-only statement.

    Args:
        cypher: The raw Cypher statement submitted by the agent.

    Returns:
        The stripped statement, if validation passes.

    Raises:
        QueryNotAllowedError: If the statement is empty, contains multiple statements, or
            uses a forbidden write clause.
    """
    stripped = cypher.strip().rstrip(";").strip()
    if not stripped:
        raise QueryNotAllowedError("Empty Cypher query.")
    if ";" in stripped:
        raise QueryNotAllowedError("Only single Cypher statements are allowed: remove the ';'.")
    tokens = {t.upper() for t in _TOKEN_RE.findall(stripped)}
    forbidden = _WRITE_CLAUSES & tokens
    if forbidden:
        raise QueryNotAllowedError(
            f"Write clause not allowed: {', '.join(sorted(forbidden))}. "
            "Only read queries are allowed (MATCH/RETURN/WITH/CALL ... YIELD)."
        )
    return stripped
quote_label
quote_label(label: str) -> str

Escapes a Cypher label/identifier (doubled backticks).

Parameters:

Name Type Description Default
label str

The raw label/identifier to escape.

required

Returns:

Type Description
str

The label wrapped in backticks, with any internal backtick doubled.

Prevents a label containing a backtick from escaping the quoting in (n:\`label\`) and injecting Cypher.

Source code in src/deep_db_agents/dialects/neo4j/tools.py
62
63
64
65
66
67
68
69
70
71
72
73
74
def quote_label(label: str) -> str:
    """Escapes a Cypher label/identifier (doubled backticks).

    Args:
        label: The raw label/identifier to escape.

    Returns:
        The label wrapped in backticks, with any internal backtick doubled.

    Prevents a label containing a backtick from escaping the quoting in
    ``(n:\\`label\\`)`` and injecting Cypher.
    """
    return "`" + label.replace("`", "``") + "`"
records_to_table
records_to_table(keys: list[str], records: list[dict]) -> tuple[list[str], list[list[Any]]]

Converts Cypher records into (columns, rows) for file materialization.

Parameters:

Name Type Description Default
keys list[str]

The result column names declared by the Cypher query.

required
records list[dict]

The record data, one dict per row.

required

Returns:

Type Description
tuple[list[str], list[list[Any]]]

A tuple of (column names, row values).

Unlike documents, the columns are the keys declared by the Cypher result (not the union of keys); nodes/relationships/nested values are serialized.

Source code in src/deep_db_agents/dialects/neo4j/tools.py
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def records_to_table(keys: list[str], records: list[dict]) -> tuple[list[str], list[list[Any]]]:
    """Converts Cypher records into (columns, rows) for file materialization.

    Args:
        keys: The result column names declared by the Cypher query.
        records: The record data, one dict per row.

    Returns:
        A tuple of (column names, row values).

    Unlike documents, the columns are the ``keys`` declared by the Cypher result (not the
    union of keys); nodes/relationships/nested values are serialized.
    """
    rows = [[scalar_json(rec.get(k)) for k in keys] for rec in records]
    return list(keys), rows

opensearch

OpenSearch dialect.

OpenSearchDialect

OpenSearchDialect()

Bases: SearchDialect

Agent specialized on OpenSearch.

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.

Source code in src/deep_db_agents/dialects/search_base.py
171
172
173
174
175
def __init__(self) -> None:
    """Initializes the dialect with an empty, lazily-populated client cache."""
    # The ES/OS client internally manages a thread-safe connection pool: it is reused
    # for the whole lifetime of the agent instead of being recreated on every tool call.
    self._client_cache = LazyClient()
system_prompt
system_prompt() -> str

Return the OpenSearch-specific system prompt.

Returns:

Name Type Description
str str

The system prompt text for the OpenSearch agent.

Source code in src/deep_db_agents/dialects/opensearch/dialect.py
31
32
33
34
35
36
37
def system_prompt(self) -> str:
    """Return the OpenSearch-specific system prompt.

    Returns:
        str: The system prompt text for the OpenSearch agent.
    """
    return OPENSEARCH_SYSTEM_PROMPT

dialect

Complete OpenSearch dialect (official opensearch-py driver).

Reuses the entire tool logic from :mod:..search_base (shared by Elasticsearch and OpenSearch); this module only provides the driver-specific connection opening.

OpenSearchDialect
OpenSearchDialect()

Bases: SearchDialect

Agent specialized on OpenSearch.

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.

Source code in src/deep_db_agents/dialects/search_base.py
171
172
173
174
175
def __init__(self) -> None:
    """Initializes the dialect with an empty, lazily-populated client cache."""
    # The ES/OS client internally manages a thread-safe connection pool: it is reused
    # for the whole lifetime of the agent instead of being recreated on every tool call.
    self._client_cache = LazyClient()
system_prompt
system_prompt() -> str

Return the OpenSearch-specific system prompt.

Returns:

Name Type Description
str str

The system prompt text for the OpenSearch agent.

Source code in src/deep_db_agents/dialects/opensearch/dialect.py
31
32
33
34
35
36
37
def system_prompt(self) -> str:
    """Return the OpenSearch-specific system prompt.

    Returns:
        str: The system prompt text for the OpenSearch agent.
    """
    return OPENSEARCH_SYSTEM_PROMPT

prompt

Generic system prompt for the OpenSearch-specialized agent.

tools

Driver-specific helpers for OpenSearch (official opensearch-py driver).

connect
connect(conn: ConnectionConfig, *, request_timeout: int | None = None) -> Any

Open an OpenSearch client from the given credentials.

The driver is imported lazily so the package can be imported without the opensearch extra installed.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection configuration with host, port, and credentials. Recognized credential keys: use_ssl, verify_certs, ca_certs, user, password.

required
request_timeout int | None

Optional request timeout in seconds; forwarded to the client only if positive.

None

Returns:

Name Type Description
Any Any

An initialized OpenSearch client instance.

Raises:

Type Description
ImportError

If the opensearchpy package is not installed.

Source code in src/deep_db_agents/dialects/opensearch/tools.py
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
def connect(conn: ConnectionConfig, *, request_timeout: int | None = None) -> Any:
    """Open an OpenSearch client from the given credentials.

    The driver is imported lazily so the package can be imported without the
    ``opensearch`` extra installed.

    Args:
        conn: Connection configuration with host, port, and credentials. Recognized
            ``credential`` keys: ``use_ssl``, ``verify_certs``, ``ca_certs``, ``user``,
            ``password``.
        request_timeout: Optional request timeout in seconds; forwarded to the
            client only if positive.

    Returns:
        Any: An initialized ``OpenSearch`` client instance.

    Raises:
        ImportError: If the ``opensearchpy`` package is not installed.
    """
    try:
        from opensearchpy import OpenSearch
    except ImportError as exc:  # pragma: no cover - depends on the installed extra
        raise ImportError(
            "The OpenSearch dialect requires the 'opensearch' extra "
            "(pip install 'deep-db-agents[opensearch]')."
        ) from exc

    cred = conn.credential
    host = conn.host or "localhost"
    port = conn.port or 9200
    kwargs: dict[str, Any] = {
        "hosts": [{"host": host, "port": port}],
        "use_ssl": bool(cred.get("use_ssl", False)),
        "verify_certs": cred.get("verify_certs", True),
    }
    if request_timeout and request_timeout > 0:
        kwargs["timeout"] = request_timeout
    if cred.get("ca_certs"):
        kwargs["ca_certs"] = cred["ca_certs"]
    if cred.get("user"):
        kwargs["http_auth"] = (cred["user"], cred.get("password", ""))
    return OpenSearch(**kwargs)

postgres

Postgres dialect.

PostgresDialect

Bases: SqlDialect

Agent specialized on PostgreSQL.

system_prompt
system_prompt() -> str

Return the PostgreSQL-specific system prompt.

Returns:

Type Description
str

The full system prompt text for the PostgreSQL agent.

Source code in src/deep_db_agents/dialects/postgres/dialect.py
21
22
23
24
25
26
27
def system_prompt(self) -> str:
    """Return the PostgreSQL-specific system prompt.

    Returns:
        The full system prompt text for the PostgreSQL agent.
    """
    return POSTGRES_SYSTEM_PROMPT

dialect

Full PostgreSQL dialect (psycopg 3 driver).

PostgresDialect

Bases: SqlDialect

Agent specialized on PostgreSQL.

system_prompt
system_prompt() -> str

Return the PostgreSQL-specific system prompt.

Returns:

Type Description
str

The full system prompt text for the PostgreSQL agent.

Source code in src/deep_db_agents/dialects/postgres/dialect.py
21
22
23
24
25
26
27
def system_prompt(self) -> str:
    """Return the PostgreSQL-specific system prompt.

    Returns:
        The full system prompt text for the PostgreSQL agent.
    """
    return POSTGRES_SYSTEM_PROMPT

prompt

Generic system prompt for the PostgreSQL-specialized agent.

tools

PostgreSQL-specific driver helpers (psycopg 3).

connect
connect(conn: ConnectionConfig, *, network_timeout_s: int | None = None)

Open a psycopg connection from credentials, with lazy driver import.

network_timeout_s configures TCP keepalives so that a connection whose peer is dead (socket hanging, no data) is torn down within ~network_timeout_s, instead of waiting indefinitely. It should be derived from guardrails.query_timeout_s.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters (host, port, credentials, database).

required
network_timeout_s int | None

Idle time in seconds before keepalive probing starts, or None/0 to disable keepalives.

None

Returns:

Type Description

An open psycopg connection.

Raises:

Type Description
ImportError

If the psycopg driver is not installed.

Source code in src/deep_db_agents/dialects/postgres/tools.py
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def connect(conn: ConnectionConfig, *, network_timeout_s: int | None = None):
    """Open a psycopg connection from credentials, with lazy driver import.

    ``network_timeout_s`` configures TCP keepalives so that a connection whose peer is dead
    (socket hanging, no data) is torn down within ~``network_timeout_s``, instead of waiting
    indefinitely. It should be derived from ``guardrails.query_timeout_s``.

    Args:
        conn: Connection parameters (host, port, credentials, database).
        network_timeout_s: Idle time in seconds before keepalive probing starts, or
            ``None``/``0`` to disable keepalives.

    Returns:
        An open psycopg connection.

    Raises:
        ImportError: If the ``psycopg`` driver is not installed.
    """
    try:
        import psycopg
    except ImportError as exc:  # pragma: no cover - depends on the installed extra
        raise ImportError(
            "The Postgres dialect requires the 'postgres' extra "
            "(pip install 'deep-db-agents[postgres]')."
        ) from exc

    cred = conn.credential
    keepalive_kwargs: dict[str, object] = {}
    if network_timeout_s and network_timeout_s > 0:
        # Detect a dead peer within ~network_timeout_s: after network_timeout_s of
        # inactivity, send a probe every few seconds and close after a few failed probes.
        keepalive_kwargs = {
            "keepalives": 1,
            "keepalives_idle": network_timeout_s,
            "keepalives_interval": 5,
            "keepalives_count": 3,
        }
    return psycopg.connect(
        host=conn.host or "localhost",
        port=conn.port or 5432,
        user=cred.get("user"),
        password=cred.get("password"),
        dbname=conn.database,
        connect_timeout=cred.get("connect_timeout", 10),
        autocommit=True,
        **keepalive_kwargs,
    )
estimate_rows
estimate_rows(cursor, sql: str) -> int

Estimate rows via EXPLAIN (FORMAT JSON), reading Plan Rows of the root node.

Parameters:

Name Type Description Default
cursor

DB-API cursor to run the EXPLAIN on.

required
sql str

SQL query to estimate.

required

Returns:

Type Description
int

Estimated row count, or 0 if it cannot be determined.

Source code in src/deep_db_agents/dialects/postgres/tools.py
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def estimate_rows(cursor, sql: str) -> int:
    """Estimate rows via ``EXPLAIN (FORMAT JSON)``, reading ``Plan Rows`` of the root node.

    Args:
        cursor: DB-API cursor to run the ``EXPLAIN`` on.
        sql: SQL query to estimate.

    Returns:
        Estimated row count, or 0 if it cannot be determined.
    """
    cursor.execute(f"EXPLAIN (FORMAT JSON) {sql}")
    plan = cursor.fetchone()[0]
    if isinstance(plan, str):
        plan = json.loads(plan)
    try:
        return int(plan[0]["Plan"]["Plan Rows"])
    except (KeyError, IndexError, TypeError, ValueError):
        return 0

search_base

Shared base for document 'search' dialects (Elasticsearch, OpenSearch).

elasticsearch-py and opensearch-py expose the same REST API (cat.indices, indices.get_mapping, count, search), so the entire tool logic is shared here; each concrete dialect implements only the driver-specific connection opening (_connect).

Unlike the SQL dialects there is no need for a statement whitelist: the tools only use read endpoints (_search, _count, _cat, _mapping), never write or delete endpoints (_bulk, _update, _delete_by_query...) — the main guardrail is therefore the restriction on queryable indices, configured via the index key of the credentials (single name, CSV list or * pattern).

SearchDialect

SearchDialect()

Bases: DbDialect

Base for document search engine dialects (Elasticsearch, OpenSearch).

Initializes the dialect with an empty, lazily-populated client cache.

Source code in src/deep_db_agents/dialects/search_base.py
171
172
173
174
175
def __init__(self) -> None:
    """Initializes the dialect with an empty, lazily-populated client cache."""
    # The ES/OS client internally manages a thread-safe connection pool: it is reused
    # for the whole lifetime of the agent instead of being recreated on every tool call.
    self._client_cache = LazyClient()
build_tools
build_tools(conn: ConnectionConfig, guardrails: GuardrailConfig, materialize_enable: bool = False) -> Sequence[BaseTool]

Builds the LangChain tools exposed to the agent for this search dialect.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters injected into the tool closures.

required
guardrails GuardrailConfig

Guardrail configuration (limits, timeout, budget) enforced by every tool.

required
materialize_enable bool

Whether to also expose the materialize_query tool.

False

Returns:

Type Description
Sequence[BaseTool]

The sequence of tools: list_indices, describe_index, count_documents,

Sequence[BaseTool]

sample_documents, run_query, search_query, aggregate and, if

Sequence[BaseTool]

enabled, materialize_query.

Source code in src/deep_db_agents/dialects/search_base.py
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
def build_tools(
    self,
    conn: ConnectionConfig,
    guardrails: GuardrailConfig,
    materialize_enable: bool = False,
) -> Sequence[BaseTool]:
    """Builds the LangChain tools exposed to the agent for this search dialect.

    Args:
        conn: Connection parameters injected into the tool closures.
        guardrails: Guardrail configuration (limits, timeout, budget) enforced by every
            tool.
        materialize_enable: Whether to also expose the ``materialize_query`` tool.

    Returns:
        The sequence of tools: ``list_indices``, ``describe_index``, ``count_documents``,
        ``sample_documents``, ``run_query``, ``search_query``, ``aggregate`` and, if
        enabled, ``materialize_query``.
    """
    metrics = getattr(self, "_metrics", None)
    budget = SessionBudget(guardrails.row_budget, metrics=metrics)
    allowed = allowed_index_pattern(conn)
    timeout = guardrails.query_timeout_s

    def _paginated_search(
        target: str, q: dict, page: int, page_size: int | None
    ) -> tuple[list[dict], int]:
        limit = guardrails.clamp_limit(page_size)
        offset = max(page, 0) * limit
        with self._client(conn, guardrails) as client:
            estimate = client.count(index=target, body={"query": q}, request_timeout=timeout)[
                "count"
            ]
            guardrails.check_estimate(estimate, metrics)
            res = client.search(
                index=target,
                body={"query": q, "from": offset, "size": limit},
                request_timeout=timeout,
            )
        return res.get("hits", {}).get("hits", []), limit

    @tool
    def list_indices() -> str:
        """Lists the indices visible to this agent (within the configured index scope)."""
        try:
            with self._client(conn, guardrails) as client:
                stats = client.cat.indices(
                    index=allowed, format="json", request_timeout=timeout
                )
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=allowed, what="index listing")
        if not stats:
            return f"No index found for the allowed pattern {allowed!r}."
        lines = [f"{s.get('index')}: {s.get('docs.count', '?')} documents" for s in stats]
        return "Indices:\n" + "\n".join(lines)

    @tool
    def describe_index(index: str | None = None) -> str:
        """Reports the field mapping (names and types) of one or more indices.

        ``index`` can be a single index name, a CSV list, or a wildcard pattern; it must
        stay within the configured index scope.
        """
        try:
            target = resolve_index(index, allowed)
            with self._client(conn, guardrails) as client:
                mapping = client.indices.get_mapping(index=target, request_timeout=timeout)
        except QueryNotAllowedError as exc:  # noqa: BLE001 - scope violation -> feedback to the agent
            return format_query_error(exc, query=index, what="index mapping")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=index, what="index mapping")
        lines = []
        for idx_name, idx_mapping in mapping.items():
            props = idx_mapping.get("mappings", {}).get("properties", {})
            fields = ", ".join(f"{f}:{d.get('type', '?')}" for f, d in props.items())
            lines.append(f"{idx_name}: {fields or '(no mapped field)'}")
        return "\n".join(lines)

    @tool
    def count_documents(index: str | None = None, query: str | dict | None = None) -> str:
        """Counts the documents matching a Query DSL clause, without retrieving them.

        Use it BEFORE searching, to assess whether the result volume is manageable
        (explore before extracting). ``query`` is a string containing the JSON Query DSL
        'query' clause, e.g. '{"match": {"title": "foo"}}'; omit it to count all documents.
        """
        try:
            target = resolve_index(index, allowed)
            if isinstance(query, dict):
                q = query
            else:
                q = parse_query(query, what="query")
            with self._client(conn, guardrails) as client:
                total = client.count(index=target, body={"query": q}, request_timeout=timeout)[
                    "count"
                ]
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="count")
        return f"{target}: {total:,} documents" + (f" with query {query}" if query else "")

    @tool
    def sample_documents(index: str | None = None, limit: int = 5) -> str:
        """Returns a small sample of documents from the index scope (default 5)."""
        try:
            target = resolve_index(index, allowed)
            n = guardrails.clamp_limit(min(limit, 20))
            with self._client(conn, guardrails) as client:
                res = client.search(
                    index=target,
                    body={"query": {"match_all": {}}, "size": n},
                    request_timeout=timeout,
                )
        except QueryNotAllowedError as exc:  # noqa: BLE001 - scope violation -> feedback to the agent
            return format_query_error(exc, query=index, what="search")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=target, what="search")
        docs = _hits_to_docs(res.get("hits", {}).get("hits", []))
        budget.charge(len(docs))
        return f"Sample of {target} ({len(docs)} documents): {docs}"

    @tool
    def run_query(
        index: str | None = None,
        query: str | dict | None = None,
        page: int = 0,
        page_size: int | None = None,
    ) -> str:
        """Runs a read-only Query DSL search with forced limit and pagination.

        ``query`` is a string containing the JSON Query DSL 'query' clause (e.g.
        '{"bool": {"must": [{"match": {"status": "ok"}}]}}'). Before extraction the number
        of matching documents is estimated and the search is blocked if it exceeds the
        threshold. Prefer pushing aggregations into the engine (terms/metric aggregations,
        `_count`) instead of downloading documents and aggregating them by hand.
        """
        try:
            target = resolve_index(index, allowed)
            if isinstance(query, dict):
                q = query
            else:
                q = parse_query(query, what="query")
            hits, limit = _paginated_search(target, q, page, page_size)
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="search")
        budget.charge(len(hits))
        docs = _hits_to_docs(hits)
        more = (
            f" More documents available: ask for page={page + 1}." if len(hits) == limit else ""
        )
        return f"{len(docs)} documents (page {page}, page_size {limit}).{more}\n{docs}"

    @tool
    def search_query(
        index: str | None = None,
        query_string: str = "",
        page: int = 0,
        page_size: int | None = None,
    ) -> str:
        """Runs a simplified Lucene-like text search (the `query_string` DSL clause).

        Use this for natural, search-engine-style queries (e.g. 'status:active AND
        price:[10 TO 50]') instead of building a full Query DSL clause by hand. Forced
        limit and pagination apply, just like `run_query`.
        """
        try:
            target = resolve_index(index, allowed)
            q = {"query_string": {"query": query_string}}
            hits, limit = _paginated_search(target, q, page, page_size)
        except QueryNotAllowedError as exc:  # noqa: BLE001 - scope violation -> feedback to the agent
            return format_query_error(exc, query=query_string, what="query_string search")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=query_string, what="query_string search")
        budget.charge(len(hits))
        docs = _hits_to_docs(hits)
        more = (
            f" More documents available: ask for page={page + 1}." if len(hits) == limit else ""
        )
        return f"{len(docs)} documents (page {page}, page_size {limit}).{more}\n{docs}"

    @tool
    def aggregate(
        runtime: ToolRuntime,
        index: str | None = None,
        aggs: str | dict = "",
        query: str | dict | None = None,
    ) -> str:
        """Runs one or more Query DSL aggregations in the engine, without returning documents.

        This is the PREFERRED way to summarize data (counts per group, averages, min/max,
        cardinality...): the engine does the work and you get back only compact results,
        not raw documents. Pass ``aggs`` as a JSON object mapping each aggregation name to
        its definition, e.g. '{"by_carrier": {"terms": {"field": "Carrier", "size": 20}},
        "avg_price": {"avg": {"field": "AvgTicketPrice"}}}'. An optional ``query`` (Query DSL
        clause) restricts the documents the aggregations run over; omit it to aggregate all
        documents. ``size`` is forced to 0, so no documents are downloaded. If the aggregation
        result is too large for the context (e.g. thousands of `terms` buckets) it is saved to
        file and only a preview is returned: reduce the buckets `size` to get it inline.
        """
        try:
            target = resolve_index(index, allowed)
            a = parse_aggs(aggs)
            if isinstance(query, dict):
                q = query
            else:
                q = parse_query(query, what="query")
            with self._client(conn, guardrails) as client:
                res = client.search(
                    index=target,
                    body={"query": q, "size": 0, "aggs": a},
                    request_timeout=timeout,
                )
        except QueryNotAllowedError as exc:  # noqa: BLE001 - parsing error -> feedback to the agent
            return format_query_error(exc, query=str(aggs), what="aggregation")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=str(aggs), what="aggregation")
        aggregations = res.get("aggregations", {})
        payload = json.dumps(aggregations, default=str)
        if len(payload) <= MAX_AGG_INLINE_CHARS:
            return f"Aggregations on {target}:\n{payload}"
        # Output past the threshold: like the other tools, don't pour it into the
        # context. If a backend is available, materialize it to file (raw JSON, because
        # aggregations are nested and not tabular like materialize_result); otherwise
        # truncate with a preview and invite the agent to reduce the buckets.
        preview = payload[:MAX_AGG_PREVIEW_CHARS]
        be_uuid = runtime.config.get("configurable", {}).get("be_uuid")
        be = BERegistry().get(be_uuid) if be_uuid is not None else None
        if be is not None:
            filename = f"aggregations_{uuid.uuid4().hex[:8]}.json"
            write_res = be.write(filename, payload)
            if not write_res.error:
                return (
                    f"Aggregations on {target}: result too large "
                    f"({len(payload):,} characters) saved to {filename}.\n"
                    f"Preview (first {len(preview):,} characters):\n{preview}\n"
                    "Reduce the buckets' 'size' or refine the aggregation to get it "
                    "fully inline."
                )
        return (
            f"Aggregations on {target}: result too large "
            f"({len(payload):,} characters), truncated to the first {len(preview):,}.\n"
            f"{preview}\n"
            "Reduce the buckets' 'size' or refine the aggregation for a complete output."
        )

    @tool
    def materialize_query(
        runtime: ToolRuntime,
        index: str | None = None,
        query: str | dict | None = None,
        fmt: str = "csv",
        filename: str | None = None,
    ) -> str:
        """Runs a Query DSL search and saves the result to file (Parquet/CSV).

        Returns ONLY metadata: path, columns, preview and statistics. Use it for analysis
        or charts on large volumes. Nested values are serialized.
        """
        ber = BERegistry()
        be_uuid = runtime.config.get("configurable", {}).get("be_uuid")
        if be_uuid is None or (be := ber.get(be_uuid)) is None:
            return "Error: cannot use this tool, backend missing."
        try:
            target = resolve_index(index, allowed)
            if isinstance(query, dict):
                q = query
            else:
                q = parse_query(query, what="query")
            with self._client(conn, guardrails) as client:
                res = client.search(
                    index=target,
                    body={"query": q, "size": guardrails.hard_max_rows},
                    request_timeout=timeout,
                )
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="search")
        hits = res.get("hits", {}).get("hits", [])
        budget.charge(len(hits))
        columns, rows = hits_to_table(hits)
        result = materialize_result(columns, rows, fmt=fmt, filename=filename, backend=be)
        return result.to_summary()

    tools_list = [
        list_indices,
        describe_index,
        count_documents,
        sample_documents,
        run_query,
        search_query,
        aggregate,
    ]
    if materialize_enable:
        tools_list.append(materialize_query)
    return tools_list
close
close() -> None

Closes the reused ES/OS client. Call this at the end of the agent's lifetime.

Source code in src/deep_db_agents/dialects/search_base.py
189
190
191
def close(self) -> None:
    """Closes the reused ES/OS client. Call this at the end of the agent's lifetime."""
    self._client_cache.close()

allowed_index_pattern

allowed_index_pattern(conn: ConnectionConfig) -> str

Returns the index pattern/CSV the agent may operate on, from credentials (default *).

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters whose credential["index"] is read.

required

Returns:

Type Description
str

The configured index pattern/CSV, or "*" if unset.

Source code in src/deep_db_agents/dialects/search_base.py
45
46
47
48
49
50
51
52
53
54
55
def allowed_index_pattern(conn: ConnectionConfig) -> str:
    """Returns the index pattern/CSV the agent may operate on, from credentials (default ``*``).

    Args:
        conn: Connection parameters whose ``credential["index"]`` is read.

    Returns:
        The configured index pattern/CSV, or ``"*"`` if unset.
    """
    pattern = conn.credential.get("index")
    return pattern.strip() if isinstance(pattern, str) and pattern.strip() else "*"

hits_to_table

hits_to_table(hits: list[dict]) -> tuple[list[str], list[list[Any]]]

Converts search hits into (columns, rows) for materialization to file.

Parameters:

Name Type Description Default
hits list[dict]

The raw list of hit objects returned by a search.

required

Returns:

Type Description
list[str]

A tuple of (column names, row values). Columns are the union of _source keys

list[list[Any]]

plus _id/_index; nested values are serialized.

Source code in src/deep_db_agents/dialects/search_base.py
141
142
143
144
145
146
147
148
149
150
151
def hits_to_table(hits: list[dict]) -> tuple[list[str], list[list[Any]]]:
    """Converts search hits into (columns, rows) for materialization to file.

    Args:
        hits: The raw list of hit objects returned by a search.

    Returns:
        A tuple of (column names, row values). Columns are the union of ``_source`` keys
        plus ``_id``/``_index``; nested values are serialized.
    """
    return docs_to_table(_hits_to_docs(hits))

parse_aggs

parse_aggs(value: str | dict) -> dict

Decodes the aggs clause (JSON), which must be a non-empty object.

Parameters:

Name Type Description Default
value str | dict

The aggs clause, either already a dict or a raw JSON string.

required

Returns:

Type Description
dict

The decoded aggs clause as a dict.

Raises:

Type Description
QueryNotAllowedError

If value is empty, not valid JSON, or not a non-empty JSON object.

Source code in src/deep_db_agents/dialects/search_base.py
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def parse_aggs(value: str | dict) -> dict:
    """Decodes the ``aggs`` clause (JSON), which must be a non-empty object.

    Args:
        value: The ``aggs`` clause, either already a dict or a raw JSON string.

    Returns:
        The decoded ``aggs`` clause as a dict.

    Raises:
        QueryNotAllowedError: If ``value`` is empty, not valid JSON, or not a non-empty
            JSON object.
    """
    if isinstance(value, dict):
        parsed = value
    else:
        if value is None or not value.strip():
            raise QueryNotAllowedError("The 'aggs' clause is required and cannot be empty.")
        try:
            parsed = json.loads(value)
        except ValueError as exc:
            raise QueryNotAllowedError(f"Invalid aggs JSON: {exc}") from exc
    if not isinstance(parsed, dict) or not parsed:
        raise QueryNotAllowedError(
            "aggs must be a non-empty JSON object (Query DSL 'aggs' clause)."
        )
    return parsed

parse_query

parse_query(text: str | None, *, what: str = 'query') -> dict

Decodes the Query DSL clause (JSON), or returns match_all if absent.

Parameters:

Name Type Description Default
text str | None

The raw JSON text of the Query DSL clause, or None.

required
what str

Label used in the error message if parsing fails.

'query'

Returns:

Type Description
dict

The decoded query clause as a dict, or {"match_all": {}} if text is empty.

Raises:

Type Description
QueryNotAllowedError

If text is not valid JSON or is not a JSON object.

Source code in src/deep_db_agents/dialects/search_base.py
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def parse_query(text: str | None, *, what: str = "query") -> dict:
    """Decodes the Query DSL clause (JSON), or returns ``match_all`` if absent.

    Args:
        text: The raw JSON text of the Query DSL clause, or ``None``.
        what: Label used in the error message if parsing fails.

    Returns:
        The decoded query clause as a dict, or ``{"match_all": {}}`` if ``text`` is empty.

    Raises:
        QueryNotAllowedError: If ``text`` is not valid JSON or is not a JSON object.
    """
    if text is None or not text.strip():
        return {"match_all": {}}
    try:
        parsed = json.loads(text)
    except ValueError as exc:
        raise QueryNotAllowedError(f"Invalid {what} JSON: {exc}") from exc
    if not isinstance(parsed, dict):
        raise QueryNotAllowedError(f"{what} must be a JSON object (Query DSL clause).")
    return parsed

resolve_index

resolve_index(requested: str | None, allowed: str) -> str

Validates the requested index against the allowed pattern, raising on violation.

allowed can be a single name, a CSV list or a * pattern. If requested is None, allowed is used directly (pattern/CSV expansion is left to the driver). If requested is given, every CSV element of it must match at least one allowed pattern, otherwise the request is rejected.

Parameters:

Name Type Description Default
requested str | None

The index/indices requested by the agent, or None for the default scope.

required
allowed str

The configured allowed index pattern/CSV.

required

Returns:

Type Description
str

The index/indices to actually query.

Raises:

Type Description
QueryNotAllowedError

If a requested index does not match any allowed pattern.

Source code in src/deep_db_agents/dialects/search_base.py
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def resolve_index(requested: str | None, allowed: str) -> str:
    """Validates the requested index against the allowed pattern, raising on violation.

    ``allowed`` can be a single name, a CSV list or a ``*`` pattern. If ``requested`` is
    ``None``, ``allowed`` is used directly (pattern/CSV expansion is left to the driver). If
    ``requested`` is given, every CSV element of it must match at least one allowed pattern,
    otherwise the request is rejected.

    Args:
        requested: The index/indices requested by the agent, or ``None`` for the default scope.
        allowed: The configured allowed index pattern/CSV.

    Returns:
        The index/indices to actually query.

    Raises:
        QueryNotAllowedError: If a requested index does not match any allowed pattern.
    """
    if allowed in ("*", ""):
        return requested or "*"
    if requested is None:
        return allowed
    allowed_parts = [p.strip() for p in allowed.split(",") if p.strip()]
    requested_parts = [p.strip() for p in requested.split(",") if p.strip()]
    for part in requested_parts:
        if not any(fnmatch.fnmatchcase(part, pat) for pat in allowed_parts):
            raise QueryNotAllowedError(f"Index {part!r} not allowed. Allowed indices: {allowed}.")
    return requested

sql_base

Shared base for SQL dialects (MySQL, Postgres).

Implements the SQL tools exactly once — schema exploration, counting, sampling, query execution with guardrails and materialization — delegating the driver-specific part (connection, row estimation via EXPLAIN, timeout, quoting) to abstract methods. The SQL drivers used (pymysql, psycopg) are both DB-API 2.0, so the execution logic is shared.

FileSqlDialect

Bases: SqlDialect

Base for file-based SQL dialects (SQLite, DuckDB).

These engines have no server-side statement_timeout: the execution timeout is enforced by a watchdog that calls connection.interrupt() from a separate thread past query_timeout_s, so a long or hung query gets interrupted instead of freezing the whole tool node when tool calls run in parallel.

is_directory staticmethod
is_directory(conn: ConnectionConfig) -> bool

Checks whether the URL points to a folder (trailing slash or existing directory).

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters whose path is inspected.

required

Returns:

Type Description
bool

True if the URL path is a directory (or has a trailing slash), False otherwise.

Source code in src/deep_db_agents/dialects/sql_base.py
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
@staticmethod
def is_directory(conn: ConnectionConfig) -> bool:
    """Checks whether the URL points to a folder (trailing slash or existing directory).

    Args:
        conn: Connection parameters whose ``path`` is inspected.

    Returns:
        True if the URL path is a directory (or has a trailing slash), False otherwise.
    """
    raw = conn.path or ""
    if raw == ":memory:":
        return False
    if raw.endswith("/"):
        return True
    return os.path.isdir(FileSqlDialect.resolve_path(raw))
resolve_path staticmethod
resolve_path(path: str | None) -> str

Resolves the DB path relative to the working dir (:memory: passes through).

Parameters:

Name Type Description Default
path str | None

The raw path from the connection URL, or None.

required

Returns:

Type Description
str

The absolute, expanded path, or ":memory:" unchanged.

Source code in src/deep_db_agents/dialects/sql_base.py
482
483
484
485
486
487
488
489
490
491
492
493
494
@staticmethod
def resolve_path(path: str | None) -> str:
    """Resolves the DB path relative to the working dir (``:memory:`` passes through).

    Args:
        path: The raw path from the connection URL, or ``None``.

    Returns:
        The absolute, expanded path, or ``":memory:"`` unchanged.
    """
    if not path or path == ":memory:":
        return ":memory:"
    return os.path.abspath(os.path.expanduser(path))

SqlDialect

Bases: DbDialect

Abstract dialect for relational databases with a DB-API 2.0 interface.

build_tools
build_tools(conn: ConnectionConfig, guardrails: GuardrailConfig, materialize_enable: bool = False) -> Sequence[BaseTool]

Builds the LangChain tools exposed to the agent for this SQL dialect.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters injected into the tool closures.

required
guardrails GuardrailConfig

Guardrail configuration (limits, timeout, allowed statements, budget) enforced by every tool.

required
materialize_enable bool

Whether to also expose the materialize_query tool.

False

Returns:

Type Description
Sequence[BaseTool]

The sequence of tools: list_tables, describe_table, count_rows,

Sequence[BaseTool]

sample_rows, run_query and, if enabled, materialize_query.

Source code in src/deep_db_agents/dialects/sql_base.py
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def build_tools(
    self,
    conn: ConnectionConfig,
    guardrails: GuardrailConfig,
    materialize_enable: bool = False,
) -> Sequence[BaseTool]:
    """Builds the LangChain tools exposed to the agent for this SQL dialect.

    Args:
        conn: Connection parameters injected into the tool closures.
        guardrails: Guardrail configuration (limits, timeout, allowed statements, budget)
            enforced by every tool.
        materialize_enable: Whether to also expose the ``materialize_query`` tool.

    Returns:
        The sequence of tools: ``list_tables``, ``describe_table``, ``count_rows``,
        ``sample_rows``, ``run_query`` and, if enabled, ``materialize_query``.
    """
    metrics = getattr(self, "_metrics", None)
    budget = SessionBudget(guardrails.row_budget, metrics=metrics)

    @tool
    def list_tables() -> str:
        """Lists the tables available in the current database."""
        sql = self._list_tables_sql(conn)
        try:
            with self._cursor(conn, guardrails) as cur:
                cur.execute(sql)
                tables = [row[0] for row in cur.fetchall()]
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, what="table listing")
        return "Tables: " + (", ".join(tables) if tables else "(none)")

    @tool
    def describe_table(table: str) -> str:
        """Describes the columns (name and type) of a table."""
        try:
            self._quote_ident(table)
            sql, params = self._describe_table_sql(table)
            with self._cursor(conn, guardrails) as cur:
                cur.execute(sql, params)
                cols = self._columns(cur)
                rows = cur.fetchall()
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=table, what="table description")
        lines = [", ".join(f"{c}={v}" for c, v in zip(cols, r, strict=False)) for r in rows]
        return f"Columns of {table}:\n" + ("\n".join(lines) or "(none)")

    @tool
    def count_rows(table: str, where: str | None = None) -> str:
        """Counts the rows of a table with an optional WHERE filter.

        Use this tool BEFORE extracting data, to decide whether the query is manageable
        or needs to be refined/aggregated (explore before extracting).
        """
        qtable = self._quote_ident(table)
        sql = f"SELECT COUNT(*) FROM {qtable}"
        try:
            if where:
                sql += f" WHERE {_validate_where(where)}"
            with self._cursor(conn, guardrails) as cur:
                cur.execute(sql)
                total = cur.fetchone()[0]
        except QueryNotAllowedError as exc:  # noqa: BLE001 - filter not allowed -> feedback
            return format_query_error(exc, query=where, what="WHERE filter")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=sql, what="count")
        return f"{table}: {total:,} rows" + (f" with WHERE {where}" if where else "")

    @tool
    def sample_rows(table: str, limit: int = 5) -> str:
        """Returns a small sample of rows from a table (default 5)."""
        qtable = self._quote_ident(table)
        n = guardrails.clamp_limit(min(limit, 20))
        try:
            with self._cursor(conn, guardrails) as cur:
                cur.execute(f"SELECT * FROM {qtable} LIMIT {n}")
                cols = self._columns(cur)
                rows = cur.fetchall()
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=table, what="sample")
        budget.charge(len(rows))
        data = [dict(zip(cols, r, strict=False)) for r in rows]
        return f"Sample of {table} ({len(data)} rows): {data}"

    @tool
    def run_query(sql: str, page: int = 0, page_size: int | None = None) -> str:
        """Runs a read-only query (SELECT) with forced LIMIT and pagination.

        Only SELECT/WITH are allowed. A maximum LIMIT that cannot be bypassed is applied,
        and the query is blocked if EXPLAIN estimates too many rows. For large datasets to
        analyze, use `materialize_query` instead. Prefer already-aggregated queries.
        """

        try:
            stmt = _ensure_single_select(sql, guardrails.allowed_statements)
            limit = guardrails.clamp_limit(page_size)
            offset = max(page, 0) * limit
            paginated = f"SELECT * FROM ({stmt}) AS _q LIMIT {limit} OFFSET {offset}"
            with self._cursor(conn, guardrails) as cur:
                guardrails.check_estimate(self._estimate_rows(cur, stmt), metrics)
                cur.execute(paginated)
                cols = self._columns(cur)
                rows = cur.fetchall()
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=sql, what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=sql)
        budget.charge(len(rows))
        data = [dict(zip(cols, r, strict=False)) for r in rows]
        more = f" More rows available: ask for page={page + 1}." if len(rows) == limit else ""
        return f"{len(rows)} rows (page {page}, page_size {limit}).{more}\n{data}"

    @tool
    def materialize_query(
        runtime: ToolRuntime, sql: str, fmt: str = "csv", filename: str | None = None
    ) -> str:
        """Runs a SELECT and saves the entire result to file (Parquet/CSV).

        Returns ONLY metadata: file path, columns, preview and statistics.
        Use this tool when the data is needed for analysis or charts but is too much to
        read in chat. The maximum LIMIT still applies as a safety ceiling.
        """
        ber = BERegistry()
        be_uuid = runtime.config.get("configurable", {}).get("be_uuid")
        if be_uuid is None or (be := ber.get(be_uuid)) is None:
            return "Error: cannot use this tool, backend missing."
        try:
            stmt = _ensure_single_select(sql, guardrails.allowed_statements)
            capped = f"SELECT * FROM ({stmt}) AS _q LIMIT {guardrails.hard_max_rows}"
            with self._cursor(conn, guardrails) as cur:
                guardrails.check_estimate(self._estimate_rows(cur, stmt), metrics)
                cur.execute(capped)
                cols = self._columns(cur)
                rows = cur.fetchall()
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=sql, what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=sql)
        budget.charge(len(rows))
        result = materialize_result(cols, rows, fmt=fmt, filename=filename, backend=be)
        return result.to_summary()

    tool_list = [list_tables, describe_table, count_rows, sample_rows, run_query]
    if materialize_enable:
        tool_list.append(materialize_query)
    return tool_list

sqlite

SQLite dialect.

SQLiteDialect

Bases: FileSqlDialect

Agent specialized on SQLite.

system_prompt
system_prompt() -> str

Returns the SQLite-specific system prompt.

Returns:

Type Description
str

The system prompt text for the SQLite agent.

Source code in src/deep_db_agents/dialects/sqlite/dialect.py
32
33
34
35
36
37
38
def system_prompt(self) -> str:
    """Returns the SQLite-specific system prompt.

    Returns:
        The system prompt text for the SQLite agent.
    """
    return SQLITE_SYSTEM_PROMPT

dialect

Full SQLite dialect (stdlib sqlite3 module).

SQLite is a file-based SQL database: the URL sqlite://<path> indicates the file (SQLAlchemy style: sqlite:///rel.db relative to the working dir, sqlite:////abs.db absolute, sqlite://:memory: in memory). It reuses all of SqlDialect's logic (DB-API 2.0); the execution timeout is enforced by FileSqlDialect's watchdog because SQLite has no native statement_timeout. File connections are opened read-only (mode=ro), a defense-in-depth measure on top of the statement whitelist.

Warning: with :memory: every tool call opens a brand-new, empty database (the connection is per-call): useful only for smoke tests, not for data that must survive across calls.

SQLiteDialect

Bases: FileSqlDialect

Agent specialized on SQLite.

system_prompt
system_prompt() -> str

Returns the SQLite-specific system prompt.

Returns:

Type Description
str

The system prompt text for the SQLite agent.

Source code in src/deep_db_agents/dialects/sqlite/dialect.py
32
33
34
35
36
37
38
def system_prompt(self) -> str:
    """Returns the SQLite-specific system prompt.

    Returns:
        The system prompt text for the SQLite agent.
    """
    return SQLITE_SYSTEM_PROMPT

prompt

Generic system prompt for the SQLite-specialized agent.

tools

Driver-specific helpers for SQLite (stdlib sqlite3 module).

connect
connect(path: str, *, busy_timeout_s: int = 10, read_only: bool = False)

Opens a SQLite connection.

busy_timeout_s is the wait timeout on a locked database (the timeout parameter of sqlite3.connect); query execution duration is bounded separately by a watchdog calling interrupt() (see FileSqlDialect). check_same_thread=False allows the watchdog thread to call interrupt() on the connection.

read_only=True opens the file in read-only mode (URI ?mode=ro), so that even a bypass of the keyword whitelist still cannot write to the database (defense in depth). This does not apply to :memory:, which remains writable.

Parameters:

Name Type Description Default
path str

Path to the SQLite database file, or :memory: for an in-memory database.

required
busy_timeout_s int

Seconds to wait on a locked database before giving up.

10
read_only bool

If True, opens the file in read-only mode via a file: URI.

False

Returns:

Type Description

An open sqlite3.Connection.

Source code in src/deep_db_agents/dialects/sqlite/tools.py
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
def connect(path: str, *, busy_timeout_s: int = 10, read_only: bool = False):
    """Opens a SQLite connection.

    ``busy_timeout_s`` is the wait timeout on a locked database (the ``timeout`` parameter of
    ``sqlite3.connect``); query execution duration is bounded separately by a watchdog calling
    ``interrupt()`` (see ``FileSqlDialect``). ``check_same_thread=False`` allows the watchdog
    thread to call ``interrupt()`` on the connection.

    ``read_only=True`` opens the file in read-only mode (URI ``?mode=ro``), so that even a
    bypass of the keyword whitelist still cannot write to the database (defense in depth).
    This does not apply to ``:memory:``, which remains writable.

    Args:
        path: Path to the SQLite database file, or ``:memory:`` for an in-memory database.
        busy_timeout_s: Seconds to wait on a locked database before giving up.
        read_only: If True, opens the file in read-only mode via a ``file:`` URI.

    Returns:
        An open ``sqlite3.Connection``.
    """
    import sqlite3

    if read_only and path != ":memory:":
        uri = f"file:{pathname2url(path)}?mode=ro"
        return sqlite3.connect(uri, timeout=busy_timeout_s, check_same_thread=False, uri=True)
    return sqlite3.connect(path, timeout=busy_timeout_s, check_same_thread=False)

SQL base

deep_db_agents.dialects.sql_base

Shared base for SQL dialects (MySQL, Postgres).

Implements the SQL tools exactly once — schema exploration, counting, sampling, query execution with guardrails and materialization — delegating the driver-specific part (connection, row estimation via EXPLAIN, timeout, quoting) to abstract methods. The SQL drivers used (pymysql, psycopg) are both DB-API 2.0, so the execution logic is shared.

FileSqlDialect

Bases: SqlDialect

Base for file-based SQL dialects (SQLite, DuckDB).

These engines have no server-side statement_timeout: the execution timeout is enforced by a watchdog that calls connection.interrupt() from a separate thread past query_timeout_s, so a long or hung query gets interrupted instead of freezing the whole tool node when tool calls run in parallel.

is_directory staticmethod

is_directory(conn: ConnectionConfig) -> bool

Checks whether the URL points to a folder (trailing slash or existing directory).

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters whose path is inspected.

required

Returns:

Type Description
bool

True if the URL path is a directory (or has a trailing slash), False otherwise.

Source code in src/deep_db_agents/dialects/sql_base.py
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
@staticmethod
def is_directory(conn: ConnectionConfig) -> bool:
    """Checks whether the URL points to a folder (trailing slash or existing directory).

    Args:
        conn: Connection parameters whose ``path`` is inspected.

    Returns:
        True if the URL path is a directory (or has a trailing slash), False otherwise.
    """
    raw = conn.path or ""
    if raw == ":memory:":
        return False
    if raw.endswith("/"):
        return True
    return os.path.isdir(FileSqlDialect.resolve_path(raw))

resolve_path staticmethod

resolve_path(path: str | None) -> str

Resolves the DB path relative to the working dir (:memory: passes through).

Parameters:

Name Type Description Default
path str | None

The raw path from the connection URL, or None.

required

Returns:

Type Description
str

The absolute, expanded path, or ":memory:" unchanged.

Source code in src/deep_db_agents/dialects/sql_base.py
482
483
484
485
486
487
488
489
490
491
492
493
494
@staticmethod
def resolve_path(path: str | None) -> str:
    """Resolves the DB path relative to the working dir (``:memory:`` passes through).

    Args:
        path: The raw path from the connection URL, or ``None``.

    Returns:
        The absolute, expanded path, or ``":memory:"`` unchanged.
    """
    if not path or path == ":memory:":
        return ":memory:"
    return os.path.abspath(os.path.expanduser(path))

SqlDialect

Bases: DbDialect

Abstract dialect for relational databases with a DB-API 2.0 interface.

build_tools

build_tools(conn: ConnectionConfig, guardrails: GuardrailConfig, materialize_enable: bool = False) -> Sequence[BaseTool]

Builds the LangChain tools exposed to the agent for this SQL dialect.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters injected into the tool closures.

required
guardrails GuardrailConfig

Guardrail configuration (limits, timeout, allowed statements, budget) enforced by every tool.

required
materialize_enable bool

Whether to also expose the materialize_query tool.

False

Returns:

Type Description
Sequence[BaseTool]

The sequence of tools: list_tables, describe_table, count_rows,

Sequence[BaseTool]

sample_rows, run_query and, if enabled, materialize_query.

Source code in src/deep_db_agents/dialects/sql_base.py
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def build_tools(
    self,
    conn: ConnectionConfig,
    guardrails: GuardrailConfig,
    materialize_enable: bool = False,
) -> Sequence[BaseTool]:
    """Builds the LangChain tools exposed to the agent for this SQL dialect.

    Args:
        conn: Connection parameters injected into the tool closures.
        guardrails: Guardrail configuration (limits, timeout, allowed statements, budget)
            enforced by every tool.
        materialize_enable: Whether to also expose the ``materialize_query`` tool.

    Returns:
        The sequence of tools: ``list_tables``, ``describe_table``, ``count_rows``,
        ``sample_rows``, ``run_query`` and, if enabled, ``materialize_query``.
    """
    metrics = getattr(self, "_metrics", None)
    budget = SessionBudget(guardrails.row_budget, metrics=metrics)

    @tool
    def list_tables() -> str:
        """Lists the tables available in the current database."""
        sql = self._list_tables_sql(conn)
        try:
            with self._cursor(conn, guardrails) as cur:
                cur.execute(sql)
                tables = [row[0] for row in cur.fetchall()]
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, what="table listing")
        return "Tables: " + (", ".join(tables) if tables else "(none)")

    @tool
    def describe_table(table: str) -> str:
        """Describes the columns (name and type) of a table."""
        try:
            self._quote_ident(table)
            sql, params = self._describe_table_sql(table)
            with self._cursor(conn, guardrails) as cur:
                cur.execute(sql, params)
                cols = self._columns(cur)
                rows = cur.fetchall()
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=table, what="table description")
        lines = [", ".join(f"{c}={v}" for c, v in zip(cols, r, strict=False)) for r in rows]
        return f"Columns of {table}:\n" + ("\n".join(lines) or "(none)")

    @tool
    def count_rows(table: str, where: str | None = None) -> str:
        """Counts the rows of a table with an optional WHERE filter.

        Use this tool BEFORE extracting data, to decide whether the query is manageable
        or needs to be refined/aggregated (explore before extracting).
        """
        qtable = self._quote_ident(table)
        sql = f"SELECT COUNT(*) FROM {qtable}"
        try:
            if where:
                sql += f" WHERE {_validate_where(where)}"
            with self._cursor(conn, guardrails) as cur:
                cur.execute(sql)
                total = cur.fetchone()[0]
        except QueryNotAllowedError as exc:  # noqa: BLE001 - filter not allowed -> feedback
            return format_query_error(exc, query=where, what="WHERE filter")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=sql, what="count")
        return f"{table}: {total:,} rows" + (f" with WHERE {where}" if where else "")

    @tool
    def sample_rows(table: str, limit: int = 5) -> str:
        """Returns a small sample of rows from a table (default 5)."""
        qtable = self._quote_ident(table)
        n = guardrails.clamp_limit(min(limit, 20))
        try:
            with self._cursor(conn, guardrails) as cur:
                cur.execute(f"SELECT * FROM {qtable} LIMIT {n}")
                cols = self._columns(cur)
                rows = cur.fetchall()
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=table, what="sample")
        budget.charge(len(rows))
        data = [dict(zip(cols, r, strict=False)) for r in rows]
        return f"Sample of {table} ({len(data)} rows): {data}"

    @tool
    def run_query(sql: str, page: int = 0, page_size: int | None = None) -> str:
        """Runs a read-only query (SELECT) with forced LIMIT and pagination.

        Only SELECT/WITH are allowed. A maximum LIMIT that cannot be bypassed is applied,
        and the query is blocked if EXPLAIN estimates too many rows. For large datasets to
        analyze, use `materialize_query` instead. Prefer already-aggregated queries.
        """

        try:
            stmt = _ensure_single_select(sql, guardrails.allowed_statements)
            limit = guardrails.clamp_limit(page_size)
            offset = max(page, 0) * limit
            paginated = f"SELECT * FROM ({stmt}) AS _q LIMIT {limit} OFFSET {offset}"
            with self._cursor(conn, guardrails) as cur:
                guardrails.check_estimate(self._estimate_rows(cur, stmt), metrics)
                cur.execute(paginated)
                cols = self._columns(cur)
                rows = cur.fetchall()
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=sql, what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=sql)
        budget.charge(len(rows))
        data = [dict(zip(cols, r, strict=False)) for r in rows]
        more = f" More rows available: ask for page={page + 1}." if len(rows) == limit else ""
        return f"{len(rows)} rows (page {page}, page_size {limit}).{more}\n{data}"

    @tool
    def materialize_query(
        runtime: ToolRuntime, sql: str, fmt: str = "csv", filename: str | None = None
    ) -> str:
        """Runs a SELECT and saves the entire result to file (Parquet/CSV).

        Returns ONLY metadata: file path, columns, preview and statistics.
        Use this tool when the data is needed for analysis or charts but is too much to
        read in chat. The maximum LIMIT still applies as a safety ceiling.
        """
        ber = BERegistry()
        be_uuid = runtime.config.get("configurable", {}).get("be_uuid")
        if be_uuid is None or (be := ber.get(be_uuid)) is None:
            return "Error: cannot use this tool, backend missing."
        try:
            stmt = _ensure_single_select(sql, guardrails.allowed_statements)
            capped = f"SELECT * FROM ({stmt}) AS _q LIMIT {guardrails.hard_max_rows}"
            with self._cursor(conn, guardrails) as cur:
                guardrails.check_estimate(self._estimate_rows(cur, stmt), metrics)
                cur.execute(capped)
                cols = self._columns(cur)
                rows = cur.fetchall()
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=sql, what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=sql)
        budget.charge(len(rows))
        result = materialize_result(cols, rows, fmt=fmt, filename=filename, backend=be)
        return result.to_summary()

    tool_list = [list_tables, describe_table, count_rows, sample_rows, run_query]
    if materialize_enable:
        tool_list.append(materialize_query)
    return tool_list

Search base

deep_db_agents.dialects.search_base

Shared base for document 'search' dialects (Elasticsearch, OpenSearch).

elasticsearch-py and opensearch-py expose the same REST API (cat.indices, indices.get_mapping, count, search), so the entire tool logic is shared here; each concrete dialect implements only the driver-specific connection opening (_connect).

Unlike the SQL dialects there is no need for a statement whitelist: the tools only use read endpoints (_search, _count, _cat, _mapping), never write or delete endpoints (_bulk, _update, _delete_by_query...) — the main guardrail is therefore the restriction on queryable indices, configured via the index key of the credentials (single name, CSV list or * pattern).

SearchDialect

SearchDialect()

Bases: DbDialect

Base for document search engine dialects (Elasticsearch, OpenSearch).

Initializes the dialect with an empty, lazily-populated client cache.

Source code in src/deep_db_agents/dialects/search_base.py
171
172
173
174
175
def __init__(self) -> None:
    """Initializes the dialect with an empty, lazily-populated client cache."""
    # The ES/OS client internally manages a thread-safe connection pool: it is reused
    # for the whole lifetime of the agent instead of being recreated on every tool call.
    self._client_cache = LazyClient()

build_tools

build_tools(conn: ConnectionConfig, guardrails: GuardrailConfig, materialize_enable: bool = False) -> Sequence[BaseTool]

Builds the LangChain tools exposed to the agent for this search dialect.

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters injected into the tool closures.

required
guardrails GuardrailConfig

Guardrail configuration (limits, timeout, budget) enforced by every tool.

required
materialize_enable bool

Whether to also expose the materialize_query tool.

False

Returns:

Type Description
Sequence[BaseTool]

The sequence of tools: list_indices, describe_index, count_documents,

Sequence[BaseTool]

sample_documents, run_query, search_query, aggregate and, if

Sequence[BaseTool]

enabled, materialize_query.

Source code in src/deep_db_agents/dialects/search_base.py
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
def build_tools(
    self,
    conn: ConnectionConfig,
    guardrails: GuardrailConfig,
    materialize_enable: bool = False,
) -> Sequence[BaseTool]:
    """Builds the LangChain tools exposed to the agent for this search dialect.

    Args:
        conn: Connection parameters injected into the tool closures.
        guardrails: Guardrail configuration (limits, timeout, budget) enforced by every
            tool.
        materialize_enable: Whether to also expose the ``materialize_query`` tool.

    Returns:
        The sequence of tools: ``list_indices``, ``describe_index``, ``count_documents``,
        ``sample_documents``, ``run_query``, ``search_query``, ``aggregate`` and, if
        enabled, ``materialize_query``.
    """
    metrics = getattr(self, "_metrics", None)
    budget = SessionBudget(guardrails.row_budget, metrics=metrics)
    allowed = allowed_index_pattern(conn)
    timeout = guardrails.query_timeout_s

    def _paginated_search(
        target: str, q: dict, page: int, page_size: int | None
    ) -> tuple[list[dict], int]:
        limit = guardrails.clamp_limit(page_size)
        offset = max(page, 0) * limit
        with self._client(conn, guardrails) as client:
            estimate = client.count(index=target, body={"query": q}, request_timeout=timeout)[
                "count"
            ]
            guardrails.check_estimate(estimate, metrics)
            res = client.search(
                index=target,
                body={"query": q, "from": offset, "size": limit},
                request_timeout=timeout,
            )
        return res.get("hits", {}).get("hits", []), limit

    @tool
    def list_indices() -> str:
        """Lists the indices visible to this agent (within the configured index scope)."""
        try:
            with self._client(conn, guardrails) as client:
                stats = client.cat.indices(
                    index=allowed, format="json", request_timeout=timeout
                )
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=allowed, what="index listing")
        if not stats:
            return f"No index found for the allowed pattern {allowed!r}."
        lines = [f"{s.get('index')}: {s.get('docs.count', '?')} documents" for s in stats]
        return "Indices:\n" + "\n".join(lines)

    @tool
    def describe_index(index: str | None = None) -> str:
        """Reports the field mapping (names and types) of one or more indices.

        ``index`` can be a single index name, a CSV list, or a wildcard pattern; it must
        stay within the configured index scope.
        """
        try:
            target = resolve_index(index, allowed)
            with self._client(conn, guardrails) as client:
                mapping = client.indices.get_mapping(index=target, request_timeout=timeout)
        except QueryNotAllowedError as exc:  # noqa: BLE001 - scope violation -> feedback to the agent
            return format_query_error(exc, query=index, what="index mapping")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=index, what="index mapping")
        lines = []
        for idx_name, idx_mapping in mapping.items():
            props = idx_mapping.get("mappings", {}).get("properties", {})
            fields = ", ".join(f"{f}:{d.get('type', '?')}" for f, d in props.items())
            lines.append(f"{idx_name}: {fields or '(no mapped field)'}")
        return "\n".join(lines)

    @tool
    def count_documents(index: str | None = None, query: str | dict | None = None) -> str:
        """Counts the documents matching a Query DSL clause, without retrieving them.

        Use it BEFORE searching, to assess whether the result volume is manageable
        (explore before extracting). ``query`` is a string containing the JSON Query DSL
        'query' clause, e.g. '{"match": {"title": "foo"}}'; omit it to count all documents.
        """
        try:
            target = resolve_index(index, allowed)
            if isinstance(query, dict):
                q = query
            else:
                q = parse_query(query, what="query")
            with self._client(conn, guardrails) as client:
                total = client.count(index=target, body={"query": q}, request_timeout=timeout)[
                    "count"
                ]
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="count")
        return f"{target}: {total:,} documents" + (f" with query {query}" if query else "")

    @tool
    def sample_documents(index: str | None = None, limit: int = 5) -> str:
        """Returns a small sample of documents from the index scope (default 5)."""
        try:
            target = resolve_index(index, allowed)
            n = guardrails.clamp_limit(min(limit, 20))
            with self._client(conn, guardrails) as client:
                res = client.search(
                    index=target,
                    body={"query": {"match_all": {}}, "size": n},
                    request_timeout=timeout,
                )
        except QueryNotAllowedError as exc:  # noqa: BLE001 - scope violation -> feedback to the agent
            return format_query_error(exc, query=index, what="search")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=target, what="search")
        docs = _hits_to_docs(res.get("hits", {}).get("hits", []))
        budget.charge(len(docs))
        return f"Sample of {target} ({len(docs)} documents): {docs}"

    @tool
    def run_query(
        index: str | None = None,
        query: str | dict | None = None,
        page: int = 0,
        page_size: int | None = None,
    ) -> str:
        """Runs a read-only Query DSL search with forced limit and pagination.

        ``query`` is a string containing the JSON Query DSL 'query' clause (e.g.
        '{"bool": {"must": [{"match": {"status": "ok"}}]}}'). Before extraction the number
        of matching documents is estimated and the search is blocked if it exceeds the
        threshold. Prefer pushing aggregations into the engine (terms/metric aggregations,
        `_count`) instead of downloading documents and aggregating them by hand.
        """
        try:
            target = resolve_index(index, allowed)
            if isinstance(query, dict):
                q = query
            else:
                q = parse_query(query, what="query")
            hits, limit = _paginated_search(target, q, page, page_size)
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="search")
        budget.charge(len(hits))
        docs = _hits_to_docs(hits)
        more = (
            f" More documents available: ask for page={page + 1}." if len(hits) == limit else ""
        )
        return f"{len(docs)} documents (page {page}, page_size {limit}).{more}\n{docs}"

    @tool
    def search_query(
        index: str | None = None,
        query_string: str = "",
        page: int = 0,
        page_size: int | None = None,
    ) -> str:
        """Runs a simplified Lucene-like text search (the `query_string` DSL clause).

        Use this for natural, search-engine-style queries (e.g. 'status:active AND
        price:[10 TO 50]') instead of building a full Query DSL clause by hand. Forced
        limit and pagination apply, just like `run_query`.
        """
        try:
            target = resolve_index(index, allowed)
            q = {"query_string": {"query": query_string}}
            hits, limit = _paginated_search(target, q, page, page_size)
        except QueryNotAllowedError as exc:  # noqa: BLE001 - scope violation -> feedback to the agent
            return format_query_error(exc, query=query_string, what="query_string search")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=query_string, what="query_string search")
        budget.charge(len(hits))
        docs = _hits_to_docs(hits)
        more = (
            f" More documents available: ask for page={page + 1}." if len(hits) == limit else ""
        )
        return f"{len(docs)} documents (page {page}, page_size {limit}).{more}\n{docs}"

    @tool
    def aggregate(
        runtime: ToolRuntime,
        index: str | None = None,
        aggs: str | dict = "",
        query: str | dict | None = None,
    ) -> str:
        """Runs one or more Query DSL aggregations in the engine, without returning documents.

        This is the PREFERRED way to summarize data (counts per group, averages, min/max,
        cardinality...): the engine does the work and you get back only compact results,
        not raw documents. Pass ``aggs`` as a JSON object mapping each aggregation name to
        its definition, e.g. '{"by_carrier": {"terms": {"field": "Carrier", "size": 20}},
        "avg_price": {"avg": {"field": "AvgTicketPrice"}}}'. An optional ``query`` (Query DSL
        clause) restricts the documents the aggregations run over; omit it to aggregate all
        documents. ``size`` is forced to 0, so no documents are downloaded. If the aggregation
        result is too large for the context (e.g. thousands of `terms` buckets) it is saved to
        file and only a preview is returned: reduce the buckets `size` to get it inline.
        """
        try:
            target = resolve_index(index, allowed)
            a = parse_aggs(aggs)
            if isinstance(query, dict):
                q = query
            else:
                q = parse_query(query, what="query")
            with self._client(conn, guardrails) as client:
                res = client.search(
                    index=target,
                    body={"query": q, "size": 0, "aggs": a},
                    request_timeout=timeout,
                )
        except QueryNotAllowedError as exc:  # noqa: BLE001 - parsing error -> feedback to the agent
            return format_query_error(exc, query=str(aggs), what="aggregation")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=str(aggs), what="aggregation")
        aggregations = res.get("aggregations", {})
        payload = json.dumps(aggregations, default=str)
        if len(payload) <= MAX_AGG_INLINE_CHARS:
            return f"Aggregations on {target}:\n{payload}"
        # Output past the threshold: like the other tools, don't pour it into the
        # context. If a backend is available, materialize it to file (raw JSON, because
        # aggregations are nested and not tabular like materialize_result); otherwise
        # truncate with a preview and invite the agent to reduce the buckets.
        preview = payload[:MAX_AGG_PREVIEW_CHARS]
        be_uuid = runtime.config.get("configurable", {}).get("be_uuid")
        be = BERegistry().get(be_uuid) if be_uuid is not None else None
        if be is not None:
            filename = f"aggregations_{uuid.uuid4().hex[:8]}.json"
            write_res = be.write(filename, payload)
            if not write_res.error:
                return (
                    f"Aggregations on {target}: result too large "
                    f"({len(payload):,} characters) saved to {filename}.\n"
                    f"Preview (first {len(preview):,} characters):\n{preview}\n"
                    "Reduce the buckets' 'size' or refine the aggregation to get it "
                    "fully inline."
                )
        return (
            f"Aggregations on {target}: result too large "
            f"({len(payload):,} characters), truncated to the first {len(preview):,}.\n"
            f"{preview}\n"
            "Reduce the buckets' 'size' or refine the aggregation for a complete output."
        )

    @tool
    def materialize_query(
        runtime: ToolRuntime,
        index: str | None = None,
        query: str | dict | None = None,
        fmt: str = "csv",
        filename: str | None = None,
    ) -> str:
        """Runs a Query DSL search and saves the result to file (Parquet/CSV).

        Returns ONLY metadata: path, columns, preview and statistics. Use it for analysis
        or charts on large volumes. Nested values are serialized.
        """
        ber = BERegistry()
        be_uuid = runtime.config.get("configurable", {}).get("be_uuid")
        if be_uuid is None or (be := ber.get(be_uuid)) is None:
            return "Error: cannot use this tool, backend missing."
        try:
            target = resolve_index(index, allowed)
            if isinstance(query, dict):
                q = query
            else:
                q = parse_query(query, what="query")
            with self._client(conn, guardrails) as client:
                res = client.search(
                    index=target,
                    body={"query": q, "size": guardrails.hard_max_rows},
                    request_timeout=timeout,
                )
        except QueryNotAllowedError as exc:  # noqa: BLE001 - query parsing error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="query")
        except (DeepDbAgentError, ImportError):
            raise
        except Exception as exc:  # noqa: BLE001 - driver error -> feedback to the agent
            return format_query_error(exc, query=str(query), what="search")
        hits = res.get("hits", {}).get("hits", [])
        budget.charge(len(hits))
        columns, rows = hits_to_table(hits)
        result = materialize_result(columns, rows, fmt=fmt, filename=filename, backend=be)
        return result.to_summary()

    tools_list = [
        list_indices,
        describe_index,
        count_documents,
        sample_documents,
        run_query,
        search_query,
        aggregate,
    ]
    if materialize_enable:
        tools_list.append(materialize_query)
    return tools_list

close

close() -> None

Closes the reused ES/OS client. Call this at the end of the agent's lifetime.

Source code in src/deep_db_agents/dialects/search_base.py
189
190
191
def close(self) -> None:
    """Closes the reused ES/OS client. Call this at the end of the agent's lifetime."""
    self._client_cache.close()

allowed_index_pattern

allowed_index_pattern(conn: ConnectionConfig) -> str

Returns the index pattern/CSV the agent may operate on, from credentials (default *).

Parameters:

Name Type Description Default
conn ConnectionConfig

Connection parameters whose credential["index"] is read.

required

Returns:

Type Description
str

The configured index pattern/CSV, or "*" if unset.

Source code in src/deep_db_agents/dialects/search_base.py
45
46
47
48
49
50
51
52
53
54
55
def allowed_index_pattern(conn: ConnectionConfig) -> str:
    """Returns the index pattern/CSV the agent may operate on, from credentials (default ``*``).

    Args:
        conn: Connection parameters whose ``credential["index"]`` is read.

    Returns:
        The configured index pattern/CSV, or ``"*"`` if unset.
    """
    pattern = conn.credential.get("index")
    return pattern.strip() if isinstance(pattern, str) and pattern.strip() else "*"

hits_to_table

hits_to_table(hits: list[dict]) -> tuple[list[str], list[list[Any]]]

Converts search hits into (columns, rows) for materialization to file.

Parameters:

Name Type Description Default
hits list[dict]

The raw list of hit objects returned by a search.

required

Returns:

Type Description
list[str]

A tuple of (column names, row values). Columns are the union of _source keys

list[list[Any]]

plus _id/_index; nested values are serialized.

Source code in src/deep_db_agents/dialects/search_base.py
141
142
143
144
145
146
147
148
149
150
151
def hits_to_table(hits: list[dict]) -> tuple[list[str], list[list[Any]]]:
    """Converts search hits into (columns, rows) for materialization to file.

    Args:
        hits: The raw list of hit objects returned by a search.

    Returns:
        A tuple of (column names, row values). Columns are the union of ``_source`` keys
        plus ``_id``/``_index``; nested values are serialized.
    """
    return docs_to_table(_hits_to_docs(hits))

parse_aggs

parse_aggs(value: str | dict) -> dict

Decodes the aggs clause (JSON), which must be a non-empty object.

Parameters:

Name Type Description Default
value str | dict

The aggs clause, either already a dict or a raw JSON string.

required

Returns:

Type Description
dict

The decoded aggs clause as a dict.

Raises:

Type Description
QueryNotAllowedError

If value is empty, not valid JSON, or not a non-empty JSON object.

Source code in src/deep_db_agents/dialects/search_base.py
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def parse_aggs(value: str | dict) -> dict:
    """Decodes the ``aggs`` clause (JSON), which must be a non-empty object.

    Args:
        value: The ``aggs`` clause, either already a dict or a raw JSON string.

    Returns:
        The decoded ``aggs`` clause as a dict.

    Raises:
        QueryNotAllowedError: If ``value`` is empty, not valid JSON, or not a non-empty
            JSON object.
    """
    if isinstance(value, dict):
        parsed = value
    else:
        if value is None or not value.strip():
            raise QueryNotAllowedError("The 'aggs' clause is required and cannot be empty.")
        try:
            parsed = json.loads(value)
        except ValueError as exc:
            raise QueryNotAllowedError(f"Invalid aggs JSON: {exc}") from exc
    if not isinstance(parsed, dict) or not parsed:
        raise QueryNotAllowedError(
            "aggs must be a non-empty JSON object (Query DSL 'aggs' clause)."
        )
    return parsed

parse_query

parse_query(text: str | None, *, what: str = 'query') -> dict

Decodes the Query DSL clause (JSON), or returns match_all if absent.

Parameters:

Name Type Description Default
text str | None

The raw JSON text of the Query DSL clause, or None.

required
what str

Label used in the error message if parsing fails.

'query'

Returns:

Type Description
dict

The decoded query clause as a dict, or {"match_all": {}} if text is empty.

Raises:

Type Description
QueryNotAllowedError

If text is not valid JSON or is not a JSON object.

Source code in src/deep_db_agents/dialects/search_base.py
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def parse_query(text: str | None, *, what: str = "query") -> dict:
    """Decodes the Query DSL clause (JSON), or returns ``match_all`` if absent.

    Args:
        text: The raw JSON text of the Query DSL clause, or ``None``.
        what: Label used in the error message if parsing fails.

    Returns:
        The decoded query clause as a dict, or ``{"match_all": {}}`` if ``text`` is empty.

    Raises:
        QueryNotAllowedError: If ``text`` is not valid JSON or is not a JSON object.
    """
    if text is None or not text.strip():
        return {"match_all": {}}
    try:
        parsed = json.loads(text)
    except ValueError as exc:
        raise QueryNotAllowedError(f"Invalid {what} JSON: {exc}") from exc
    if not isinstance(parsed, dict):
        raise QueryNotAllowedError(f"{what} must be a JSON object (Query DSL clause).")
    return parsed

resolve_index

resolve_index(requested: str | None, allowed: str) -> str

Validates the requested index against the allowed pattern, raising on violation.

allowed can be a single name, a CSV list or a * pattern. If requested is None, allowed is used directly (pattern/CSV expansion is left to the driver). If requested is given, every CSV element of it must match at least one allowed pattern, otherwise the request is rejected.

Parameters:

Name Type Description Default
requested str | None

The index/indices requested by the agent, or None for the default scope.

required
allowed str

The configured allowed index pattern/CSV.

required

Returns:

Type Description
str

The index/indices to actually query.

Raises:

Type Description
QueryNotAllowedError

If a requested index does not match any allowed pattern.

Source code in src/deep_db_agents/dialects/search_base.py
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def resolve_index(requested: str | None, allowed: str) -> str:
    """Validates the requested index against the allowed pattern, raising on violation.

    ``allowed`` can be a single name, a CSV list or a ``*`` pattern. If ``requested`` is
    ``None``, ``allowed`` is used directly (pattern/CSV expansion is left to the driver). If
    ``requested`` is given, every CSV element of it must match at least one allowed pattern,
    otherwise the request is rejected.

    Args:
        requested: The index/indices requested by the agent, or ``None`` for the default scope.
        allowed: The configured allowed index pattern/CSV.

    Returns:
        The index/indices to actually query.

    Raises:
        QueryNotAllowedError: If a requested index does not match any allowed pattern.
    """
    if allowed in ("*", ""):
        return requested or "*"
    if requested is None:
        return allowed
    allowed_parts = [p.strip() for p in allowed.split(",") if p.strip()]
    requested_parts = [p.strip() for p in requested.split(",") if p.strip()]
    for part in requested_parts:
        if not any(fnmatch.fnmatchcase(part, pat) for pat in allowed_parts):
            raise QueryNotAllowedError(f"Index {part!r} not allowed. Allowed indices: {allowed}.")
    return requested