Adding more routes

main/api.py

@@ -1,10 +1,11 @@
 import httpx
-from typing import Optional, AsyncIterator, Dict, Any, Iterator
+from typing import Optional, AsyncIterator, Dict, Any, Iterator, List
 import logging
 import asyncio
 from litserve import LitAPI
 from pydantic import BaseModel
 
+
 class GenerationResponse(BaseModel):
     generated_text: str
 
@@ -62,6 +63,107 @@ class InferenceApi(LitAPI):
             response = await self.generate_response(x, **kwargs)
             yield response
 
+    async def generate_embedding(self, text: str) -> List[float]:
+        """Generate embedding vector from input text."""
+        self.logger.debug(f"Forwarding embedding request for text: {text[:50]}...")
+
+        try:
+            async with await self._get_client() as client:
+                response = await client.post(
+                    self._get_endpoint('embedding'),
+                    json={"text": text}
+                )
+                response.raise_for_status()
+                data = response.json()
+                return data["embedding"]
+
+        except Exception as e:
+            self.logger.error(f"Error in generate_embedding: {str(e)}")
+            raise
+
+    async def check_system_status(self) -> Dict[str, Any]:
+        """Check system status of the LLM Server."""
+        self.logger.debug("Checking system status...")
+
+        try:
+            async with await self._get_client() as client:
+                response = await client.get(
+                    self._get_endpoint('system_status')
+                )
+                response.raise_for_status()
+                return response.json()
+
+        except Exception as e:
+            self.logger.error(f"Error in check_system_status: {str(e)}")
+            raise
+
+    async def download_model(self, model_name: Optional[str] = None) -> Dict[str, str]:
+        """Download model files from the LLM Server."""
+        self.logger.debug(f"Forwarding model download request for: {model_name or 'default model'}")
+
+        try:
+            async with await self._get_client() as client:
+                response = await client.post(
+                    self._get_endpoint('model_download'),
+                    params={"model_name": model_name} if model_name else None
+                )
+                response.raise_for_status()
+                return response.json()
+
+        except Exception as e:
+            self.logger.error(f"Error in download_model: {str(e)}")
+            raise
+
+    async def validate_system(self) -> Dict[str, Any]:
+        """Validate system configuration and setup."""
+        self.logger.debug("Validating system configuration...")
+
+        try:
+            async with await self._get_client() as client:
+                response = await client.get(
+                    self._get_endpoint('system_validate')
+                )
+                response.raise_for_status()
+                return response.json()
+
+        except Exception as e:
+            self.logger.error(f"Error in validate_system: {str(e)}")
+            raise
+
+    async def initialize_model(self, model_name: Optional[str] = None) -> Dict[str, Any]:
+        """Initialize specified model or default model."""
+        self.logger.debug(f"Initializing model: {model_name or 'default'}")
+
+        try:
+            async with await self._get_client() as client:
+                response = await client.post(
+                    self._get_endpoint('model_initialize'),
+                    json={"model_name": model_name} if model_name else {}
+                )
+                response.raise_for_status()
+                return response.json()
+
+        except Exception as e:
+            self.logger.error(f"Error in initialize_model: {str(e)}")
+            raise
+
+    async def initialize_embedding_model(self, model_name: Optional[str] = None) -> Dict[str, Any]:
+        """Initialize embedding model."""
+        self.logger.debug(f"Initializing embedding model: {model_name or 'default'}")
+
+        try:
+            async with await self._get_client() as client:
+                response = await client.post(
+                    self._get_endpoint('model_initialize_embedding'),
+                    json={"model_name": model_name} if model_name else {}
+                )
+                response.raise_for_status()
+                return response.json()
+
+        except Exception as e:
+            self.logger.error(f"Error in initialize_embedding_model: {str(e)}")
+            raise
+
     def decode_request(self, request: Any, **kwargs) -> str:
         """Convert the request payload to input format."""
         if isinstance(request, dict) and "prompt" in request:

main/config.yaml

@@ -15,4 +15,5 @@ llm_server:
     system_status: "/system/status"
     system_validate: "/system/validate"
     model_initialize: "/model/initialize"
-    model_initialize_embedding: "/model/initialize/embedding"
+    model_initialize_embedding: "/model/initialize/embedding"
+    model_download: "/model/download"

main/routes.py

@@ -17,7 +17,7 @@ from .schemas import (
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
-api = None
+api = InferenceApi()
 
 async def init_router(inference_api: InferenceApi):
     """Initialize router with an already setup API instance"""
@@ -174,6 +174,25 @@ async def initialize_embedding_model(model_name: Optional[str] = None):
         logger.error(f"Error initializing embedding model: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
 
+@router.post("/model/download",
+             summary="Download default or specified model",
+             description="Downloads model files. Uses default model from config if none specified.")
+async def download_model(model_name: Optional[str] = None):
+    """Download model files to local storage"""
+    try:
+        # Use model name from config if none provided
+        model_to_download = model_name or config["model"]["defaults"]["model_name"]
+        logger.info(f"Received request to download model: {model_to_download}")
+
+        result = await api.download_model(model_to_download)
+        logger.info(f"Successfully downloaded model: {model_to_download}")
+
+        return result
+
+    except Exception as e:
+        logger.error(f"Error downloading model: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
 @router.on_event("shutdown")
 async def shutdown_event():
     """Clean up resources on shutdown"""