annie-sdk 0.2.0__tar.gz

annie_sdk-0.2.0/.gitignore

@@ -0,0 +1,5 @@
+.venv/
+__pycache__/
+*.egg-info/
+dist/
+build/

annie_sdk-0.2.0/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.4
+Name: annie-sdk
+Version: 0.2.0
+Summary: Python SDK for querying databases using natural language via the Annie API
+Requires-Python: >=3.10
+Requires-Dist: httpx>=0.28.0
+Requires-Dist: pydantic>=2.0.0
+Provides-Extra: all
+Requires-Dist: mysql-connector-python>=9.0.0; extra == 'all'
+Requires-Dist: psycopg2-binary>=2.9.0; extra == 'all'
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == 'dev'
+Requires-Dist: ruff>=0.8.0; extra == 'dev'
+Provides-Extra: mysql
+Requires-Dist: mysql-connector-python>=9.0.0; extra == 'mysql'
+Provides-Extra: postgres
+Requires-Dist: psycopg2-binary>=2.9.0; extra == 'postgres'

annie_sdk-0.2.0/annie_sdk/__init__.py

@@ -0,0 +1,142 @@
+"""
+Annie SDK
+
+A Python SDK for querying databases using natural language.
+The SDK connects to your database through connectors, sends your
+query and schema context to the Annie API, and executes the
+resulting SQL locally against your database.
+
+Your data never leaves your infrastructure — only schema metadata
+is sent to the API.
+
+Quick start:
+    ```python
+    from annie_sdk import Agent, PostgresConnector, ConnectorTable, ConnectorTableColumn
+
+    connector = PostgresConnector(
+        connection_string="postgresql://user:pass@localhost:5432/mydb",
+        tables=[
+            ConnectorTable(
+                name="orders",
+                description="Customer orders",
+                columns=[
+                    ConnectorTableColumn(name="id", type="integer"),
+                    ConnectorTableColumn(name="amount", type="decimal"),
+                    ConnectorTableColumn(name="status", type="string"),
+                ],
+            ),
+        ],
+    )
+
+    with Agent(connector=connector, api_key="your-api-key") as agent:
+        response = agent.run("Show me top 10 orders by amount")
+        print(response.sql)
+        print(response.data)
+    ```
+"""
+
+# Agent
+from .agent import Agent
+
+# Connectors
+from .connectors.base import (
+    BaseConnector,
+    ConnectorRelationship,
+    ConnectorTable,
+    ConnectorTableColumn,
+)
+from .connectors.mock import MockConnector
+from .connectors.mysql import MySQLConnector
+from .connectors.postgres import PostgresConnector
+
+# Exceptions
+from .exceptions import (
+    AnnieError,
+    APIError,
+    AuthenticationError,
+    ConnectionError,
+    RateLimitError,
+    TimeoutError,
+    ValidationError,
+)
+
+# Models (DSL)
+from .models import (
+    AggregationFunction,
+    AggregationSpec,
+    BySpec,
+    DataModel,
+    DateBinning,
+    DateBinningInterval,
+    FilterCondition,
+    FilterStep,
+    LimitStep,
+    NumberBinning,
+    NumberBinningStrategy,
+    Operator,
+    OrderByColumn,
+    OrderByStep,
+    SelectStep,
+    SortDirection,
+    SummarizeStep,
+    VisualizationType,
+)
+
+# Response
+from .response import (
+    AgentResponse,
+    ChartContent,
+    ChartFormat,
+    DataContent,
+    ErrorContent,
+    ResponseType,
+    TextContent,
+)
+
+__all__ = [
+    # Agent
+    "Agent",
+    # Connectors
+    "BaseConnector",
+    "ConnectorRelationship",
+    "ConnectorTable",
+    "ConnectorTableColumn",
+    "MockConnector",
+    "MySQLConnector",
+    "PostgresConnector",
+    # Exceptions
+    "AnnieError",
+    "APIError",
+    "AuthenticationError",
+    "ConnectionError",
+    "RateLimitError",
+    "TimeoutError",
+    "ValidationError",
+    # Models
+    "AggregationFunction",
+    "AggregationSpec",
+    "BySpec",
+    "DataModel",
+    "DateBinning",
+    "DateBinningInterval",
+    "FilterCondition",
+    "FilterStep",
+    "LimitStep",
+    "NumberBinning",
+    "NumberBinningStrategy",
+    "Operator",
+    "OrderByColumn",
+    "OrderByStep",
+    "SelectStep",
+    "SortDirection",
+    "SummarizeStep",
+    "VisualizationType",
+    # Response
+    "AgentResponse",
+    "ChartContent",
+    "ChartFormat",
+    "DataContent",
+    "ErrorContent",
+    "ResponseType",
+    "TextContent",
+]

annie_sdk-0.2.0/annie_sdk/agent.py

@@ -0,0 +1,410 @@
+"""
+Annie SDK Agent
+
+The main entry point for the Annie SDK. The Agent handles communication
+with the Annie API and local SQL execution via connectors.
+
+Flow:
+1. User calls agent.run("natural language query")
+2. Agent sends NL query + schema context to Annie API
+3. Annie API returns DSL + SQL in the connector's dialect
+4. Agent executes SQL locally via the connector
+5. Results are formatted as AgentResponse
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import random
+from decimal import Decimal
+from typing import Any
+
+import httpx
+
+from .connectors.base import BaseConnector
+from .exceptions import APIError, AuthenticationError, RateLimitError, TimeoutError, ValidationError
+from .response import AgentResponse, ChartFormat
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_API_URL = "https://api.pandas-ai.com"
+DEFAULT_TIMEOUT = 60
+
+DEFAULT_DENIAL_MESSAGES = [
+    "I can only help with questions about your data. Try asking something like 'Show revenue by month'.",
+    "That doesn't seem to be a data question. Try asking about your data, like 'Top 10 customers by sales'.",
+    "I'm a data assistant. Ask me about your data — for example, 'What's the average order value?'",
+]
+
+
+class Agent:
+    """Annie SDK Agent.
+
+    Handles natural language queries against your database using the
+    Annie API for NL-to-SQL conversion and local execution for data privacy.
+
+    Args:
+        connector: Database connector instance
+        api_key: Annie API key (or set ANNIE_API_KEY env var)
+        api_url: Annie API URL (or set ANNIE_API_URL env var)
+        denial_messages: Custom messages for non-data queries
+
+    Example:
+        ```python
+        from sdk import Agent, PostgresConnector, ConnectorTable
+
+        connector = PostgresConnector(
+            connection_string="postgresql://user:pass@localhost/db",
+            tables=[ConnectorTable(name="orders", description="Customer orders")],
+        )
+
+        with Agent(connector=connector, api_key="your-key") as agent:
+            response = agent.run("Show me top 10 orders by amount")
+            print(response.data)
+        ```
+    """
+
+    def __init__(
+        self,
+        connector: BaseConnector,
+        api_key: str | None = None,
+        api_url: str | None = None,
+        denial_messages: list[str] | None = None,
+    ):
+        # Resolve API key
+        self._api_key = api_key or os.environ.get("ANNIE_API_KEY")
+        if not self._api_key:
+            raise AuthenticationError(
+                "API key is required. Pass api_key parameter or set ANNIE_API_KEY environment variable."
+            )
+
+        # Resolve API URL
+        self._api_url = (
+            api_url or os.environ.get("ANNIE_API_URL") or DEFAULT_API_URL
+        ).rstrip("/")
+
+        self._connector = connector
+        self._denial_messages = denial_messages or DEFAULT_DENIAL_MESSAGES
+
+        # Validate connector has tables
+        if not connector.tables:
+            raise ValidationError(
+                "Connector must have at least one table configured",
+                field="connector.tables",
+            )
+
+        # Initialize HTTP client
+        self._client = httpx.Client(
+            base_url=self._api_url,
+            headers={
+                "Authorization": f"Bearer {self._api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=DEFAULT_TIMEOUT,
+        )
+
+    def run(self, query: str, *, explain: bool = False) -> AgentResponse:
+        """Run a natural language query against your database.
+
+        Args:
+            query: Natural language query (e.g., "Show me top 10 customers by revenue")
+            explain: If True, also generate AI insights about the results
+
+        Returns:
+            AgentResponse with data, SQL, DSL, and optional insights
+        """
+        if not query or not query.strip():
+            return AgentResponse.from_error("Query cannot be empty")
+
+        # 1. Build schema context from connector
+        schema_context = self._connector.get_schema_context()
+
+        # 2. Call Annie API: NL → DSL + SQL
+        api_response = self._call_query_api(query, schema_context)
+
+        # 3. Handle non-data responses (guardrail denials)
+        if api_response.get("type") == "text":
+            message = api_response.get("message") or random.choice(self._denial_messages)
+            response = AgentResponse(
+                success=True,
+                visualization="table",
+            )
+            response.add_text(message)
+            return response
+
+        # 4. Handle API errors
+        if not api_response.get("success", True):
+            error_msg = api_response.get("message", "API returned an error")
+            return AgentResponse.from_error(error_msg)
+
+        sql = api_response.get("sql")
+        dsl = api_response.get("dsl")
+        visualization = api_response.get("visualization", "table")
+
+        if not sql:
+            return AgentResponse.from_error("API did not return SQL")
+
+        # 5. Execute SQL locally via connector
+        try:
+            if not self._connector.is_connected():
+                self._connector.connect()
+
+            results = self._connector.execute(sql)
+        except Exception as e:
+            response = AgentResponse(
+                sql=sql,
+                dsl=dsl,
+                visualization=visualization,
+                success=False,
+            )
+            response.add_error(f"Query execution failed: {e}")
+            return response
+
+        # 6. Format results as AgentResponse
+        response = AgentResponse(
+            sql=sql,
+            dsl=dsl,
+            visualization=visualization,
+            success=True,
+        )
+
+        if results:
+            columns = list(results[0].keys())
+            rows = [list(row.values()) for row in results]
+            response.add_data(columns=columns, rows=rows)
+        else:
+            response.add_data(columns=[], rows=[])
+
+        # 6b. Generate chart from results
+        if results and visualization not in ("table", "kpi"):
+            chart_config = self._generate_chart(results, dsl or {}, visualization)
+            if chart_config:
+                response.add_chart(
+                    content=chart_config["content"],
+                    format=ChartFormat(chart_config.get("format", "chartjs")),
+                    title=chart_config.get("title"),
+                )
+
+        # 7. Optionally get AI explanation
+        if explain and results:
+            try:
+                explanation = self._call_explain_api(
+                    query=query,
+                    results=results[:50],
+                    sql=sql,
+                    dsl=dsl,
+                    schema=schema_context,
+                )
+                if explanation:
+                    response.add_text(explanation)
+            except Exception as e:
+                logger.warning(f"Failed to get explanation: {e}")
+                # Don't fail the whole response if explain fails
+
+        return response
+
+    def close(self) -> None:
+        """Close the agent and its connector."""
+        try:
+            self._connector.disconnect()
+        except Exception:
+            pass
+        try:
+            self._client.close()
+        except Exception:
+            pass
+
+    # =========================================================================
+    # Context manager
+    # =========================================================================
+
+    def __enter__(self) -> Agent:
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
+    # =========================================================================
+    # Private methods
+    # =========================================================================
+
+    def _call_query_api(
+        self,
+        query: str,
+        schema: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Call the Annie /sdk/query endpoint."""
+        payload = {
+            "query": query,
+            "schema": schema,
+            "dialect": self._connector.dialect,
+        }
+
+        try:
+            response = self._client.post("/sdk/query", json=payload)
+        except httpx.TimeoutException as e:
+            raise TimeoutError(f"API request timed out: {e}")
+        except httpx.ConnectError as e:
+            raise APIError(f"Failed to connect to Annie API at {self._api_url}: {e}")
+
+        if response.status_code == 401:
+            raise AuthenticationError("Invalid API key")
+        if response.status_code == 429:
+            retry_after = response.headers.get("retry-after")
+            raise RateLimitError(
+                "Rate limit exceeded",
+                retry_after=int(retry_after) if retry_after else None,
+            )
+        if response.status_code >= 400:
+            raise APIError(
+                f"Annie API error: {response.status_code}",
+                status_code=response.status_code,
+            )
+
+        try:
+            return response.json()
+        except Exception as e:
+            body_preview = response.text[:200] if response.text else "(empty)"
+            raise APIError(
+                f"Failed to parse API response (status {response.status_code}): {e}. "
+                f"Response: {body_preview}"
+            )
+
+    def _call_explain_api(
+        self,
+        query: str,
+        results: list[dict[str, Any]],
+        sql: str | None,
+        dsl: dict[str, Any] | None,
+        schema: dict[str, Any] | None,
+    ) -> str | None:
+        """Call the Annie /sdk/explain endpoint."""
+        payload = {
+            "query": query,
+            "results": results,
+            "sql": sql,
+            "dsl": dsl,
+            "schema": schema,
+        }
+
+        try:
+            response = self._client.post("/sdk/explain", json=payload)
+
+            if response.status_code != 200:
+                logger.warning(f"Explain API returned {response.status_code}")
+                return None
+
+            data = response.json()
+            return data.get("explanation") if data.get("success") else None
+
+        except Exception as e:
+            logger.warning(f"Explain API call failed: {e}")
+            return None
+
+    def _generate_chart(
+        self,
+        results: list[dict[str, Any]],
+        dsl: dict[str, Any],
+        visualization: str,
+    ) -> dict[str, Any] | None:
+        """Generate a ChartJS configuration from query results.
+
+        Creates a ChartJS-compatible configuration based on the data structure
+        and the visualization type returned by the API.
+
+        Args:
+            results: Query results as list of dicts
+            dsl: The DSL used for the query
+            visualization: Visualization type from the API (bar, line, pie, etc.)
+
+        Returns:
+            Chart configuration dict or None if not applicable
+        """
+        if not results or len(results) < 2:
+            return None
+
+        # Find label column (first string-typed) and value columns (numeric).
+        # Scan up to 5 rows to handle cases where the first row has NULLs.
+        columns = list(results[0].keys())
+        label_col = None
+        value_cols = []
+
+        for col in columns:
+            # Find the first non-None value in up to 5 rows
+            sample_val = None
+            for row in results[:5]:
+                v = row.get(col)
+                if v is not None:
+                    sample_val = v
+                    break
+
+            if sample_val is not None and isinstance(sample_val, str) and label_col is None:
+                label_col = col
+            elif isinstance(sample_val, (int, float, Decimal)):
+                value_cols.append(col)
+
+        if not label_col or not value_cols:
+            return None
+
+        # Extract data
+        labels = [str(row.get(label_col) or "") for row in results]
+
+        colors = [
+            "rgba(54, 162, 235, 0.8)",
+            "rgba(255, 99, 132, 0.8)",
+            "rgba(255, 206, 86, 0.8)",
+            "rgba(75, 192, 192, 0.8)",
+            "rgba(153, 102, 255, 0.8)",
+            "rgba(255, 159, 64, 0.8)",
+            "rgba(199, 199, 199, 0.8)",
+            "rgba(83, 102, 255, 0.8)",
+        ]
+
+        chart_type = visualization if visualization in ("bar", "line", "pie", "doughnut") else "bar"
+
+        datasets = []
+        for i, val_col in enumerate(value_cols[:2]):  # Limit to 2 datasets
+            values = [float(row.get(val_col) or 0) for row in results]
+            dataset: dict[str, Any] = {
+                "label": val_col.replace("_", " ").title(),
+                "data": values,
+                "borderWidth": 1,
+            }
+
+            if chart_type in ("pie", "doughnut"):
+                # Pie/doughnut: different color per slice
+                dataset["backgroundColor"] = colors[: len(labels)]
+            elif len(value_cols) == 1:
+                # Single dataset: different color per bar
+                dataset["backgroundColor"] = colors[: len(labels)]
+            else:
+                # Multiple datasets: same color per dataset
+                dataset["backgroundColor"] = colors[i % len(colors)]
+                dataset["borderColor"] = colors[i % len(colors)].replace("0.8", "1")
+
+            datasets.append(dataset)
+
+        title = (
+            f"{value_cols[0].replace('_', ' ').title()} by "
+            f"{label_col.replace('_', ' ').title()}"
+        )
+
+        return {
+            "format": "chartjs",
+            "title": title,
+            "content": {
+                "type": chart_type,
+                "data": {
+                    "labels": labels,
+                    "datasets": datasets,
+                },
+                "options": {
+                    "responsive": True,
+                    "plugins": {
+                        "legend": {"position": "top"},
+                        "title": {"display": True, "text": title},
+                    },
+                },
+            },
+        }

annie_sdk-0.2.0/annie_sdk/connectors/__init__.py

@@ -0,0 +1,16 @@
+"""Annie SDK Connectors."""
+
+from .base import BaseConnector, ConnectorRelationship, ConnectorTable, ConnectorTableColumn
+from .mock import MockConnector
+from .mysql import MySQLConnector
+from .postgres import PostgresConnector
+
+__all__ = [
+    "BaseConnector",
+    "ConnectorRelationship",
+    "ConnectorTable",
+    "ConnectorTableColumn",
+    "MockConnector",
+    "MySQLConnector",
+    "PostgresConnector",
+]