Update tools.py

tools.py

@@ -58,39 +58,156 @@ def _require_results(session) -> None:
 # Model lifecycle
 # ---------------------------------------------------------------------------
 
-def upload_model(inp_content: str, filename: str = "model.inp") -> dict:
-    """Upload an EPA SWMM .inp model (raw text or base64) and create a session.
+MAX_MODEL_BYTES = 20 * 1024 * 1024
+
+
+def _decode_model_payload(
+    inp_content: str,
+    filename: str = "model.inp",
+    encoding: str = "auto",
+) -> tuple[str, str]:
+    """Normalize raw text, base64, or a JSON-wrapped SWMM model payload."""
+    if not isinstance(inp_content, str) or not inp_content.strip():
+        raise ValueError("inp_content is required and must be a non-empty string.")
+
+    value = inp_content.strip()
+    resolved_filename = str(filename or "model.inp").strip() or "model.inp"
+    selected_encoding = str(encoding or "auto").strip().lower()
+
+    # Accept a complete JSON object passed as the inp_content value:
+    # {"filename": "...", "inp_content": "...", "encoding": "base64|raw|auto"}
+    if value.startswith("{"):
+        try:
+            payload = json.loads(value)
+        except json.JSONDecodeError:
+            payload = None
+
+        if isinstance(payload, dict):
+            if "inp_content" not in payload:
+                raise ValueError("JSON upload payload must contain an 'inp_content' field.")
+            value = str(payload.get("inp_content") or "").strip()
+            resolved_filename = (
+                str(payload.get("filename") or "").strip()
+                or resolved_filename
+            )
+            selected_encoding = str(
+                payload.get("encoding") or selected_encoding
+            ).strip().lower()
+
+    if not value:
+        raise ValueError("The upload payload contains no SWMM model content.")
+
+    if selected_encoding in {"raw", "text"}:
+        model_text = value
+
+    elif selected_encoding == "base64":
+        try:
+            decoded = base64.b64decode(value, validate=True)
+        except (binascii.Error, ValueError) as exc:
+            raise ValueError("Invalid base64 SWMM model content.") from exc
+        try:
+            model_text = decoded.decode("utf-8-sig")
+        except UnicodeDecodeError:
+            model_text = decoded.decode("cp1252")
+
+    elif selected_encoding == "auto":
+        # Raw INP text normally exposes a section header near the beginning.
+        if "[" in value[:2000]:
+            model_text = value
+        else:
+            try:
+                decoded = base64.b64decode(value, validate=True)
+                try:
+                    candidate = decoded.decode("utf-8-sig")
+                except UnicodeDecodeError:
+                    candidate = decoded.decode("cp1252")
+            except (binascii.Error, ValueError, UnicodeDecodeError):
+                candidate = ""
+
+            model_text = candidate if "[" in candidate[:2000] else value
+
+    else:
+        raise ValueError("encoding must be one of: auto, raw, text, or base64.")
+
+    encoded_size = len(model_text.encode("utf-8"))
+    if encoded_size > MAX_MODEL_BYTES:
+        raise ValueError(
+            f"SWMM model exceeds the {MAX_MODEL_BYTES // (1024 * 1024)} MB upload limit."
+        )
+
+    upper = model_text.upper()
+    if "[OPTIONS]" not in upper and "[JUNCTIONS]" not in upper:
+        raise ValueError(
+            "Content does not look like a SWMM .inp file "
+            "(no [OPTIONS] or [JUNCTIONS] section)."
+        )
+
+    safe_name = Path(resolved_filename).name or "model.inp"
+    if not safe_name.lower().endswith(".inp"):
+        safe_name = f"{Path(safe_name).stem}.inp"
+
+    return model_text, safe_name
+
+
+def upload_model(
+    inp_content: str,
+    filename: str = "model.inp",
+    encoding: str = "auto",
+) -> dict:
+    """Upload an EPA SWMM model as raw text, base64, or a JSON-wrapped payload.
+
+    JSON payload format:
+    {"filename": "model.inp", "inp_content": "...", "encoding": "auto|raw|base64"}
 
     Returns a session_id used by every other tool, plus element counts.
     """
-    text = inp_content
-    if "[" not in inp_content[:2000]:  # likely base64
-        try:
-            text = base64.b64decode(inp_content, validate=True).decode("utf-8", errors="replace")
-        except (binascii.Error, ValueError):
-            pass
-    if "[OPTIONS]" not in text.upper() and "[JUNCTIONS]" not in text.upper():
-        raise ValueError("Content does not look like a SWMM .inp file (no [OPTIONS]/[JUNCTIONS] section).")
+    model_text, safe_name = _decode_model_payload(
+        inp_content=inp_content,
+        filename=filename,
+        encoding=encoding,
+    )
+
     session = STORE.create()
-    safe_name = Path(filename).name or "model.inp"
-    if not safe_name.lower().endswith(".inp"):
-        safe_name += ".inp"
-    inp_path = session.workdir / safe_name
-    inp_path.write_text(text, encoding="utf-8")
-    sections = mp.parse_inp_sections(str(inp_path))
-    session.data.update({"filename": safe_name, "inp_path": str(inp_path), "sections": sections})
-    counts = {name: len(rows) for name, rows in sections.items()
-              if name in ("JUNCTIONS", "OUTFALLS", "STORAGE", "CONDUITS", "PUMPS", "WEIRS",
-                          "ORIFICES", "OUTLETS", "SUBCATCHMENTS", "RAINGAGES", "TIMESERIES")}
-    gages = [row[0] for row in sections.get("RAINGAGES", []) if row]
-    return {
-        "session_id": session.id,
-        "filename": safe_name,
-        "element_counts": counts,
-        "rain_gages": gages,
-        "design_event_inference": infer_design_event(gages) if gages else None,
-        "next_step": "Call run_simulation with this session_id.",
-    }
+    try:
+        inp_path = session.workdir / safe_name
+        inp_path.write_text(model_text, encoding="utf-8", newline="\n")
+
+        sections = mp.parse_inp_sections(str(inp_path))
+        if not sections:
+            raise ValueError("The uploaded file could not be parsed into SWMM sections.")
+
+        session.data.update({
+            "filename": safe_name,
+            "inp_path": str(inp_path),
+            "sections": sections,
+            "upload_encoding": encoding,
+            "model_size_bytes": inp_path.stat().st_size,
+        })
+
+        counts = {
+            name: len(rows)
+            for name, rows in sections.items()
+            if name in (
+                "JUNCTIONS", "OUTFALLS", "STORAGE", "CONDUITS", "PUMPS",
+                "WEIRS", "ORIFICES", "OUTLETS", "SUBCATCHMENTS",
+                "RAINGAGES", "TIMESERIES",
+            )
+        }
+        gages = [row[0] for row in sections.get("RAINGAGES", []) if row]
+
+        return {
+            "session_id": session.id,
+            "filename": safe_name,
+            "model_size_bytes": inp_path.stat().st_size,
+            "element_counts": counts,
+            "rain_gages": gages,
+            "design_event_inference": infer_design_event(gages) if gages else None,
+            "next_step": "Call run_simulation with this session_id.",
+        }
+
+    except Exception:
+        STORE.drop(session.id)
+        raise
 
 
 def run_simulation(session_id: str) -> dict:
@@ -535,4 +652,4 @@ TOOL_REGISTRY: dict[str, Callable[..., dict]] = {
         calgary_screening, preliminary_design_review, get_reconciliation,
         run_scenario, attach_figure, set_report_details, generate_report,
     ]
-}
+}