supernote 0.13.6__py3-none-any.whl

supernote/server/services/processor_modules/summary.py

@@ -0,0 +1,302 @@
+import json
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Optional
+
+from mashumaro.jsonschema import build_json_schema
+from mashumaro.mixins.json import DataClassJSONMixin
+from sqlalchemy import select
+
+from supernote.models.summary import (
+    METADATA_SEGMENTS,
+    AddSummaryDTO,
+    UpdateSummaryDTO,
+)
+from supernote.server.config import ServerConfig
+from supernote.server.db.models.file import UserFileDO
+from supernote.server.db.models.note_processing import NotePageContentDO
+from supernote.server.db.models.user import UserDO
+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.services.summary import SummaryService
+from supernote.server.utils.note_content import format_page_metadata
+from supernote.server.utils.paths import get_summary_id, get_transcript_id
+
+logger = logging.getLogger(__name__)
+
+
+# Define structured output schema
+@dataclass
+class SummarySegment(DataClassJSONMixin):
+    date_range: str = field(
+        metadata={
+            "description": "The date range covered by this segment (e.g., '2023-10-27', 'Week of Oct 27')."
+        }
+    )
+    summary: str = field(
+        metadata={
+            "description": "A concise summary of the events, tasks, and notes for this period."
+        }
+    )
+    extracted_dates: List[str] = field(
+        metadata={
+            "description": "List of specific dates derived from the content in ISO 8601 format (YYYY-MM-DD)."
+        }
+    )
+    page_refs: List[int] = field(
+        metadata={
+            "description": "List of 1-indexed page numbers typically found in the text as '--- Page X ---'."
+        }
+    )
+
+
+@dataclass
+class SummaryResponse(DataClassJSONMixin):
+    segments: List[SummarySegment] = field(
+        metadata={
+            "description": "List of summary segments extracted from the transcript."
+        }
+    )
+
+
+class SummaryModule(ProcessorModule):
+    """Module responsible for aggregating OCR text and generating summaries for a note."""
+
+    def __init__(
+        self,
+        file_service: FileService,
+        config: ServerConfig,
+        gemini_service: GeminiService,
+        summary_service: SummaryService,
+    ) -> None:
+        self.file_service = file_service
+        self.config = config
+        self.gemini_service = gemini_service
+        self.summary_service = summary_service
+
+    @property
+    def name(self) -> str:
+        return "SummaryModule"
+
+    @property
+    def task_type(self) -> str:
+        return "SUMMARY_GENERATION"
+
+    async def run_if_needed(
+        self,
+        file_id: int,
+        session_manager: DatabaseSessionManager,
+        page_index: Optional[int] = None,
+        page_id: Optional[str] = None,
+    ) -> bool:
+        """Determines if summary generation is needed."""
+        # Summary is a file-level task (global), not page-level.
+        if page_index is not 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):
+            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:
+        """
+        Generates an AI summary for the given file.
+
+        1. Aggregates all OCR text for the file.
+        2. Sends to Gemini for summarization and date extraction.
+        3. Stores the result as a new Summary.
+        """
+        logger.info(f"Starting summary generation for file_id={file_id}")
+
+        async with session_manager.session() as session:
+            # 1. Fetch File Info
+            file_do = await session.get(UserFileDO, file_id)
+            if not file_do:
+                logger.error(f"File {file_id} not found.")
+                return
+
+            user: Optional[UserDO] = await session.get(UserDO, file_do.user_id)
+            if not user or not user.email:
+                logger.error(f"User for file {file_id} not found.")
+                return
+            user_email = user.email
+
+            # 2. Key Generation
+            file_basis = file_do.storage_key or str(file_do.id)
+            summary_uuid = get_summary_id(file_basis)
+            transcript_uuid = get_transcript_id(file_basis)
+
+            # 3. Aggregate OCR Text
+            # We want to format the text with page markers so the model can cite them.
+            stmt = (
+                select(NotePageContentDO)
+                .where(NotePageContentDO.file_id == file_id)
+                .order_by(NotePageContentDO.page_index)
+            )
+            result = await session.execute(stmt)
+            pages = result.scalars().all()
+
+        if not pages:
+            logger.info(f"No OCR pages found for file {file_id}. Skipping summary.")
+            return
+
+        # Format transcript with Page markers and inferred dates
+        text_parts = []
+        for p in pages:
+            if p.text_content:
+                metadata = format_page_metadata(
+                    page_index=p.page_index,
+                    page_id=p.page_id or "",
+                    file_name=file_do.file_name,
+                    notebook_create_time=file_do.create_time,
+                    include_section_divider=True,
+                )
+                text_parts.append(f"{metadata}\n{p.text_content}")
+
+        full_text = "\n\n".join(text_parts)
+
+        # 4. Generate Transcript Summary (Preserve existing functionality)
+        # Store the raw aggregated text as a 'transcript' summary type first
+        # This is a good baseline to have.
+        await self._upsert_summary(
+            user_email,
+            AddSummaryDTO(
+                file_id=file_id,
+                unique_identifier=transcript_uuid,
+                content=full_text,
+                data_source="OCR",
+                source_path=file_do.file_name,
+            ),
+        )
+
+        # 5. Generate AI Summary using Gemini
+        # Determine prompt based on filename/type
+        custom_type = Path(file_do.file_name).stem.lower()
+
+        # Load Prompt using specialized logic: Common + (Custom OR Default)
+        prompt_template = PROMPT_LOADER.get_prompt(
+            PromptId.SUMMARY_GENERATION, custom_type=custom_type
+        )
+        prompt = f"{prompt_template}\n\nTRANSCRIPT:\n{full_text}"
+
+        try:
+            response = await self.gemini_service.generate_content(
+                model=self.config.gemini_ocr_model,
+                contents=prompt,
+                config={
+                    "response_mime_type": "application/json",
+                    "response_json_schema": build_json_schema(
+                        SummaryResponse
+                    ).to_dict(),
+                },
+            )
+        except Exception as e:
+            logger.error(f"Failed to generate AI summary for file {file_id}: {e}")
+            return
+
+        # Parse JSON response
+        ai_summary = "No summary generated."
+        metadata_str = None
+
+        if response.text:
+            try:
+                data = json.loads(response.text)
+                segments_data = data.get("segments", [])
+
+                # Format segments into Markdown and collect metadata
+                summary_parts = []
+                all_extracted_dates = []
+
+                # Check for page citations and build clean segment list for metadata
+                clean_segments = []
+
+                for seg in segments_data:
+                    date_range = seg.get("date_range", "Unknown Date")
+                    text = seg.get("summary", "")
+                    dates = seg.get("extracted_dates", [])
+                    page_refs = seg.get("page_refs", [])
+
+                    # Markdown Formatting
+                    header = f"## {date_range}"
+                    # Add page citations to markdown too for readability
+                    if page_refs:
+                        # e.g. (Pages 45, 46)
+                        pages_str = ", ".join(str(p) for p in page_refs)
+                        header += f" (Pages {pages_str})"
+
+                    summary_parts.append(f"{header}\n{text}")
+                    all_extracted_dates.extend(dates)
+
+                    clean_segments.append(
+                        {
+                            "date_range": date_range,
+                            "summary": text,  # Optional: could truncate or omit to save space
+                            "extracted_dates": dates,
+                            "page_refs": page_refs,
+                        }
+                    )
+
+                if summary_parts:
+                    ai_summary = "\n\n".join(summary_parts)
+
+                # Construct rich metadata for retrieval
+                if clean_segments:
+                    meta_obj = {METADATA_SEGMENTS: clean_segments}
+                    metadata_str = json.dumps(meta_obj)
+                    logger.info(
+                        f"Generated metadata for file {file_id}: segments={len(clean_segments)}"
+                    )
+
+            except json.JSONDecodeError:
+                logger.error(f"Failed to parse JSON response for file {file_id}")
+                ai_summary = response.text
+
+        await self._upsert_summary(
+            user_email,
+            AddSummaryDTO(
+                file_id=file_id,
+                unique_identifier=summary_uuid,
+                content=ai_summary,
+                data_source="GEMINI",
+                source_path=file_do.file_name,
+                metadata=metadata_str,
+            ),
+        )
+
+    async def _upsert_summary(self, user_email: str, dto: AddSummaryDTO) -> None:
+        """Helper to either add or update a summary based on its unique identifier."""
+        if not dto.unique_identifier:
+            logger.error("Cannot upsert summary without a unique identifier")
+            return
+
+        existing = await self.summary_service.get_summary_by_uuid(
+            user_email, dto.unique_identifier
+        )
+        if existing and existing.id is not None:
+            await self.summary_service.update_summary(
+                user_email,
+                UpdateSummaryDTO(
+                    id=existing.id,
+                    content=dto.content,
+                    data_source=dto.data_source,
+                    source_path=dto.source_path,
+                    metadata=dto.metadata,
+                ),
+            )
+        else:
+            await self.summary_service.add_summary(user_email, dto)

supernote/server/services/prompt_loader.py

@@ -0,0 +1,133 @@
+import enum
+import importlib.resources
+import logging
+from importlib.resources.abc import Traversable
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+RESOURCES_DIR = importlib.resources.files("supernote.server") / "resources" / "prompts"
+
+
+class PromptId(str, enum.Enum):
+    OCR_TRANSCRIPTION = "ocr_transcription"
+    SUMMARY_GENERATION = "summary_generation"
+
+
+CATEGORY_MAP = {
+    "ocr": PromptId.OCR_TRANSCRIPTION.value,
+    "summary": PromptId.SUMMARY_GENERATION.value,
+}
+COMMON = "common"
+DEFAULT = "default"
+
+
+class PromptLoader:
+    """Service to load and manage externalized prompts."""
+
+    def __init__(
+        self, resources_dir: Optional[Union[Path, Traversable]] = None
+    ) -> None:
+        self.resources_dir = resources_dir or RESOURCES_DIR
+        # Map: prompt_id -> (type -> prompt_text)
+        # type can be "common", "default", or specific custom types like "monthly"
+        self.prompts: Dict[str, Dict[str, str]] = {}
+        self._load_prompts()
+
+    def _load_prompts(self) -> None:
+        """Load all prompts from the resources directory into memory."""
+        if not self.resources_dir.is_dir():
+            logger.warning(f"Prompts directory not found at {self.resources_dir}")
+            return
+
+        # Initialize dicts for known prompt IDs
+        for pid in PromptId:
+            self.prompts[pid.value] = {}
+
+        try:
+            # Map categories to PromptId
+            for category, prompt_id in CATEGORY_MAP.items():
+                category_dir = self.resources_dir / category
+                if not category_dir.is_dir():
+                    continue
+
+                if prompt_id not in self.prompts:
+                    self.prompts[prompt_id] = {}
+
+                # 1. Load Common Prompts (Always On)
+                common_dir = category_dir / COMMON
+                if common_dir.is_dir():
+                    common_text = self._read_prompts_from_dir(common_dir)
+                    if common_text:
+                        self.prompts[prompt_id][COMMON] = common_text
+
+                # 2. Load Default Prompts
+                default_dir = category_dir / DEFAULT
+                if default_dir.is_dir():
+                    default_text = self._read_prompts_from_dir(default_dir)
+                    if default_text:
+                        self.prompts[prompt_id][DEFAULT] = default_text
+
+                # 3. Load Custom Prompts (sub-directories)
+                for item in category_dir.iterdir():
+                    if item.is_dir() and item.name not in [COMMON, DEFAULT]:
+                        custom_type = item.name.lower()
+                        custom_text = self._read_prompts_from_dir(item)
+                        if custom_text:
+                            self.prompts[prompt_id][custom_type] = custom_text
+
+            logger.info(f"Loaded prompts from {self.resources_dir}")
+
+        except Exception as e:
+            logger.error(f"Failed to load prompts from {self.resources_dir}: {e}")
+
+    def _read_prompts_from_dir(self, directory: Union[Path, Traversable]) -> str:
+        """Read and concatenate all .md files in a directory."""
+        prompts = []
+        # Use iterdir instead of glob for Traversable compatibility
+        files = sorted(
+            [f for f in directory.iterdir() if f.name.endswith(".md")],
+            key=lambda x: x.name,
+        )
+        for file_path in files:
+            if file_path.is_file():
+                prompts.append(file_path.read_text(encoding="utf-8").strip())
+        return "\n\n".join(prompts)
+
+    def get_prompt(self, prompt_id: str, custom_type: Optional[str] = None) -> str:
+        """Retrieve a prompt by its ID, optionally overridden by a custom type.
+
+        Logic: Common + (Custom if exists else Default)
+        """
+        if prompt_id not in self.prompts:
+            raise ValueError(f"Prompt ID '{prompt_id}' not found.")
+
+        type_map = self.prompts[prompt_id]
+
+        parts = []
+
+        # 1. Add Common
+        if COMMON in type_map:
+            parts.append(type_map[COMMON])
+
+        # 2. Add Specific (Custom or Default)
+        specific_prompt = None
+        if custom_type and custom_type in type_map:
+            logger.info(f"Using custom prompt '{custom_type}' for {prompt_id}")
+            specific_prompt = type_map[custom_type]
+        elif DEFAULT in type_map:
+            specific_prompt = type_map[DEFAULT]
+
+        if specific_prompt:
+            parts.append(specific_prompt)
+        elif not parts:
+            # If no common and no specific/default, that's an issue
+            raise ValueError(
+                f"No prompt content found for '{prompt_id}' (Custom: {custom_type})"
+            )
+
+        return "\n\n".join(parts)
+
+
+PROMPT_LOADER = PromptLoader()

supernote/server/services/schedule.py

@@ -0,0 +1,161 @@
+import logging
+from typing import Any, Optional
+
+from sqlalchemy import delete, select, update
+
+from supernote.server.db.models.schedule import ScheduleTaskDO, ScheduleTaskGroupDO
+from supernote.server.db.session import DatabaseSessionManager
+
+logger = logging.getLogger(__name__)
+
+MAX_TITLE_LENGTH = 255
+MAX_DETAIL_LENGTH = 1 * 1024 * 1024  # 1MB
+
+
+class ScheduleService:
+    """Schedule service."""
+
+    def __init__(self, session_manager: DatabaseSessionManager):
+        """Initialize the schedule service."""
+        self.session_manager = session_manager
+
+    # Task Group Operations
+
+    async def create_group(self, user_id: int, title: str) -> ScheduleTaskGroupDO:
+        """Create a new task group."""
+        if len(title) > MAX_TITLE_LENGTH:
+            raise ValueError("Title is too long")
+        async with self.session_manager.session() as session:
+            group = ScheduleTaskGroupDO(user_id=user_id, title=title)
+            session.add(group)
+            # Flush to generate ID
+            await session.flush()
+            # Commit to persist
+            await session.commit()
+            # Refresh to get defaults if any
+            await session.refresh(group)
+            return group
+
+    async def list_groups(self, user_id: int) -> list[ScheduleTaskGroupDO]:
+        """List all task groups for a user."""
+        async with self.session_manager.session() as session:
+            stmt = (
+                select(ScheduleTaskGroupDO)
+                .where(ScheduleTaskGroupDO.user_id == user_id)
+                .order_by(ScheduleTaskGroupDO.create_time.desc())
+            )
+            result = await session.execute(stmt)
+            return list(result.scalars().all())
+
+    async def delete_group(self, user_id: int, group_id: int) -> bool:
+        """Delete a task group. Returns True if found and deleted."""
+        # TODO: Handle cascade delete of tasks?
+        # For now, simplistic delete.
+        async with self.session_manager.session() as session:
+            stmt = delete(ScheduleTaskGroupDO).where(
+                ScheduleTaskGroupDO.user_id == user_id,
+                ScheduleTaskGroupDO.task_list_id == group_id,
+            )
+            result = await session.execute(stmt)
+            await session.commit()
+            return bool(result.rowcount > 0)  # type: ignore[attr-defined]
+
+    # Task Operations
+
+    async def create_task(
+        self,
+        user_id: int,
+        group_id: int,
+        title: str,
+        detail: str = "",
+        status: str = "needsAction",
+        importance: str | None = None,
+        due_time: int | None = None,
+        recurrence: str | None = None,
+        is_reminder_on: bool = False,
+    ) -> ScheduleTaskDO:
+        """Create a new task."""
+        if len(title) > MAX_TITLE_LENGTH:
+            raise ValueError("Title is too long")
+        if len(detail) > MAX_DETAIL_LENGTH:
+            raise ValueError("Detail is too long")
+        async with self.session_manager.session() as session:
+            task = ScheduleTaskDO(
+                user_id=user_id,
+                task_list_id=group_id,
+                title=title,
+                detail=detail,
+                status=status,
+                importance=importance,
+                due_time=due_time,
+                recurrence=recurrence,
+                is_reminder_on=is_reminder_on,
+            )
+            session.add(task)
+            await session.flush()
+            await session.commit()
+            await session.refresh(task)
+            return task
+
+    async def list_tasks(
+        self,
+        user_id: int,
+        group_id: int | None = None,
+    ) -> list[ScheduleTaskDO]:
+        """List tasks for a user, optionally filtered by group."""
+        async with self.session_manager.session() as session:
+            query = select(ScheduleTaskDO).where(ScheduleTaskDO.user_id == user_id)
+            if group_id is not None:
+                query = query.where(ScheduleTaskDO.task_list_id == group_id)
+
+            query = query.order_by(ScheduleTaskDO.create_time.desc())
+            result = await session.execute(query)
+            return list(result.scalars().all())
+
+    async def update_task(
+        self, user_id: int, task_id: int, **kwargs: Any
+    ) -> Optional[ScheduleTaskDO]:
+        """Update a task."""
+        # Clean kwargs to only allow update of specific fields?
+        # For simplicity, we assume caller passes valid fields that match DO columns.
+        allowed_fields = {
+            "title",
+            "detail",
+            "status",
+            "importance",
+            "due_time",
+            "completed_time",
+            "recurrence",
+            "is_reminder_on",
+            "task_list_id",
+        }
+        updates = {k: v for k, v in kwargs.items() if k in allowed_fields}
+
+        async with self.session_manager.session() as session:
+            stmt = (
+                update(ScheduleTaskDO)
+                .where(
+                    ScheduleTaskDO.user_id == user_id, ScheduleTaskDO.task_id == task_id
+                )
+                .values(**updates)
+                .execution_options(synchronize_session="fetch")
+            )
+            await session.execute(stmt)
+            await session.commit()
+
+            # Retrieve updated
+            stmt_get = select(ScheduleTaskDO).where(
+                ScheduleTaskDO.user_id == user_id, ScheduleTaskDO.task_id == task_id
+            )
+            result = await session.execute(stmt_get)
+            return result.scalar_one_or_none()
+
+    async def delete_task(self, user_id: int, task_id: int) -> bool:
+        """Delete a task."""
+        async with self.session_manager.session() as session:
+            stmt = delete(ScheduleTaskDO).where(
+                ScheduleTaskDO.user_id == user_id, ScheduleTaskDO.task_id == task_id
+            )
+            result = await session.execute(stmt)
+            await session.commit()
+            return bool(result.rowcount > 0)  # type: ignore[attr-defined]