97 lines
4.7 KiB
Python
97 lines
4.7 KiB
Python
"""Side-effect-limited adapters for existing domain calculations."""
|
|
|
|
import asyncio
|
|
from threading import BoundedSemaphore
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
from app.mcp.scenario_service import SCENARIO_TOOLS, calculate_scenario
|
|
from app.mcp.errors import CalculationCapacityExceeded
|
|
|
|
# These services synchronously read SCADA/model series or write online results.
|
|
# Topology-only isolation and model-based sensor placement need only BizDB.
|
|
TIMESCALE_CALCULATIONS = SCENARIO_TOOLS | frozenset({
|
|
"run_burst_detection", "run_burst_location", "run_leakage_identification",
|
|
})
|
|
|
|
# A timed-out HTTP request cannot kill a native algorithm thread. Keep its slot
|
|
# until the calculation really exits, so retries cannot accumulate CPU jobs.
|
|
_slots = BoundedSemaphore(2)
|
|
|
|
|
|
def _calculate(name: str, parameters: dict[str, Any], network: str, username: str) -> dict[str, Any]:
|
|
if not _slots.acquire(blocking=False):
|
|
raise CalculationCapacityExceeded()
|
|
try:
|
|
if name in SCENARIO_TOOLS:
|
|
return calculate_scenario(name, parameters, network=network, username=username)
|
|
if name == "calculate_demand_allocation":
|
|
from app.services import demand_allocation
|
|
target = parameters.pop("target")
|
|
limit, offset = parameters.pop("limit"), parameters.pop("offset")
|
|
demand = parameters["demand"]
|
|
if target == "nodes":
|
|
result = demand_allocation.calculate_demand_to_nodes(network, demand, parameters["node_ids"])
|
|
elif target == "region":
|
|
result = demand_allocation.calculate_demand_to_region(network, demand, parameters["region_id"])
|
|
else:
|
|
result = demand_allocation.calculate_demand_to_network(network, demand)
|
|
keys = sorted(result)
|
|
return {"items": [{"node_id": key, "demand": result[key]} for key in keys[offset:offset + limit]],
|
|
"total": len(keys), "limit": limit, "offset": offset,
|
|
"has_more": offset + limit < len(keys), "persisted": False}
|
|
if name == "optimize_sensor_placement":
|
|
from app.services.sensor_placement import optimize_sensor_placement_by_kmeans, optimize_sensor_placement_by_sensitivity
|
|
parameters = dict(parameters)
|
|
method = parameters.pop("method")
|
|
optimizer = optimize_sensor_placement_by_sensitivity if method == "sensitivity" else optimize_sensor_placement_by_kmeans
|
|
return optimizer(project_code=network, run_name="MCP preview", created_by=username,
|
|
persist_results=False, **parameters)
|
|
if name == "run_leakage_identification":
|
|
from app.services.dma_leakage_estimation import run_leakage_identification
|
|
|
|
duration = (parameters["scada_end"] - parameters["scada_start"]).total_seconds() / 3600
|
|
return run_leakage_identification(
|
|
network=network, username=username, scheme_name=None,
|
|
duration=duration, q_sum_unit="m3/s", output_flow_unit="m3/s",
|
|
n_workers=1, pop_size=50, max_gen=100, **parameters,
|
|
)
|
|
if name == "run_burst_detection":
|
|
from app.services.burst_detection import run_burst_detection
|
|
|
|
return run_burst_detection(network=network, username=username, scheme_name=None, **parameters)
|
|
if name == "run_burst_location":
|
|
from app.services.burst_location import run_burst_location_by_network
|
|
|
|
return run_burst_location_by_network(
|
|
network=network, username=username, scheme_name=None,
|
|
final_candidates_csv_path=None, **parameters,
|
|
)
|
|
if name == "analyze_valve_isolation":
|
|
from app.services.valve_isolation import analyze_valve_isolation
|
|
|
|
return analyze_valve_isolation(network=network, **parameters)
|
|
raise ValueError("Calculation is not registered")
|
|
finally:
|
|
_slots.release()
|
|
|
|
|
|
async def calculate(name: str, args: BaseModel, *, network: str, username: str) -> dict[str, Any]:
|
|
# to_thread copies the validated project's routing context. Never accept
|
|
# database routes, actor identity, arbitrary measurements or paths from MCP.
|
|
return await asyncio.to_thread(
|
|
_calculate, name, args.model_dump(exclude={"context_handle", "timeout_seconds"}, exclude_none=True), network, username,
|
|
)
|
|
|
|
|
|
async def bounded_calculation(function, *args, **kwargs):
|
|
"""Use the same native-work capacity for persistent calculation tools."""
|
|
def run():
|
|
if not _slots.acquire(blocking=False):
|
|
raise CalculationCapacityExceeded()
|
|
try:
|
|
return function(*args, **kwargs)
|
|
finally:
|
|
_slots.release()
|
|
return await asyncio.to_thread(run)
|