Update private_gpt/launcher.py

private_gpt/launcher.py

@@ -7,7 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.openapi.utils import get_openapi
 from injector import Injector
 from fastapi import APIRouter
-#from fastapi import JSONResponse
+
 
 from private_gpt.paths import docs_path
 from private_gpt.server.chat.chat_router import chat_router
@@ -20,9 +20,6 @@ from private_gpt.server.utils import authentication
 from private_gpt.settings.settings import Settings
 from private_gpt.components.llm.llm_component import LLMComponent
 
-#databse for user authentication integration
-#from private_gpt.server.utils import models 
-#from private_gpt.server.utils.database import engine, SessionLocal
 from typing import Annotated
 from sqlalchemy.orm import Session
 from private_gpt.server.utils.authentication import get_current_user
@@ -82,28 +79,43 @@ def create_app(root_injector: Injector) -> FastAPI:
             request.state.injector = root_injector
 
         app = FastAPI(dependencies=[Depends(bind_injector_to_request)])
-        
-
         model_router = APIRouter(prefix ='/v1',dependencies=[Depends(get_current_user)])
 
+        #global hardcoded dictionary for LLMs
+        models_data = [
+                {"id": 1, "name": "gpt-3.5-turbo", "access": ["user", "admin"]},
+                {"id": 2, "name": "gpt-4", "access": ["admin"]},
+                {"id": 3, "name": "mistral-7B", "access": ["user", "admin"]},
+                
+            ]
+
         @model_router.post("/switch_model")
         async def switch_model(
             new_model: str, current_user: dict = Depends(get_current_user)
         ):
-            # Check if the user has the "admin" role
-            if "admin" not in current_user.get("role", []):
+            # Check if the user has either "admin" or "user" role
+            if "user" not in current_user.get("role", []):
+                raise HTTPException(
+                    status_code=status.HTTP_403_FORBIDDEN,
+                    detail="You are not authorized to use this API.",
+                )
+
+            # Retrieve model information and validate access
+            model_info = next((m for m in models_data if m["name"] == new_model), None)
+            if not model_info or current_user.get("role") not in model_info["access"]:
                 raise HTTPException(
                     status_code=status.HTTP_403_FORBIDDEN,
-                    detail="You are not an admin and cannot use this API.",
+                    detail="You are not authorized to access this model.",
                 )
 
-            #logic to switch the LLM model based on the user's request
-            
+            # Switch the model using the LLMComponent
             llm_component = root_injector.get(LLMComponent)
             llm_component.switch_model(new_model, settings=settings)
 
+            # Return a success message
             return {"message": f"Model switched to {new_model}"}
 
+
         # Define a new APIRouter for the model_list
         model_list_router = APIRouter(prefix="/v1", dependencies=[Depends(get_current_user)])    
 
@@ -112,24 +124,7 @@ def create_app(root_injector: Injector) -> FastAPI:
             """
             Get a list of models with their details.
             """
-            # In this example, hardcoding some sample model data
-            models_data = [
-                {"id": 1, "name": "gpt-3.5-turbo", "access": ["user", "admin"]},
-                {"id": 2, "name": "gpt-4", "access": ["admin"]},
-                # Add more models as needed
-            ]
-
-            return models_data
-
-        # @model_router.post("/switch_model")
-        # async def switch_model(new_model: str):
-        #     # Implement logic to switch the LLM model based on the user's request
-        #     # Example: Change the LLM model in LLMComponent based on the new_model parameter
-        #     llm_component = root_injector.get(LLMComponent)
-        #     llm_component.switch_model(new_model, settings=settings)
-
-        #     return {"message": f"Model switched to {new_model}"}
-
+            return models_data  # Directly return the global variable
 
         def custom_openapi() -> dict[str, Any]:
             if app.openapi_schema:
@@ -160,17 +155,6 @@ def create_app(root_injector: Injector) -> FastAPI:
 
         app.openapi = custom_openapi  # type: ignore[method-assign]
 
-        #models.Base.metadata.create_all(bind=engine)
-
-        # def get_db():
-        #     db = SessionLocal()
-        #     try:
-        #         yield db
-        #     finally:
-        #         db.close()
-
-        #db_dependency = Annotated[Session, Depends(get_db)]
-        #user_dependency = Annotated[dict, Depends(get_current_user)]
 
         @app.get("/v1/me", status_code=status.HTTP_200_OK)
         async def user(current_user: dict = Depends(get_current_user)):
@@ -209,4 +193,4 @@ def create_app(root_injector: Injector) -> FastAPI:
             ui = root_injector.get(PrivateGptUi)
             ui.mount_in_app(app, settings.ui.path)
 
-        return app
+        return app