♻️ refactor me model module.

Author: jiangpeiling

backend/apps/me_model_managment_app.py

@@ -1,117 +1,77 @@
-import asyncio
+from http import HTTPStatus
 
-import aiohttp
 from fastapi import APIRouter, Query
+from fastapi.responses import JSONResponse
 
-from consts.const import MODEL_ENGINE_APIKEY, MODEL_ENGINE_HOST
-from consts.model import ModelConnectStatusEnum, ModelResponse
-from services.model_health_service import check_me_model_connectivity
+from consts.model import ModelResponse, ModelConnectStatusEnum
+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
 
 router = APIRouter(prefix="/me")
 
 
-@router.get("/model/list", response_model=ModelResponse)
+@router.get("/model/list")
 async def get_me_models(
         type: str = Query(
             default="", description="Model type: embed/chat/rerank"),
         timeout: int = Query(
             default=2, description="Request timeout in seconds")
 ):
     try:
-        headers = {
-            'Authorization': f'Bearer {MODEL_ENGINE_APIKEY}',
-        }
-
-        async with aiohttp.ClientSession(
-            timeout=aiohttp.ClientTimeout(total=timeout),
-            connector=aiohttp.TCPConnector(verify_ssl=False)
-        ) as session:
-            async with session.get(
-                f"{MODEL_ENGINE_HOST}/open/router/v1/models",
-                headers=headers
-            ) as response:
-                response.raise_for_status()
-                result_data = await response.json()
-                result: list = result_data['data']
-
-        # Type filtering
-        filtered_result = []
-        if type:
-            for data in result:
-                if data['type'] == type:
-                    filtered_result.append(data)
-            if not filtered_result:
-                result_types = set(data['type'] for data in result)
-                return ModelResponse(
-                    code=404,
-                    message=f"No models found with type '{type}'. Available types: {result_types}",
-                    data=[]
-                )
-        else:
-            filtered_result = result
-
-        return ModelResponse(
-            code=200,
-            message="Successfully retrieved",
-            data=filtered_result
-        )
-
-    except asyncio.TimeoutError:
-        return ModelResponse(
-            code=408,
-            message="Request timeout",
-            data=[]
+        filtered_result = await get_me_models_impl(timeout, type)
+        return JSONResponse(
+            status_code=HTTPStatus.OK,
+            content={
+                "code": HTTPStatus.OK,
+                "message": "Successfully retrieved",
+                "data": filtered_result
+            }
         )
     except Exception as e:
-        return ModelResponse(
-            code=500,
-            message=f"Failed to get model list: {str(e)}",
-            data=[]
+        return JSONResponse(
+            status_code=HTTPStatus.OK,
+            content={
+                "code": HTTPStatus.INTERNAL_SERVER_ERROR,
+                "message": f"Failed to get model list: {str(e)}",
+                "data": None
+            }
         )
 
 
-@router.get("/healthcheck", response_model=ModelResponse)
+@router.get("/healthcheck")
 async def check_me_connectivity(timeout: int = Query(default=2, description="Timeout in seconds")):
     try:
-        headers = {'Authorization': f'Bearer {MODEL_ENGINE_APIKEY}'}
-
-        async with aiohttp.ClientSession(
-            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
-                ) as response:
-                    if response.status == 200:
-                        return ModelResponse(
-                            code=200,
-                            message="Connection successful",
-                            data={"status": "Connected", "desc": "Connection successful",
-                                  "connect_status": ModelConnectStatusEnum.AVAILABLE.value}
-                        )
-                    else:
-                        return ModelResponse(
-                            code=response.status,
-                            message=f"Connection failed, error code: {response.status}",
-                            data={"status": "Disconnected", "desc": f"Connection failed, error code: {response.status}",
-                                  "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value}
-                        )
-            except asyncio.TimeoutError:
-                return ModelResponse(
-                    code=408,
-                    message="Connection timeout",
-                    data={"status": "Disconnected", "desc": "Connection timeout",
-                          "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value}
-                )
-
+        message, code, status_data = await check_me_connectivity_impl(timeout)
+        if code == HTTPStatus.OK:
+            return JSONResponse(
+                status_code=HTTPStatus.OK,
+                content={
+                    "code": HTTPStatus.OK,
+                    "message": message,
+                    "data": status_data
+                }
+            )
+        else:
+            return JSONResponse(
+                status_code=HTTPStatus.OK,
+                content={
+                    "code": code,
+                    "message": message,
+                    "data": status_data
+                }
+            )
     except Exception as e:
-        return ModelResponse(
-            code=500,
-            message=f"Unknown error occurred: {str(e)}",
-            data={"status": "Disconnected", "desc": f"Unknown error occurred: {str(e)}",
-                  "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value}
+        return JSONResponse(
+            status_code=HTTPStatus.OK,
+            content={
+                "code": HTTPStatus.INTERNAL_SERVER_ERROR,
+                "message": f"Connection failed: {str(e)}",
+                "data": {
+                    "status": "Disconnected",
+                    "desc": f"Connection failed: {str(e)}",
+                    "connect_status": ModelConnectStatusEnum.UNAVAILABLE.value
+                }
+            }
         )
 
 

backend/consts/exceptions.py

@@ -27,7 +27,7 @@ class MemoryPreparationException(Exception):
     """Raised when memory preprocessing or retrieval fails prior to agent run."""
     pass
 
-  
+
 class MCPConnectionError(Exception):
     """Raised when MCP connection fails."""
     pass
@@ -50,4 +50,8 @@ class IncorrectInviteCodeException(Exception):
 
 class UserRegistrationException(Exception):
     """Raised when user registration fails."""
+    pass
+
+
+class TimeoutException(Exception):
     pass

backend/services/me_model_management_service.py

@@ -0,0 +1,54 @@
+import asyncio
+from typing import List, Dict
+
+import aiohttp
+from onnx.reference.ops.op_optional import Optional
+
+from consts.const import MODEL_ENGINE_APIKEY, MODEL_ENGINE_HOST
+from consts.exceptions import TimeoutException
+
+
+async def get_me_models_impl(timeout: int = 2, type: Optional[str] = None) -> List[Dict]:
+    """
+    Fetches a list of models from the model engine API, optionally filtering by type.
+    Parameters:
+        timeout (float): The total timeout for the request in seconds.
+        type (str or None): The type of model to filter for. If None, returns all models.
+    Returns:
+        list: A list of model data dictionaries, filtered by the specified type if provided.
+    Raises:
+        TimeoutException: If the request times out.
+        Exception: If no models of the specified type are found, or if another request error occurs.
+    """
+    try:
+        headers = {
+            'Authorization': f'Bearer {MODEL_ENGINE_APIKEY}',
+        }
+        async with aiohttp.ClientSession(
+                timeout=aiohttp.ClientTimeout(total=timeout),
+                connector=aiohttp.TCPConnector(verify_ssl=False)
+        ) as session:
+            async with session.get(
+                    f"{MODEL_ENGINE_HOST}/open/router/v1/models",
+                    headers=headers
+            ) as response:
+                response.raise_for_status()
+                result_data = await response.json()
+                result: list = result_data['data']
+        # Type filtering
+        filtered_result = []
+        if type:
+            for data in result:
+                if data['type'] == type:
+                    filtered_result.append(data)
+            if not filtered_result:
+                result_types = set(data['type'] for data in result)
+                raise Exception(
+                    f"No models found with type '{type}'. Available types: {result_types}.")
+        else:
+            filtered_result = result
+        return filtered_result
+    except asyncio.TimeoutError:
+        raise TimeoutException("Request timeout.")
+    except Exception as e:
+        raise Exception(f"Request error: {e}.")