Metadata-Version: 2.4
Name: datryn
Version: 0.1.1
Summary: Relational databases as first-class tools for LLM agents
Author: datryn contributors
License: MIT
License-File: LICENSE
Keywords: agent,ai,database,langchain,llm,mcp,sql,text-to-sql
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Database
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Python: >=3.11
Requires-Dist: rank-bm25>=0.2
Requires-Dist: sqlalchemy>=2.0
Requires-Dist: sqlglot>=25.0
Provides-Extra: all
Requires-Dist: anthropic>=0.40; extra == 'all'
Requires-Dist: langchain-core>=0.3; extra == 'all'
Requires-Dist: langgraph>=0.1; extra == 'all'
Requires-Dist: mcp>=1.0; extra == 'all'
Requires-Dist: openai>=1.50; extra == 'all'
Requires-Dist: psycopg2-binary>=2.9; extra == 'all'
Requires-Dist: sentence-transformers>=2.0; extra == 'all'
Provides-Extra: anthropic
Requires-Dist: anthropic>=0.40; extra == 'anthropic'
Provides-Extra: dev
Requires-Dist: mypy>=1.11; extra == 'dev'
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Provides-Extra: langchain
Requires-Dist: langchain-core>=0.3; extra == 'langchain'
Provides-Extra: langgraph
Requires-Dist: langgraph>=0.1; extra == 'langgraph'
Provides-Extra: mcp
Requires-Dist: mcp>=1.0; extra == 'mcp'
Provides-Extra: openai
Requires-Dist: openai>=1.50; extra == 'openai'
Provides-Extra: postgres
Requires-Dist: psycopg2-binary>=2.9; extra == 'postgres'
Provides-Extra: sentence-transformers
Requires-Dist: sentence-transformers>=2.0; extra == 'sentence-transformers'
Description-Content-Type: text/markdown

# datryn

[![PyPI](https://img.shields.io/pypi/v/datryn)](https://pypi.org/project/datryn/)
[![Python](https://img.shields.io/pypi/pyversions/datryn)](https://pypi.org/project/datryn/)
[![License: MIT](https://img.shields.io/badge/license-MIT-blue)](LICENSE)

**Relational databases as first-class tools for LLM agents.**

datryn translates natural language questions into safe, validated SQL — and executes them. It wires schema introspection, context-aware table selection, LLM generation, multi-layer validation, execution, and audit logging behind a single callable object.

```python
from datryn import DatrynTool, AnthropicLLM

with DatrynTool.from_uri("postgresql://user:pass@localhost/mydb", llm=AnthropicLLM()) as tool:
    result = tool("Which products had the highest revenue last quarter?")
    print(result.answer)
```

---

## How it works

Each call to `tool(question)` runs a six-stage pipeline:

```
question
  │
  ▼
[1] Schema introspection (once, then cached)
    SQLAlchemy inspector reads all tables, columns, PKs, FKs, and sample values.
    Builds a BM25 index (default) or pre-computes embeddings (if a semantic
    context provider is configured).
  │
  ▼
[2] Context selection
    Selects the most relevant tables for the question — via BM25 keyword search
    by default, or a pluggable ContextProvider (LLM-based, embedding-based, etc.).
    FK expansion adds related tables by following foreign keys (up to 3 hops).
  │
  ▼
[3] SQL generation
    LLM receives only the relevant schema fragment + the question.
    Returns a SELECT statement.
  │
  ▼
[4] Validation (multi-layer)
    - Read-only enforcement (no INSERT / UPDATE / DELETE / DROP / ...)
    - Allowed tables whitelist
    - Denied column patterns (password, token, api_key, ssn, ...)
    - Dangerous function block (SLEEP, BENCHMARK, pg_sleep, ...)
    - Subquery depth limit
    - Schema-aware hallucination detection (table-alias aware, scope-aware)
    On failure → correction prompt sent back to LLM (up to max_retries).
  │
  ▼
[5] Execution
    Row limit enforced in Python (cannot be bypassed by SQL LIMIT).
    Statement timeout set at the driver level.
  │
  ▼
[6] Formatting + Audit
    Result formatted as Markdown / JSON / CSV / natural text.
    Every call recorded in the audit log (SQLite, JSONL, or any SQLAlchemy DB).
```

---

## Installation

**Core + PostgreSQL + Anthropic:**
```bash
pip install datryn[postgres,anthropic]
```

**With OpenAI:**
```bash
pip install datryn[postgres,openai]
```

**All extras:**
```bash
pip install datryn[all]
```

---

## Quick start

```python
import os
from datryn import DatrynTool, AnthropicLLM

with DatrynTool.from_uri(
    uri=os.environ["DATABASE_URL"],
    llm=AnthropicLLM(),          # reads ANTHROPIC_API_KEY from env
    dialect="postgresql",
    read_only=True,              # blocks INSERT / UPDATE / DELETE (default)
    row_limit=200,
) as tool:
    result = tool("How many orders were placed this month?")

if result.success:
    print(result.answer)         # formatted Markdown table
    print(result.sql)            # generated SQL
    print(result.row_count)      # number of rows returned
    print(result.truncated)      # True if result was cut at row_limit
    print(result.metadata)       # model, duration_ms, attempts, input_tokens, output_tokens, validation_warnings
else:
    print(result.error)          # human-readable error
    print(result.error_code)     # DatrynErrorCode enum value
```

---

## Configuration

```python
DatrynTool.from_uri(
    uri="postgresql://...",
    llm=AnthropicLLM(model="claude-haiku-4-5"),

    # SQL dialect injected into every prompt
    dialect="postgresql",          # default

    # Security
    read_only=True,                # block all writes (default: True)
    allowed_tables=["orders", "products", "customers"],  # whitelist; None = all
    denied_columns={               # block specific columns per table
        "customers": ["password_hash", "ssn"],
    },

    # Execution limits
    row_limit=500,                 # max rows returned (default: 500)
    timeout_seconds=30,            # query timeout (default: 30)

    # Generation
    max_retries=2,                 # correction attempts on validation failure

    # Context selection
    max_tables=10,                 # max tables included in each prompt (default: 10)
    include_sample_values=True,    # include sample values in schema context (default: True)
                                   # set False for sensitive data or faster startup
    annotations={                  # domain synonyms added to BM25 index per table
        "tbl_rev":   ["revenue", "income", "sales"],
        "stk_movmt": ["inventory", "stock", "warehouse"],
    },
    context_provider=None,         # semantic provider (LLMContextProvider, OpenAIEmbeddingProvider, ...)

    # Audit
    audit_log=SQLiteAuditLogger(), # default: ~/.datryn/audit.db

    # Traceability
    session_id="session-abc",
    user_id="user-123",
)
```

---

## LLM adapters

### Anthropic
```python
from datryn import AnthropicLLM
llm = AnthropicLLM(model="claude-haiku-4-5")    # reads ANTHROPIC_API_KEY
```

### OpenAI
```python
from datryn import OpenAILLM
llm = OpenAILLM(model="gpt-5.4-nano")            # reads OPENAI_API_KEY
```

### Custom adapter
```python
from datryn.llm.base import BaseLLM
from datryn.llm.models import LLMResponse

class MyLLM(BaseLLM):
    @property
    def model_name(self) -> str:
        return "my-model"

    def complete(self, prompt: str, system: str | None = None) -> LLMResponse:
        response = my_api.generate(prompt, system=system)
        return LLMResponse(
            text=response.text,
            input_tokens=response.usage.input,   # use 0 if your API doesn't report usage
            output_tokens=response.usage.output,
        )
```

---

## Audit logging

Every call is logged — the question, generated SQL, validation result, execution result, model used, and duration.

```python
from datryn import SQLiteAuditLogger, SQLAlchemyAuditLogger, JSONLineAuditLogger, NoOpAuditLogger

# Default: local SQLite file at ~/.datryn/audit.db
audit = SQLiteAuditLogger()
audit = SQLiteAuditLogger("./logs/audit.db")

# Any SQLAlchemy-compatible database
audit = SQLAlchemyAuditLogger("postgresql://audit-user:pw@host/audit_db")

# Same database being queried — enables querying the audit log via datryn itself
audit = SQLAlchemyAuditLogger(os.environ["DATABASE_URL"])

# JSONL file — one JSON object per line
audit = JSONLineAuditLogger("./logs/queries.jsonl")

# Disabled
audit = NoOpAuditLogger()

tool = DatrynTool.from_uri(uri, llm=llm, audit_log=audit)
```

When the audit database is the same as the queried database, you can query your own audit log:
```python
result = tool("Which questions triggered validation errors this week?")
result = tool("What is the average query duration per model?")
```

---

## Integrations

### LangChain
```python
tool.warm_up()   # build schema index before the agent's first call

# Single database
lc_tool = tool.as_langchain_tool()

# Multiple databases in the same agent — give each tool a distinct name
sales_tool   = sales_datryn.as_langchain_tool(
    name="query_sales",
    description="Query the sales database (orders, customers, revenue)",
)
inventory_tool = inventory_datryn.as_langchain_tool(
    name="query_inventory",
    description="Query the inventory database (products, stock, warehouses)",
)

# show_sql=True (default): the generated SQL is appended to the answer so the
# agent can understand which filters and joins were applied.
# Set show_sql=False to return only the formatted result.
lc_tool = tool.as_langchain_tool(show_sql=False)
```

### LangGraph
```python
tool.warm_up()   # build schema index before the graph's first run

node_fn = tool.as_langgraph_node()             # show_sql=True by default
node_fn = tool.as_langgraph_node(show_sql=False)  # omit SQL from state["answer"]

graph = StateGraph(State)
graph.add_node("query_db", node_fn)
# Node reads state["question"], writes state["answer"] and state["datryn_result"]
# state["datryn_result"].sql is always available to downstream nodes
```

### MCP (Model Context Protocol)
```python
tool.warm_up()   # build schema index before accepting connections

# Single database
server = tool.as_mcp_server()
server.start()  # blocking, stdio transport

# Multiple databases — give each server and tool a distinct name
server = tool.as_mcp_server(
    server_name="sales-db",
    tool_name="query_sales",
    tool_description="Query the sales database (orders, customers, revenue)",
)
server.start()
```

Register in Claude Desktop (`claude_desktop_config.json`):
```json
{
  "mcpServers": {
    "sales-db": {
      "command": "C:\\apps\\myproject\\.venv\\Scripts\\python.exe",
      "args": ["C:\\apps\\myproject\\mcp_sales.py"],
      "env": {
        "ANTHROPIC_API_KEY": "sk-...",
        "DATABASE_URL": "postgresql://user:password@localhost:5432/sales"
      }
    },
    "inventory-db": {
      "command": "C:\\apps\\myproject\\.venv\\Scripts\\python.exe",
      "args": ["C:\\apps\\myproject\\mcp_inventory.py"],
      "env": {
        "ANTHROPIC_API_KEY": "sk-...",
        "DATABASE_URL": "postgresql://user:password@localhost:5432/inventory"
      }
    }
  }
}
```

> **Important:** always use absolute paths. Claude Desktop runs with `CWD = C:\Windows\System32` on Windows, so relative paths will fail.

---

## CLI

```bash
# Single query
datryn query --uri postgresql://... --llm anthropic "How many active users?"

# Show generated SQL
datryn query --uri postgresql://... --show-sql "Top 10 products by revenue"

# MCP server (stdio)
datryn serve --uri postgresql://... --llm anthropic

# OpenAI instead of Anthropic
datryn query --uri postgresql://... --llm openai --model gpt-5.4-nano "Total orders today"
```

---

## Custom validator

```python
from datryn import SQLValidator
from datryn.validation.models import ValidationResult

class MaxTablesValidator(SQLValidator):
    def __init__(self, max_tables: int = 3, **kwargs):
        super().__init__(**kwargs)
        self._max_tables = max_tables

    def validate(self, sql: str, known_columns=None) -> ValidationResult:
        result = super().validate(sql, known_columns=known_columns)
        if result.valid and len(result.referenced_tables) > self._max_tables:
            result.valid = False
            result.reasons.append(
                f"Query touches {len(result.referenced_tables)} tables "
                f"(max allowed: {self._max_tables})"
            )
        return result

tool = DatrynTool.from_uri(uri, llm=llm, validator=MaxTablesValidator(max_tables=3))
```

---

## Schema management

The schema index is built once on the first query and cached in memory.

Call `warm_up()` before starting an agent or server to move that cost to startup time, so all queries respond at the same speed:

```python
tool.warm_up()
```

Call `refresh_schema()` after database migrations to rebuild the index:

```python
tool.refresh_schema()
```

---

## Output formats

```python
from datryn import ResultFormatter

tool = DatrynTool.from_uri(uri, llm=llm, formatter=ResultFormatter(format="json"))
# format options: "markdown" (default), "json", "csv", "natural"

# The summary header (*N row(s) returned in Xms*) is included by default —
# useful when the answer is fed directly to an LLM agent.
# Set show_header=False when the application displays row count and timing itself.
tool = DatrynTool.from_uri(uri, llm=llm, formatter=ResultFormatter(show_header=False))
```

---

## Improving context quality

datryn selects which tables to include in each prompt using BM25 keyword matching against
a document built from table names, column names, and — when available — table and column
comments. The more descriptive the schema, the better the table selection.

### Database comments (zero config)

Table and column comments are indexed automatically. Adding them once is the most
effective way to improve accuracy for schemas with non-obvious naming:

```sql
-- PostgreSQL
COMMENT ON TABLE  orders      IS 'customer purchase orders — source of revenue data';
COMMENT ON COLUMN orders.amt  IS 'total order value in USD';

-- MySQL / MariaDB
ALTER TABLE orders COMMENT = 'customer purchase orders — source of revenue data';
```

datryn re-reads comments every time the schema index is rebuilt (first call or
`tool.refresh_schema()`), so changes take effect without code changes.

### Annotations (explicit synonyms)

For schemas without comments, or when domain vocabulary differs from column names,
use `annotations` to map business terms to specific tables:

```python
tool = DatrynTool.from_uri(
    uri="postgresql://...",
    llm=llm,
    annotations={
        "tbl_rev":   ["revenue", "income", "sales"],
        "stk_movmt": ["inventory", "stock", "warehouse"],
    },
)
```

Annotations are additive — they work alongside database comments, not instead of them.

### Semantic context providers

BM25 is the default selection strategy and works well when queries and schema names
are in the same language. For multilingual queries or highly abbreviated schemas,
use a semantic context provider:

```python
from datryn import DatrynTool, AnthropicLLM, LLMContextProvider
from datryn import OpenAIEmbeddingProvider, SentenceTransformerProvider

llm = AnthropicLLM()

# Uses the configured LLM for table selection — no extra dependencies
tool = DatrynTool.from_uri(uri, llm=llm,
    context_provider=LLMContextProvider(llm))

# Table docs pre-computed at build time; question embedded per query (one small API call)
tool = DatrynTool.from_uri(uri, llm=llm,
    context_provider=OpenAIEmbeddingProvider())

# Local embeddings, no API key required
tool = DatrynTool.from_uri(uri, llm=llm,
    context_provider=SentenceTransformerProvider())
```

| Provider | Extra dependency | API key | Per-query cost |
|---|---|---|---|
| `LLMContextProvider` | none | none (uses configured LLM) | one small LLM call |
| `OpenAIEmbeddingProvider` | `openai` (already optional) | OpenAI or compatible | one embedding call (question only) |
| `SentenceTransformerProvider` | `sentence-transformers` | none (runs locally) | none (pre-computed) |

`OpenAIEmbeddingProvider` is compatible with any OpenAI-compatible endpoint (Ollama,
Azure OpenAI, etc.) via the `base_url` / `api_key` kwargs.

---

## Security model

| Layer | What it blocks |
|-------|---------------|
| Read-only mode | `INSERT`, `UPDATE`, `DELETE`, `DROP`, `TRUNCATE`, `ALTER`, `CREATE`, `GRANT`, ... |
| Allowed tables | Any table not in the explicit whitelist |
| Denied columns | `password`, `token`, `api_key`, `ssn`, `credit_card`, `cvv`, and custom patterns |
| Dangerous functions | `SLEEP()`, `pg_sleep()`, `BENCHMARK()`, `xp_cmdshell()`, `sys_exec()`, ... |
| Subquery depth | Configurable limit (default: 3 levels) |
| Hallucination detection | Qualified (`c.name`) and unqualified (`name`) column references validated against the actual schema, per scope |
| Row limit | Python slicing enforced for row-returning queries; aggregations (scalar and GROUP BY without ranking intent) are exempt from the `LIMIT` prompt instruction |
| Query timeout | Set at driver level per dialect (PostgreSQL `statement_timeout`, MySQL `MAX_EXECUTION_TIME`) |

---

## License

MIT
