hatidata-agent 0.1.0__tar.gz

hatidata_agent-0.1.0/PKG-INFO

@@ -0,0 +1,183 @@
+Metadata-Version: 2.4
+Name: hatidata-agent
+Version: 0.1.0
+Summary: HatiData Agent SDK — RAM for Agents
+Author-email: Marviy Pte Ltd <eng@hatidata.com>
+License: Apache-2.0
+Project-URL: Homepage, https://hatidata.com
+Project-URL: Documentation, https://docs.hatiosai.com/hatidata
+Project-URL: Repository, https://github.com/HatiOS-AI/HatiData-Core
+Project-URL: Changelog, https://github.com/HatiOS-AI/HatiData-Core/releases
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Database
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: psycopg2-binary>=2.9
+Provides-Extra: async
+Requires-Dist: asyncpg>=0.29; extra == "async"
+Provides-Extra: langchain
+Requires-Dist: langchain-community>=0.2; extra == "langchain"
+Provides-Extra: mcp
+Requires-Dist: mcp>=1.0; extra == "mcp"
+Provides-Extra: all
+Requires-Dist: asyncpg>=0.29; extra == "all"
+Requires-Dist: langchain-community>=0.2; extra == "all"
+Requires-Dist: mcp>=1.0; extra == "all"
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.23; extra == "dev"
+
+# HatiData Agent SDK
+
+[![PyPI version](https://img.shields.io/pypi/v/hatidata-agent.svg)](https://pypi.org/project/hatidata-agent/)
+[![Python versions](https://img.shields.io/pypi/pyversions/hatidata-agent.svg)](https://pypi.org/project/hatidata-agent/)
+[![License](https://img.shields.io/pypi/l/hatidata-agent.svg)](https://github.com/HatiOS-AI/HatiData-Core/blob/main/LICENSE)
+
+**RAM for Agents** — Python SDK for AI agents to query HatiData's in-VPC data warehouse with sub-10ms latency via Postgres wire protocol.
+
+## Installation
+
+```bash
+pip install hatidata-agent
+
+# With async support
+pip install hatidata-agent[async]
+
+# With LangChain support
+pip install hatidata-agent[langchain]
+
+# With MCP server
+pip install hatidata-agent[mcp]
+
+# Everything
+pip install hatidata-agent[all]
+```
+
+## Quick Start
+
+```python
+from hatidata_agent import HatiDataAgent
+
+agent = HatiDataAgent(
+    host="proxy.internal",
+    port=5439,
+    agent_id="my-agent",
+    framework="langchain",
+)
+
+# Simple query
+rows = agent.query("SELECT * FROM customers WHERE status = 'active' LIMIT 10")
+print(rows)  # [{"id": 1, "name": "Acme Corp", ...}, ...]
+```
+
+## Features
+
+- **Sub-10ms query latency** -- In-VPC execution, no data leaves your network
+- **Postgres wire protocol** -- Works with any Postgres client library
+- **Reasoning chain tracking** -- Multi-step audit trails for agent workflows
+- **RAG context retrieval** -- Full-text and vector similarity search
+- **LangChain integration** -- Drop-in `SQLDatabase` replacement for LangChain agents
+- **MCP server** -- Expose HatiData as tools for Claude and other MCP-compatible agents
+- **Agent identification** -- Per-agent billing, priority scheduling, and audit via Postgres startup parameters
+- **Snowflake SQL compatible** -- Bring existing queries without rewrites
+
+## Reasoning Chain Tracking
+
+Track multi-step reasoning chains for audit and debugging:
+
+```python
+with agent.reasoning_chain("req-001") as chain:
+    # Step 0: Discover tables
+    tables = chain.query("SELECT table_name FROM information_schema.tables")
+
+    # Step 1: Get relevant data
+    customers = chain.query("SELECT * FROM customers WHERE tier = 'enterprise'", step=1)
+
+    # Step 2: Aggregate
+    revenue = chain.query("SELECT SUM(revenue) FROM orders WHERE customer_id IN (...)", step=2)
+```
+
+## RAG Context Retrieval
+
+```python
+# Full-text search
+context = agent.get_context("customers", "enterprise accounts in US", top_k=5)
+
+# Vector similarity search
+context = agent.get_rag_context("docs", "embedding", query_vector, top_k=10)
+```
+
+## LangChain Integration
+
+```python
+from hatidata_agent.langchain import HatiDataSQLDatabase
+from langchain.agents import create_sql_agent
+
+db = HatiDataSQLDatabase(
+    host="proxy.internal",
+    port=5439,
+    agent_id="sql-agent-1",
+)
+
+agent = create_sql_agent(llm=llm, db=db, verbose=True)
+result = agent.run("How many enterprise customers do we have?")
+```
+
+## MCP Server
+
+Run as an MCP server for Claude and other MCP-compatible agents:
+
+```bash
+# Start the MCP server
+hatidata-mcp-server --host proxy.internal --port 5439
+
+# Or add to Claude Code's MCP config:
+# ~/.claude/mcp.json
+{
+  "mcpServers": {
+    "hatidata": {
+      "command": "hatidata-mcp-server",
+      "args": ["--host", "proxy.internal", "--port", "5439"]
+    }
+  }
+}
+```
+
+### MCP Tools
+
+| Tool | Description |
+|------|-------------|
+| `query` | Execute SQL and return JSON results |
+| `list_tables` | List available tables |
+| `describe_table` | Get table schema |
+| `get_context` | RAG context retrieval via full-text search |
+
+## Agent Identification
+
+The SDK automatically identifies agents via Postgres startup parameters:
+
+| Parameter | Purpose |
+|-----------|---------|
+| `hatidata_agent_id` | Unique agent identifier |
+| `hatidata_framework` | AI framework (langchain, crewai, autogen, etc.) |
+| `hatidata_priority` | Scheduling priority (low, normal, high, critical) |
+| `hatidata_request_id` | Request/reasoning chain ID |
+| `hatidata_reasoning_step` | Step number within a chain |
+
+These enable per-agent billing, priority scheduling, audit trails, and the Agent Tax Report showing savings vs legacy cloud warehouses.
+
+## Documentation
+
+Full documentation is available at [docs.hatiosai.com/hatidata](https://docs.hatiosai.com/hatidata).
+
+## License
+
+Apache License 2.0. Copyright (c) Marviy Pte Ltd. See [LICENSE](https://github.com/HatiOS-AI/HatiData-Core/blob/main/LICENSE) for details.

hatidata_agent-0.1.0/README.md

@@ -0,0 +1,146 @@
+# HatiData Agent SDK
+
+[![PyPI version](https://img.shields.io/pypi/v/hatidata-agent.svg)](https://pypi.org/project/hatidata-agent/)
+[![Python versions](https://img.shields.io/pypi/pyversions/hatidata-agent.svg)](https://pypi.org/project/hatidata-agent/)
+[![License](https://img.shields.io/pypi/l/hatidata-agent.svg)](https://github.com/HatiOS-AI/HatiData-Core/blob/main/LICENSE)
+
+**RAM for Agents** — Python SDK for AI agents to query HatiData's in-VPC data warehouse with sub-10ms latency via Postgres wire protocol.
+
+## Installation
+
+```bash
+pip install hatidata-agent
+
+# With async support
+pip install hatidata-agent[async]
+
+# With LangChain support
+pip install hatidata-agent[langchain]
+
+# With MCP server
+pip install hatidata-agent[mcp]
+
+# Everything
+pip install hatidata-agent[all]
+```
+
+## Quick Start
+
+```python
+from hatidata_agent import HatiDataAgent
+
+agent = HatiDataAgent(
+    host="proxy.internal",
+    port=5439,
+    agent_id="my-agent",
+    framework="langchain",
+)
+
+# Simple query
+rows = agent.query("SELECT * FROM customers WHERE status = 'active' LIMIT 10")
+print(rows)  # [{"id": 1, "name": "Acme Corp", ...}, ...]
+```
+
+## Features
+
+- **Sub-10ms query latency** -- In-VPC execution, no data leaves your network
+- **Postgres wire protocol** -- Works with any Postgres client library
+- **Reasoning chain tracking** -- Multi-step audit trails for agent workflows
+- **RAG context retrieval** -- Full-text and vector similarity search
+- **LangChain integration** -- Drop-in `SQLDatabase` replacement for LangChain agents
+- **MCP server** -- Expose HatiData as tools for Claude and other MCP-compatible agents
+- **Agent identification** -- Per-agent billing, priority scheduling, and audit via Postgres startup parameters
+- **Snowflake SQL compatible** -- Bring existing queries without rewrites
+
+## Reasoning Chain Tracking
+
+Track multi-step reasoning chains for audit and debugging:
+
+```python
+with agent.reasoning_chain("req-001") as chain:
+    # Step 0: Discover tables
+    tables = chain.query("SELECT table_name FROM information_schema.tables")
+
+    # Step 1: Get relevant data
+    customers = chain.query("SELECT * FROM customers WHERE tier = 'enterprise'", step=1)
+
+    # Step 2: Aggregate
+    revenue = chain.query("SELECT SUM(revenue) FROM orders WHERE customer_id IN (...)", step=2)
+```
+
+## RAG Context Retrieval
+
+```python
+# Full-text search
+context = agent.get_context("customers", "enterprise accounts in US", top_k=5)
+
+# Vector similarity search
+context = agent.get_rag_context("docs", "embedding", query_vector, top_k=10)
+```
+
+## LangChain Integration
+
+```python
+from hatidata_agent.langchain import HatiDataSQLDatabase
+from langchain.agents import create_sql_agent
+
+db = HatiDataSQLDatabase(
+    host="proxy.internal",
+    port=5439,
+    agent_id="sql-agent-1",
+)
+
+agent = create_sql_agent(llm=llm, db=db, verbose=True)
+result = agent.run("How many enterprise customers do we have?")
+```
+
+## MCP Server
+
+Run as an MCP server for Claude and other MCP-compatible agents:
+
+```bash
+# Start the MCP server
+hatidata-mcp-server --host proxy.internal --port 5439
+
+# Or add to Claude Code's MCP config:
+# ~/.claude/mcp.json
+{
+  "mcpServers": {
+    "hatidata": {
+      "command": "hatidata-mcp-server",
+      "args": ["--host", "proxy.internal", "--port", "5439"]
+    }
+  }
+}
+```
+
+### MCP Tools
+
+| Tool | Description |
+|------|-------------|
+| `query` | Execute SQL and return JSON results |
+| `list_tables` | List available tables |
+| `describe_table` | Get table schema |
+| `get_context` | RAG context retrieval via full-text search |
+
+## Agent Identification
+
+The SDK automatically identifies agents via Postgres startup parameters:
+
+| Parameter | Purpose |
+|-----------|---------|
+| `hatidata_agent_id` | Unique agent identifier |
+| `hatidata_framework` | AI framework (langchain, crewai, autogen, etc.) |
+| `hatidata_priority` | Scheduling priority (low, normal, high, critical) |
+| `hatidata_request_id` | Request/reasoning chain ID |
+| `hatidata_reasoning_step` | Step number within a chain |
+
+These enable per-agent billing, priority scheduling, audit trails, and the Agent Tax Report showing savings vs legacy cloud warehouses.
+
+## Documentation
+
+Full documentation is available at [docs.hatiosai.com/hatidata](https://docs.hatiosai.com/hatidata).
+
+## License
+
+Apache License 2.0. Copyright (c) Marviy Pte Ltd. See [LICENSE](https://github.com/HatiOS-AI/HatiData-Core/blob/main/LICENSE) for details.

hatidata_agent-0.1.0/hatidata_agent/__init__.py

@@ -0,0 +1,23 @@
+"""HatiData Agent SDK — Sub-10ms SQL for AI agents.
+
+Provides agent-aware database access to HatiData's in-VPC data warehouse
+via the Postgres wire protocol. Agents identify themselves through startup
+parameters, enabling per-agent billing, scheduling, and audit trails.
+
+Quick start::
+
+    from hatidata_agent import HatiDataAgent
+
+    agent = HatiDataAgent(
+        host="localhost",
+        port=5439,
+        agent_id="my-agent",
+        framework="langchain",
+    )
+    rows = agent.query("SELECT * FROM customers WHERE status = 'active' LIMIT 10")
+"""
+
+from hatidata_agent.client import HatiDataAgent
+
+__version__ = "0.1.0"
+__all__ = ["HatiDataAgent"]

hatidata_agent-0.1.0/hatidata_agent/client.py

@@ -0,0 +1,283 @@
+"""HatiData agent client — Postgres wire protocol with agent startup params.
+
+Usage::
+
+    from hatidata_agent import HatiDataAgent
+
+    agent = HatiDataAgent(
+        host="proxy.internal",
+        port=5439,
+        agent_id="data-analyst-agent",
+        framework="langchain",
+        database="analytics",
+    )
+
+    # Simple query
+    rows = agent.query("SELECT COUNT(*) FROM orders")
+
+    # With reasoning chain tracking
+    with agent.reasoning_chain("req-001") as chain:
+        tables = chain.query("SELECT table_name FROM information_schema.tables")
+        data = chain.query("SELECT * FROM customers LIMIT 10", step=1)
+        summary = chain.query("SELECT AVG(revenue) FROM orders", step=2)
+
+    # RAG context retrieval
+    context = agent.get_context("customers", "enterprise accounts", top_k=5)
+"""
+
+from __future__ import annotations
+
+import uuid
+from contextlib import contextmanager
+from typing import Any, Generator, Optional
+
+import psycopg2
+import psycopg2.extras
+
+
+class HatiDataAgent:
+    """Agent-aware HatiData client using psycopg2.
+
+    Connects via Postgres wire protocol and identifies itself through
+    startup parameters that the proxy reads for billing, scheduling,
+    and audit.
+
+    Args:
+        host: Proxy hostname.
+        port: Proxy port (default 5439).
+        agent_id: Unique identifier for this agent instance.
+        framework: AI framework name (langchain, crewai, autogen, etc.).
+        database: Database name (default "hatidata").
+        user: Username (default "agent").
+        password: Password (default empty for dev mode).
+        priority: Query priority (low, normal, high, critical).
+        connect_timeout: Connection timeout in seconds.
+    """
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 5439,
+        agent_id: str = "",
+        framework: str = "custom",
+        database: str = "hatidata",
+        user: str = "agent",
+        password: str = "",
+        priority: str = "normal",
+        connect_timeout: int = 10,
+    ):
+        self.host = host
+        self.port = port
+        self.agent_id = agent_id or f"agent-{uuid.uuid4().hex[:8]}"
+        self.framework = framework
+        self.database = database
+        self.user = user
+        self.password = password
+        self.priority = priority
+        self.connect_timeout = connect_timeout
+        self._conn: Optional[psycopg2.extensions.connection] = None
+
+    def _get_connection(self) -> psycopg2.extensions.connection:
+        """Get or create a connection with agent startup params."""
+        if self._conn is None or self._conn.closed:
+            # Agent identification via startup parameters.
+            # These are read by the HatiData proxy during connection setup.
+            options = (
+                f"-c hatidata_agent_id={self.agent_id} "
+                f"-c hatidata_framework={self.framework} "
+                f"-c hatidata_priority={self.priority}"
+            )
+
+            self._conn = psycopg2.connect(
+                host=self.host,
+                port=self.port,
+                dbname=self.database,
+                user=self.user,
+                password=self.password,
+                options=options,
+                connect_timeout=self.connect_timeout,
+                application_name=f"hatidata-agent/{self.framework}",
+            )
+            self._conn.autocommit = True
+
+        return self._conn
+
+    def query(
+        self,
+        sql: str,
+        params: Optional[tuple[Any, ...]] = None,
+        request_id: Optional[str] = None,
+    ) -> list[dict[str, Any]]:
+        """Execute a SQL query and return results as a list of dicts.
+
+        Args:
+            sql: SQL query string.
+            params: Optional query parameters (for parameterized queries).
+            request_id: Optional request ID for tracing.
+
+        Returns:
+            List of row dicts with column names as keys.
+        """
+        conn = self._get_connection()
+
+        # Set request_id for this query if provided.
+        if request_id:
+            with conn.cursor() as cur:
+                cur.execute(f"SET hatidata_request_id = '{request_id}'")
+
+        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
+            cur.execute(sql, params)
+            if cur.description:
+                return [dict(row) for row in cur.fetchall()]
+            return []
+
+    def execute(self, sql: str, params: Optional[tuple[Any, ...]] = None) -> int:
+        """Execute a SQL statement (INSERT, UPDATE, DELETE) and return row count.
+
+        Args:
+            sql: SQL statement.
+            params: Optional query parameters.
+
+        Returns:
+            Number of affected rows.
+        """
+        conn = self._get_connection()
+        with conn.cursor() as cur:
+            cur.execute(sql, params)
+            return cur.rowcount
+
+    def get_context(
+        self,
+        table: str,
+        search_query: str,
+        top_k: int = 10,
+    ) -> list[dict[str, Any]]:
+        """Retrieve RAG context using HatiData's built-in context function.
+
+        This uses the `hatidata_context()` UDF which the proxy rewrites
+        to an efficient full-text search subquery.
+
+        Args:
+            table: Table name to search.
+            search_query: Natural language search query.
+            top_k: Maximum number of results.
+
+        Returns:
+            List of matching rows as dicts.
+        """
+        sql = f"SELECT * FROM hatidata_context('{table}', '{search_query}', {top_k})"
+        return self.query(sql)
+
+    def get_rag_context(
+        self,
+        table: str,
+        embedding_col: str,
+        vector: list[float],
+        top_k: int = 10,
+    ) -> list[dict[str, Any]]:
+        """Retrieve RAG context via vector similarity search.
+
+        Args:
+            table: Table name containing embeddings.
+            embedding_col: Column name with embedding vectors.
+            vector: Query embedding vector.
+            top_k: Maximum number of results.
+
+        Returns:
+            List of matching rows ordered by similarity.
+        """
+        vector_str = f"[{', '.join(str(v) for v in vector)}]"
+        sql = (
+            f"SELECT * FROM hatidata_rag_context('{table}', '{embedding_col}', "
+            f"{vector_str}, {top_k})"
+        )
+        return self.query(sql)
+
+    @contextmanager
+    def reasoning_chain(
+        self,
+        request_id: Optional[str] = None,
+        parent_request_id: Optional[str] = None,
+    ) -> Generator[ReasoningChain, None, None]:
+        """Context manager for tracking a multi-step reasoning chain.
+
+        Usage::
+
+            with agent.reasoning_chain("req-001") as chain:
+                tables = chain.query("SELECT table_name FROM information_schema.tables")
+                data = chain.query("SELECT * FROM orders LIMIT 10", step=1)
+
+        Args:
+            request_id: Unique ID for this reasoning chain.
+            parent_request_id: Parent chain ID for nested reasoning.
+        """
+        rid = request_id or f"req-{uuid.uuid4().hex[:12]}"
+        chain = ReasoningChain(self, rid, parent_request_id)
+        try:
+            yield chain
+        finally:
+            chain._finalize()
+
+    def close(self) -> None:
+        """Close the underlying database connection."""
+        if self._conn and not self._conn.closed:
+            self._conn.close()
+            self._conn = None
+
+    def __enter__(self) -> HatiDataAgent:
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self.close()
+
+
+class ReasoningChain:
+    """Tracks a multi-step reasoning chain with auto-incrementing steps."""
+
+    def __init__(
+        self,
+        agent: HatiDataAgent,
+        request_id: str,
+        parent_request_id: Optional[str] = None,
+    ):
+        self._agent = agent
+        self.request_id = request_id
+        self.parent_request_id = parent_request_id
+        self._step = 0
+
+    def query(
+        self,
+        sql: str,
+        params: Optional[tuple[Any, ...]] = None,
+        step: Optional[int] = None,
+    ) -> list[dict[str, Any]]:
+        """Execute a query as part of this reasoning chain.
+
+        Args:
+            sql: SQL query.
+            params: Optional query parameters.
+            step: Explicit step number (auto-increments if not provided).
+        """
+        current_step = step if step is not None else self._step
+        self._step = current_step + 1
+
+        conn = self._agent._get_connection()
+        with conn.cursor() as cur:
+            cur.execute(f"SET hatidata_request_id = '{self.request_id}'")
+            cur.execute(f"SET hatidata_reasoning_step = '{current_step}'")
+            if self.parent_request_id:
+                cur.execute(
+                    f"SET hatidata_parent_request_id = '{self.parent_request_id}'"
+                )
+
+        return self._agent.query(sql, params)
+
+    def _finalize(self) -> None:
+        """Reset session vars after chain completes."""
+        try:
+            conn = self._agent._get_connection()
+            with conn.cursor() as cur:
+                cur.execute("SET hatidata_request_id = ''")
+                cur.execute("SET hatidata_reasoning_step = ''")
+        except Exception:
+            pass

hatidata_agent-0.1.0/hatidata_agent/langchain.py

@@ -0,0 +1,169 @@
+"""LangChain integration for HatiData agent queries.
+
+Provides a LangChain-compatible SQLDatabase wrapper that automatically
+identifies itself as a LangChain agent, enabling per-agent billing,
+scheduling, and audit in HatiData.
+
+Usage::
+
+    from hatidata_agent.langchain import HatiDataSQLDatabase
+    from langchain.agents import create_sql_agent
+
+    db = HatiDataSQLDatabase(
+        host="proxy.internal",
+        port=5439,
+        agent_id="sql-agent-1",
+    )
+
+    agent = create_sql_agent(llm=llm, db=db, verbose=True)
+    result = agent.run("How many active enterprise customers do we have?")
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Sequence
+
+from hatidata_agent.client import HatiDataAgent
+
+
+class HatiDataSQLDatabase:
+    """LangChain-compatible SQL database wrapper for HatiData.
+
+    Implements the interface expected by LangChain's SQL agent tools:
+    - `run(command)` — Execute SQL and return string result
+    - `get_usable_table_names()` — List available tables
+    - `get_table_info(table_names)` — Get DDL/schema for tables
+    - `dialect` — Database dialect identifier
+
+    This wrapper automatically sets `framework="langchain"` for
+    agent identification in HatiData's billing and audit systems.
+    """
+
+    dialect = "postgresql"
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 5439,
+        agent_id: str = "langchain-agent",
+        database: str = "hatidata",
+        user: str = "agent",
+        password: str = "",
+        priority: str = "normal",
+        include_tables: Optional[Sequence[str]] = None,
+        sample_rows_in_table_info: int = 3,
+    ):
+        self._agent = HatiDataAgent(
+            host=host,
+            port=port,
+            agent_id=agent_id,
+            framework="langchain",
+            database=database,
+            user=user,
+            password=password,
+            priority=priority,
+        )
+        self._include_tables = set(include_tables) if include_tables else None
+        self._sample_rows = sample_rows_in_table_info
+
+    def run(self, command: str, fetch: str = "all") -> str:
+        """Execute SQL and return results as a string.
+
+        Args:
+            command: SQL query or statement.
+            fetch: "all" to fetch all rows, "one" for single row.
+
+        Returns:
+            String representation of the results.
+        """
+        rows = self._agent.query(command)
+        if not rows:
+            return ""
+        if fetch == "one" and rows:
+            return str(rows[0])
+        return str(rows)
+
+    def run_no_throw(self, command: str, fetch: str = "all") -> str:
+        """Execute SQL, returning error message on failure instead of raising."""
+        try:
+            return self.run(command, fetch)
+        except Exception as e:
+            return f"Error: {e}"
+
+    def get_usable_table_names(self) -> list[str]:
+        """Return list of table names available for querying."""
+        rows = self._agent.query(
+            "SELECT table_name FROM information_schema.tables "
+            "WHERE table_schema = 'main' ORDER BY table_name"
+        )
+        tables = [row["table_name"] for row in rows]
+
+        if self._include_tables:
+            tables = [t for t in tables if t in self._include_tables]
+
+        return tables
+
+    def get_table_info(self, table_names: Optional[list[str]] = None) -> str:
+        """Get CREATE TABLE DDL and sample rows for the given tables.
+
+        Args:
+            table_names: Tables to describe. If None, describes all usable tables.
+
+        Returns:
+            Multi-line string with DDL and sample rows for each table.
+        """
+        if table_names is None:
+            table_names = self.get_usable_table_names()
+
+        info_parts = []
+        for table in table_names:
+            # Get column info.
+            cols = self._agent.query(
+                f"SELECT column_name, data_type, is_nullable "
+                f"FROM information_schema.columns "
+                f"WHERE table_name = '{table}' ORDER BY ordinal_position"
+            )
+
+            # Build CREATE TABLE statement.
+            col_defs = []
+            for col in cols:
+                nullable = "" if col.get("is_nullable") == "YES" else " NOT NULL"
+                col_defs.append(
+                    f"  {col['column_name']} {col['data_type']}{nullable}"
+                )
+
+            ddl = f"CREATE TABLE {table} (\n" + ",\n".join(col_defs) + "\n);"
+
+            # Sample rows.
+            sample = ""
+            if self._sample_rows > 0:
+                rows = self._agent.query(
+                    f"SELECT * FROM {table} LIMIT {self._sample_rows}"
+                )
+                if rows:
+                    sample = (
+                        f"\n\n/*\n{self._sample_rows} rows from {table} table:\n"
+                        + "\t".join(rows[0].keys())
+                        + "\n"
+                    )
+                    for row in rows:
+                        sample += "\t".join(str(v) for v in row.values()) + "\n"
+                    sample += "*/"
+
+            info_parts.append(f"{ddl}{sample}")
+
+        return "\n\n".join(info_parts)
+
+    def get_table_info_no_throw(
+        self, table_names: Optional[list[str]] = None
+    ) -> str:
+        """Get table info, returning error message on failure."""
+        try:
+            return self.get_table_info(table_names)
+        except Exception as e:
+            return f"Error: {e}"
+
+    @property
+    def table_info(self) -> str:
+        """Property for LangChain compatibility."""
+        return self.get_table_info()

hatidata_agent-0.1.0/hatidata_agent/mcp_server.py

@@ -0,0 +1,268 @@
+"""MCP server for HatiData agent access.
+
+Exposes HatiData as an MCP (Model Context Protocol) server, allowing
+Claude and other MCP-compatible agents to query the data warehouse
+directly. Tools provided:
+
+- `query` — Execute SQL and return results
+- `list_tables` — List available tables
+- `describe_table` — Get schema for a table
+- `get_context` — RAG context retrieval via full-text search
+
+Usage::
+
+    # Start the MCP server (stdio transport)
+    hatidata-mcp-server --host localhost --port 5439
+
+    # Or in Python
+    from hatidata_agent.mcp_server import create_server
+    server = create_server(host="localhost", port=5439)
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from typing import Any
+
+from hatidata_agent.client import HatiDataAgent
+
+
+def create_tools(agent: HatiDataAgent) -> list[dict[str, Any]]:
+    """Create MCP tool definitions."""
+    return [
+        {
+            "name": "query",
+            "description": (
+                "Execute a SQL query against the HatiData warehouse. "
+                "Returns results as a JSON array of objects. "
+                "Supports Snowflake-compatible SQL syntax."
+            ),
+            "inputSchema": {
+                "type": "object",
+                "properties": {
+                    "sql": {
+                        "type": "string",
+                        "description": "The SQL query to execute",
+                    },
+                },
+                "required": ["sql"],
+            },
+        },
+        {
+            "name": "list_tables",
+            "description": "List all available tables in the database.",
+            "inputSchema": {
+                "type": "object",
+                "properties": {},
+            },
+        },
+        {
+            "name": "describe_table",
+            "description": "Get the schema (columns, types) for a specific table.",
+            "inputSchema": {
+                "type": "object",
+                "properties": {
+                    "table_name": {
+                        "type": "string",
+                        "description": "Name of the table to describe",
+                    },
+                },
+                "required": ["table_name"],
+            },
+        },
+        {
+            "name": "get_context",
+            "description": (
+                "Retrieve relevant rows from a table using full-text search. "
+                "Useful for RAG context retrieval."
+            ),
+            "inputSchema": {
+                "type": "object",
+                "properties": {
+                    "table": {
+                        "type": "string",
+                        "description": "Table to search",
+                    },
+                    "search_query": {
+                        "type": "string",
+                        "description": "Natural language search query",
+                    },
+                    "top_k": {
+                        "type": "integer",
+                        "description": "Maximum results (default 10)",
+                        "default": 10,
+                    },
+                },
+                "required": ["table", "search_query"],
+            },
+        },
+    ]
+
+
+def handle_tool_call(
+    agent: HatiDataAgent, tool_name: str, arguments: dict[str, Any]
+) -> dict[str, Any]:
+    """Handle an MCP tool call and return the result."""
+    try:
+        if tool_name == "query":
+            rows = agent.query(arguments["sql"])
+            return {
+                "content": [
+                    {"type": "text", "text": json.dumps(rows, default=str)},
+                ],
+            }
+
+        elif tool_name == "list_tables":
+            rows = agent.query(
+                "SELECT table_name FROM information_schema.tables "
+                "WHERE table_schema = 'main' ORDER BY table_name"
+            )
+            tables = [row["table_name"] for row in rows]
+            return {
+                "content": [
+                    {"type": "text", "text": json.dumps(tables)},
+                ],
+            }
+
+        elif tool_name == "describe_table":
+            table = arguments["table_name"]
+            rows = agent.query(
+                f"SELECT column_name, data_type, is_nullable "
+                f"FROM information_schema.columns "
+                f"WHERE table_name = '{table}' ORDER BY ordinal_position"
+            )
+            return {
+                "content": [
+                    {"type": "text", "text": json.dumps(rows, default=str)},
+                ],
+            }
+
+        elif tool_name == "get_context":
+            rows = agent.get_context(
+                table=arguments["table"],
+                search_query=arguments["search_query"],
+                top_k=arguments.get("top_k", 10),
+            )
+            return {
+                "content": [
+                    {"type": "text", "text": json.dumps(rows, default=str)},
+                ],
+            }
+
+        else:
+            return {
+                "content": [
+                    {"type": "text", "text": f"Unknown tool: {tool_name}"},
+                ],
+                "isError": True,
+            }
+
+    except Exception as e:
+        return {
+            "content": [
+                {"type": "text", "text": f"Error: {e}"},
+            ],
+            "isError": True,
+        }
+
+
+def run_stdio_server(agent: HatiDataAgent) -> None:
+    """Run MCP server over stdio transport (JSON-RPC)."""
+    tools = create_tools(agent)
+
+    # Server info
+    server_info = {
+        "name": "hatidata",
+        "version": "0.1.0",
+        "capabilities": {
+            "tools": {},
+        },
+    }
+
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+
+        try:
+            request = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+
+        method = request.get("method", "")
+        req_id = request.get("id")
+
+        if method == "initialize":
+            response = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "result": {
+                    "protocolVersion": "2024-11-05",
+                    "serverInfo": server_info,
+                    "capabilities": server_info["capabilities"],
+                },
+            }
+
+        elif method == "tools/list":
+            response = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "result": {"tools": tools},
+            }
+
+        elif method == "tools/call":
+            params = request.get("params", {})
+            tool_name = params.get("name", "")
+            arguments = params.get("arguments", {})
+            result = handle_tool_call(agent, tool_name, arguments)
+            response = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "result": result,
+            }
+
+        elif method == "notifications/initialized":
+            continue  # No response needed for notifications.
+
+        else:
+            response = {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "error": {
+                    "code": -32601,
+                    "message": f"Method not found: {method}",
+                },
+            }
+
+        sys.stdout.write(json.dumps(response) + "\n")
+        sys.stdout.flush()
+
+
+def main() -> None:
+    """Entry point for the MCP server CLI."""
+    parser = argparse.ArgumentParser(description="HatiData MCP Server")
+    parser.add_argument("--host", default="localhost", help="Proxy host")
+    parser.add_argument("--port", type=int, default=5439, help="Proxy port")
+    parser.add_argument("--agent-id", default="mcp-agent", help="Agent ID")
+    parser.add_argument("--database", default="hatidata", help="Database")
+    parser.add_argument("--user", default="agent", help="Username")
+    parser.add_argument("--password", default="", help="Password")
+    args = parser.parse_args()
+
+    agent = HatiDataAgent(
+        host=args.host,
+        port=args.port,
+        agent_id=args.agent_id,
+        framework="anthropic",
+        database=args.database,
+        user=args.user,
+        password=args.password,
+    )
+
+    run_stdio_server(agent)
+
+
+if __name__ == "__main__":
+    main()

hatidata_agent-0.1.0/hatidata_agent.egg-info/PKG-INFO

@@ -0,0 +1,183 @@
+Metadata-Version: 2.4
+Name: hatidata-agent
+Version: 0.1.0
+Summary: HatiData Agent SDK — RAM for Agents
+Author-email: Marviy Pte Ltd <eng@hatidata.com>
+License: Apache-2.0
+Project-URL: Homepage, https://hatidata.com
+Project-URL: Documentation, https://docs.hatiosai.com/hatidata
+Project-URL: Repository, https://github.com/HatiOS-AI/HatiData-Core
+Project-URL: Changelog, https://github.com/HatiOS-AI/HatiData-Core/releases
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Database
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: psycopg2-binary>=2.9
+Provides-Extra: async
+Requires-Dist: asyncpg>=0.29; extra == "async"
+Provides-Extra: langchain
+Requires-Dist: langchain-community>=0.2; extra == "langchain"
+Provides-Extra: mcp
+Requires-Dist: mcp>=1.0; extra == "mcp"
+Provides-Extra: all
+Requires-Dist: asyncpg>=0.29; extra == "all"
+Requires-Dist: langchain-community>=0.2; extra == "all"
+Requires-Dist: mcp>=1.0; extra == "all"
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.23; extra == "dev"
+
+# HatiData Agent SDK
+
+[![PyPI version](https://img.shields.io/pypi/v/hatidata-agent.svg)](https://pypi.org/project/hatidata-agent/)
+[![Python versions](https://img.shields.io/pypi/pyversions/hatidata-agent.svg)](https://pypi.org/project/hatidata-agent/)
+[![License](https://img.shields.io/pypi/l/hatidata-agent.svg)](https://github.com/HatiOS-AI/HatiData-Core/blob/main/LICENSE)
+
+**RAM for Agents** — Python SDK for AI agents to query HatiData's in-VPC data warehouse with sub-10ms latency via Postgres wire protocol.
+
+## Installation
+
+```bash
+pip install hatidata-agent
+
+# With async support
+pip install hatidata-agent[async]
+
+# With LangChain support
+pip install hatidata-agent[langchain]
+
+# With MCP server
+pip install hatidata-agent[mcp]
+
+# Everything
+pip install hatidata-agent[all]
+```
+
+## Quick Start
+
+```python
+from hatidata_agent import HatiDataAgent
+
+agent = HatiDataAgent(
+    host="proxy.internal",
+    port=5439,
+    agent_id="my-agent",
+    framework="langchain",
+)
+
+# Simple query
+rows = agent.query("SELECT * FROM customers WHERE status = 'active' LIMIT 10")
+print(rows)  # [{"id": 1, "name": "Acme Corp", ...}, ...]
+```
+
+## Features
+
+- **Sub-10ms query latency** -- In-VPC execution, no data leaves your network
+- **Postgres wire protocol** -- Works with any Postgres client library
+- **Reasoning chain tracking** -- Multi-step audit trails for agent workflows
+- **RAG context retrieval** -- Full-text and vector similarity search
+- **LangChain integration** -- Drop-in `SQLDatabase` replacement for LangChain agents
+- **MCP server** -- Expose HatiData as tools for Claude and other MCP-compatible agents
+- **Agent identification** -- Per-agent billing, priority scheduling, and audit via Postgres startup parameters
+- **Snowflake SQL compatible** -- Bring existing queries without rewrites
+
+## Reasoning Chain Tracking
+
+Track multi-step reasoning chains for audit and debugging:
+
+```python
+with agent.reasoning_chain("req-001") as chain:
+    # Step 0: Discover tables
+    tables = chain.query("SELECT table_name FROM information_schema.tables")
+
+    # Step 1: Get relevant data
+    customers = chain.query("SELECT * FROM customers WHERE tier = 'enterprise'", step=1)
+
+    # Step 2: Aggregate
+    revenue = chain.query("SELECT SUM(revenue) FROM orders WHERE customer_id IN (...)", step=2)
+```
+
+## RAG Context Retrieval
+
+```python
+# Full-text search
+context = agent.get_context("customers", "enterprise accounts in US", top_k=5)
+
+# Vector similarity search
+context = agent.get_rag_context("docs", "embedding", query_vector, top_k=10)
+```
+
+## LangChain Integration
+
+```python
+from hatidata_agent.langchain import HatiDataSQLDatabase
+from langchain.agents import create_sql_agent
+
+db = HatiDataSQLDatabase(
+    host="proxy.internal",
+    port=5439,
+    agent_id="sql-agent-1",
+)
+
+agent = create_sql_agent(llm=llm, db=db, verbose=True)
+result = agent.run("How many enterprise customers do we have?")
+```
+
+## MCP Server
+
+Run as an MCP server for Claude and other MCP-compatible agents:
+
+```bash
+# Start the MCP server
+hatidata-mcp-server --host proxy.internal --port 5439
+
+# Or add to Claude Code's MCP config:
+# ~/.claude/mcp.json
+{
+  "mcpServers": {
+    "hatidata": {
+      "command": "hatidata-mcp-server",
+      "args": ["--host", "proxy.internal", "--port", "5439"]
+    }
+  }
+}
+```
+
+### MCP Tools
+
+| Tool | Description |
+|------|-------------|
+| `query` | Execute SQL and return JSON results |
+| `list_tables` | List available tables |
+| `describe_table` | Get table schema |
+| `get_context` | RAG context retrieval via full-text search |
+
+## Agent Identification
+
+The SDK automatically identifies agents via Postgres startup parameters:
+
+| Parameter | Purpose |
+|-----------|---------|
+| `hatidata_agent_id` | Unique agent identifier |
+| `hatidata_framework` | AI framework (langchain, crewai, autogen, etc.) |
+| `hatidata_priority` | Scheduling priority (low, normal, high, critical) |
+| `hatidata_request_id` | Request/reasoning chain ID |
+| `hatidata_reasoning_step` | Step number within a chain |
+
+These enable per-agent billing, priority scheduling, audit trails, and the Agent Tax Report showing savings vs legacy cloud warehouses.
+
+## Documentation
+
+Full documentation is available at [docs.hatiosai.com/hatidata](https://docs.hatiosai.com/hatidata).
+
+## License
+
+Apache License 2.0. Copyright (c) Marviy Pte Ltd. See [LICENSE](https://github.com/HatiOS-AI/HatiData-Core/blob/main/LICENSE) for details.

hatidata_agent-0.1.0/hatidata_agent.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+README.md
+pyproject.toml
+hatidata_agent/__init__.py
+hatidata_agent/client.py
+hatidata_agent/langchain.py
+hatidata_agent/mcp_server.py
+hatidata_agent.egg-info/PKG-INFO
+hatidata_agent.egg-info/SOURCES.txt
+hatidata_agent.egg-info/dependency_links.txt
+hatidata_agent.egg-info/entry_points.txt
+hatidata_agent.egg-info/requires.txt
+hatidata_agent.egg-info/top_level.txt

hatidata_agent-0.1.0/hatidata_agent.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

hatidata_agent-0.1.0/hatidata_agent.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+hatidata-mcp-server = hatidata_agent.mcp_server:main

hatidata_agent-0.1.0/hatidata_agent.egg-info/requires.txt

@@ -0,0 +1,19 @@
+psycopg2-binary>=2.9
+
+[all]
+asyncpg>=0.29
+langchain-community>=0.2
+mcp>=1.0
+
+[async]
+asyncpg>=0.29
+
+[dev]
+pytest>=8
+pytest-asyncio>=0.23
+
+[langchain]
+langchain-community>=0.2
+
+[mcp]
+mcp>=1.0

hatidata_agent-0.1.0/hatidata_agent.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+dist
+hatidata_agent

hatidata_agent-0.1.0/pyproject.toml

@@ -0,0 +1,55 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "hatidata-agent"
+version = "0.1.0"
+description = "HatiData Agent SDK — RAM for Agents"
+readme = "README.md"
+license = {text = "Apache-2.0"}
+requires-python = ">=3.9"
+authors = [
+    {name = "Marviy Pte Ltd", email = "eng@hatidata.com"},
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Database",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "psycopg2-binary>=2.9",
+]
+
+[project.optional-dependencies]
+async = ["asyncpg>=0.29"]
+langchain = ["langchain-community>=0.2"]
+mcp = ["mcp>=1.0"]
+all = [
+    "asyncpg>=0.29",
+    "langchain-community>=0.2",
+    "mcp>=1.0",
+]
+dev = [
+    "pytest>=8",
+    "pytest-asyncio>=0.23",
+]
+
+[project.urls]
+Homepage = "https://hatidata.com"
+Documentation = "https://docs.hatiosai.com/hatidata"
+Repository = "https://github.com/HatiOS-AI/HatiData-Core"
+Changelog = "https://github.com/HatiOS-AI/HatiData-Core/releases"
+
+[project.scripts]
+hatidata-mcp-server = "hatidata_agent.mcp_server:main"
+
+[tool.setuptools.packages.find]
+where = ["."]

hatidata_agent-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+