225 lines
12 KiB
Python
225 lines
12 KiB
Python
"""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)
|