feat: add project-scoped operations
This commit is contained in:
@@ -0,0 +1,224 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user