♻️ refactor me model module.

Author: jiangpeiling

backend/apps/me_model_managment_app.py

@@ -3,7 +3,6 @@
 from fastapi import APIRouter, Query
 from fastapi.responses import JSONResponse
 
-from consts.model import ModelResponse
 from services.me_model_management_service import get_me_models_impl
 from services.model_health_service import check_me_model_connectivity, check_me_connectivity_impl
 
@@ -20,7 +19,6 @@ async def get_me_models(
     """
     Get list of models from model engine API
     """
-    # Call service function to get business logic result
     code, message, data = await get_me_models_impl(timeout=timeout, type=type)
     return JSONResponse(
         status_code=HTTPStatus.OK,
@@ -37,7 +35,6 @@ async def check_me_connectivity(timeout: int = Query(default=2, description="Tim
     """
     Health check from model engine API
     """
-    # Call service function to health check
     code, message, data = await check_me_connectivity_impl(timeout)
     return JSONResponse(
         status_code=HTTPStatus.OK,
@@ -49,8 +46,15 @@ async def check_me_connectivity(timeout: int = Query(default=2, description="Tim
     )
 
 
-@router.get("/model/healthcheck", response_model=ModelResponse)
+@router.get("/model/healthcheck")
 async def check_me_model_healthcheck(
         model_name: str = Query(..., description="Model name to check")
 ):
-    return await check_me_model_connectivity(model_name)
+    code, message, data = await check_me_model_connectivity(model_name)
+    return JSONResponse(
+        status_code=HTTPStatus.OK,
+        content={
+        "code": code,
+        "message": message,
+        "data": data
+    })

backend/services/model_health_service.py

@@ -25,7 +25,6 @@ async def _embedding_dimension_check(
         model_type: str,
         model_base_url: str,
         model_api_key: str):
-
     # Test connectivity based on different model types
     if model_type == "embedding":
         embedding = await OpenAICompatibleEmbedding(
@@ -50,11 +49,11 @@ async def _embedding_dimension_check(
 
 
 async def _perform_connectivity_check(
-    model_name: str,
-    model_type: str,
-    model_base_url: str,
-    model_api_key: str,
-    embedding_dim: int = 1024
+        model_name: str,
+        model_type: str,
+        model_base_url: str,
+        model_api_key: str,
+        embedding_dim: int = 1024
 ) -> bool:
     """
     Perform specific model connectivity check
@@ -183,8 +182,9 @@ async def check_me_model_connectivity(model_name: str):
             model_data = next(
                 (item for item in result if item['id'] == model_name), None)
             if not model_data:
-                return ModelResponse(code=404, message="Specified model not found",
-                                     data={"connectivity": False, "message": "Specified model not found", "connect_status": ""})
+                return HTTPStatus.NOT_FOUND, "Specified model not found", {"connectivity": False,
+                                                                           "message": "Specified model not found",
+                                                                           "connect_status": ""}
 
             model_type = model_data['type']
 
@@ -205,27 +205,28 @@ async def check_me_model_connectivity(model_name: str):
                     json=payload
                 )
             else:
-                return ModelResponse(code=400, message=f"Health check not supported for {model_type} type models",
-                                     data={"connectivity": False, "message": f"Health check not supported for {model_type} type models",
-                                           "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value})
+                return HTTPStatus.BAD_REQUEST, f"Health check not supported for {model_type} type models", {
+                    "connectivity": False, "message": f"Health check not supported for {model_type} type models",
+                    "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value}
 
             status_code = api_response.status_code
             response_text = api_response.text
 
-            if status_code == 200:
+            if status_code == HTTPStatus.OK:
                 connect_status = ModelConnectStatusEnum.AVAILABLE.value
-                return ModelResponse(code=200, message=f"Model {model_name} responded normally",
-                                     data={"connectivity": True, "message": f"Model {model_name} responded normally", "connect_status": connect_status})
+                return HTTPStatus.OK, f"Model {model_name} responded normally", {"connectivity": True,
+                                                                                 "message": f"Model {model_name} responded normally",
+                                                                                 "connect_status": connect_status}
             else:
                 connect_status = ModelConnectStatusEnum.UNAVAILABLE.value
-                return ModelResponse(code=status_code, message=f"Model {model_name} response failed",
-                                     data={"connectivity": False, "message": f"Model {model_name} response failed: {response_text}",
-                                           "connect_status": connect_status})
+                return status_code, f"Model {model_name} response failed", {"connectivity": False,
+                                                                            "message": f"Model {model_name} response failed: {response_text}",
+                                                                            "connect_status": connect_status}
 
     except Exception as e:
-        return ModelResponse(code=500, message=f"Unknown error occurred: {str(e)}",
-                             data={"connectivity": False, "message": f"Unknown error occurred: {str(e)}",
-                                   "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value})
+        return HTTPStatus.INTERNAL_SERVER_ERROR, f"Unknown error occurred: {str(e)}", {"connectivity": False,
+                                                                                       "message": f"Unknown error occurred: {str(e)}",
+                                                                                       "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value}
 
 
 async def check_me_connectivity_impl(timeout: int):
@@ -240,13 +241,13 @@ async def check_me_connectivity_impl(timeout: int):
         headers = {'Authorization': f'Bearer {MODEL_ENGINE_APIKEY}'}
 
         async with aiohttp.ClientSession(
-            timeout=aiohttp.ClientTimeout(total=timeout),
-            connector=aiohttp.TCPConnector(ssl=False)
+                timeout=aiohttp.ClientTimeout(total=timeout),
+                connector=aiohttp.TCPConnector(ssl=False)
         ) as session:
             try:
                 async with session.get(
-                    f"{MODEL_ENGINE_HOST}/open/router/v1/models",
-                    headers=headers
+                        f"{MODEL_ENGINE_HOST}/open/router/v1/models",
+                        headers=headers
                 ) as response:
                     if response.status == HTTPStatus.OK:
                         return (
@@ -282,10 +283,10 @@ async def check_me_connectivity_impl(timeout: int):
     except Exception as e:
         return (
             HTTPStatus.INTERNAL_SERVER_ERROR,
-            f"Connection failed: {str(e)}",
+            f"Unknown error occurred: {str(e)}",
             {
                 "status": "Disconnected",
-                "desc": f"Connection failed: {str(e)}",
+                "desc": f"Unknown error occurred: {str(e)}",
                 "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value
             }
         )

test/backend/app/test_me_model_managment_app.py

@@ -220,8 +220,7 @@ async def mock_list_models_not_found(type: str = Query(None)):
 
         # Filter models if type is provided
         if type:
-            filtered_models = [
-                model for model in all_models if model["type"] == type]
+            filtered_models = [model for model in all_models if model["type"] == type]
 
             # Return 404 if no models found with this type
             if not filtered_models:

test/backend/services/test_model_health_service.py

@@ -130,8 +130,7 @@ async def test_perform_connectivity_check_embedding():
     # Setup
     with mock.patch("backend.services.model_health_service.OpenAICompatibleEmbedding") as mock_embedding:
         mock_embedding_instance = mock.MagicMock()
-        mock_embedding_instance.dimension_check = mock.AsyncMock(return_value=[
-                                                                 1])
+        mock_embedding_instance.dimension_check = mock.AsyncMock(return_value=[1])
         mock_embedding.return_value = mock_embedding_instance
 
         # Execute
@@ -192,8 +191,7 @@ async def test_perform_connectivity_check_llm():
         mock_observer.return_value = mock_observer_instance
 
         mock_model_instance = mock.MagicMock()
-        mock_model_instance.check_connectivity = mock.AsyncMock(
-            return_value=True)
+        mock_model_instance.check_connectivity = mock.AsyncMock(return_value=True)
         mock_model.return_value = mock_model_instance
 
         # Execute
@@ -224,8 +222,7 @@ async def test_perform_connectivity_check_vlm():
         mock_observer.return_value = mock_observer_instance
 
         mock_model_instance = mock.MagicMock()
-        mock_model_instance.check_connectivity = mock.AsyncMock(
-            return_value=True)
+        mock_model_instance.check_connectivity = mock.AsyncMock(return_value=True)
         mock_model.return_value = mock_model_instance
 
         # Execute
@@ -520,12 +517,12 @@ async def test_check_me_model_connectivity_llm_success():
         mock_client_instance.post.return_value = mock_test_response
 
         # Execute
-        response = await check_me_model_connectivity("gpt-4")
+        code, _, data = await check_me_model_connectivity("gpt-4")
 
         # Assert
-        assert response.code == 200
-        assert response.data["connectivity"] is True
-        assert response.data["connect_status"] == "available"
+        assert code == 200
+        assert data["connectivity"] is True
+        assert data["connect_status"] == "available"
 
         # Verify API calls
         mock_client_instance.get.assert_called_once_with(
@@ -573,12 +570,12 @@ async def test_check_me_model_connectivity_embedding_success():
         mock_client_instance.post.return_value = mock_test_response
 
         # Execute
-        response = await check_me_model_connectivity("text-embedding-ada-002")
+        code, _, data = await check_me_model_connectivity("text-embedding-ada-002")
 
         # Assert
-        assert response.code == 200
-        assert response.data["connectivity"] is True
-        assert response.data["connect_status"] == "available"
+        assert code == 200
+        assert data["connectivity"] is True
+        assert data["connect_status"] == "available"
 
         # Verify API calls
         mock_client_instance.get.assert_called_once_with(
@@ -614,12 +611,12 @@ async def test_check_me_model_connectivity_model_not_found():
         mock_client_instance.get.return_value = mock_response_obj
 
         # Execute
-        response = await check_me_model_connectivity("nonexistent-model")
+        code, _, data = await check_me_model_connectivity("nonexistent-model")
 
         # Assert
-        assert response.code == 404
-        assert response.data["connectivity"] is False
-        assert response.data["message"] == "Specified model not found"
+        assert code == 404
+        assert data["connectivity"] is False
+        assert data["message"] == "Specified model not found"
 
 
 @pytest.mark.asyncio
@@ -647,13 +644,13 @@ async def test_check_me_model_connectivity_unsupported_type():
         mock_client_instance.get.return_value = mock_response_obj
 
         # Execute
-        response = await check_me_model_connectivity("unsupported-model")
+        code, _, data = await check_me_model_connectivity("unsupported-model")
 
         # Assert
-        assert response.code == 400
-        assert response.data["connectivity"] is False
-        assert response.data["connect_status"] == "unavailable"
-        assert "Health check not supported" in response.data["message"]
+        assert code == 400
+        assert data["connectivity"] is False
+        assert data["connect_status"] == "unavailable"
+        assert "Health check not supported" in data["message"]
 
 
 @pytest.mark.asyncio
@@ -687,13 +684,13 @@ async def test_check_me_model_connectivity_api_error():
         mock_client_instance.post.return_value = mock_test_response
 
         # Execute
-        response = await check_me_model_connectivity("gpt-4")
+        code, _, data = await check_me_model_connectivity("gpt-4")
 
         # Assert
-        assert response.code == 500
-        assert response.data["connectivity"] is False
-        assert response.data["connect_status"] == "unavailable"
-        assert "response failed" in response.data["message"]
+        assert code == 500
+        assert data["connectivity"] is False
+        assert data["connect_status"] == "unavailable"
+        assert "response failed" in data["message"]
 
 
 @pytest.mark.asyncio
@@ -711,13 +708,13 @@ async def test_check_me_model_connectivity_exception():
             "Connection error")
 
         # Execute
-        response = await check_me_model_connectivity("gpt-4")
+        code, _, data = await check_me_model_connectivity("gpt-4")
 
         # Assert
-        assert response.code == 500
-        assert response.data["connectivity"] is False
-        assert response.data["connect_status"] == "unavailable"
-        assert "Unknown error" in response.data["message"]
+        assert code == 500
+        assert data["connectivity"] is False
+        assert data["connect_status"] == "unavailable"
+        assert "Unknown error" in data["message"]
 
 
 @pytest.mark.asyncio