feat: add LlamaIndex tools for Hotdata managed databases

Author: eddietejeda

.gitignore

@@ -0,0 +1,23 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+
+# Virtual environments
+.env
+.venv
+env/
+venv/
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+
+# Packaging
+*.egg-info/
+dist/
+build/
+
+.DS_Store

README.md

@@ -0,0 +1,60 @@
+# hotdata-llamaindex
+
+LlamaIndex tools for [Hotdata](https://hotdata.dev), built on **hotdata-runtime**.
+
+## Features
+
+- **SQL tool** — run workspace SQL and return JSON rows for agents
+- **Managed database tools** — list, create, and load parquet into Hotdata-owned catalogs (replaces legacy dataset uploads)
+
+## Install
+
+```bash
+pip install hotdata-llamaindex
+```
+
+Requires `HOTDATA_API_KEY`. Optionally set `HOTDATA_WORKSPACE`, `HOTDATA_API_URL`, or `HOTDATA_SANDBOX`.
+
+## Usage
+
+```python
+import hotdata_llamaindex as hli
+
+client = hli.from_env()
+tools = hli.make_hotdata_tools(client)
+
+for tool in tools:
+    print(tool.metadata.name, tool.metadata.description)
+```
+
+Managed database example:
+
+```python
+tools = {tool.metadata.name: tool for tool in hli.make_hotdata_tools(client)}
+
+tools["hotdata_create_managed_database"].call(
+    name="sales",
+    schema_name="public",
+    tables="orders",
+)
+
+tools["hotdata_load_managed_table"].call(
+    database="sales",
+    table="orders",
+    file="/path/to/orders.parquet",
+)
+```
+
+## Examples
+
+```bash
+uv run python examples/llamaindex_basic.py
+uv run python examples/llamaindex_managed_db.py
+```
+
+## Development
+
+```bash
+uv sync --locked
+uv run pytest
+```

examples/llamaindex_basic.py

@@ -0,0 +1,21 @@
+"""Minimal LlamaIndex tool usage with hotdata-llamaindex."""
+
+import hotdata_llamaindex as hli
+
+
+def main() -> None:
+    client = hli.from_env()
+    tools = hli.make_hotdata_tools(client)
+    by_name = {tool.metadata.name: tool for tool in tools}
+
+    sql_tool = by_name["hotdata_execute_sql"]
+    print(sql_tool.call(sql="SELECT 1 AS ok"))
+
+    list_tool = by_name["hotdata_list_managed_databases"]
+    print(list_tool.call())
+
+    client.close()
+
+
+if __name__ == "__main__":
+    main()

examples/llamaindex_managed_db.py

@@ -0,0 +1,34 @@
+"""Managed database tools for LlamaIndex agents."""
+
+import hotdata_llamaindex as hli
+
+
+def main() -> None:
+    client = hli.from_env()
+    tools = hli.make_hotdata_tools(client)
+    by_name = {tool.metadata.name: tool for tool in tools}
+
+    create = by_name["hotdata_create_managed_database"]
+    print(
+        create.call(
+            name="demo_sales",
+            schema_name="public",
+            tables="orders\ncustomers",
+        )
+    )
+
+    load = by_name["hotdata_load_managed_table"]
+    print(
+        load.call(
+            database="demo_sales",
+            table="orders",
+            file="/path/to/orders.parquet",
+            schema_name="public",
+        )
+    )
+
+    client.close()
+
+
+if __name__ == "__main__":
+    main()

hotdata_llamaindex/__init__.py

@@ -0,0 +1,37 @@
+"""LlamaIndex tools for Hotdata runtime."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("hotdata-llamaindex")
+except PackageNotFoundError:
+    __version__ = "0.0.0+unknown"
+
+from hotdata_runtime import HotdataClient, QueryResult, from_env
+from hotdata_llamaindex.databases import (
+    create_managed_database,
+    list_managed_databases_json,
+    load_managed_table,
+    load_result_summary,
+    managed_database_summary,
+)
+from hotdata_llamaindex.tools import (
+    execute_sql_json,
+    make_hotdata_tools,
+    result_rows_for_llm,
+)
+
+__all__ = [
+    "__version__",
+    "HotdataClient",
+    "QueryResult",
+    "create_managed_database",
+    "execute_sql_json",
+    "from_env",
+    "list_managed_databases_json",
+    "load_managed_table",
+    "load_result_summary",
+    "make_hotdata_tools",
+    "managed_database_summary",
+    "result_rows_for_llm",
+]

hotdata_llamaindex/databases.py

@@ -0,0 +1,60 @@
+"""Managed database helpers for LlamaIndex agents."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from hotdata_runtime import (
+    DEFAULT_SCHEMA,
+    HotdataClient,
+    LoadManagedTableResult,
+    ManagedDatabase,
+)
+
+
+def list_managed_databases_json(client: HotdataClient) -> str:
+    rows = [
+        {
+            "name": db.name,
+            "id": db.id,
+            "sql_prefix": f"{db.name}.{{schema}}.{{table}}",
+        }
+        for db in client.list_managed_databases()
+    ]
+    return json.dumps(rows, indent=2)
+
+
+def create_managed_database(
+    client: HotdataClient,
+    *,
+    name: str,
+    schema: str = DEFAULT_SCHEMA,
+    tables: list[str] | None = None,
+) -> ManagedDatabase:
+    return client.create_managed_database(name, schema=schema, tables=tables)
+
+
+def load_managed_table(
+    client: HotdataClient,
+    *,
+    database: str,
+    table: str,
+    file: str,
+    schema: str = DEFAULT_SCHEMA,
+) -> LoadManagedTableResult:
+    return client.load_managed_table(database, table, schema=schema, file=file)
+
+
+def managed_database_summary(db: ManagedDatabase) -> dict[str, str]:
+    return {"id": db.id, "name": db.name, "source_type": db.source_type}
+
+
+def load_result_summary(result: LoadManagedTableResult) -> dict[str, Any]:
+    return {
+        "connection_id": result.connection_id,
+        "schema_name": result.schema_name,
+        "table_name": result.table_name,
+        "row_count": result.row_count,
+        "full_name": result.full_name,
+    }

hotdata_llamaindex/tools.py

@@ -0,0 +1,97 @@
+"""LlamaIndex tools built on hotdata-runtime."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from llama_index.core.tools import FunctionTool
+
+from hotdata_runtime import DEFAULT_SCHEMA, HotdataClient, QueryResult
+
+from hotdata_llamaindex.databases import (
+    create_managed_database,
+    list_managed_databases_json,
+    load_managed_table,
+    load_result_summary,
+    managed_database_summary,
+)
+
+
+def result_rows_for_llm(result: QueryResult, *, max_rows: int = 20) -> list[dict[str, Any]]:
+    return result.to_records(max_rows=max_rows)
+
+
+def execute_sql_json(client: HotdataClient, sql: str, *, max_rows: int = 100) -> str:
+    result = client.execute_sql(sql)
+    payload = {
+        "metadata": result.metadata_dict(),
+        "rows": result.to_records(max_rows=max_rows),
+    }
+    return json.dumps(payload, indent=2)
+
+
+def make_hotdata_tools(
+    client: HotdataClient,
+    *,
+    max_rows: int = 100,
+) -> list[FunctionTool]:
+    """Return LlamaIndex tools for SQL and managed database workflows."""
+
+    def hotdata_execute_sql(sql: str) -> str:
+        """Run SQL against the Hotdata workspace and return JSON rows."""
+        return execute_sql_json(client, sql, max_rows=max_rows)
+
+    def hotdata_list_managed_databases() -> str:
+        """List Hotdata-managed databases in the workspace."""
+        return list_managed_databases_json(client)
+
+    def hotdata_create_managed_database(
+        name: str,
+        schema_name: str = DEFAULT_SCHEMA,
+        tables: str = "",
+    ) -> str:
+        """Create a Hotdata-managed database and optionally declare tables (one per line)."""
+        table_names = [line.strip() for line in tables.splitlines() if line.strip()]
+        db = create_managed_database(
+            client,
+            name=name,
+            schema=schema_name or DEFAULT_SCHEMA,
+            tables=table_names or None,
+        )
+        return json.dumps(managed_database_summary(db), indent=2)
+
+    def hotdata_load_managed_table(
+        database: str,
+        table: str,
+        file: str,
+        schema_name: str = DEFAULT_SCHEMA,
+    ) -> str:
+        """Load a local parquet file into a declared managed table."""
+        loaded = load_managed_table(
+            client,
+            database=database,
+            table=table,
+            file=file,
+            schema=schema_name or DEFAULT_SCHEMA,
+        )
+        return json.dumps(load_result_summary(loaded), indent=2)
+
+    return [
+        FunctionTool.from_defaults(
+            fn=hotdata_execute_sql,
+            name="hotdata_execute_sql",
+        ),
+        FunctionTool.from_defaults(
+            fn=hotdata_list_managed_databases,
+            name="hotdata_list_managed_databases",
+        ),
+        FunctionTool.from_defaults(
+            fn=hotdata_create_managed_database,
+            name="hotdata_create_managed_database",
+        ),
+        FunctionTool.from_defaults(
+            fn=hotdata_load_managed_table,
+            name="hotdata_load_managed_table",
+        ),
+    ]

pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "hotdata-llamaindex"
+version = "0.1.0"
+description = "LlamaIndex tools for Hotdata runtime"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+dependencies = [
+  "hotdata-runtime>=0.1.1",
+  "llama-index-core>=0.12.0",
+]
+
+[dependency-groups]
+dev = [
+  "pytest>=8.0",
+]
+
+[tool.uv]
+default-groups = ["dev"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["hotdata_llamaindex"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

tests/conftest.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from hotdata_runtime import QueryResult
+
+
+@pytest.fixture
+def sample_result() -> QueryResult:
+    return QueryResult(
+        columns=["n"],
+        rows=[[1], [2]],
+        row_count=2,
+        result_id="res_1",
+        query_run_id="run_1",
+        execution_time_ms=12,
+        warning=None,
+        error_message=None,
+    )
+
+
+@pytest.fixture
+def mock_client(sample_result: QueryResult):
+    client = MagicMock()
+    client.workspace_id = "ws_test"
+    client.execute_sql = MagicMock(return_value=sample_result)
+    client.list_managed_databases = MagicMock(return_value=[])
+    return client