supernote 0.13.6__py3-none-any.whl

supernote/server/services/processor_modules/gemini_embedding.py

@@ -0,0 +1,112 @@
+import json
+import logging
+from typing import Optional
+
+from supernote.server.config import ServerConfig
+from supernote.server.db.session import DatabaseSessionManager
+from supernote.server.services.file import FileService
+from supernote.server.services.gemini import GeminiService
+from supernote.server.services.processor_modules import ProcessorModule
+from supernote.server.utils.note_content import get_page_content_by_id
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiEmbeddingModule(ProcessorModule):
+    """Module responsible for generating embeddings for note pages using Gemini."""
+
+    def __init__(
+        self,
+        file_service: FileService,
+        config: ServerConfig,
+        gemini_service: GeminiService,
+    ) -> None:
+        self.file_service = file_service
+        self.config = config
+        self.gemini_service = gemini_service
+
+    @property
+    def name(self) -> str:
+        return "GeminiEmbeddingModule"
+
+    @property
+    def task_type(self) -> str:
+        return "EMBEDDING_GENERATION"
+
+    async def run_if_needed(
+        self,
+        file_id: int,
+        session_manager: DatabaseSessionManager,
+        page_index: Optional[int] = None,
+        page_id: Optional[str] = None,
+    ) -> bool:
+        if not page_id:
+            return False
+
+        if not self.gemini_service.is_configured:
+            return False
+
+        if not await super().run_if_needed(
+            file_id, session_manager, page_index, page_id
+        ):
+            return False
+
+        async with session_manager.session() as session:
+            # Check Prerequisites (Text Content must exist)
+            content = await get_page_content_by_id(session, file_id, page_id)
+
+            if not content or not content.text_content:
+                # Dependency not met yet (OCR not done or empty page)
+                return False
+
+        return True
+
+    async def process(
+        self,
+        file_id: int,
+        session_manager: DatabaseSessionManager,
+        page_index: Optional[int] = None,
+        page_id: Optional[str] = None,
+        **kwargs: object,
+    ) -> None:
+        if not page_id:
+            return
+
+        # Get Text Content
+        text_content = ""
+        async with session_manager.session() as session:
+            content = await get_page_content_by_id(session, file_id, page_id)
+
+            if not content or not content.text_content:
+                logger.warning(
+                    f"No text content found for embedding: file {file_id} page {page_id} (idx {page_index})"
+                )
+                return
+            text_content = content.text_content
+
+        # Call Gemini API
+        if not self.gemini_service.is_configured:
+            raise ValueError("Gemini API key not configured")
+
+        model_id = self.config.gemini_embedding_model
+        response = await self.gemini_service.embed_content(
+            model=model_id,
+            contents=text_content,
+        )
+
+        if not response.embeddings:
+            raise ValueError("No embeddings returned from Gemini API")
+
+        # Assuming single embedding for the whole text block for now
+        embedding_values = response.embeddings[0].values
+        embedding_json = json.dumps(embedding_values)
+
+        # Save Result
+        async with session_manager.session() as session:
+            content = await get_page_content_by_id(session, file_id, page_id)
+
+            if content:
+                content.embedding = embedding_json
+            await session.commit()
+
+        logger.info(f"Completed Embedding for file {file_id} page {page_index}")

supernote/server/services/processor_modules/gemini_ocr.py

@@ -0,0 +1,149 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+from google.genai import types
+
+from supernote.server.config import ServerConfig
+from supernote.server.constants import CACHE_BUCKET
+from supernote.server.db.models.file import UserFileDO
+from supernote.server.db.session import DatabaseSessionManager
+from supernote.server.services.file import FileService
+from supernote.server.services.gemini import GeminiService
+from supernote.server.services.processor_modules import ProcessorModule
+from supernote.server.services.prompt_loader import PROMPT_LOADER, PromptId
+from supernote.server.utils.note_content import (
+    format_page_metadata,
+    get_page_content_by_id,
+)
+from supernote.server.utils.paths import get_page_png_path
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiOcrModule(ProcessorModule):
+    """Module responsible for extracting text from note pages using Gemini OCR."""
+
+    def __init__(
+        self,
+        file_service: FileService,
+        config: ServerConfig,
+        gemini_service: GeminiService,
+    ) -> None:
+        self.file_service = file_service
+        self.config = config
+        self.gemini_service = gemini_service
+
+    @property
+    def name(self) -> str:
+        return "GeminiOcrModule"
+
+    @property
+    def task_type(self) -> str:
+        return "OCR_EXTRACTION"
+
+    async def run_if_needed(
+        self,
+        file_id: int,
+        session_manager: DatabaseSessionManager,
+        page_index: Optional[int] = None,
+        page_id: Optional[str] = None,
+    ) -> bool:
+        if page_index is None:
+            return False
+
+        if not self.gemini_service.is_configured:
+            return False
+
+        if not await super().run_if_needed(
+            file_id, session_manager, page_index, page_id
+        ):
+            return False
+
+        if not page_id:
+            return False
+
+        # Check if PNG exists (Prerequisite)
+        png_path = get_page_png_path(file_id, page_id)
+        if not await self.file_service.blob_storage.exists(CACHE_BUCKET, png_path):
+            logger.warning(
+                f"PNG prerequisite not met for OCR of {file_id} page {page_id}"
+            )
+            return False
+
+        return True
+
+    async def process(
+        self,
+        file_id: int,
+        session_manager: DatabaseSessionManager,
+        page_index: Optional[int] = None,
+        page_id: Optional[str] = None,
+        **kwargs: object,
+    ) -> None:
+        if page_id is None:
+            logger.error(f"Page ID required for OCR processing of file {file_id}")
+            return
+
+        # Get PNG Content
+        png_path = get_page_png_path(file_id, page_id)
+        png_data = b""
+        async for chunk in self.file_service.blob_storage.get(CACHE_BUCKET, png_path):
+            png_data += chunk
+
+        # Call Gemini API
+        if not self.gemini_service.is_configured:
+            raise ValueError("Gemini API key not configured")
+
+        # Get File Info for custom prompt and metadata
+        file_name_basis: Optional[str] = None
+        file_name: Optional[str] = None
+        notebook_create_time: Optional[int] = None
+        async with session_manager.session() as session:
+            file_do = await session.get(UserFileDO, file_id)
+            if file_do:
+                file_name = file_do.file_name
+                notebook_create_time = file_do.create_time
+                file_name_basis = Path(file_name).stem.lower()
+
+        model_id = self.config.gemini_ocr_model
+        prompt = PROMPT_LOADER.get_prompt(
+            PromptId.OCR_TRANSCRIPTION, custom_type=file_name_basis
+        )
+
+        metadata_block = format_page_metadata(
+            page_index=page_index or 0,
+            page_id=page_id,
+            file_name=file_name,
+            notebook_create_time=notebook_create_time,
+            include_section_divider=True,
+        )
+        prompt = f"{metadata_block}\n\n{prompt}"
+
+        response = await self.gemini_service.generate_content(
+            model=model_id,
+            contents=[
+                types.Content(
+                    parts=[
+                        types.Part.from_text(text=prompt),
+                        types.Part.from_bytes(data=png_data, mime_type="image/png"),
+                    ]
+                )
+            ],
+            config={"media_resolution": "media_resolution_high"},  # type: ignore[arg-type]
+        )
+
+        text_content = response.text if response.text else ""
+
+        # Save Result
+        async with session_manager.session() as session:
+            content = await get_page_content_by_id(session, file_id, page_id)
+            if content:
+                content.text_content = text_content
+            else:
+                logger.warning(
+                    f"NotePageContentDO missing for {file_id} page {page_id} during OCR"
+                )
+            await session.commit()
+
+        logger.info(f"Completed OCR for file {file_id} page {page_id}")

supernote/server/services/processor_modules/page_hashing.py

@@ -0,0 +1,220 @@
+import asyncio
+import logging
+from functools import partial
+from typing import Any, Optional
+
+from sqlalchemy import delete, select
+
+from supernote.notebook.parser import parse_metadata
+from supernote.server.constants import CACHE_BUCKET, USER_DATA_BUCKET
+from supernote.server.db.models.file import UserFileDO
+from supernote.server.db.models.note_processing import NotePageContentDO, SystemTaskDO
+from supernote.server.db.session import DatabaseSessionManager
+from supernote.server.services.file import FileService
+from supernote.server.services.processor_modules import ProcessorModule
+from supernote.server.utils.hashing import get_md5_hash
+from supernote.server.utils.paths import get_page_png_path
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_helper(path: str) -> Any:
+    with open(path, "rb") as f:
+        return parse_metadata(f)  # type: ignore[arg-type]
+
+
+class PageHashingModule(ProcessorModule):
+    """Module responsible for detecting changes in .note files at the page level.
+
+    This module performs the following:
+    - Parses the binary .note file using `SupernoteParser`.
+    - Computes a unique MD5 hash for each page based on its layer metadata.
+    - Updates the `NotePageContentDO` table:
+        - Creates entries for new pages.
+        - Updates hashes for changed pages and invalidates downstream data (OCR, Embeddings) to trigger reprocessing.
+        - Removes entries for deleted pages.
+    - Updates the `SystemTaskDO` status for the 'HASHING' task (handled by base class).
+
+    This module acts as the entry point and "change detector" for the incremental processing pipeline.
+    """
+
+    def __init__(self, file_service: FileService) -> None:
+        self.file_service = file_service
+
+    @property
+    def name(self) -> str:
+        return "PageHashingModule"
+
+    @property
+    def task_type(self) -> str:
+        return "HASHING"
+
+    async def run_if_needed(
+        self,
+        file_id: int,
+        session_manager: DatabaseSessionManager,
+        page_index: Optional[int] = None,
+        page_id: Optional[str] = None,
+    ) -> bool:
+        """Hashing acts as the change detector, so it MUST run every time the file is processed."""
+        return True
+
+    async def process(
+        self,
+        file_id: int,
+        session_manager: DatabaseSessionManager,
+        page_index: Optional[int] = None,
+        page_id: Optional[str] = None,
+        **kwargs: object,
+    ) -> None:
+        """Parses the .note file, computes page hashes, and updates NotePageContentDO."""
+        logger.info(f"Starting PageHashingModule for file_id={file_id}")
+
+        # Resolve file path
+        async with session_manager.session() as session:
+            # Get UserFileDO to find owner & storage key
+            result = await session.execute(
+                select(UserFileDO).where(UserFileDO.id == file_id)
+            )
+            user_file = result.scalars().first()
+            if not user_file:
+                logger.error(f"File {file_id} not found in DB")
+                return
+
+            storage_key = user_file.storage_key
+            if not storage_key:
+                logger.error(f"File {file_id} has no storage_key")
+                return
+
+        # Construct real OS path via BlobStorage (No DB access needed here)
+        try:
+            # This assumes LocalBlobStorage. For S3, we'd need to stream,
+            # but SupernoteParser currently expects a file path.
+            abs_path = self.file_service.blob_storage.get_blob_path(
+                USER_DATA_BUCKET, storage_key
+            )
+        except Exception as e:
+            logger.error(f"Failed to resolve blob path for {file_id}: {e}")
+            return
+
+        if not abs_path.exists():
+            logger.error(f"File {abs_path} does not exist on disk")
+            return
+
+        # Parse .note file
+        try:
+            # Run parser in thread pool
+            loop = asyncio.get_running_loop()
+            metadata = await loop.run_in_executor(
+                None, partial(_parse_helper, str(abs_path))
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to parse .note file {file_id}: {e}")
+            return
+
+        # Iterate pages and update DB
+        total_pages = metadata.get_total_pages()
+        if not metadata.pages:
+            return
+
+        async with session_manager.session() as session:
+            # 1. Fetch all existing pages for this file and map by page_id
+            existing_pages = await session.execute(
+                select(NotePageContentDO).where(NotePageContentDO.file_id == file_id)
+            )
+            existing_map = {p.page_id: p for p in existing_pages.scalars().all()}
+
+            # Track which page_ids are present in the current file
+            current_page_ids = set()
+
+            for i in range(total_pages):
+                page_info = metadata.pages[i]
+                page_id = page_info.get("PAGEID")
+
+                if not page_id:
+                    logger.warning(f"Page {i} of file {file_id} has no PAGEID.")
+                    continue
+
+                current_page_ids.add(page_id)
+
+                # Canonical string representation of page metadata for hashing
+                page_hash_input = str(page_info)
+                current_hash = get_md5_hash(page_hash_input)
+
+                if page_id in existing_map:
+                    # UPDATE path
+                    row = existing_map[page_id]
+
+                    # Check for Move (Index Changed)
+                    if row.page_index != i:
+                        logger.info(
+                            f"Page {page_id} moved from {row.page_index} to {i}"
+                        )
+                        row.page_index = i
+
+                    # Check for Content Change
+                    if row.content_hash != current_hash:
+                        logger.info(
+                            f"Page {page_id} (idx {i}) changed. Resetting content."
+                        )
+                        row.content_hash = current_hash
+                        row.text_content = None  # Clear OCR
+                        row.embedding = None  # Clear Embedding
+
+                        # Invalidate downstream tasks
+                        # Note: Tasks are keyed by page_index usually?
+                        # We should probably migrate tasks to use page_id too,
+                        # but for now let's invalidate based on the NEW index
+                        # or just all tasks for this file/page logic.
+                        # Ideally, SystemTaskDO key should be page_id.
+                        # For now, let's assume keys are "page_{page_index}".
+                        # If page moved, the old task key "page_{old_idx}" is irrelevant,
+                        # and we need to trigger "page_{new_idx}".
+                        # Refactoring SystemTaskDO is out of scope?
+                        # Let's clean up tasks for this "logical" page.
+                        # Actually, keeping it simple: invalidate tasks for the CURRENT index.
+                        # Update: Use page_id for task keys now.
+                        page_task_key = f"page_{page_id}"
+                        await session.execute(
+                            delete(SystemTaskDO)
+                            .where(SystemTaskDO.file_id == file_id)
+                            .where(SystemTaskDO.key == page_task_key)
+                        )
+                else:
+                    # INSERT path
+                    logger.info(f"New page {page_id} at index {i} detected.")
+                    new_content = NotePageContentDO(
+                        file_id=file_id,
+                        page_index=i,
+                        page_id=page_id,
+                        content_hash=current_hash,
+                        text_content=None,
+                        embedding=None,
+                    )
+                    session.add(new_content)
+
+            # 2. Handle Deletions
+            # Any page_id in existing_map but NOT in current_page_ids is deleted
+            for pid, row in existing_map.items():
+                if pid not in current_page_ids:
+                    logger.info(f"Page {pid} deleted (was at {row.page_index}).")
+                    await session.delete(row)
+
+                    # Cleanup Blobs
+                    png_path = get_page_png_path(file_id, pid)
+                    try:
+                        await self.file_service.blob_storage.delete(
+                            CACHE_BUCKET, png_path
+                        )
+                    except Exception as e:
+                        logger.warning(f"Failed to delete orphan PNG {pid}: {e}")
+
+                    # Cleanup Tasks
+                    await session.execute(
+                        delete(SystemTaskDO)
+                        .where(SystemTaskDO.file_id == file_id)
+                        .where(SystemTaskDO.key == f"page_{pid}")
+                    )
+
+            await session.commit()

supernote/server/services/processor_modules/png_conversion.py

@@ -0,0 +1,107 @@
+import asyncio
+import io
+import logging
+from functools import partial
+from typing import Optional
+
+from sqlalchemy import select
+
+from supernote.notebook.converter import ImageConverter
+from supernote.notebook.parser import load_notebook
+from supernote.server.constants import CACHE_BUCKET, USER_DATA_BUCKET
+from supernote.server.db.models.file import UserFileDO
+from supernote.server.db.models.note_processing import NotePageContentDO
+from supernote.server.db.session import DatabaseSessionManager
+from supernote.server.services.file import FileService
+from supernote.server.services.processor_modules import ProcessorModule
+from supernote.server.utils.paths import get_page_png_path
+
+logger = logging.getLogger(__name__)
+
+
+def _convert_helper(path: str, page_index: int) -> bytes:
+    # Use loose policy to attempt parsing even if signature is unknown
+    notebook = load_notebook(path, policy="loose")  # type: ignore[no-untyped-call]
+    converter = ImageConverter(notebook)  # type: ignore[no-untyped-call]
+    img = converter.convert(page_index)  # type: ignore[no-untyped-call]
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
+
+
+class PngConversionModule(ProcessorModule):
+    def __init__(self, file_service: FileService) -> None:
+        self.file_service = file_service
+
+    @property
+    def name(self) -> str:
+        return "PngConversionModule"
+
+    @property
+    def task_type(self) -> str:
+        return "PNG_CONVERSION"
+
+    async def process(
+        self,
+        file_id: int,
+        session_manager: DatabaseSessionManager,
+        page_index: Optional[int] = None,
+        page_id: Optional[str] = None,
+        **kwargs: object,
+    ) -> None:
+        """
+        Converts a specific page of the .note file to PNG and saves it to blob storage.
+        """
+        if page_index is None:
+            logger.warning(
+                f"PngConversionModule requires page_index for file {file_id}"
+            )
+            return
+
+        # Resolve file path
+        async with session_manager.session() as session:
+            result = await session.execute(
+                select(UserFileDO).where(UserFileDO.id == file_id)
+            )
+            user_file = result.scalars().first()
+            if not user_file or not user_file.storage_key:
+                logger.error(f"File {file_id} not found or missing storage_key")
+                return
+            storage_key = user_file.storage_key
+
+            # Get Page ID for stable path
+            page_result = await session.execute(
+                select(NotePageContentDO.page_id).where(
+                    NotePageContentDO.file_id == file_id,
+                    NotePageContentDO.page_index == page_index,
+                )
+            )
+            page_id = page_result.scalars().first()
+            if not page_id:
+                logger.error(f"Page ID not found for {file_id} index {page_index}")
+                return
+
+        try:
+            abs_path = self.file_service.blob_storage.get_blob_path(
+                USER_DATA_BUCKET, storage_key
+            )
+        except Exception as e:
+            logger.error(f"Failed to resolve blob path for {file_id}: {e}")
+            raise
+
+        if not abs_path.exists():
+            logger.error(f"File {abs_path} does not exist on disk")
+            raise FileNotFoundError(f"File {abs_path} does not exist on disk")
+
+        # Run Conversion in Thread Pool
+        loop = asyncio.get_running_loop()
+        png_data = await loop.run_in_executor(
+            None, partial(_convert_helper, str(abs_path), page_index)
+        )
+
+        # Upload to Blob Storage
+        blob_path = get_page_png_path(file_id, page_id)
+        await self.file_service.blob_storage.put(CACHE_BUCKET, blob_path, png_data)
+        logger.info(
+            f"Successfully converted page {page_index} ({page_id}) of {file_id} to PNG"
+        )