apersona/backend/app/main.py

from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.responses import JSONResponse
import time
import logging
from contextlib import asynccontextmanager

from app.core.config import settings
from app.db.database import engine
from app.db.models import Base

# Import routers
from app.api.auth import router as auth_router
# from app.api.files import router as files_router
# from app.api.chat import router as chat_router
# from app.api.reminders import router as reminders_router
# from app.api.search import router as search_router

logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Lifespan context manager for startup and shutdown events"""
    # Startup
    logger.info("Starting aPersona backend...")
    
    # Create database tables
    Base.metadata.create_all(bind=engine)
    logger.info("Database tables created")
    
    # Initialize AI components
    try:
        from ai_core.embeddings.embedding_service import embedding_service
        from ai_core.rag.vector_store import vector_store
        from ai_core.llm.ollama_client import ollama_client
        
        # Test Ollama connection
        is_healthy = await ollama_client.check_health()
        if is_healthy:
            logger.info("Ollama connection established")
        else:
            logger.warning("Ollama service not available - some features may be limited")
        
        # Initialize vector store
        stats = vector_store.get_collection_stats()
        logger.info(f"Vector store initialized: {stats}")
        
        # Test embedding service
        embedding_info = embedding_service.get_model_info()
        logger.info(f"Embedding service ready: {embedding_info}")
        
    except Exception as e:
        logger.error(f"Failed to initialize AI components: {e}")
    
    yield
    
    # Shutdown
    logger.info("Shutting down aPersona backend...")
    try:
        await ollama_client.close()
    except:
        pass


# Create FastAPI app
app = FastAPI(
    title=settings.APP_NAME,
    version=settings.APP_VERSION,
    description="AI-powered personal assistant that works completely offline",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.BACKEND_CORS_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Add trusted host middleware for security
app.add_middleware(
    TrustedHostMiddleware,
    allowed_hosts=["localhost", "127.0.0.1", "*.localhost"]
)


# Request timing middleware
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    """Add processing time to response headers"""
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response


# Global exception handler
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    """Global exception handler for unhandled errors"""
    logger.error(f"Unhandled error for {request.url}: {exc}")
    return JSONResponse(
        status_code=500,
        content={
            "detail": "Internal server error",
            "error": str(exc) if settings.DEBUG else "An unexpected error occurred"
        }
    )


# Health check endpoint
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    try:
        from ai_core.llm.ollama_client import ollama_client
        ollama_healthy = await ollama_client.check_health()
        
        return {
            "status": "healthy",
            "app_name": settings.APP_NAME,
            "version": settings.APP_VERSION,
            "services": {
                "database": "healthy",
                "ollama": "healthy" if ollama_healthy else "unhealthy",
                "embeddings": "healthy",
                "vector_store": "healthy"
            }
        }
    except Exception as e:
        logger.error(f"Health check failed: {e}")
        return JSONResponse(
            status_code=503,
            content={
                "status": "unhealthy",
                "error": str(e)
            }
        )


# Root endpoint
@app.get("/")
async def root():
    """Root endpoint"""
    return {
        "message": f"Welcome to {settings.APP_NAME}",
        "version": settings.APP_VERSION,
        "description": "AI-powered personal assistant - fully local and private",
        "endpoints": {
            "health": "/health",
            "docs": "/docs",
            "api": settings.API_V1_STR
        }
    }


# System info endpoint
@app.get(f"{settings.API_V1_STR}/system/info")
async def get_system_info():
    """Get system information and capabilities"""
    try:
        from ai_core.embeddings.embedding_service import embedding_service
        from ai_core.rag.vector_store import vector_store
        from ai_core.llm.ollama_client import ollama_client
        
        # Get AI service information
        embedding_info = embedding_service.get_model_info()
        vector_stats = vector_store.get_collection_stats()
        available_models = await ollama_client.list_models()
        
        return {
            "app_info": {
                "name": settings.APP_NAME,
                "version": settings.APP_VERSION,
                "debug": settings.DEBUG
            },
            "ai_services": {
                "embedding_model": embedding_info,
                "vector_store": vector_stats,
                "available_llm_models": available_models,
                "current_llm_model": settings.DEFAULT_LLM_MODEL
            },
            "capabilities": {
                "file_processing": [
                    "PDF", "DOCX", "TXT", "Images (OCR)",
                    "Markdown", "PNG", "JPEG", "GIF"
                ],
                "ai_features": [
                    "Semantic search", "Auto-categorization",
                    "Smart reminders", "Personalized responses",
                    "Learning from interactions"
                ]
            }
        }
    except Exception as e:
        logger.error(f"Failed to get system info: {e}")
        raise HTTPException(status_code=500, detail="Failed to retrieve system information")


# Include API routers
app.include_router(auth_router, prefix=f"{settings.API_V1_STR}/auth", tags=["authentication"])
# app.include_router(files_router, prefix=f"{settings.API_V1_STR}/files", tags=["files"])
# app.include_router(chat_router, prefix=f"{settings.API_V1_STR}/chat", tags=["chat"])
# app.include_router(reminders_router, prefix=f"{settings.API_V1_STR}/reminders", tags=["reminders"])
# app.include_router(search_router, prefix=f"{settings.API_V1_STR}/search", tags=["search"])


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=8000,
        reload=settings.DEBUG,
        log_level="info" if not settings.DEBUG else "debug"
    )