theroyallab/tabbyAPI
1
0
Fork 0
mirror of https://github.com/theroyallab/tabbyAPI.git synced 2026-03-15 00:07:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

API: Error on invalid key permissions and cleanup format

Browse Source 
If a user requesting a model change isn't admin, error.

Better to place the load function before the generate functions.

Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
kingbri
2024-09-04 21:44:14 -04:00
parent 21f14d4318
commit 9c10789ca1
1 changed files with 42 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
endpoints/OAI/utils/completion.py
View File

@@ -108,6 +108,48 @@ async def _stream_collector(
await gen_queue.put(e)
async def load_inline_model(model_name: str, request: Request):
"""Load a model from the data.model parameter"""
# Return if the model container already exists
if model.container and model.container.model_dir.name == model_name:
return
model_config = config.model_config()
# Inline model loading isn't enabled or the user isn't an admin
if not get_key_permission(request) == "admin":
error_message = handle_request_error(
f"Unable to switch model to {model_name} because "
+ "an admin key isn't provided",
exc_info=False,
).error.message
raise HTTPException(401, error_message)
if not unwrap(model_config.get("inline_model_loading"), False):
logger.warning(
f"Unable to switch model to {model_name} because "
'"inline_model_load" is not True in config.yml.'
)
return
model_path = pathlib.Path(unwrap(model_config.get("model_dir"), "models"))
model_path = model_path / model_name
# Model path doesn't exist
if not model_path.exists():
logger.warning(
f"Could not find model path {str(model_path)}. Skipping inline model load."
)
return
# Load the model
await model.load_model(model_path)
async def stream_generate_completion(
data: CompletionRequest, request: Request, model_path: pathlib.Path
):
@@ -178,47 +220,6 @@ async def stream_generate_completion(
)
async def load_inline_model(model_name: str, request: Request):
"""Load a model from the data.model parameter"""
# Return if the model container already exists
if model.container and model.container.model_dir.name == model_name:
return
model_config = config.model_config()
# Inline model loading isn't enabled or the user isn't an admin
if not get_key_permission(request) == "admin":
logger.warning(
f"Unable to switch model to {model_name} "
"because an admin key isn't provided."
)
return
if not unwrap(model_config.get("inline_model_loading"), False):
logger.warning(
f"Unable to switch model to {model_name} because "
'"inline_model_load" is not True in config.yml.'
)
return
model_path = pathlib.Path(unwrap(model_config.get("model_dir"), "models"))
model_path = model_path / model_name
# Model path doesn't exist
if not model_path.exists():
logger.warning(
f"Could not find model path {str(model_path)}. Skipping inline model load."
)
return
# Load the model
await model.load_model(model_path)
async def generate_completion(
data: CompletionRequest, request: Request, model_path: pathlib.Path
):