src/iotdb_mcp_server/services/json_response.py - iotdb-mcp-server - Git at Google

 #
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
 #

 import json
 from contextvars import ContextVar
 from datetime import datetime, timezone
 from typing import Any, Callable, Iterable

 from mcp.types import TextContent

 from iotdb_mcp_server.result_store import FileResultStore

 _RESPONSE_VERSION = "1.0"
 _target_context: ContextVar[dict[str, Any] | None] = ContextVar(
     "iotdb_target_context", default=None
 )

 _TRUE_VALUES = {"1", "true", "yes", "on", "available", "running", "enabled"}
 _FALSE_VALUES = {"0", "false", "no", "off", "unavailable", "stopped", "disabled"}


 def set_iotdb_target_context(target: dict[str, Any] | None):
     return _target_context.set(target)


 def reset_iotdb_target_context(token) -> None:
     _target_context.reset(token)


 def _coerce_availability(value: Any) -> str | None:
     if value is None:
         return None
     if isinstance(value, bool):
         return "yes" if value else "no"
     normalized = str(value).strip().lower()
     if normalized in _TRUE_VALUES:
         return "yes"
     if normalized in _FALSE_VALUES:
         return "no"
     if normalized in {"unknown", "probe_failed", "not_probed"}:
         return "unknown"
     return None


 def _ainode_availability(target: dict[str, Any]) -> tuple[str, Any]:
     availability = _coerce_availability(target.get("ainode_available")) or "unknown"
     source = target.get("ainode_availability_source")
     if source is None:
         source = {
             "method": "SHOW AINODES",
             "ok": False,
             "error": "not_probed",
         }
     return availability, source


 def _cli_endpoint(target: dict[str, Any]) -> str:
     node_urls = target.get("node_urls")
     if isinstance(node_urls, list) and node_urls:
         first = str(node_urls[0]).strip()
         if first:
             return first
     host = str(target.get("host") or "127.0.0.1").strip()
     port = str(target.get("port") or "6667").strip()
     return f"{host}:{port}"


 def _row_count(payload: Any) -> int:
     if not isinstance(payload, dict):
         return 0
     rows = payload.get("rows")
     if isinstance(rows, list):
         return len(rows)
     for key in ("row_count", "rows_count", "count"):
         value = payload.get(key)
         if isinstance(value, int):
             return value
         if isinstance(value, str) and value.strip().isdigit():
             return int(value.strip())
     return 0


 def _iotdb_context(tool: str, payload: Any, target: dict[str, Any]) -> dict[str, Any]:
     ainode_available, ainode_source = _ainode_availability(target)
     return {
         "target_id": target.get("target_id"),
         "cli_endpoint": _cli_endpoint(target),
         "ainode_available": ainode_available,
         "ainode_availability_source": ainode_source,
         "dialect": target.get("sql_dialect"),
         "database": target.get("database"),
         "tool": tool,
         "rows": _row_count(payload),
     }


 def _render_iotdb_context(context: dict[str, Any]) -> str:
     lines = [
         "TimeSeek IoTDB context:",
         f"- target_id: {context.get('target_id') or ''}",
         f"- cli_endpoint: {context.get('cli_endpoint') or ''}",
         f"- ainode_available: {context.get('ainode_available') or 'unknown'}",
         f"- dialect: {context.get('dialect') or ''}",
         f"- database: {context.get('database') or ''}",
         f"- tool: {context.get('tool') or ''}",
         f"- rows: {context.get('rows', 0)}",
     ]
     return "\n".join(lines)


 class JsonParser:
     """Simple parser/checker for MCP tool JSON envelopes."""

     _required_keys = ("version", "tool", "ok", "timestamp", "payload")

     @staticmethod
     def check_format(obj: Any) -> bool:
         if not isinstance(obj, dict):
             return False
         for key in JsonParser._required_keys:
             if key not in obj:
                 return False
         if not isinstance(obj["version"], str):
             return False
         if not isinstance(obj["tool"], str) or not obj["tool"].strip():
             return False
         if not isinstance(obj["ok"], bool):
             return False
         if not isinstance(obj["timestamp"], str) or not obj["timestamp"].strip():
             return False
         return True

     @staticmethod
     def parse(text: str) -> dict[str, Any]:
         try:
             obj = json.loads(text)
         except json.JSONDecodeError as e:
             raise ValueError(f"Invalid JSON text: {e.msg}") from e

         if not JsonParser.check_format(obj):
             raise ValueError(
                 "Invalid wrapped response format. Expected keys: "
                 "version, tool, ok, timestamp, payload."
             )
         return obj


 def _envelope(
     tool: str, payload: Any, ok: bool = True, message: str | None = None
 ) -> dict[str, Any]:
     target = _target_context.get()
     context = _iotdb_context(tool, payload, target) if target else None
     if target:
         if isinstance(payload, dict):
             payload = dict(payload)
             payload.setdefault("iotdb_target", target)
             if context:
                 payload.setdefault("iotdb_context", context)
         else:
             payload = {
                 "value": payload,
                 "iotdb_target": target,
             }
             if context:
                 payload["iotdb_context"] = context
     obj: dict[str, Any] = {
         "version": _RESPONSE_VERSION,
         "tool": tool,
         "ok": ok,
         "timestamp": datetime.now(timezone.utc).isoformat(),
         "payload": payload,
     }
     if context:
         obj["context"] = {"iotdb": context}
     if message is not None:
         obj["message"] = message
     return obj


 def _to_text_content(obj: dict[str, Any]) -> list[TextContent]:
     if not JsonParser.check_format(obj):
         raise ValueError("JsonParser format check failed before output.")
     text = json.dumps(obj, ensure_ascii=False)
     JsonParser.parse(text)
     context = obj.get("context")
     iotdb_context = context.get("iotdb") if isinstance(context, dict) else None
     if isinstance(iotdb_context, dict):
         return [
             TextContent(type="text", text=_render_iotdb_context(iotdb_context)),
             TextContent(type="text", text=text),
         ]
     return [TextContent(type="text", text=text)]


 def payload_response(
     tool: str, payload: Any, message: str | None = None
 ) -> list[TextContent]:
     return _to_text_content(_envelope(tool=tool, payload=payload, message=message))


 def error_response(tool: str, error: str) -> list[TextContent]:
     return _to_text_content(
         _envelope(
             tool=tool,
             payload={"error": error},
             ok=False,
             message="Tool execution failed.",
         )
     )


 def text_payload_response(
     tool: str, text: str, message: str | None = None
 ) -> list[TextContent]:
     return payload_response(
         tool=tool,
         payload={"format": "text", "text": text},
         message=message,
     )


 def _csv_text(columns: list[str], rows: list[str]) -> str:
     return "\n".join([",".join(columns)] + rows)


 def csv_payload_response(
     tool: str,
     columns: list[str],
     rows: list[str],
     message: str | None = None,
     diagnostics: list[dict[str, Any]] | None = None,
 ) -> list[TextContent]:
     csv_text = _csv_text(columns, rows)
     payload: dict[str, Any] = {
         "format": "csv",
         "columns": columns,
         "rows": rows,
         "text": csv_text,
         "row_count": len(rows),
         "inline_row_count": len(rows),
         "inline_truncated": False,
     }
     if diagnostics:
         payload["diagnostics"] = diagnostics
     return payload_response(
         tool=tool,
         payload=payload,
         message=message,
     )


 def csv_result_payload_response(
     tool: str,
     columns: list[str],
     rows: Iterable[str],
     *,
     result_store: FileResultStore,
     source: dict[str, Any] | None = None,
     message: str | None = None,
     diagnostics: list[dict[str, Any]] | None = None,
     diagnostics_factory: Callable[[int], list[dict[str, Any]] | None] | None = None,
     owner_session_id: str | None = None,
     max_inline_rows: int | None = None,
     page_size_rows: int | None = None,
 ) -> list[TextContent]:
     stored = result_store.write_csv_result(
         tool=tool,
         columns=columns,
         rows=rows,
         source=source,
         owner_session_id=owner_session_id,
         preview_rows=max_inline_rows,
         page_size_rows=page_size_rows,
     )
     preview_rows = stored.preview_rows
     inline_truncated = stored.row_count > len(preview_rows)
     payload: dict[str, Any] = {
         "format": "csv",
         "columns": columns,
         "rows": preview_rows,
         "text": _csv_text(columns, preview_rows),
         "row_count": stored.row_count,
         "inline_row_count": len(preview_rows),
         "inline_truncated": inline_truncated,
         "result_id": stored.result_id,
         "owner_session_id": stored.owner_session_id,
         "next_cursor": str(len(preview_rows)) if inline_truncated else None,
         "result_store": {
             "type": "file",
             "result_id": stored.result_id,
             "owner_session_id": stored.owner_session_id,
             "page_tool": "read_result_page",
             "batch_page_tool": "read_result_pages",
             "page_size_rows": stored.page_size_rows,
             "manifest_path": stored.manifest_path,
             "shard_count": stored.shard_count,
             "byte_count": stored.byte_count,
         },
     }
     combined_diagnostics = list(diagnostics or [])
     if diagnostics_factory is not None:
         combined_diagnostics.extend(diagnostics_factory(stored.row_count) or [])
     if combined_diagnostics:
         payload["diagnostics"] = combined_diagnostics
     return payload_response(tool=tool, payload=payload, message=message)


 def sql_success_response(tool: str, sql: str) -> list[TextContent]:
     return payload_response(
         tool=tool,
         payload={"sql": sql, "result": "success"},
         message="SQL executed successfully.",
     )
	#
	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.
	#

	import json
	from contextvars import ContextVar
	from datetime import datetime, timezone
	from typing import Any, Callable, Iterable

	from mcp.types import TextContent

	from iotdb_mcp_server.result_store import FileResultStore

	_RESPONSE_VERSION = "1.0"
	_target_context: ContextVar[dict[str, Any] \| None] = ContextVar(
	"iotdb_target_context", default=None
	)

	_TRUE_VALUES = {"1", "true", "yes", "on", "available", "running", "enabled"}
	_FALSE_VALUES = {"0", "false", "no", "off", "unavailable", "stopped", "disabled"}


	def set_iotdb_target_context(target: dict[str, Any] \| None):
	return _target_context.set(target)


	def reset_iotdb_target_context(token) -> None:
	_target_context.reset(token)


	def _coerce_availability(value: Any) -> str \| None:
	if value is None:
	return None
	if isinstance(value, bool):
	return "yes" if value else "no"
	normalized = str(value).strip().lower()
	if normalized in _TRUE_VALUES:
	return "yes"
	if normalized in _FALSE_VALUES:
	return "no"
	if normalized in {"unknown", "probe_failed", "not_probed"}:
	return "unknown"
	return None


	def _ainode_availability(target: dict[str, Any]) -> tuple[str, Any]:
	availability = _coerce_availability(target.get("ainode_available")) or "unknown"
	source = target.get("ainode_availability_source")
	if source is None:
	source = {
	"method": "SHOW AINODES",
	"ok": False,
	"error": "not_probed",
	}
	return availability, source


	def _cli_endpoint(target: dict[str, Any]) -> str:
	node_urls = target.get("node_urls")
	if isinstance(node_urls, list) and node_urls:
	first = str(node_urls[0]).strip()
	if first:
	return first
	host = str(target.get("host") or "127.0.0.1").strip()
	port = str(target.get("port") or "6667").strip()
	return f"{host}:{port}"


	def _row_count(payload: Any) -> int:
	if not isinstance(payload, dict):
	return 0
	rows = payload.get("rows")
	if isinstance(rows, list):
	return len(rows)
	for key in ("row_count", "rows_count", "count"):
	value = payload.get(key)
	if isinstance(value, int):
	return value
	if isinstance(value, str) and value.strip().isdigit():
	return int(value.strip())
	return 0


	def _iotdb_context(tool: str, payload: Any, target: dict[str, Any]) -> dict[str, Any]:
	ainode_available, ainode_source = _ainode_availability(target)
	return {
	"target_id": target.get("target_id"),
	"cli_endpoint": _cli_endpoint(target),
	"ainode_available": ainode_available,
	"ainode_availability_source": ainode_source,
	"dialect": target.get("sql_dialect"),
	"database": target.get("database"),
	"tool": tool,
	"rows": _row_count(payload),
	}


	def _render_iotdb_context(context: dict[str, Any]) -> str:
	lines = [
	"TimeSeek IoTDB context:",
	f"- target_id: {context.get('target_id') or ''}",
	f"- cli_endpoint: {context.get('cli_endpoint') or ''}",
	f"- ainode_available: {context.get('ainode_available') or 'unknown'}",
	f"- dialect: {context.get('dialect') or ''}",
	f"- database: {context.get('database') or ''}",
	f"- tool: {context.get('tool') or ''}",
	f"- rows: {context.get('rows', 0)}",
	]
	return "\n".join(lines)


	class JsonParser:
	"""Simple parser/checker for MCP tool JSON envelopes."""

	_required_keys = ("version", "tool", "ok", "timestamp", "payload")

	@staticmethod
	def check_format(obj: Any) -> bool:
	if not isinstance(obj, dict):
	return False
	for key in JsonParser._required_keys:
	if key not in obj:
	return False
	if not isinstance(obj["version"], str):
	return False
	if not isinstance(obj["tool"], str) or not obj["tool"].strip():
	return False
	if not isinstance(obj["ok"], bool):
	return False
	if not isinstance(obj["timestamp"], str) or not obj["timestamp"].strip():
	return False
	return True

	@staticmethod
	def parse(text: str) -> dict[str, Any]:
	try:
	obj = json.loads(text)
	except json.JSONDecodeError as e:
	raise ValueError(f"Invalid JSON text: {e.msg}") from e

	if not JsonParser.check_format(obj):
	raise ValueError(
	"Invalid wrapped response format. Expected keys: "
	"version, tool, ok, timestamp, payload."
	)
	return obj


	def _envelope(
	tool: str, payload: Any, ok: bool = True, message: str \| None = None
	) -> dict[str, Any]:
	target = _target_context.get()
	context = _iotdb_context(tool, payload, target) if target else None
	if target:
	if isinstance(payload, dict):
	payload = dict(payload)
	payload.setdefault("iotdb_target", target)
	if context:
	payload.setdefault("iotdb_context", context)
	else:
	payload = {
	"value": payload,
	"iotdb_target": target,
	}
	if context:
	payload["iotdb_context"] = context
	obj: dict[str, Any] = {
	"version": _RESPONSE_VERSION,
	"tool": tool,
	"ok": ok,
	"timestamp": datetime.now(timezone.utc).isoformat(),
	"payload": payload,
	}
	if context:
	obj["context"] = {"iotdb": context}
	if message is not None:
	obj["message"] = message
	return obj


	def _to_text_content(obj: dict[str, Any]) -> list[TextContent]:
	if not JsonParser.check_format(obj):
	raise ValueError("JsonParser format check failed before output.")
	text = json.dumps(obj, ensure_ascii=False)
	JsonParser.parse(text)
	context = obj.get("context")
	iotdb_context = context.get("iotdb") if isinstance(context, dict) else None
	if isinstance(iotdb_context, dict):
	return [
	TextContent(type="text", text=_render_iotdb_context(iotdb_context)),
	TextContent(type="text", text=text),
	]
	return [TextContent(type="text", text=text)]


	def payload_response(
	tool: str, payload: Any, message: str \| None = None
	) -> list[TextContent]:
	return _to_text_content(_envelope(tool=tool, payload=payload, message=message))


	def error_response(tool: str, error: str) -> list[TextContent]:
	return _to_text_content(
	_envelope(
	tool=tool,
	payload={"error": error},
	ok=False,
	message="Tool execution failed.",
	)
	)


	def text_payload_response(
	tool: str, text: str, message: str \| None = None
	) -> list[TextContent]:
	return payload_response(
	tool=tool,
	payload={"format": "text", "text": text},
	message=message,
	)


	def _csv_text(columns: list[str], rows: list[str]) -> str:
	return "\n".join([",".join(columns)] + rows)


	def csv_payload_response(
	tool: str,
	columns: list[str],
	rows: list[str],
	message: str \| None = None,
	diagnostics: list[dict[str, Any]] \| None = None,
	) -> list[TextContent]:
	csv_text = _csv_text(columns, rows)
	payload: dict[str, Any] = {
	"format": "csv",
	"columns": columns,
	"rows": rows,
	"text": csv_text,
	"row_count": len(rows),
	"inline_row_count": len(rows),
	"inline_truncated": False,
	}
	if diagnostics:
	payload["diagnostics"] = diagnostics
	return payload_response(
	tool=tool,
	payload=payload,
	message=message,
	)


	def csv_result_payload_response(
	tool: str,
	columns: list[str],
	rows: Iterable[str],
	*,
	result_store: FileResultStore,
	source: dict[str, Any] \| None = None,
	message: str \| None = None,
	diagnostics: list[dict[str, Any]] \| None = None,
	diagnostics_factory: Callable[[int], list[dict[str, Any]] \| None] \| None = None,
	owner_session_id: str \| None = None,
	max_inline_rows: int \| None = None,
	page_size_rows: int \| None = None,
	) -> list[TextContent]:
	stored = result_store.write_csv_result(
	tool=tool,
	columns=columns,
	rows=rows,
	source=source,
	owner_session_id=owner_session_id,
	preview_rows=max_inline_rows,
	page_size_rows=page_size_rows,
	)
	preview_rows = stored.preview_rows
	inline_truncated = stored.row_count > len(preview_rows)
	payload: dict[str, Any] = {
	"format": "csv",
	"columns": columns,
	"rows": preview_rows,
	"text": _csv_text(columns, preview_rows),
	"row_count": stored.row_count,
	"inline_row_count": len(preview_rows),
	"inline_truncated": inline_truncated,
	"result_id": stored.result_id,
	"owner_session_id": stored.owner_session_id,
	"next_cursor": str(len(preview_rows)) if inline_truncated else None,
	"result_store": {
	"type": "file",
	"result_id": stored.result_id,
	"owner_session_id": stored.owner_session_id,
	"page_tool": "read_result_page",
	"batch_page_tool": "read_result_pages",
	"page_size_rows": stored.page_size_rows,
	"manifest_path": stored.manifest_path,
	"shard_count": stored.shard_count,
	"byte_count": stored.byte_count,
	},
	}
	combined_diagnostics = list(diagnostics or [])
	if diagnostics_factory is not None:
	combined_diagnostics.extend(diagnostics_factory(stored.row_count) or [])
	if combined_diagnostics:
	payload["diagnostics"] = combined_diagnostics
	return payload_response(tool=tool, payload=payload, message=message)


	def sql_success_response(tool: str, sql: str) -> list[TextContent]:
	return payload_response(
	tool=tool,
	payload={"sql": sql, "result": "success"},
	message="SQL executed successfully.",
	)