Files
site_aegisone/max_bot/app/file_storage.py
T

78 lines
2.9 KiB
Python

import uuid
import json
import logging
from pathlib import Path
from datetime import datetime
ALLOWED_EXTENSIONS = {".pdf", ".docx", ".xlsx", ".txt", ".jpg", ".jpeg", ".png", ".gif"}
logger = logging.getLogger(__name__)
class FileStorage:
def __init__(self, base_dir: str = "uploads/bot"):
self.base_dir = Path(base_dir)
self.base_dir.mkdir(parents=True, exist_ok=True)
async def save_file(self, file_bytes: bytes, filename: str, conversation_id: int,
file_type: str = "document", db=None) -> str:
ext = Path(filename).suffix.lower()
if ext not in ALLOWED_EXTENSIONS:
ext = ".bin"
safe_name = f"{uuid.uuid4().hex}{ext}"
conv_dir = self.base_dir / str(conversation_id)
conv_dir.mkdir(parents=True, exist_ok=True)
file_path = conv_dir / safe_name
file_path.write_bytes(file_bytes)
meta = {
"type": "conversation",
"id": conversation_id,
"filename": filename,
"file_type": file_type,
"timestamp": datetime.now().isoformat(),
"local_path": str(file_path),
}
meta_path = conv_dir / "metadata.json"
if meta_path.exists():
try:
existing = json.loads(meta_path.read_text())
if "attachments" not in existing:
existing["attachments"] = []
existing["attachments"].append({"filename": safe_name, "original": filename})
meta_path.write_text(json.dumps(existing, ensure_ascii=False, indent=2))
except Exception:
meta_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2))
else:
meta["attachments"] = [{"filename": safe_name, "original": filename}]
meta_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2))
# Queue for Yandex Disk sync (service portal picks this up)
if db is not None:
try:
from app.models.models import FileSyncQueue
remote_path = f"bot/conversations/{conversation_id}/{safe_name}"
queue = FileSyncQueue(
local_path=str(file_path),
remote_path=remote_path,
entity_type="conversation",
entity_id=conversation_id,
status="pending",
)
db.add(queue)
await db.flush()
except Exception as e:
logger.error(f"Failed to queue file sync: {e}")
return str(file_path)
def get_file_type(self, filename: str) -> str:
ext = Path(filename).suffix.lower()
if ext in {".jpg", ".jpeg", ".png", ".gif"}:
return "image"
if ext in {".pdf", ".docx", ".xlsx", ".txt"}:
return "document"
return "other"
storage = FileStorage()