140 lines
4.9 KiB
Python
140 lines
4.9 KiB
Python
"""Shared upload-policy defaults and filename classification.
|
|
|
|
The policy itself is persisted by :mod:`deerflow.persistence.upload_policy`.
|
|
This module deliberately has no FastAPI dependency so gateway uploads, IM
|
|
channels and clients can share the same defaults and classification rules.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from pathlib import Path
|
|
from threading import RLock
|
|
from typing import Any
|
|
|
|
MIB = 1024 * 1024
|
|
MAX_RUNTIME_FILES = 100
|
|
MAX_RUNTIME_TOTAL_BYTES = 200 * MIB
|
|
MAX_RUNTIME_FILE_BYTES = 100 * MIB
|
|
MAX_RUNTIME_RESULT_CHARS = 32_000
|
|
|
|
DEFAULT_UPLOAD_POLICY: dict[str, Any] = {
|
|
"max_files_per_message": 20,
|
|
"max_total_bytes": 100 * MIB,
|
|
"groups": {
|
|
"document": {
|
|
"max_bytes": 10 * MIB,
|
|
"extensions": [".txt", ".md", ".doc", ".docx", ".ppt", ".pptx", ".pdf"],
|
|
},
|
|
"image": {
|
|
"max_bytes": 10 * MIB,
|
|
"extensions": [".jpg", ".jpeg", ".png"],
|
|
},
|
|
"spreadsheet": {
|
|
"max_bytes": 5 * MIB,
|
|
"extensions": [".xls", ".xlsx", ".csv"],
|
|
},
|
|
},
|
|
"large_file": {
|
|
"threshold_bytes": 1 * MIB,
|
|
"max_result_chars": 12_000,
|
|
"hard_max_result_chars": 32_000,
|
|
},
|
|
}
|
|
|
|
_runtime_lock = RLock()
|
|
_runtime_policy: dict[str, Any] = deepcopy(DEFAULT_UPLOAD_POLICY)
|
|
|
|
|
|
def default_upload_policy() -> dict[str, Any]:
|
|
"""Return an isolated copy of the built-in policy."""
|
|
|
|
return deepcopy(DEFAULT_UPLOAD_POLICY)
|
|
|
|
|
|
def set_runtime_upload_policy(policy: dict[str, Any]) -> None:
|
|
"""Publish a validated DB policy to synchronous harness consumers."""
|
|
|
|
global _runtime_policy
|
|
safe = default_upload_policy()
|
|
try:
|
|
safe["max_files_per_message"] = min(
|
|
MAX_RUNTIME_FILES,
|
|
max(1, int(policy.get("max_files_per_message") or safe["max_files_per_message"])),
|
|
)
|
|
safe["max_total_bytes"] = min(
|
|
MAX_RUNTIME_TOTAL_BYTES,
|
|
max(1, int(policy.get("max_total_bytes") or safe["max_total_bytes"])),
|
|
)
|
|
configured_groups = policy.get("groups") or {}
|
|
for group_name, default_group in safe["groups"].items():
|
|
configured_group = configured_groups.get(group_name) or {}
|
|
default_group["max_bytes"] = min(
|
|
safe["max_total_bytes"],
|
|
MAX_RUNTIME_FILE_BYTES,
|
|
max(1, int(configured_group.get("max_bytes") or default_group["max_bytes"])),
|
|
)
|
|
if "extensions" in configured_group and isinstance(configured_group["extensions"], list):
|
|
default_group["extensions"] = [
|
|
normalized
|
|
for item in configured_group["extensions"]
|
|
if (normalized := normalize_extension(item))
|
|
]
|
|
configured_large = policy.get("large_file") or {}
|
|
large = safe["large_file"]
|
|
large["threshold_bytes"] = min(
|
|
safe["max_total_bytes"],
|
|
max(1, int(configured_large.get("threshold_bytes") or large["threshold_bytes"])),
|
|
)
|
|
large["hard_max_result_chars"] = min(
|
|
MAX_RUNTIME_RESULT_CHARS,
|
|
max(1024, int(configured_large.get("hard_max_result_chars") or large["hard_max_result_chars"])),
|
|
)
|
|
large["max_result_chars"] = min(
|
|
large["hard_max_result_chars"],
|
|
max(1024, int(configured_large.get("max_result_chars") or large["max_result_chars"])),
|
|
)
|
|
except (TypeError, ValueError):
|
|
safe = default_upload_policy()
|
|
with _runtime_lock:
|
|
_runtime_policy = safe
|
|
|
|
|
|
def get_runtime_upload_policy() -> dict[str, Any]:
|
|
"""Return the latest policy published by the gateway in this worker."""
|
|
|
|
with _runtime_lock:
|
|
return deepcopy(_runtime_policy)
|
|
|
|
|
|
def get_runtime_large_file_policy() -> dict[str, int]:
|
|
"""Return bounded-read settings without importing the application layer."""
|
|
|
|
policy = get_runtime_upload_policy()
|
|
return dict(policy.get("large_file") or DEFAULT_UPLOAD_POLICY["large_file"])
|
|
|
|
|
|
def normalize_extension(value: object) -> str:
|
|
"""Normalize a configured extension to lowercase ``.ext`` form."""
|
|
|
|
extension = str(value or "").strip().lower()
|
|
if extension and not extension.startswith("."):
|
|
extension = f".{extension}"
|
|
return extension
|
|
|
|
|
|
def classify_upload_filename(filename: str, policy: dict[str, Any]) -> tuple[str, int] | None:
|
|
"""Return ``(group_name, max_bytes)`` for *filename*, or ``None``.
|
|
|
|
Matching is case-insensitive. A defensive first-match rule is used even
|
|
though the API rejects extension overlap before a policy is stored.
|
|
"""
|
|
|
|
extension = Path(filename).suffix.lower()
|
|
for group_name in ("document", "image", "spreadsheet"):
|
|
group = (policy.get("groups") or {}).get(group_name) or {}
|
|
allowed = {normalize_extension(item) for item in group.get("extensions") or []}
|
|
if extension in allowed:
|
|
return group_name, int(group.get("max_bytes") or 0)
|
|
return None
|