"""Adapters for bounded, project-owned queries and downloadable exports.""" import asyncio import base64 import hashlib import io import json import zipfile from datetime import datetime, timezone from uuid import uuid4 from typing import Any, TYPE_CHECKING from fastapi import HTTPException from fastapi.encoders import jsonable_encoder from psycopg import sql from pydantic import BaseModel from app.infra.db.postgresql.read_queries import ELEMENT_VIEWS, read_rows from app.mcp.models import PageInfo, QueryResult, BoundContext from app.mcp.scope import QueryScope if TYPE_CHECKING: from app.mcp.catalog import QueryTool # Both function names and object identities are selected from a finite registry. CONFIG_FUNCTIONS = { "pattern": ("pattern", "id"), "curve": ("curve", "id"), "control": ("control", None), "rule": ("rule", None), "quality": ("quality", "node"), "source": ("source", "node"), "reaction": ("reaction", None), "pipe-reaction": ("pipe_reaction", "pipe"), "tank-reaction": ("tank_reaction", "tank"), "mixing": ("mixing", "tank"), "emitter": ("emitter", "junction"), "demand": ("demand", "junction"), "time": ("time", None), "energy": ("energy", None), "pump-energy": ("pump_energy", "pump"), "option": ("option_v3", None), } EXTENDED_READ_TOOLS = { "list_regions", "get_region", "get_region_nodes", "get_node_connections", "query_network_extent", "list_major_nodes", "get_model_configuration", "list_model_configurations", "search_project_places", "reverse_geocode", "get_project_info", "export_project_model", "get_sensor_placement_candidate", "export_sensor_placement", "list_task_schedules", "get_task_run_history", "get_mutation_status", } ARTIFACT_TOOLS = {"export_project_model", "export_sensor_placement"} def configuration(network: str, kind: str, object_id: str | None = None) -> Any: from app.services import tjnetwork function, identity = CONFIG_FUNCTIONS[kind] args = (network, object_id) if identity else (network,) return getattr(tjnetwork, "get_" + function)(*args) def artifact(filename: str, payload: bytes, mime_type: str) -> dict[str, Any]: if len(payload) > 5 * 1024 * 1024: raise HTTPException(413, "Export exceeds the 5 MiB compressed file budget") return {"file": { "name": filename, "mime_type": mime_type, "size": len(payload), "sha256": hashlib.sha256(payload).hexdigest(), "encoding": "base64", "data": base64.b64encode(payload).decode("ascii"), }} def project_export(network: str, args: BaseModel) -> dict[str, Any]: from app.services import tjnetwork from app.services.project_inp import temporary_project_inp output = io.BytesIO() with zipfile.ZipFile(output, "w", zipfile.ZIP_DEFLATED) as archive: if args.format == "inp": with temporary_project_inp(network, purpose="mcp_export", version=args.version) as path: archive.write(path, arcname="model.inp") else: changes = tjnetwork.export_inp(network, args.version) if changes.operations: changes.operations[0]["vertex"] = json.dumps(tjnetwork.get_all_vertices(network)) changes.operations[0]["scada"] = json.dumps(tjnetwork.get_all_scada_info(network)) archive.writestr("model.json", json.dumps(jsonable_encoder(changes.operations), ensure_ascii=False)) return artifact("model.zip", output.getvalue(), "application/zip") def sensor_export(network: str, args: BaseModel, user: Any, project: Any) -> dict[str, Any]: from app.services.sensor_placement import get_sensor_placement_run, build_sensor_placement_workbook run = get_sensor_placement_run(network, args.run_id) workbook = build_sensor_placement_workbook( network=network, scheme=run, sensor_location=run["sensor_locations"], adjustment_status={}, ) return artifact("sensor-placement.xlsx", workbook.getvalue(), "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet") async def execute_extended(scope: QueryScope, binding: BoundContext, tool: "QueryTool", args: BaseModel) -> QueryResult: name, conn, network = tool.name, scope.connection, scope.project.project_code page, gaps, truncated = None, [], False def paginate(items, *, already_paged=False): nonlocal page, truncated if already_paged: more, selected, total = len(items) > args.limit, items[:args.limit], None else: total = len(items) more = args.offset + args.limit < total selected = items[args.offset:args.offset + args.limit] page = PageInfo(limit=args.limit, offset=args.offset, total=total, has_more=more) truncated = more or args.offset > 0 return selected if name == "list_regions": rows = await read_rows(conn, "SELECT id,region_type FROM gis.regions ORDER BY id LIMIT %s OFFSET %s", (args.limit + 1, args.offset)) data = paginate(rows, already_paged=True) elif name == "get_region": rows = await read_rows(conn, "SELECT id,region_type,ST_AsGeoJSON(ST_Transform(boundary,4326))::json AS geometry FROM gis.regions WHERE id=%s", (args.region_id,)) if not rows: raise HTTPException(404, "Region not found in the current project") data = {**rows[0], "crs": "EPSG:4326"} elif name == "get_region_nodes": rows = await read_rows(conn, """WITH configured AS ( SELECT node_id FROM gis.region_nodes WHERE region_id=%s ), members AS ( SELECT node_id FROM configured UNION ALL SELECT n.node_id FROM gis.node_geometries n JOIN gis.regions r ON r.id=%s AND ST_Intersects(n.geom,r.boundary) WHERE NOT EXISTS (SELECT 1 FROM configured) ) SELECT node_id AS id FROM members ORDER BY node_id LIMIT %s OFFSET %s""", (args.region_id, args.region_id, args.limit + 1, args.offset)) data = paginate(rows, already_paged=True) gaps = ["Uses configured region membership, or geometry if none is configured; membership does not prove hydraulic connectivity"] elif name == "get_node_connections": rows = await read_rows(conn, """SELECT id,link_type,start_node_id,end_node_id FROM network.links WHERE start_node_id=%s OR end_node_id=%s ORDER BY id LIMIT %s OFFSET %s""", (args.node_id, args.node_id, args.limit + 1, args.offset)) data = paginate(rows, already_paged=True) elif name == "query_network_extent": statement = sql.SQL("""SELECT v.id,to_jsonb(v)-'geom' AS properties, ST_AsGeoJSON(ST_Transform(v.geom,4326))::json AS geometry FROM gis.{} v WHERE ST_Intersects(v.geom,ST_Transform(ST_MakeEnvelope(%s,%s,%s,%s,4326),3857)) ORDER BY v.id LIMIT %s OFFSET %s""").format(sql.Identifier(ELEMENT_VIEWS[args.element_type])) rows = await read_rows(conn, statement, (*args.bounds, args.limit + 1, args.offset)) data = {"items": paginate(rows, already_paged=True), "crs": "EPSG:4326", "source": "current_project_model"} gaps = ["Only current-project model elements with valid stored geometry are covered; missing GIS geometry is not evidence of an empty network"] elif name == "list_major_nodes": rows = await read_rows(conn, """SELECT n.id,ST_X(ST_Transform(n.geom,4326)) AS longitude, ST_Y(ST_Transform(n.geom,4326)) AS latitude FROM gis.junctions n WHERE EXISTS (SELECT 1 FROM gis.pipes p WHERE p.diameter >= %s AND (p.node1=n.id OR p.node2=n.id)) ORDER BY n.id LIMIT %s OFFSET %s""", (args.min_diameter_mm, args.limit + 1, args.offset)) data = {"items": paginate(rows, already_paged=True), "crs": "EPSG:4326"} gaps = ["Only junctions with stored geometry and qualifying model pipes are covered; no other project's GIS is queried"] elif name == "get_model_configuration": data = await asyncio.to_thread(configuration, network, args.kind, args.object_id) for key in args.data_path: if not isinstance(data, dict) or key not in data: raise HTTPException(404, "Configuration path not found") data = data[key] if isinstance(data, dict): data = {"items": paginate([{"key": key, "value": value} for key, value in sorted(data.items())]), "data_path": args.data_path} elif isinstance(data, list): data = {"items": paginate(data), "data_path": args.data_path} else: data = {"value": data, "data_path": args.data_path} elif name == "list_model_configurations": from app.services import tjnetwork rows = await asyncio.to_thread(getattr(tjnetwork, "get_" + args.kind + "s"), network) data = paginate(sorted(rows or [])) elif name in {"search_project_places", "reverse_geocode", "get_project_info"}: project = await scope.metadata.get_project_by_id(scope.project.project_id) if project is None: raise HTTPException(404, "Project not found") if name == "get_project_info": data = {key: getattr(project, key, None) for key in ("name", "code", "description", "map_extent", "status")} data["project_id"] = str(scope.project.project_id) else: from app.domain.schemas.project_map_config import ProjectMapConfig from app.domain.schemas.place_search import PlaceSearchRequest, ReverseGeocodingRequest from app.services.place_search import search_tianditu_places, PlaceSearchError from app.services.reverse_geocoding import reverse_geocode_tianditu config = ProjectMapConfig.model_validate(project.map_config or {}).place_search if config is None or not config.enabled: raise HTTPException(503, "Current project place search is not configured") await scope.metadata.session.commit() try: data = (await search_tianditu_places(PlaceSearchRequest(keyword=args.keyword), config=config) if name == "search_project_places" else await reverse_geocode_tianditu( ReverseGeocodingRequest(longitude=args.longitude, latitude=args.latitude), config=config)) truncated = bool(getattr(data, "truncated", False)) except PlaceSearchError as exc: raise HTTPException(exc.status_code, exc.detail) from exc elif name == "export_project_model": data = await asyncio.to_thread(project_export, network, args) elif name == "get_sensor_placement_candidate": from app.services.sensor_placement import get_sensor_placement_candidate data = await asyncio.to_thread(get_sensor_placement_candidate, network, args.node_id) elif name == "export_sensor_placement": data = await asyncio.to_thread(sensor_export, network, args, scope.user, scope.project) elif name == "get_mutation_status": from app.infra.db.postgresql.mcp_queries import McpQueryRepository repo = McpQueryRepository(conn, scope.project.project_id) previous = await repo.receipt("mutation:" + binding.subject + ":" + binding.native_session_id, str(args.request_id)) if previous is None: data = {"status": "not_started", "request_id": str(args.request_id)} else: output = previous["result"].get("output", {}) data = {"request_id": str(args.request_id), "status": previous["result"]["status"], "tool": previous["parameters"]["tool"], "source_ref": output.get("source_ref"), "queried_at": output.get("queried_at")} if previous["result"]["status"] == "rejected": from app.mcp.mutation_service import retryable_capacity_rejection data["retryable"] = retryable_capacity_rejection(previous["result"]) data["error"] = previous["result"].get("error") gaps = ["started does not prove success or rollback; verify target state before any new operation"] elif name in {"list_task_schedules", "get_task_run_history"}: from app.mcp.query_service import task_service service = task_service(scope) if name == "list_task_schedules": data = paginate(await service.list_schedules()) else: rows = await service.task_runs(args.target_task_id, limit=args.limit + 1, offset=args.offset) data = paginate(rows, already_paged=True) else: raise ValueError("Unknown extended query") return QueryResult(source=name, source_ref="mcp:" + uuid4().hex, queried_at=datetime.now(timezone.utc), semantics=tool.semantics, data=jsonable_encoder(data), page=page, truncated=truncated, data_gaps=gaps)