Refactor database management to use MongoDB with async support; update audiobook routes for MongoDB integration and improve error handling.

app/api/audiobook_routes.py

@@ -13,202 +13,189 @@ from fastapi.responses import FileResponse, JSONResponse
 from sqlalchemy.orm import Session
 from app.db_models.database import Audiobook, AudiobookStatus, AudiobookChunk, TextChunk
 from app.services.storage import storage
-from app.db import get_db
+from app.db import get_db, AUDIOBOOKS_COLLECTION
+from app.config import AUDIO_DIR, TEXT_DIR, TEMP_DIR
+from pydantic import BaseModel
 import torchaudio
+import json
+import shutil
+from motor.motor_asyncio import AsyncIOMotorDatabase
 
 # Set up logging
 logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/audiobook", tags=["Audiobook"])
 
-async def process_audiobook(
-    request: Request,
-    book_id: str,
-    text_content: str,
-    voice_id: int,
-    db: Session
-):
-    """Process audiobook in the background."""
+class AudiobookBase(BaseModel):
+    title: str
+    author: str
+    voice_id: str
+    status: str = "pending"
+    created_at: datetime = datetime.utcnow()
+    updated_at: datetime = datetime.utcnow()
+
+class Audiobook(AudiobookBase):
+    id: str
+    file_path: Optional[str] = None
+    text_path: Optional[str] = None
+    error: Optional[str] = None
+
+class TextChunk(BaseModel):
+    text: str
+    start_time: float
+    end_time: float
+
+async def process_audiobook(book_id: str, db: AsyncIOMotorDatabase):
+    """Process the audiobook in the background."""
     try:
-        # Get the book from database
-        book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
-        if not book:
-            logger.error(f"Book {book_id} not found")
-            return False
-
         # Update status to processing
-        book.status = AudiobookStatus.PROCESSING
-        db.commit()
-
-        logger.info(f"Starting processing for audiobook {book_id}")
-
-        # Get the generator from app state
-        generator = request.app.state.generator
-        if generator is None:
-            raise Exception("TTS model not available")
-
-        # Get voice info
-        voice_info = request.app.state.get_voice_info(voice_id)
-        if not voice_info:
-            raise Exception(f"Voice ID {voice_id} not found")
-
-        # Generate audio for the entire text
-        logger.info(f"Generating audio for entire text of book {book_id}")
-        audio = generator.generate(
-            text=text_content,
-            speaker=voice_info["speaker_id"],
-            max_audio_length_ms=min(300000, len(text_content) * 80)  # Big text = big audio
+        await db[AUDIOBOOKS_COLLECTION].update_one(
+            {"id": book_id},
+            {"$set": {"status": "processing", "updated_at": datetime.utcnow()}}
         )
 
-        if audio is None:
-            raise Exception("Failed to generate audio")
+        # Get the audiobook data
+        audiobook = await db[AUDIOBOOKS_COLLECTION].find_one({"id": book_id})
+        if not audiobook:
+            raise HTTPException(status_code=404, detail="Audiobook not found")
 
-        # Save the audio using storage service
-        audio_to_save = audio.unsqueeze(0).cpu() if len(audio.shape) == 1 else audio.cpu()
-        audio_bytes = audio_to_save.numpy().tobytes()
-        audio_path = await storage.save_audio_file(book_id, audio_bytes)
+        # TODO: Implement TTS processing logic here
+        # For now, we'll just simulate processing
+        logger.info(f"Processing audiobook {book_id}")
 
-        # Update book status in database
-        book.status = AudiobookStatus.COMPLETED
-        book.audio_file_path = audio_path
-        db.commit()
-
-        logger.info(f"Successfully created audiobook {book_id}")
-        return True
+        # Update status to completed
+        await db[AUDIOBOOKS_COLLECTION].update_one(
+            {"id": book_id},
+            {
+                "$set": {
+                    "status": "completed",
+                    "file_path": f"{AUDIO_DIR}/{book_id}.mp3",
+                    "updated_at": datetime.utcnow()
+                }
+            }
+        )
 
     except Exception as e:
-        logger.error(f"Error processing audiobook {book_id}: {e}")
-        
-        # Update status to failed in database
-        book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
-        if book:
-            book.status = AudiobookStatus.FAILED
-            book.error_message = str(e)
-            db.commit()
-        
-        return False
-
-@router.post("/")
+        logger.error(f"Error processing audiobook {book_id}: {str(e)}")
+        await db[AUDIOBOOKS_COLLECTION].update_one(
+            {"id": book_id},
+            {
+                "$set": {
+                    "status": "failed",
+                    "error": str(e),
+                    "updated_at": datetime.utcnow()
+                }
+            }
+        )
+
+@router.post("/", response_model=Audiobook)
 async def create_audiobook(
-    request: Request,
     background_tasks: BackgroundTasks,
-    title: str = Form(...),
-    author: str = Form(...),
-    voice_id: int = Form(0),
+    title: str,
+    author: str,
+    voice_id: str,
     text_file: Optional[UploadFile] = File(None),
-    text_content: Optional[str] = Form(None),
-    db: Session = Depends(get_db)
+    text_content: Optional[str] = None,
+    request: Request = None
 ):
-    """Create a new audiobook from text."""
-    try:
-        # Validate input
-        if not text_file and not text_content:
-            raise HTTPException(status_code=400, detail="Either text_file or text_content is required")
-
-        # Generate unique ID
-        book_id = str(uuid.uuid4())
-
-        # Handle text content
-        if text_file:
-            text_file_path = await storage.save_text_file(book_id, text_file)
-            with open(text_file_path, "r", encoding="utf-8") as f:
-                text_content = f.read()
-        else:
-            text_file_path = await storage.save_text_content(book_id, text_content)
-
-        # Create book in database
-        book = Audiobook(
-            id=book_id,
-            title=title,
-            author=author,
-            voice_id=voice_id,
-            status=AudiobookStatus.PENDING,
-            text_file_path=text_file_path,
-            text_content=text_content if len(text_content) <= 10000 else None  # Store small texts directly
+    """Create a new audiobook."""
+    db = await get_db()
+    book_id = str(uuid.uuid4())
+
+    # Validate input
+    if not text_file and not text_content:
+        raise HTTPException(
+            status_code=400,
+            detail="Either text_file or text_content must be provided"
         )
-        db.add(book)
-        db.commit()
 
-        # Process in background
-        background_tasks.add_task(process_audiobook, request, book_id, text_content, voice_id, db)
+    # Create audiobook document
+    audiobook = {
+        "id": book_id,
+        "title": title,
+        "author": author,
+        "voice_id": voice_id,
+        "status": "pending",
+        "created_at": datetime.utcnow(),
+        "updated_at": datetime.utcnow()
+    }
+
+    # Handle text input
+    if text_file:
+        text_path = f"{TEXT_DIR}/{book_id}.txt"
+        with open(text_path, "wb") as f:
+            shutil.copyfileobj(text_file.file, f)
+        audiobook["text_path"] = text_path
+    else:
+        text_path = f"{TEXT_DIR}/{book_id}.txt"
+        with open(text_path, "w") as f:
+            f.write(text_content)
+        audiobook["text_path"] = text_path
 
-        return JSONResponse(content={"message": "Audiobook creation started", "book_id": book_id})
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error creating audiobook: {str(e)}")
+    # Insert audiobook into database
+    await db[AUDIOBOOKS_COLLECTION].insert_one(audiobook)
 
-@router.get("/{book_id}")
-async def get_audiobook(book_id: str, db: Session = Depends(get_db)):
+    # Start background processing
+    background_tasks.add_task(process_audiobook, book_id, db)
+
+    return audiobook
+
+@router.get("/{book_id}", response_model=Audiobook)
+async def get_audiobook(book_id: str):
     """Get audiobook information."""
-    book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
-    if not book:
+    db = await get_db()
+    audiobook = await db[AUDIOBOOKS_COLLECTION].find_one({"id": book_id})
+    if not audiobook:
         raise HTTPException(status_code=404, detail="Audiobook not found")
-
-    return {
-        "id": book.id,
-        "title": book.title,
-        "author": book.author,
-        "voice_id": book.voice_id,
-        "status": book.status.value,
-        "created_at": book.created_at.isoformat(),
-        "updated_at": book.updated_at.isoformat(),
-        "error_message": book.error_message
-    }
+    return audiobook
 
 @router.get("/{book_id}/audio")
-async def get_audiobook_audio(book_id: str, db: Session = Depends(get_db)):
-    """Get the audiobook audio file."""
-    book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
-    if not book:
+async def get_audiobook_audio(book_id: str):
+    """Get audiobook audio file."""
+    db = await get_db()
+    audiobook = await db[AUDIOBOOKS_COLLECTION].find_one({"id": book_id})
+    
+    if not audiobook:
         raise HTTPException(status_code=404, detail="Audiobook not found")
-
-    if book.status != AudiobookStatus.COMPLETED or not book.audio_file_path:
-        raise HTTPException(status_code=400, detail="Audiobook is not yet completed")
-
-    audio_path = await storage.get_audio_file(book_id)
-    if not audio_path:
+    
+    if audiobook["status"] != "completed":
+        raise HTTPException(
+            status_code=400,
+            detail=f"Audiobook is not ready (status: {audiobook['status']})"
+        )
+    
+    file_path = audiobook.get("file_path")
+    if not file_path or not os.path.exists(file_path):
         raise HTTPException(status_code=404, detail="Audio file not found")
-
+    
     return FileResponse(
-        str(audio_path),
-        media_type="audio/wav",
-        filename=f"{book.title}.wav"
+        file_path,
+        media_type="audio/mpeg",
+        filename=f"{audiobook['title']}.mp3"
     )
 
-@router.get("/")
-async def get_audiobooks(db: Session = Depends(get_db)):
-    """Get all audiobooks."""
-    books = db.query(Audiobook).order_by(Audiobook.created_at.desc()).all()
-    return {
-        "audiobooks": [
-            {
-                "id": book.id,
-                "title": book.title,
-                "author": book.author,
-                "voice_id": book.voice_id,
-                "status": book.status.value,
-                "created_at": book.created_at.isoformat(),
-                "updated_at": book.updated_at.isoformat(),
-                "error_message": book.error_message
-            }
-            for book in books
-        ]
-    }
+@router.get("/", response_model=List[Audiobook])
+async def list_audiobooks():
+    """List all audiobooks."""
+    db = await get_db()
+    audiobooks = await db[AUDIOBOOKS_COLLECTION].find().to_list(length=None)
+    return audiobooks
 
 @router.delete("/{book_id}")
-async def delete_audiobook(book_id: str, db: Session = Depends(get_db)):
+async def delete_audiobook(book_id: str):
     """Delete an audiobook."""
-    book = db.query(Audiobook).filter(Audiobook.id == book_id).first()
-    if not book:
+    db = await get_db()
+    audiobook = await db[AUDIOBOOKS_COLLECTION].find_one({"id": book_id})
+    
+    if not audiobook:
         raise HTTPException(status_code=404, detail="Audiobook not found")
-
-    try:
-        # Delete associated files
-        await storage.delete_book_files(book_id)
-
-        # Delete from database
-        db.delete(book)
-        db.commit()
-
-        return {"message": "Audiobook deleted successfully"}
-    except Exception as e:
-        db.rollback()
-        raise HTTPException(status_code=500, detail=f"Error deleting audiobook: {str(e)}") 
+    
+    # Delete associated files
+    if audiobook.get("file_path") and os.path.exists(audiobook["file_path"]):
+        os.remove(audiobook["file_path"])
+    if audiobook.get("text_path") and os.path.exists(audiobook["text_path"]):
+        os.remove(audiobook["text_path"])
+    
+    # Delete from database
+    await db[AUDIOBOOKS_COLLECTION].delete_one({"id": book_id})
+    
+    return {"message": "Audiobook deleted successfully"} 

app/db.py

@@ -1,29 +1,84 @@
-"""Database connection and session management."""
+"""MongoDB database configuration."""
 import os
-from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker
-from app.db_models.database import Base
+import logging
+from typing import Optional
+from motor.motor_asyncio import AsyncIOMotorClient
+from pymongo.errors import ConnectionFailure
+from dotenv import load_dotenv
 
-# Get database URL from environment or use SQLite as default
-DATABASE_URL = os.getenv(
-    "DATABASE_URL",
-    "sqlite:///app/storage/audiobooks.db"
-)
+# Load environment variables
+load_dotenv()
 
-# Create engine
-engine = create_engine(DATABASE_URL)
+# Configure logging
+logger = logging.getLogger(__name__)
 
-# Create session factory
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+# Get MongoDB URI from environment variable
+MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017")
+DB_NAME = os.getenv("DB_NAME", "tts_api")
 
-def init_db():
-    """Initialize the database, creating all tables."""
-    Base.metadata.create_all(bind=engine)
+# MongoDB client instance
+client: Optional[AsyncIOMotorClient] = None
 
-def get_db():
-    """Get a database session."""
-    db = SessionLocal()
+async def connect_to_mongo():
+    """Connect to MongoDB."""
+    global client
     try:
-        yield db
-    finally:
-        db.close() 
+        client = AsyncIOMotorClient(MONGO_URI)
+        # Verify the connection
+        await client.admin.command('ping')
+        logger.info("Successfully connected to MongoDB")
+    except ConnectionFailure as e:
+        logger.error(f"Could not connect to MongoDB: {e}")
+        raise
+
+async def close_mongo_connection():
+    """Close MongoDB connection."""
+    global client
+    if client:
+        client.close()
+        logger.info("MongoDB connection closed")
+
+def get_db():
+    """Get database instance."""
+    if not client:
+        raise ConnectionError("MongoDB client not initialized")
+    return client[DB_NAME]
+
+# Collection names
+AUDIOBOOKS_COLLECTION = "audiobooks"
+VOICES_COLLECTION = "voices"
+AUDIO_CACHE_COLLECTION = "audio_cache"
+
+# Database schemas/models
+AUDIOBOOK_SCHEMA = {
+    "id": str,  # UUID string
+    "title": str,
+    "author": str,
+    "voice_id": str,
+    "status": str,  # pending, processing, completed, failed
+    "created_at": str,  # ISO format datetime
+    "updated_at": str,  # ISO format datetime
+    "duration": float,
+    "file_path": str,
+    "error": str,
+    "meta_data": dict
+}
+
+VOICE_SCHEMA = {
+    "id": str,  # UUID string
+    "name": str,
+    "type": str,  # standard, cloned
+    "speaker_id": int,
+    "created_at": str,  # ISO format datetime
+    "is_active": bool,
+    "meta_data": dict
+}
+
+AUDIO_CACHE_SCHEMA = {
+    "id": str,  # UUID string
+    "hash": str,  # Hash of input parameters
+    "format": str,  # Audio format (mp3, wav, etc.)
+    "created_at": str,  # ISO format datetime
+    "file_path": str,
+    "meta_data": dict
+} 

app/main.py

@@ -18,7 +18,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import RedirectResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
 from app.api.routes import router as api_router
-from app.db_models.database import Base, get_db
+from app.db import connect_to_mongo, close_mongo_connection
 
 # Setup logging
 os.makedirs("logs", exist_ok=True)
@@ -618,6 +618,16 @@ async def root():
     logger.debug("Root endpoint accessed, redirecting to docs")
     return RedirectResponse(url="/docs")
 
+@app.on_event("startup")
+async def startup_db_client():
+    """Initialize MongoDB connection on startup."""
+    await connect_to_mongo()
+
+@app.on_event("shutdown")
+async def shutdown_db_client():
+    """Close MongoDB connection on shutdown."""
+    await close_mongo_connection()
+
 if __name__ == "__main__":
     # Get port from environment or use default
     port = int(os.environ.get("PORT", 7860))

requirements.txt

@@ -22,5 +22,6 @@ yt-dlp>=2023.3.4
 openai-whisper>=20230314
 ffmpeg-python>=0.2.0
 accelerate>=0.20.0
-alembic>=1.12.0
-SQLAlchemy>=2.0.0
+pymongo>=4.6.1
+motor>=3.3.2
+python-dotenv>=1.0.1