Spaces:
Running on Zero
Running on Zero
| """One upload -> course, kind, clean title, duplicate checks, stored. Used by the admin upload handler.""" | |
| import asyncio | |
| import hashlib | |
| import logging | |
| import os | |
| import re | |
| from dataclasses import dataclass, field | |
| from typing import Optional | |
| from aiogram import Bot | |
| from aiogram.types import BufferedInputFile | |
| from sqlalchemy.ext.asyncio import AsyncSession | |
| from app.config import get_settings | |
| from app.db.models import Course, Material | |
| from app.db.repos import library as repo | |
| from app.services import classifier as C | |
| from app.services import library as lib | |
| from app.services import watermark as wm | |
| from app.services.doc_analyzer import DocInfo, analyze | |
| log = logging.getLogger(__name__) | |
| BOT_API_DOWNLOAD_LIMIT = 20 * 1024 * 1024 # a bot can only download files up to 20 MB | |
| class Outcome: | |
| status: str # saved | review | near_duplicate | duplicate | |
| material: Optional[Material] = None | |
| course: Optional[Course] = None | |
| kind: str = "note" | |
| title: str = "" | |
| why: str = "" | |
| notes: list[str] = field(default_factory=list) | |
| suggestions: list = field(default_factory=list) | |
| duplicate_of: Optional[Material] = None | |
| file_name: Optional[str] = None | |
| async def _download(bot: Bot, file_id: str) -> Optional[bytes]: | |
| try: | |
| buf = await asyncio.wait_for(bot.download(file_id), timeout=45) | |
| return buf.read() if buf else None | |
| except Exception as exc: | |
| log.info("Download for analysis failed: %r", exc) | |
| return None | |
| async def ingest_file( | |
| bot: Bot, | |
| session: AsyncSession, | |
| *, | |
| uploader_id: int, | |
| file_type: str, | |
| file_id: str, | |
| unique_id: str, | |
| file_name: Optional[str], | |
| size: Optional[int], | |
| caption: Optional[str], | |
| forced_course: Optional[Course] = None, | |
| forced_kind: Optional[str] = None, | |
| forced_title: Optional[str] = None, | |
| backup_channel: int = 0, | |
| message=None, | |
| stamp: bool = True, | |
| ) -> Outcome: | |
| # 1) exact same Telegram file | |
| existing = await repo.find_by_unique_id(session, unique_id) | |
| if existing: | |
| return Outcome("duplicate", duplicate_of=existing, why="same Telegram file", file_name=file_name) | |
| # 2) read the content (only possible up to 20 MB) -> hash, text, pages | |
| content_hash, info, data = None, DocInfo(), None | |
| if file_type == "document" and size is not None and size <= BOT_API_DOWNLOAD_LIMIT: | |
| data = await _download(bot, file_id) | |
| if data: | |
| content_hash = hashlib.sha256(data).hexdigest() | |
| twin = await repo.find_by_hash(session, content_hash) | |
| if twin: | |
| return Outcome("duplicate", duplicate_of=twin, why="identical content, different name", file_name=file_name) | |
| info = await asyncio.to_thread(analyze, file_name, data) | |
| ext = os.path.splitext(file_name or "")[1].lower() | |
| notes: list[str] = [] | |
| # 3) which course? | |
| suggestions: list = [] | |
| confident = True | |
| course = forced_course | |
| if course is None: | |
| courses = await lib.all_courses_cached(session) | |
| clean_caption = re.sub(r"(https?://\S+|t\.me/\S+|@\w+)", " ", caption or "") | |
| sources = [(file_name or "", "name"), (clean_caption, "text"), (info.text[:3000], "text")] | |
| best, confident, suggestions = C.decide(C.rank_courses(courses, sources)) | |
| course = best if confident else None | |
| # 4) which kind? | |
| if forced_kind: | |
| kind, why = forced_kind, "set by you" | |
| else: | |
| kind, why = C.detect_kind(file_name or "", caption or "", info.text, info.pages, size, ext) | |
| # 5) a clean title (+ year / Mid / Final for exams) | |
| year = C.detect_year(file_name or "", caption or "", info.text[:1500]) | |
| stage = C.exam_stage(file_name or "", caption or "") if kind == "exam" else None | |
| course_name = course.name if course else "Unsorted" | |
| title = forced_title or C.clean_title(file_name, caption, course_name, kind, year, stage) | |
| # 6) near-duplicate: same text, different bytes (re-scan, re-export, different compression) | |
| sig = C.simhash64(info.text) | |
| near, near_distance = None, 99 | |
| if sig is not None: | |
| for mid, other_sig, other_title in await repo.all_sigs(session): | |
| distance = C.hamming(sig, other_sig) | |
| if distance <= C.NEAR_DUPLICATE_BITS and distance < near_distance: | |
| near, near_distance = await repo.get_material(session, mid), distance | |
| if near is not None and near_distance == 0 and info.pages is not None and near.pages == info.pages: | |
| return Outcome("duplicate", duplicate_of=near, file_name=file_name, | |
| why=f"same text and same page count ({info.pages} p), different file bytes") | |
| # 7) same title, DIFFERENT file in this course/kind -> keep both, tell them apart | |
| if course is not None: | |
| taken = {C.norm_title_key(t) for t in await repo.titles_in_course(session, course.id, kind)} | |
| if C.norm_title_key(title) in taken: | |
| new_title = C.disambiguate(title, taken, year=year, pages=info.pages, size=size) | |
| notes.append(f"same title as an existing different file -> saved as “{new_title}”") | |
| title = new_title | |
| # 7b) stamp PDFs with the official channel link; students receive the STAMPED copy | |
| stored_file_id, orig_file_id, stamped_ok = file_id, None, False | |
| library_chat = library_msg = None | |
| if stamp and get_settings().watermark_enabled and ext == ".pdf" and file_type == "document": | |
| if data is None: | |
| big = size is not None and size > BOT_API_DOWNLOAD_LIMIT | |
| notes.append("not watermarked: over 20 MB, the bot cannot read it" if big else "not watermarked: could not read the file") | |
| else: | |
| stamped, reason = await wm.stamp(data, wm.current_style()) | |
| if stamped is None: | |
| notes.append(f"not watermarked: {reason}") | |
| else: | |
| target = backup_channel or uploader_id # the library channel doubles as the backup | |
| try: | |
| sent = await bot.send_document( | |
| target, BufferedInputFile(stamped, filename=file_name or "document.pdf"), | |
| caption=f"📌 {title}"[:200], disable_notification=True, | |
| ) | |
| stored_file_id, orig_file_id, stamped_ok = sent.document.file_id, file_id, True | |
| if backup_channel: | |
| library_chat, library_msg = backup_channel, sent.message_id | |
| except Exception as exc: | |
| log.warning("Could not upload the stamped copy: %r", exc) | |
| notes.append("could not upload the stamped copy: original kept") | |
| status = "saved" | |
| if near is not None: | |
| status = "near_duplicate" | |
| elif course is None: | |
| status = "review" | |
| material = await repo.add_material( | |
| session, | |
| course_id=course.id if course else None, | |
| kind=kind, | |
| title=title[:250], | |
| file_id=stored_file_id, | |
| orig_file_id=orig_file_id, | |
| watermarked=stamped_ok, | |
| library_chat_id=library_chat, | |
| library_message_id=library_msg, | |
| file_type=file_type, | |
| file_unique_id=unique_id, | |
| uploader_id=uploader_id, | |
| content_hash=content_hash, | |
| text_sig=sig, | |
| pages=info.pages, | |
| size=size, | |
| source_name=(file_name or "")[:250] or None, | |
| status="ok" if status == "saved" else "review", | |
| ) | |
| if backup_channel and message is not None and not stamped_ok: | |
| try: | |
| copied = await message.copy_to(backup_channel) | |
| material.library_chat_id, material.library_message_id = backup_channel, copied.message_id | |
| except Exception as exc: # the backup copy is optional | |
| log.warning("Library backup copy failed: %s", exc) | |
| return Outcome( | |
| status, material=material, course=course, kind=kind, title=material.title, why=why, | |
| notes=notes, suggestions=suggestions, duplicate_of=near, file_name=file_name, | |
| ) | |
| async def stamp_existing(bot: Bot, material: Material, target_chat: int) -> tuple[bool, str]: | |
| """Stamp a PDF that is already stored (a student's file at approval time). (ok, reason when not ok)""" | |
| data = await _download(bot, material.orig_file_id or material.file_id) | |
| if not data: | |
| return False, "could not download" | |
| stamped, reason = await wm.stamp(data, wm.current_style()) | |
| if stamped is None: | |
| return False, reason | |
| try: | |
| sent = await bot.send_document( | |
| target_chat, BufferedInputFile(stamped, filename=material.source_name or "document.pdf"), | |
| caption=f"📌 {material.title}"[:200], disable_notification=True, | |
| ) | |
| except Exception as exc: | |
| return False, f"upload failed: {type(exc).__name__}" | |
| material.orig_file_id = material.orig_file_id or material.file_id | |
| material.file_id, material.watermarked = sent.document.file_id, True | |
| return True, "" | |