File size: 3,397 Bytes
0231daa
 
fea62df
0231daa
 
 
fea62df
0231daa
 
 
 
fea62df
0231daa
 
 
90528a8
0231daa
fea62df
 
 
 
 
0231daa
3ab8523
0231daa
 
 
 
 
3ab8523
0231daa
 
3ab8523
0231daa
3ab8523
0231daa
 
3ab8523
fea62df
0231daa
 
 
3ab8523
0231daa
 
 
 
 
3ab8523
0231daa
3ab8523
 
0231daa
fea62df
3ab8523
fea62df
3ab8523
0231daa
3ab8523
0231daa
 
 
 
 
 
 
3ab8523
0231daa
fea62df
 
 
0231daa
 
3ab8523
fea62df
0231daa
fea62df
0231daa
3ab8523
0231daa
fea62df
 
3ab8523
0231daa
 
fea62df
0231daa
fea62df
 
 
3ab8523
fea62df
3ab8523
0231daa
 
 
 
 
 
 
fea62df
0231daa
3ab8523
80db4a8
0231daa
90528a8
0231daa
3ab8523
0231daa
fea62df
 
0231daa
 
fea62df
 
0231daa
 
3ab8523
0231daa
3ab8523
0231daa
 
 
 
 
 
3ab8523
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""
Unified Embedding API - Main Application Entry Point

A unified API for dense and sparse embeddings with support for
multiple models from Hugging Face.
"""

from contextlib import asynccontextmanager
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from loguru import logger

from src.config.settings import get_settings
from src.core.manager import ModelManager
from src.api import dependencies
from src.api.routers import embedding, model_list, health, rerank
from src.utils.logger import setup_logger


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Application lifespan manager.

    Handles startup and shutdown events:
    - Startup: Initialize models and cache
    - Shutdown: Cleanup resources
    """
    settings = get_settings()

    # Setup logging
    setup_logger(
        level=settings.LOG_LEVEL, log_file=settings.LOG_FILE, log_dir=settings.LOG_DIR
    )

    logger.info(f"Starting {settings.APP_NAME} v{settings.VERSION}")
    logger.info(f"Environment: {settings.ENVIRONMENT}")

    try:
        logger.info("Initializing model manager...")
        manager = ModelManager(settings.MODEL_CONFIG_PATH)
        dependencies.set_model_manager(manager)

        if settings.PRELOAD_MODELS:
            logger.info("Preloading models...")
            manager.preload_all_models()
        else:
            logger.info("Model preloading disabled (lazy loading)")

        logger.success("Startup complete! API is ready to serve requests.")

    except Exception:
        logger.exception("Failed to initialize application")
        raise

    yield

    logger.info("Shutting down...")

    try:
        manager = dependencies.get_model_manager()
        if manager:
            manager.unload_all_models()
            logger.info("All models unloaded")
    except Exception as e:
        logger.error(f"Error during shutdown: {e}")

    logger.info("Shutdown complete")


def create_app() -> FastAPI:
    """
    Create and configure the FastAPI application.

    Returns:
        Configured FastAPI application instance
    """
    settings = get_settings()

    temp_manager = ModelManager(settings.MODEL_CONFIG_PATH)
    api_description = temp_manager.generate_api_description()
    del temp_manager

    app = FastAPI(
        title=settings.APP_NAME,
        description=api_description,
        version=settings.VERSION,
        lifespan=lifespan,
        docs_url="/docs",
        redoc_url="/redoc",
        debug=settings.DEBUG,
    )

    if settings.CORS_ENABLED:
        app.add_middleware(
            CORSMiddleware,
            allow_origins=settings.CORS_ORIGINS,
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )
        logger.info(f"CORS enabled for origins: {settings.CORS_ORIGINS}")

    app.include_router(health.router)
    app.include_router(embedding.router, prefix="/api/v1")
    app.include_router(rerank.router, prefix="/api/v1")
    app.include_router(model_list.router, prefix="/api/v1")

    return app


# Create app instance
app = create_app()


if __name__ == "__main__":
    import uvicorn

    settings = get_settings()

    uvicorn.run(
        "app:app",
        host=settings.HOST,
        port=settings.PORT,
        reload=settings.RELOAD,
        workers=settings.WORKERS,
        log_level=settings.LOG_LEVEL.lower(),
    )