"""BitableTool — Agent tool for bitable data ingestion and CRUD via HTTP. Implements KTD5 (REST API boundary even when co-deployed) and KTD11 (internal service token auth). The tool uses ``httpx.AsyncClient`` to call the bitable REST API; it never imports BitableService directly. Actions: create_table, import_excel, import_database, collect_api, upsert_records, query_records, create_view, update_view, update_field, delete_view. Batch chunking: upsert and import operations send at most ``BATCH_SIZE`` records per HTTP request. On partial failure, the result includes ``successful_count`` and ``resume_from`` for breakpoint continuation. """ from __future__ import annotations import asyncio import logging import httpx from agentkit.bitable.ingestion.excel import ParsedSheet, parse_excel, parse_excel_url from agentkit.bitable.ingestion.database import import_table as import_db_table from agentkit.bitable.ingestion.api_collector import transform_records from agentkit.tools.base import Tool logger = logging.getLogger(__name__) BATCH_SIZE = 500 class BitableTool(Tool): """Agent tool for bitable operations via REST API. Args: base_url: Bitable API base URL (e.g. ``http://localhost:18001/api/v1/bitable``). internal_token: Service token for KTD11 auth. If ``None``, requests go unauthenticated (will fail if the server requires auth). """ def __init__(self, base_url: str, internal_token: str | None = None) -> None: super().__init__( name="bitable", description=( "Create and manage bitable (multi-dimensional spreadsheet) tables, " "ingest data from Excel files, databases, or API responses, and " "query records. Actions: create_table, import_excel, " "import_database, collect_api, upsert_records, query_records, " "create_view, update_view, update_field, delete_view." ), input_schema={ "type": "object", "properties": { "action": { "type": "string", "enum": [ "create_table", "import_excel", "import_database", "collect_api", "upsert_records", "query_records", "create_view", "update_view", "update_field", "delete_view", ], "description": "Bitable operation to perform.", }, "table_name": { "type": "string", "description": "Name for the new bitable table (create_table, import_excel, import_database).", }, "description": { "type": "string", "description": "Table description (create_table).", }, "file_path": { "type": "string", "description": "Path to .xlsx file (import_excel).", }, "file_url": { "type": "string", "description": "URL to download .xlsx file (import_excel).", }, "connection_string": { "type": "string", "description": "Database connection string (import_database).", }, "table_names": { "type": "array", "items": {"type": "string"}, "description": "Source table names to import (import_database).", }, "table_id": { "type": "string", "description": "Target bitable table ID (collect_api, upsert_records, query_records, create_view).", }, "records": { "type": "array", "description": "Records to write (collect_api, upsert_records).", }, "field_mapping": { "type": "object", "description": "Mapping {source_key: bitable_field_id} (collect_api).", }, "primary_key_field_id": { "type": "string", "description": "Field ID of the primary key (upsert_records, collect_api).", }, "resume_from": { "type": "integer", "description": "Skip this many records before resuming a failed batch (upsert_records, collect_api).", }, "cursor": { "type": "string", "description": "Pagination cursor (query_records).", }, "limit": { "type": "integer", "description": "Max records to return (query_records).", }, "view_id": { "type": "string", "description": "View ID (update_view, delete_view).", }, "view_type": { "type": "string", "enum": ["grid", "kanban", "gantt", "gallery", "form"], "description": "View type (create_view). Defaults to grid.", }, "field_id": { "type": "string", "description": "Field ID (update_field).", }, "name": { "type": "string", "description": "Name for a view or field (create_view, update_view, update_field).", }, "type": { "type": "string", "description": "Field type for update_field (e.g. text, number, date).", }, "config": { "type": "object", "description": "View/field config dict (create_view, update_view, update_field).", }, }, "required": ["action"], }, ) self._base_url = base_url.rstrip("/") self._internal_token = internal_token self._client: httpx.AsyncClient | None = None async def _get_client(self) -> httpx.AsyncClient: if self._client is None or self._client.is_closed: headers: dict[str, str] = {} if self._internal_token: headers["X-Internal-Token"] = self._internal_token self._client = httpx.AsyncClient( base_url=self._base_url, headers=headers, timeout=60.0, ) return self._client async def close(self) -> None: if self._client is not None and not self._client.is_closed: await self._client.aclose() async def execute(self, **kwargs) -> dict[str, object]: action = kwargs.get("action") handlers = { "create_table": self._create_table, "import_excel": self._import_excel, "import_database": self._import_database, "collect_api": self._collect_api, "upsert_records": self._upsert_records, "query_records": self._query_records, "create_view": self._create_view, "update_view": self._update_view, "update_field": self._update_field, "delete_view": self._delete_view, } handler = handlers.get(action) if handler is None: return {"success": False, "error": f"Unknown action: {action!r}"} try: return await handler(**kwargs) except httpx.HTTPStatusError as e: return { "success": False, "error": f"Bitable API error {e.response.status_code}: {e.response.text[:500]}", } except httpx.ConnectError as e: return {"success": False, "error": f"Cannot connect to bitable API: {e}"} except Exception as e: return {"success": False, "error": f"{action} failed: {e}"} # ------------------------------------------------------------------ # create_table # ------------------------------------------------------------------ async def _create_table(self, **kwargs) -> dict[str, object]: table_name = kwargs.get("table_name") if not table_name: return {"success": False, "error": "Missing required field: table_name"} client = await self._get_client() resp = await client.post( "/tables", json={"name": table_name, "description": kwargs.get("description", "")}, ) resp.raise_for_status() data = resp.json() return {"success": True, "table": data["table"]} # ------------------------------------------------------------------ # import_excel # ------------------------------------------------------------------ async def _import_excel(self, **kwargs) -> dict[str, object]: file_path = kwargs.get("file_path") file_url = kwargs.get("file_url") if not file_path and not file_url: return {"success": False, "error": "Either file_path or file_url is required"} # Parse Excel — offload sync I/O to thread pool (P2 #21-23). if file_path: sheets = await asyncio.to_thread(parse_excel, file_path) else: sheets = await asyncio.to_thread(parse_excel_url, file_url) if not sheets: return {"success": False, "error": "Excel file has no sheets with data"} results: list[dict[str, object]] = [] for sheet in sheets: result = await self._import_sheet(sheet) results.append(result) return {"success": True, "sheets": results} async def _import_sheet(self, sheet: ParsedSheet) -> dict[str, object]: """Create a bitable table from a parsed sheet and upsert all rows.""" client = await self._get_client() # 1. Create table resp = await client.post("/tables", json={"name": sheet.name}) resp.raise_for_status() table_id = resp.json()["table"]["id"] # 2. Create fields field_name_to_id: dict[str, str] = {} for col_name, field_type in zip(sheet.columns, sheet.field_types): resp = await client.post( f"/tables/{table_id}/fields", json={"name": col_name, "field_type": field_type, "owner": "agent"}, ) resp.raise_for_status() field_id = resp.json()["field"]["id"] field_name_to_id[col_name] = field_id # 3. Map record keys to field IDs and batch upsert mapped_records = [ {field_name_to_id[k]: v for k, v in rec.items() if k in field_name_to_id} for rec in sheet.records ] if not mapped_records: return { "table_id": table_id, "table_name": sheet.name, "field_count": len(field_name_to_id), "record_count": 0, } # Use first field as PK fallback (import_excel doesn't require a PK) # If no PK is set, upsert won't work — use create_records instead upsert_result = await self._batch_create_records(table_id, mapped_records) return { "table_id": table_id, "table_name": sheet.name, "field_count": len(field_name_to_id), "record_count": upsert_result["successful_count"], **upsert_result, } async def _batch_create_records( self, table_id: str, records: list[dict[str, object]] ) -> dict[str, object]: """Create records in batches via POST /tables/{id}/records.""" client = await self._get_client() total = len(records) successful = 0 errors: list[dict[str, object]] = [] for start in range(0, total, BATCH_SIZE): batch = records[start : start + BATCH_SIZE] try: resp = await client.post( f"/tables/{table_id}/records", json={"records": batch}, ) resp.raise_for_status() successful += len(batch) except httpx.HTTPStatusError as e: errors.append( { "batch_start": start, "batch_size": len(batch), "status": e.response.status_code, "error": e.response.text[:300], } ) break # stop on first failure return { "successful_count": successful, "total": total, "resume_from": successful, **({"errors": errors} if errors else {}), } # ------------------------------------------------------------------ # import_database # ------------------------------------------------------------------ async def _import_database(self, **kwargs) -> dict[str, object]: conn_str = kwargs.get("connection_string") table_names = kwargs.get("table_names") if not conn_str: return {"success": False, "error": "Missing required field: connection_string"} if not table_names: return {"success": False, "error": "Missing required field: table_names"} results: list[dict[str, object]] = [] for src_table in table_names: try: # Offload sync DB reflection to thread pool (P2 #21-23). reflected = await asyncio.to_thread(import_db_table, conn_str, src_table) result = await self._import_reflected_table(reflected) results.append(result) except ConnectionError as e: return {"success": False, "error": str(e), "imported": results} except Exception as e: results.append({"table_name": src_table, "success": False, "error": str(e)}) return {"success": True, "tables": results} async def _import_reflected_table(self, reflected: dict[str, object]) -> dict[str, object]: """Create a bitable table from reflected DB data and upsert rows.""" client = await self._get_client() table_name = reflected["table_name"] # 1. Create table resp = await client.post("/tables", json={"name": table_name}) resp.raise_for_status() table_id = resp.json()["table"]["id"] # 2. Create fields field_name_to_id: dict[str, str] = {} pk_field_id: str | None = None for fdef in reflected["fields"]: resp = await client.post( f"/tables/{table_id}/fields", json={ "name": fdef["name"], "field_type": fdef["field_type"], "owner": "agent", }, ) resp.raise_for_status() fid = resp.json()["field"]["id"] field_name_to_id[fdef["name"]] = fid if fdef.get("is_primary_key"): pk_field_id = fid # 3. Set primary key if pk_field_id: await client.patch("/tables/" + table_id, json={"primary_key_field_id": pk_field_id}) # 4. Map and upsert records mapped = [ {field_name_to_id[k]: v for k, v in rec.items() if k in field_name_to_id} for rec in reflected["records"] ] if not mapped: return { "table_id": table_id, "table_name": table_name, "record_count": 0, "success": True, } if pk_field_id: upsert = await self._batch_upsert(table_id, mapped, pk_field_id) else: upsert = await self._batch_create_records(table_id, mapped) return { "table_id": table_id, "table_name": table_name, "record_count": upsert["successful_count"], "success": True, **upsert, } # ------------------------------------------------------------------ # collect_api # ------------------------------------------------------------------ async def _collect_api(self, **kwargs) -> dict[str, object]: table_id = kwargs.get("table_id") records = kwargs.get("records") field_mapping = kwargs.get("field_mapping") pk_field_id = kwargs.get("primary_key_field_id") resume_from = kwargs.get("resume_from", 0) if not table_id: return {"success": False, "error": "Missing required field: table_id"} if not records: return {"success": False, "error": "Missing required field: records"} if not field_mapping: return {"success": False, "error": "Missing required field: field_mapping"} if not pk_field_id: return {"success": False, "error": "Missing required field: primary_key_field_id"} transformed = transform_records(records, field_mapping) if resume_from > 0: transformed = transformed[resume_from:] result = await self._batch_upsert(table_id, transformed, pk_field_id) return {"success": True, **result} # ------------------------------------------------------------------ # upsert_records # ------------------------------------------------------------------ async def _upsert_records(self, **kwargs) -> dict[str, object]: table_id = kwargs.get("table_id") records = kwargs.get("records") pk_field_id = kwargs.get("primary_key_field_id") resume_from = kwargs.get("resume_from", 0) if not table_id: return {"success": False, "error": "Missing required field: table_id"} if not records: return {"success": False, "error": "Missing required field: records"} if not pk_field_id: return {"success": False, "error": "Missing required field: primary_key_field_id"} batch = records[resume_from:] if resume_from > 0 else records result = await self._batch_upsert(table_id, batch, pk_field_id) return {"success": True, **result} async def _batch_upsert( self, table_id: str, records: list[dict[str, object]], pk_field_id: str ) -> dict[str, object]: """Upsert records in batches of BATCH_SIZE via POST /tables/{id}/upsert.""" client = await self._get_client() total = len(records) successful = 0 errors: list[dict[str, object]] = [] for start in range(0, total, BATCH_SIZE): batch = records[start : start + BATCH_SIZE] try: resp = await client.post( f"/tables/{table_id}/upsert", json={ "records": batch, "primary_key_field_id": pk_field_id, }, ) resp.raise_for_status() data = resp.json() successful += data.get("inserted", 0) + data.get("updated", 0) except httpx.HTTPStatusError as e: errors.append( { "batch_start": start, "batch_size": len(batch), "status": e.response.status_code, "error": e.response.text[:300], } ) break return { "successful_count": successful, "total": total, "resume_from": successful, **({"errors": errors} if errors else {}), } # ------------------------------------------------------------------ # query_records # ------------------------------------------------------------------ async def _query_records(self, **kwargs) -> dict[str, object]: table_id = kwargs.get("table_id") if not table_id: return {"success": False, "error": "Missing required field: table_id"} client = await self._get_client() params: dict[str, object] = {} if kwargs.get("cursor"): params["cursor"] = kwargs["cursor"] if kwargs.get("limit"): params["limit"] = kwargs["limit"] resp = await client.get(f"/tables/{table_id}/records", params=params) resp.raise_for_status() data = resp.json() return { "success": True, "records": data["records"], "next_cursor": data.get("next_cursor"), } # ------------------------------------------------------------------ # View & field CRUD (U6: agent parity with REST endpoints) # ------------------------------------------------------------------ async def _create_view(self, **kwargs) -> dict[str, object]: table_id = kwargs.get("table_id") name = kwargs.get("name") if not table_id: return {"success": False, "error": "Missing required field: table_id"} if not name: return {"success": False, "error": "Missing required field: name"} view_type = kwargs.get("view_type") or "grid" config = kwargs.get("config") or {} client = await self._get_client() resp = await client.post( f"/tables/{table_id}/views", json={"name": name, "view_type": view_type, "config": config}, ) resp.raise_for_status() return {"success": True, "view": resp.json()["view"]} async def _update_view(self, **kwargs) -> dict[str, object]: view_id = kwargs.get("view_id") if not view_id: return {"success": False, "error": "Missing required field: view_id"} payload: dict[str, object] = {} if kwargs.get("name") is not None: payload["name"] = kwargs["name"] if kwargs.get("config") is not None: payload["config"] = kwargs["config"] client = await self._get_client() resp = await client.patch(f"/views/{view_id}", json=payload) resp.raise_for_status() return {"success": True, "view": resp.json()["view"]} async def _update_field(self, **kwargs) -> dict[str, object]: field_id = kwargs.get("field_id") if not field_id: return {"success": False, "error": "Missing required field: field_id"} payload: dict[str, object] = {} if kwargs.get("name") is not None: payload["name"] = kwargs["name"] # ponytail: PATCH /fields/{id} is backed by UpdateFieldRequest which # currently accepts only name + config. `type` is forwarded as # `field_type` so the tool is forward-compatible once the request # model adds it; today it is silently ignored by Pydantic (extra=ignore). if kwargs.get("type") is not None: payload["field_type"] = kwargs["type"] if kwargs.get("config") is not None: payload["config"] = kwargs["config"] client = await self._get_client() resp = await client.patch(f"/fields/{field_id}", json=payload) resp.raise_for_status() return {"success": True, "field": resp.json()["field"]} async def _delete_view(self, **kwargs) -> dict[str, object]: view_id = kwargs.get("view_id") if not view_id: return {"success": False, "error": "Missing required field: view_id"} client = await self._get_client() resp = await client.delete(f"/views/{view_id}") resp.raise_for_status() # 204 No Content has an empty body; report a stable success shape. return {"success": True, "deleted": True}