API Server Implementation¶

@app.get("/models", response_model=ModelsResponse)
async def get_available_models() -> ModelsResponse:
    """
    Get a list of available pre-trained and fine-tuned models.

    This endpoint returns information about models that can be used with the
    search endpoints. It includes:

    1. Pre-trained models categorized by type (e.g., Semantic Search, Question Answering)
    2. Fine-tuned models specifically trained for anime/manga search

    Returns:
        ModelsResponse: Available models and their descriptions

        - **models**: Dictionary of model categories and available pre-trained models
        - **fine_tuned**: Dictionary of fine-tuned model names and their paths

    Example:
        ```bash
        curl -X GET "http://localhost:8000/models"
        ```

    Note:
        Fine-tuned models are located in the `model/fine-tuned` directory.
        The API will only list models that have a valid configuration file.
    """
    # Convert Mapping to Dict to satisfy the type checker
    models = dict(BaseSearchModel.list_available_models())
    fine_tuned = BaseSearchModel.list_fine_tuned_models()

    return ModelsResponse(models=models, fine_tuned=fine_tuned)

def get_or_create_model(
    dataset_type: str,
    model_name: str,
    device: Optional[str] = None,
    include_light_novels: bool = False,
) -> BaseSearchModel:
    """
    Get a cached model or create a new one if not already cached.

    This function manages the model cache to avoid reloading models for each request.
    It handles device selection, CUDA availability checking, and model initialization.

    Args:
        dataset_type: The type of dataset to use ('anime' or 'manga')
        model_name: The name or path of the model to use
        device: Device to run the model on ('cpu', 'cuda', 'cuda:0', etc.)
            If None, automatically selects the best available device
        include_light_novels: Whether to include light novels in manga search results
            Only relevant for manga dataset_type

    Returns:
        BaseSearchModel: An initialized search model ready for queries

    Raises:
        ValueError: If the model or dataset cannot be loaded
        RuntimeError: If there are issues initializing the model

    Note:
        If CUDA is requested but not available, it will automatically
        fall back to CPU with a warning.
    """
    # Check if CUDA is available when 'cuda' is requested
    import torch  # pylint: disable=import-outside-toplevel

    cuda_requested = device is not None and "cuda" in device
    cuda_available = torch.cuda.is_available()

    # Force CPU if CUDA is requested but not available
    if cuda_requested and not cuda_available:
        logger.warning("CUDA was requested but is not available. Falling back to CPU.")
        selected_device = "cpu"
    else:
        # Use the specified device or auto-detect
        selected_device = get_device(device)

    # Create a unique key for this configuration
    key = f"{dataset_type}_{model_name}_{selected_device}_{include_light_novels}"

    if key not in model_cache:
        logger.info("Creating new model: %s on device: %s", key, selected_device)
        model = get_search_model(
            dataset_type=dataset_type,
            model_name=model_name,
            device=selected_device,
            include_light_novels=include_light_novels,
        )

        # The model's device is already set in its constructor
        model_cache[key] = model

    return model_cache[key]

@app.get("/", response_model=HealthResponse)
async def health_check() -> HealthResponse:
    """
    Check if the API is running and ready to handle requests.

    This endpoint verifies that the API server is operational and provides
    information about the status of different components:

    1. Whether the API server itself is running
    2. Whether each model type (anime, manga) can be loaded
    3. Whether CUDA is available for GPU acceleration

    Returns:
        HealthResponse: The health status of the API

        - **status**: "healthy" if critical components are working, "degraded" otherwise
        - **models_loaded**: Dictionary indicating which models loaded successfully
        - **cuda_available**: Boolean indicating if CUDA is available for GPU acceleration

    Example:
        ```bash
        curl -X GET "http://localhost:8000/"
        ```

    Note:
        This endpoint intentionally uses CPU for model loading checks to avoid
        GPU memory issues during health checking.
    """
    # Check CUDA availability
    import torch  # pylint: disable=import-outside-toplevel

    cuda_available = torch.cuda.is_available()

    # Check if models can be loaded
    models_loaded = {
        "anime": False,
        "manga": False,
    }

    try:
        # Try on CPU to avoid GPU memory issues during health check
        get_or_create_model(
            "anime", "cross-encoder/ms-marco-MiniLM-L-6-v2", device="cpu"
        )
        models_loaded["anime"] = True
    except (ImportError, ValueError, RuntimeError, FileNotFoundError) as e:
        logger.error("Error loading anime model: %s", str(e))

    try:
        # Try on CPU to avoid GPU memory issues during health check
        get_or_create_model(
            "manga", "cross-encoder/ms-marco-MiniLM-L-6-v2", device="cpu"
        )
        models_loaded["manga"] = True
    except (ImportError, ValueError, RuntimeError, FileNotFoundError) as e:
        logger.error("Error loading manga model: %s", str(e))

    return HealthResponse(
        status="healthy" if any(models_loaded.values()) else "degraded",
        models_loaded=models_loaded,
        cuda_available=cuda_available,
    )

@restricted_app.get("/models", response_model=ModelsResponse)
async def restricted_get_models():
    """
    List available models endpoint for the restricted API mode.

    This endpoint returns information about models that can be used with
    the search endpoints in restricted mode.

    Returns:
        ModelsResponse: Available models and their descriptions
    """
    return await get_available_models()

@restricted_app.get("/", response_model=HealthResponse)
async def restricted_health_check():
    """
    Health check endpoint for the restricted API mode.

    This endpoint verifies that the API server is operational in restricted mode
    and provides information about the status of different components.

    Returns:
        HealthResponse: The health status of the API
    """
    return await health_check()

@restricted_app.post("/search/anime", response_model=SearchResponse)
async def restricted_search_anime(*fn_args, **kwargs):
    """
    Search for anime endpoint for the restricted API mode.

    This endpoint performs semantic search against the anime dataset
    using the specified model in restricted mode.

    Parameters are the same as the regular search_anime endpoint.

    Returns:
        SearchResponse: The search results with relevant anime matches
    """
    return await search_anime(*fn_args, **kwargs)

@restricted_app.post("/search/manga", response_model=SearchResponse)
async def restricted_search_manga(*fn_args, **kwargs):
    """
    Search for manga endpoint for the restricted API mode.

    This endpoint performs semantic search against the manga dataset
    using the specified model in restricted mode.

    Parameters are the same as the regular search_manga endpoint.

    Returns:
        SearchResponse: The search results with relevant manga matches
    """
    return await search_manga(*fn_args, **kwargs)

@app.post("/search/anime", response_model=SearchResponse)
async def search_anime(
    request: SearchRequest,
    model_name: str = Query(
        "cross-encoder/ms-marco-MiniLM-L-6-v2", description="Model name or path"
    ),
    device: Optional[str] = Query(
        None,
        description="Device to run the model on ('cpu', 'cuda', 'cuda:0', etc.). "
        "If not specified, uses the best available device.",
    ),
) -> SearchResponse:
    """
    Search for anime matching the provided description.

    This endpoint performs semantic search against the anime dataset using
    the specified model, returning the most relevant matches sorted by score.

    ## Parameters

    - **request**: The search request body containing:
        - **query**: The search query text describing the anime
        - **num_results**: Number of results to return (default: 5, max: 100)
        - **batch_size**: Batch size for processing (default: 32)

    - **model_name**: The model to use for search (query parameter)
        - Can be a pre-trained model name or path to a fine-tuned model
        - Default: "cross-encoder/ms-marco-MiniLM-L-6-v2"

    - **device**: The device to run the model on (query parameter)
        - Options: 'cpu', 'cuda', 'cuda:0', etc.
        - If not specified, uses the best available device

    ## Returns

    - **results**: List of anime matching the query, sorted by relevance
    - **execution_time_ms**: Time taken to execute the search in milliseconds
    - **device_used**: The device used for computation (e.g., 'cpu', 'cuda')

    ## Example

    ```bash
    curl -X POST "http://localhost:8000/search/anime?device=cuda" \\
      -H "Content-Type: application/json" \\
      -d '{"query": "A story about robots and AI"}'
    ```

    ## Notes

    - For optimal performance on large queries, use GPU acceleration with `device=cuda`
    - Model caching is used to avoid reloading models between requests
    - Results include truncated synopses; full content is available in the dataset
    """
    import time  # pylint: disable=import-outside-toplevel

    try:
        # Get the search model
        start_time = time.time()
        search_model = get_or_create_model("anime", model_name, device=device)

        # Perform the search
        results = search_model.search(
            query=request.query,
            num_results=request.num_results,
            batch_size=request.batch_size,
        )

        # Convert to response format
        execution_time_ms = (time.time() - start_time) * 1000
        return SearchResponse(
            results=[SearchResult(**result) for result in results],
            execution_time_ms=execution_time_ms,
            device_used=search_model.device,
        )
    except (ImportError, ValueError, RuntimeError, FileNotFoundError) as e:
        logger.error("Error in anime search: %s", str(e), exc_info=True)
        raise HTTPException(
            status_code=500, detail=f"Error performing search: {str(e)}"
        ) from e

@app.post("/search/manga", response_model=SearchResponse)
async def search_manga(
    request: SearchRequest,
    model_name: str = Query(
        "cross-encoder/ms-marco-MiniLM-L-6-v2", description="Model name or path"
    ),
    include_light_novels: bool = Query(
        False, description="Whether to include light novels in search results"
    ),
    device: Optional[str] = Query(
        None,
        description="Device to run the model on ('cpu', 'cuda', 'cuda:0', etc.). "
        "If not specified, uses the best available device.",
    ),
) -> SearchResponse:
    """
    Search for manga matching the provided description.

    This endpoint performs semantic search against the manga dataset using
    the specified model, returning the most relevant matches sorted by score.

    ## Parameters

    - **request**: The search request body containing:
        - **query**: The search query text describing the manga
        - **num_results**: Number of results to return (default: 5, max: 100)
        - **batch_size**: Batch size for processing (default: 32)

    - **model_name**: The model to use for search (query parameter)
        - Can be a pre-trained model name or path to a fine-tuned model
        - Default: "cross-encoder/ms-marco-MiniLM-L-6-v2"

    - **include_light_novels**: Whether to include light novels in results (query parameter)
        - Default: false

    - **device**: The device to run the model on (query parameter)
        - Options: 'cpu', 'cuda', 'cuda:0', etc.
        - If not specified, uses the best available device

    ## Returns

    - **results**: List of manga matching the query, sorted by relevance
    - **execution_time_ms**: Time taken to execute the search in milliseconds
    - **device_used**: The device used for computation (e.g., 'cpu', 'cuda')

    ## Example

    ```bash
    curl -X POST "http://localhost:8000/search/manga?include_light_novels=true&device=cuda" \\
      -H "Content-Type: application/json" \\
      -d '{"query": "A fantasy adventure in a magical world", "num_results": 10}'
    ```

    ## Notes

    - Use `include_light_novels=true` to include light novels in search results
    - For optimal performance on large queries, use GPU acceleration with `device=cuda`
    - Model caching is used to avoid reloading models between requests
    - Results include truncated synopses; full content is available in the dataset
    """
    import time  # pylint: disable=import-outside-toplevel

    try:
        # Get the search model
        start_time = time.time()
        search_model = get_or_create_model(
            "manga",
            model_name,
            device=device,
            include_light_novels=include_light_novels,
        )

        # Perform the search
        results = search_model.search(
            query=request.query,
            num_results=request.num_results,
            batch_size=request.batch_size,
        )

        # Convert to response format
        execution_time_ms = (time.time() - start_time) * 1000
        return SearchResponse(
            results=[SearchResult(**result) for result in results],
            execution_time_ms=execution_time_ms,
            device_used=search_model.device,
        )
    except (ImportError, ValueError, RuntimeError, FileNotFoundError) as e:
        logger.error("Error in manga search: %s", str(e), exc_info=True)
        raise HTTPException(
            status_code=500, detail=f"Error performing search: {str(e)}"
        ) from e