""" rest_api.py REST / OpenAPI surface for Hydraulic Solver Teaching Mode. Run: CLIENT_API_KEY=your-secret uvicorn rest_api:app --host 0.0.0.0 --port 8000 OpenAPI: /openapi.json Docs: /docs """ from __future__ import annotations import os from typing import List, Optional from fastapi import Depends, FastAPI, Header, HTTPException from pydantic import BaseModel, Field import hydraulic_core as hc API_KEY = os.getenv("CLIENT_API_KEY", "") app = FastAPI( title="Hydraulic Solver Teaching Mode API", version="1.0.0", description="REST/OpenAPI surface exposing the same deterministic tools as the MCP server.", ) def require_auth(authorization: Optional[str] = Header(default=None)) -> None: if not API_KEY: return expected = f"Bearer {API_KEY}" if authorization != expected: raise HTTPException(status_code=401, detail="Missing or invalid bearer token.") class LoadNetworkRequest(BaseModel): inp_text: str title: str = "uploaded_network" class SinglePipeRequest(BaseModel): unit_system: str = "MKS / L/s" length: float diameter: float c_hw: float = 120.0 target_headloss: float initial_flow: float max_iter: int = 25 tolerance: float = 0.0001 class HardyLoopRequest(BaseModel): unit_system: str = "MKS / L/s" flows: List[float] lengths: List[float] diameters: List[float] c_values: List[float] max_iter: int = 25 tolerance: float = 0.00001 class TwoLoopRequest(BaseModel): unit_system: str = "MKS / L/s" flows: List[float] = Field(..., description="Seven initial pipe flows P1..P7.") common_length: float common_diameter: float c_hw: float = 120.0 max_iter: int = 25 tolerance: float = 0.00001 class ThreeReservoirRequest(BaseModel): unit_system: str = "MKS / L/s" reservoir_heads: List[float] demand: float initial_head: float lengths: List[float] diameters: List[float] c_values: List[float] max_iter: int = 25 tolerance: float = 0.0001 class PDDRequest(BaseModel): unit_system: str = "MKS / L/s" required_demand: float available_pressure: float minimum_pressure: float required_pressure: float exponent: float = 0.5 class TankEPSRequest(BaseModel): unit_system: str = "MKS / L/s" diameter: float initial_level: float min_level: float max_level: float timestep_hr: float = 1.0 inflows: List[float] outflows: List[float] class ValveRequest(BaseModel): unit_system: str = "MKS / L/s" valve_type: str = "PRV" upstream_head: float setting: float flow: float diameter: float minor_loss_k: float = 0.0 class PumpRequest(BaseModel): unit_system: str = "MKS / L/s" shutoff_head: float design_flow: float static_head: float pump_curve_k: float system_curve_k: float class PressureZoneRequest(BaseModel): unit_system: str = "MKS / L/s" source_head: float prv_setting: float node_elevations: List[float] demand_multiplier: float = 1.0 min_pressure: float = 14.0 max_pressure: float = 56.0 class LeakageRequest(BaseModel): unit_system: str = "MKS / L/s" average_pressure: float authorized_demand: float leakage_coefficient: float pressure_exponent: float = 1.0 class WaterAgeRequest(BaseModel): unit_system: str = "MKS / L/s" pipe_volume: float tank_volume: float demand: float dead_end_factor: float = 2.0 class ChlorineRequest(BaseModel): initial_chlorine_mg_l: float bulk_decay_per_day: float travel_time_hours: float wall_decay_factor: float = 0.0 @app.get("/health") def health(): return {"ok": True, "service": "hydraulic-solver-teaching-mode-api"} @app.post("/sessions", dependencies=[Depends(require_auth)]) def create_session(req: LoadNetworkRequest): return hc.load_network(req.inp_text, req.title) @app.get("/sessions/{session_id}", dependencies=[Depends(require_auth)]) def get_session_summary(session_id: str): return hc.network_summary(session_id) @app.delete("/sessions/{session_id}", dependencies=[Depends(require_auth)]) def delete_session(session_id: str): return hc.close_session(session_id) @app.post("/tools/solve_single_pipe", dependencies=[Depends(require_auth)]) def solve_single_pipe(req: SinglePipeRequest): return hc.solve_single_pipe(**req.model_dump()) @app.post("/tools/solve_hardy_cross_loop", dependencies=[Depends(require_auth)]) def solve_hardy_cross_loop(req: HardyLoopRequest): return hc.solve_hardy_cross_loop(**req.model_dump()) @app.post("/tools/solve_two_loop_hardy_cross", dependencies=[Depends(require_auth)]) def solve_two_loop_hardy_cross(req: TwoLoopRequest): return hc.solve_two_loop_hardy_cross(**req.model_dump()) @app.post("/tools/solve_three_reservoir", dependencies=[Depends(require_auth)]) def solve_three_reservoir(req: ThreeReservoirRequest): return hc.solve_three_reservoir(**req.model_dump()) @app.post("/tools/solve_pdd_demand", dependencies=[Depends(require_auth)]) def solve_pdd_demand(req: PDDRequest): return hc.solve_pdd_demand(**req.model_dump()) @app.post("/tools/simulate_tank_eps", dependencies=[Depends(require_auth)]) def simulate_tank_eps(req: TankEPSRequest): return hc.simulate_tank_eps(**req.model_dump()) @app.post("/tools/evaluate_valve_behavior", dependencies=[Depends(require_auth)]) def evaluate_valve_behavior(req: ValveRequest): return hc.evaluate_valve_behavior(**req.model_dump()) @app.post("/tools/solve_pump_operating_point", dependencies=[Depends(require_auth)]) def solve_pump_operating_point(req: PumpRequest): return hc.solve_pump_operating_point(**req.model_dump()) @app.post("/tools/pressure_zone_analysis", dependencies=[Depends(require_auth)]) def pressure_zone_analysis(req: PressureZoneRequest): return hc.pressure_zone_analysis(**req.model_dump()) @app.post("/tools/leakage_nrw_analysis", dependencies=[Depends(require_auth)]) def leakage_nrw_analysis(req: LeakageRequest): return hc.leakage_nrw_analysis(**req.model_dump()) @app.post("/tools/water_age_analysis", dependencies=[Depends(require_auth)]) def water_age_analysis(req: WaterAgeRequest): return hc.water_age_analysis(**req.model_dump()) @app.post("/tools/chlorine_decay_analysis", dependencies=[Depends(require_auth)]) def chlorine_decay_analysis(req: ChlorineRequest): return hc.chlorine_decay_analysis(**req.model_dump()) @app.get("/tools/generate_epanet_validation_inp", dependencies=[Depends(require_auth)]) def generate_epanet_validation_inp(case: str = "three_reservoir_mks"): return hc.generate_epanet_validation_inp(case)