diff --git a/src/app/main.py b/src/app/main.py index 0d09438..b944c19 100644 --- a/src/app/main.py +++ b/src/app/main.py @@ -23,7 +23,7 @@ RerankRequest, RerankResponse, ) -from .models import get_model +from .models import get_model as get_model, get_model_or_400 from .config import ( EMBEDDING_MODELS, RERANK_MODELS, @@ -170,16 +170,7 @@ def _get_model_or_400(model_name: str, model_type: str) -> Any: Helper for backwards compatibility with legacy tests calling _get_model_or_400. """ supported_models = EMBEDDING_MODELS if model_type == "embedding" else RERANK_MODELS - if model_name not in supported_models: - raise HTTPException( - status_code=400, - detail=f"Model '{model_name}' not found for {model_type}s.", - ) - - try: - return get_model(model_name) - except ValueError as e: - raise HTTPException(status_code=400, detail=str(e)) + return get_model_or_400(model_name, supported_models, model_type) def _determine_ruri_prefix(request: EmbeddingRequest) -> str: diff --git a/src/app/models.py b/src/app/models.py index 492495c..e30b007 100644 --- a/src/app/models.py +++ b/src/app/models.py @@ -120,3 +120,33 @@ def get_model(model_name: str, device: str | None = None): _model_cache[model_name] = model logging.info(f"Model '{model_name}' loaded successfully.") return model + + +def get_model_or_400( + model_name: str, valid_models: Any, model_type: str = "model" +) -> Any: + """ + Validates model_name against valid_models and fetches model instance via get_model. + Raises HTTPException(400) if model is invalid or if get_model raises a ValueError. + """ + from fastapi import HTTPException + + if model_name not in valid_models: + suffix = "s" if not model_type.endswith("s") else "" + raise HTTPException( + status_code=400, + detail=f"Model '{model_name}' not found for {model_type}{suffix}.", + ) + + try: + # Dynamically import main to support patches on app.main.get_model in tests + try: + import app.main as main_mod + + fetch_func = getattr(main_mod, "get_model", get_model) + except Exception: + fetch_func = get_model + + return fetch_func(model_name) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) diff --git a/src/app/services/embedding.py b/src/app/services/embedding.py index a93eb19..c74f253 100644 --- a/src/app/services/embedding.py +++ b/src/app/services/embedding.py @@ -7,6 +7,7 @@ from .base import BaseEmbeddingService from ..image_utils import load_image_from_source +from ..models import get_model_or_400 from ..schemas import ( EmbeddingRequest, EmbeddingResponse, @@ -143,17 +144,7 @@ def _tokenize_and_truncate_embeddings( def _get_model_or_400(model_name: str) -> Any: - from app.main import get_model - - if model_name not in EMBEDDING_MODELS: - raise HTTPException( - status_code=400, - detail=f"Model '{model_name}' not found for embeddings.", - ) - try: - return get_model(model_name) - except ValueError as e: - raise HTTPException(status_code=400, detail=str(e)) + return get_model_or_400(model_name, EMBEDDING_MODELS, "embedding") class EmbeddingService(BaseEmbeddingService): diff --git a/src/app/services/rerank.py b/src/app/services/rerank.py index 105212f..6aab873 100644 --- a/src/app/services/rerank.py +++ b/src/app/services/rerank.py @@ -3,6 +3,7 @@ from fastapi import HTTPException from .base import BaseRerankService +from ..models import get_model_or_400 from ..schemas import RerankRequest, RerankResponse, RerankData, Usage from ..config import RERANK_MODELS @@ -47,17 +48,7 @@ def _sort_and_format_rerank_results( def _get_model_or_400(model_name: str) -> Any: - from app.main import get_model - - if model_name not in RERANK_MODELS: - raise HTTPException( - status_code=400, - detail=f"Model '{model_name}' not found for reranks.", - ) - try: - return get_model(model_name) - except ValueError as e: - raise HTTPException(status_code=400, detail=str(e)) + return get_model_or_400(model_name, RERANK_MODELS, "rerank") class RerankService(BaseRerankService):