sandbox-notebookllamalm/src/notebookllama/notebook_manager.py

"""Notebook management helper functions for multi-document notebooks"""

from typing import List, Optional, Dict, Any
from database import get_db, Notebook, NotebookDocument, Document, DocumentSummary, User, ChatSession
from sqlalchemy.orm import Session
from datetime import datetime


def create_notebook(user_id: int, name: str, description: Optional[str] = None) -> Optional[Notebook]:
    """Create a new notebook (collection of documents)"""
    db = get_db()
    try:
        notebook = Notebook(
            user_id=user_id,
            name=name,
            description=description
        )
        db.add(notebook)
        db.commit()
        db.refresh(notebook)
        return notebook
    except Exception as e:
        db.rollback()
        print(f"Error creating notebook: {e}")
        return None
    finally:
        db.close()


def get_user_notebooks(user_id: int) -> List[Notebook]:
    """Get all notebooks for a user"""
    db = get_db()
    try:
        notebooks = db.query(Notebook).filter(Notebook.user_id == user_id).order_by(Notebook.updated_at.desc()).all()
        return notebooks
    finally:
        db.close()


def get_notebook_by_id(notebook_id: int) -> Optional[Notebook]:
    """Get a notebook by ID"""
    db = get_db()
    try:
        notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()
        return notebook
    finally:
        db.close()


def update_notebook(notebook_id: int, name: Optional[str] = None, description: Optional[str] = None) -> Optional[Notebook]:
    """Update notebook details"""
    db = get_db()
    try:
        notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()
        if notebook:
            if name:
                notebook.name = name
            if description is not None:
                notebook.description = description
            notebook.updated_at = datetime.utcnow()
            db.commit()
            db.refresh(notebook)
        return notebook
    except Exception as e:
        db.rollback()
        print(f"Error updating notebook: {e}")
        return None
    finally:
        db.close()


def delete_notebook(notebook_id: int) -> bool:
    """Delete a notebook"""
    db = get_db()
    try:
        notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()
        if notebook:
            db.delete(notebook)
            db.commit()
            return True
        return False
    except Exception as e:
        db.rollback()
        print(f"Error deleting notebook: {e}")
        return False
    finally:
        db.close()


def add_document_to_notebook(notebook_id: int, document_id: int, order: int = 0) -> Optional[NotebookDocument]:
    """Add a document to a notebook"""
    db = get_db()
    try:
        # Check if already exists
        existing = db.query(NotebookDocument).filter(
            NotebookDocument.notebook_id == notebook_id,
            NotebookDocument.document_id == document_id
        ).first()

        if existing:
            return existing

        link = NotebookDocument(
            notebook_id=notebook_id,
            document_id=document_id,
            display_order=order
        )
        db.add(link)

        # Update notebook's updated_at
        notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()
        if notebook:
            notebook.updated_at = datetime.utcnow()

        db.commit()
        db.refresh(link)
        return link
    except Exception as e:
        db.rollback()
        print(f"Error adding document to notebook: {e}")
        return None
    finally:
        db.close()


def remove_document_from_notebook(notebook_id: int, document_id: int) -> bool:
    """Remove a document from a notebook"""
    db = get_db()
    try:
        link = db.query(NotebookDocument).filter(
            NotebookDocument.notebook_id == notebook_id,
            NotebookDocument.document_id == document_id
        ).first()

        if link:
            db.delete(link)

            # Update notebook's updated_at
            notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()
            if notebook:
                notebook.updated_at = datetime.utcnow()

            db.commit()
            return True
        return False
    except Exception as e:
        db.rollback()
        print(f"Error removing document from notebook: {e}")
        return False
    finally:
        db.close()


def get_notebook_documents(notebook_id: int) -> List[Document]:
    """Get all documents in a notebook"""
    db = get_db()
    try:
        links = db.query(NotebookDocument).filter(
            NotebookDocument.notebook_id == notebook_id
        ).order_by(NotebookDocument.display_order).all()

        documents = [link.document for link in links]
        return documents
    finally:
        db.close()


def get_document_notebooks(document_id: int) -> List[Notebook]:
    """Get all notebooks containing a document"""
    db = get_db()
    try:
        links = db.query(NotebookDocument).filter(
            NotebookDocument.document_id == document_id
        ).all()

        notebooks = [link.notebook for link in links]
        return notebooks
    finally:
        db.close()


def get_notebook_document_count(notebook_id: int) -> int:
    """Get count of documents in a notebook"""
    db = get_db()
    try:
        count = db.query(NotebookDocument).filter(
            NotebookDocument.notebook_id == notebook_id
        ).count()
        return count
    finally:
        db.close()


def save_notebook_podcast(notebook_id: int, podcast_path: str) -> Optional[Notebook]:
    """Save podcast file path to notebook"""
    db = get_db()
    try:
        notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()
        if notebook:
            notebook.podcast_path = podcast_path
            notebook.podcast_generated_at = datetime.utcnow()
            db.commit()
            db.refresh(notebook)
        return notebook
    except Exception as e:
        db.rollback()
        print(f"Error saving podcast: {e}")
        return None
    finally:
        db.close()


def create_notebook_chat_session(user_id: int, notebook_id: int, title: Optional[str] = None) -> Optional[ChatSession]:
    """Create a chat session for a notebook"""
    db = get_db()
    try:
        session = ChatSession(
            user_id=user_id,
            notebook_id=notebook_id,
            title=title or f"Chat - {datetime.utcnow().strftime('%Y-%m-%d %H:%M')}"
        )
        db.add(session)
        db.commit()
        db.refresh(session)
        return session
    except Exception as e:
        db.rollback()
        print(f"Error creating chat session: {e}")
        return None
    finally:
        db.close()


def get_notebook_chat_sessions(notebook_id: int) -> List[ChatSession]:
    """Get all chat sessions for a notebook"""
    db = get_db()
    try:
        sessions = db.query(ChatSession).filter(
            ChatSession.notebook_id == notebook_id
        ).order_by(ChatSession.updated_at.desc()).all()
        return sessions
    finally:
        db.close()