# Annie SDK

Python SDK for querying databases using natural language via the Annie API.

## Overview

The Annie SDK lets you ask questions about your data in plain English. It connects to your database, sends your query and schema metadata to the Annie API, and executes the generated SQL locally — **your data never leaves your infrastructure**.

```
Natural language query
        ↓
    Annie API        ← only schema metadata is sent
        ↓
    DSL + SQL
        ↓
  Local execution    ← SQL runs on YOUR database
        ↓
  AgentResponse      ← data, charts, insights
```

## Installation

```bash
pip install annie-sdk
```

Install with database-specific drivers:

```bash
# PostgreSQL
pip install annie-sdk[postgres]

# MySQL
pip install annie-sdk[mysql]

# All drivers
pip install annie-sdk[all]
```

**Requirements:** Python 3.10+

## Quick start

```python
from annie_sdk import Agent, PostgresConnector, ConnectorTable

connector = PostgresConnector(
    connection_string="postgresql://user:pass@localhost:5432/mydb",
    tables=[
        ConnectorTable(name="orders", description="Customer orders"),
        ConnectorTable(name="products", description="Product catalog"),
    ],
)

with Agent(connector=connector, api_key="your-api-key") as agent:
    response = agent.run("Show me top 10 orders by amount")

    print(response.sql)          # Generated SQL
    print(response.data)         # DataContent with columns + rows
    print(response.chart)        # ChartJS configuration (if applicable)
```

## Connectors

Connectors handle the database connection and SQL execution. They provide schema metadata to the Annie API and run the generated SQL locally.

### PostgresConnector

```python
from annie_sdk import PostgresConnector, ConnectorTable

# Using a connection string
connector = PostgresConnector(
    connection_string="postgresql://user:pass@localhost:5432/mydb",
    tables=[ConnectorTable(name="orders")],
)

# Using individual parameters
connector = PostgresConnector(
    host="localhost",
    port=5432,
    database="mydb",
    user="user",
    password="pass",
    schema="public",  # defaults to 'public'
    tables=[ConnectorTable(name="orders")],
)
```

Requires: `psycopg2-binary` (installed with `pip install annie-sdk[postgres]`)

### MySQLConnector

```python
from annie_sdk import MySQLConnector, ConnectorTable

# Using a connection string
connector = MySQLConnector(
    connection_string="mysql://user:pass@localhost:3306/mydb",
    tables=[ConnectorTable(name="orders")],
)

# Using individual parameters
connector = MySQLConnector(
    host="localhost",
    port=3306,
    database="mydb",
    user="user",
    password="pass",
    tables=[ConnectorTable(name="orders")],
)
```

Requires: `mysql-connector-python` (installed with `pip install annie-sdk[mysql]`)

### MockConnector

For testing without a real database:

```python
from annie_sdk import MockConnector, ConnectorTable

connector = MockConnector(
    dialect_name="postgres",
    tables=[ConnectorTable(name="orders")],
    mock_results=[
        {"id": 1, "amount": 100.0},
        {"id": 2, "amount": 200.0},
    ],
)

# After executing queries, inspect what was run
connector.execute("SELECT * FROM orders")
print(connector.last_query)        # "SELECT * FROM orders"
print(connector.executed_queries)   # ["SELECT * FROM orders"]

# Change results for the next query
connector.set_results([{"total": 300.0}])
```

### Custom connectors

Extend `BaseConnector` to support other databases:

```python
from annie_sdk import BaseConnector, ConnectorTableColumn

class BigQueryConnector(BaseConnector):
    @property
    def dialect(self) -> str:
        return "bigquery"

    def connect(self) -> None: ...
    def disconnect(self) -> None: ...
    def execute(self, sql: str) -> list[dict]: ...
    def discover_tables(self) -> list[str]: ...
    def discover_columns(self, table_name: str) -> list[ConnectorTableColumn]: ...
```

## Table definitions

### ConnectorTable

```python
from annie_sdk import ConnectorTable, ConnectorTableColumn

# Minimal — columns will be auto-discovered from the database
table = ConnectorTable(name="orders")

# With description — helps the AI understand your data
table = ConnectorTable(
    name="orders",
    description="Customer orders with payment info",
)

# With explicit columns — skips auto-discovery
table = ConnectorTable(
    name="orders",
    columns=[
        ConnectorTableColumn(name="id", type="integer"),
        ConnectorTableColumn(name="customer_name", type="string"),
        ConnectorTableColumn(name="amount", type="decimal"),
        ConnectorTableColumn(
            name="status",
            type="string",
            description="Order status: active, completed, cancelled",
        ),
        ConnectorTableColumn(name="created_at", type="timestamp"),
    ],
)
```

### Auto-discovery

When `columns` is not provided, the SDK connects to your database and discovers the column names and types automatically. Discovered columns are cached for the lifetime of the connector.

### Column types

Supported type strings: `integer`, `float`, `decimal`, `boolean`, `string`, `date`, `time`, `timestamp`, `binary`, `json`, `uuid`.

## Relationships

Define relationships between tables to enable cross-table queries with automatic JOINs:

```python
from annie_sdk import ConnectorRelationship, ConnectorTable, PostgresConnector

connector = PostgresConnector(
    connection_string="postgresql://...",
    tables=[
        ConnectorTable(name="orders"),
        ConnectorTable(name="products"),
    ],
    relationships=[
        ConnectorRelationship(
            primary_table="orders",
            primary_column="product_id",
            foreign_table="products",
            foreign_column="id",
            description="Each order references a product",
        ),
    ],
)

with Agent(connector=connector, api_key="your-key") as agent:
    # This query spans two tables — the SDK handles the JOIN automatically
    response = agent.run("Show total revenue by product category")
    print(response.sql)
    # SELECT products.category, SUM(orders.amount) ...
    # FROM orders LEFT JOIN products ON orders.product_id = products.id
    # GROUP BY products.category
```

### ConnectorRelationship fields

| Field | Required | Description |
|-------|----------|-------------|
| `primary_table` | Yes | Table containing the foreign key |
| `primary_column` | Yes | Foreign key column |
| `foreign_table` | Yes | Referenced table |
| `foreign_column` | Yes | Referenced column (usually the primary key) |
| `description` | No | Helps the AI understand the relationship |

## Agent

### Initialization

```python
from annie_sdk import Agent

agent = Agent(
    connector=connector,
    api_key="your-api-key",        # or set ANNIE_API_KEY env var
    api_url="https://chat.pandas-ai.com", # or set ANNIE_API_URL env var (this is the default)
)
```

### Running queries

```python
# Basic query
response = agent.run("Show me top 10 customers by revenue")

# With AI explanation of the results
response = agent.run("Revenue trend by month", explain=True)
if response.text:
    print(response.text)  # AI-generated insight
```

### Context manager

```python
with Agent(connector=connector, api_key="your-key") as agent:
    response = agent.run("...")
# Connector and HTTP client are automatically closed
```

### Non-data queries

If a query is not data-related (e.g., "What's the weather?"), the agent returns a text response with a helpful message. You can customize these:

```python
agent = Agent(
    connector=connector,
    api_key="key",
    denial_messages=[
        "I can only answer questions about your data.",
        "Try asking something about your database.",
    ],
)
```

## Response types

### AgentResponse

The main response container. Contains multiple content items plus metadata.

```python
response = agent.run("Show revenue by month")

# Metadata
response.success        # bool — whether the query succeeded
response.sql            # str — generated SQL
response.dsl            # dict — query DSL
response.visualization  # str — recommended chart type ("bar", "line", "pie", "table", etc.)

# Content accessors (return first item or None)
response.text           # str | None — first text/insight
response.data           # DataContent | None — first data result
response.chart          # ChartContent | None — first chart config
response.error          # ErrorContent | None — first error

# Multiple items
response.texts          # list[str] — all text contents
response.charts         # list[ChartContent] — all charts
response.all_data       # list[DataContent] — all data items

# Iteration
for item in response:
    print(item.type)    # "text", "data", "chart", or "error"

# Serialization
response.to_dict()      # dict
response.to_json()      # JSON string
str(response)           # JSON string (same as to_json)
```

### DataContent

```python
data = response.data

data.columns    # list[str] — column names
data.rows       # list[list] — row data
data.row_count  # int — number of rows

# Convert to list of dicts
records = data.to_records()
# [{"name": "Alice", "amount": 100}, {"name": "Bob", "amount": 200}]

# Serialization
data.to_json()  # JSON string
str(data)       # JSON string
```

### ChartContent

Chart configurations are ChartJS-compatible and JSON-safe — `str()` and `print()` always produce valid JSON with lowercase `true`/`false`/`null` (never Python's `True`/`False`/`None`).

```python
chart = response.chart

chart.content   # dict — ChartJS configuration
chart.format    # ChartFormat.CHARTJS
chart.title     # str | None — chart title

# Safe for JavaScript consumption
print(chart)            # valid JSON string
print(chart.content)    # valid JSON string (nested dicts are JSON-safe too)
chart.to_json()         # JSON string with full metadata
```

Example `chart.content`:

```json
{
  "type": "bar",
  "data": {
    "labels": ["Jan", "Feb", "Mar"],
    "datasets": [{
      "label": "Revenue",
      "data": [1000, 1500, 1200],
      "backgroundColor": ["rgba(54, 162, 235, 0.8)", ...]
    }]
  },
  "options": {
    "responsive": true,
    "plugins": {
      "legend": {"position": "top"},
      "title": {"display": true, "text": "Revenue By Month"}
    }
  }
}
```

### TextContent

```python
text = response.text  # str — the text content
```

### ErrorContent

```python
if response.error:
    print(response.error.content)  # str — error message
    print(response.error.code)     # str | None — error code
```

### Factory methods

```python
# Create an error response directly
response = AgentResponse.from_error("Something went wrong", code="ERR_001")
```

## Error handling

All SDK exceptions inherit from `AnnieError`:

```python
from annie_sdk import (
    AnnieError,
    APIError,
    AuthenticationError,
    ConnectionError,
    RateLimitError,
    TimeoutError,
    ValidationError,
)

try:
    response = agent.run("Show me revenue")
except AuthenticationError:
    print("Invalid API key")
except RateLimitError as e:
    print(f"Rate limited. Retry after {e.retry_after}s")
except TimeoutError:
    print("Request timed out")
except APIError as e:
    print(f"API error {e.status_code}: {e.message}")
except AnnieError as e:
    print(f"SDK error: {e.message}")
```

### Exception reference

| Exception | Attributes | Description |
|-----------|-----------|-------------|
| `AnnieError` | `message` | Base exception |
| `APIError` | `message`, `status_code`, `code`, `response` | API communication errors |
| `AuthenticationError` | `message` | Invalid or missing API key |
| `ValidationError` | `message`, `field` | Invalid configuration |
| `ConnectionError` | `message`, `original_error` | Database connection failures |
| `TimeoutError` | `message` | Request timeouts |
| `RateLimitError` | `message`, `retry_after` | Rate limit exceeded |

## API reference

### Exports

| Name | Type | Description |
|------|------|-------------|
| `Agent` | class | Main entry point for NL queries |
| `PostgresConnector` | class | PostgreSQL connector |
| `MySQLConnector` | class | MySQL connector |
| `MockConnector` | class | Testing connector |
| `BaseConnector` | class | Abstract base for custom connectors |
| `ConnectorTable` | model | Table definition |
| `ConnectorTableColumn` | model | Column definition |
| `ConnectorRelationship` | model | JOIN relationship |
| `AgentResponse` | class | Query response container |
| `DataContent` | class | Data rows and columns |
| `ChartContent` | class | ChartJS configuration |
| `TextContent` | class | Text/insight content |
| `ErrorContent` | class | Error content |
| `ChartFormat` | enum | Chart format (CHARTJS) |
| `ResponseType` | enum | Content type (TEXT, DATA, CHART, ERROR) |
| `DataModel` | model | Query transformation steps |
| `FilterStep` | model | WHERE conditions |
| `SummarizeStep` | model | GROUP BY with aggregations |
| `OrderByStep` | model | ORDER BY clause |
| `LimitStep` | model | LIMIT clause |
| `SelectStep` | model | SELECT columns |
| `Operator` | enum | Filter operators (equals, gt, contains, ...) |
| `AggregationFunction` | enum | Aggregation functions (count, sum, avg, ...) |
| `SortDirection` | enum | Sort direction (asc, desc) |
| `VisualizationType` | enum | Chart types (bar, line, pie, ...) |

## Publishing

The SDK is published to [PyPI](https://pypi.org/project/annie-sdk/) automatically via GitHub Actions when a new release is created.

### Release process

1. Update `version` in `sdk/pyproject.toml`
2. Commit and merge to `main`
3. Create a GitHub release with tag `sdk-v<version>` (e.g., `sdk-v0.2.0a1`)
4. The `publish-sdk.yml` workflow builds and publishes to PyPI automatically

### Versioning

The SDK follows [PEP 440](https://peps.python.org/pep-0440/) versioning:

- `0.2.0` — stable release
- `0.2.0a1` — alpha pre-release
- `0.2.0b1` — beta pre-release
- `0.2.0rc1` — release candidate

### How it works

The GitHub Actions workflow (`.github/workflows/publish-sdk.yml`) uses [PyPI trusted publishing](https://docs.pypi.org/trusted-publishers/) (OIDC) — no API tokens or secrets are needed. When a release with an `sdk-v*` tag is published, the workflow:

1. Checks out the code
2. Builds the wheel and source distribution using `python -m build` (hatchling backend)
3. Publishes to PyPI via the official `pypa/gh-action-pypi-publish` action

### Setup (one-time)

For this to work, two things must be configured:

1. **GitHub environment** — a `pypi` environment must exist in the repo settings (Settings > Environments)
2. **PyPI trusted publisher** — the `annie-sdk` package on pypi.org must have a trusted publisher configured with:
   - **Owner**: the GitHub org
   - **Repository**: `annie`
   - **Workflow name**: `publish-sdk.yml`
   - **Environment name**: `pypi`

If the package doesn't exist on PyPI yet, use a [pending publisher](https://docs.pypi.org/trusted-publishers/creating-a-project-through-oidc/) to reserve the name and pre-authorize the workflow.
