fixed multiple issues

This commit is contained in:
Saifeddine ALOUI 2024-02-29 13:57:45 +01:00
parent c7abab41ee
commit 5989ddbc7c
4 changed files with 88 additions and 73 deletions

View File

@ -34,6 +34,7 @@ from enum import Enum
from lollms.utilities import trace_exception
from tqdm import tqdm
import sys
__author__ = "parisneo"
__github__ = "https://github.com/ParisNeo/lollms_bindings_zoo"
@ -108,6 +109,77 @@ class LLMBinding:
for models_folder in self.models_folders:
models_folder.mkdir(parents=True, exist_ok=True)
def searchModelPath(self, model_name:str):
model_path=None
for mn in self.models_folders:
if mn.name in model_name.lower():
if mn.name == "ggml":
try:
idx = model_name.index("-GGML")
models=[m for m in mn.iterdir() if model_name[:idx].lower() in m.name.lower()]
model_path = mn/models[0].name
except:
model_path = mn/model_name
elif mn.name == "gguf":
try:
idx = model_name.index("-GGUF")
models=[m for m in mn.iterdir() if model_name[:idx].lower() in m.name.lower()]
model_path = mn/models[0].name
except:
model_path = mn/model_name
else:
model_path = mn/model_name
break
if model_path is None:
model_path = self.models_folders[0]/model_name
return model_path
def download_model(self, url, binding, callback = None):
model_name = url.split("/")[-1]
folder_path = binding.searchModelPath(model_name)
model_full_path = folder_path
if binding is not None and hasattr(binding,'download_model'):
binding.download_model(url, model_full_path, callback)
else:
# Check if file already exists in folder
if model_full_path.exists():
print("File already exists in folder")
else:
# Create folder if it doesn't exist
folder_path.mkdir(parents=True, exist_ok=True)
progress_bar = tqdm(total=100, unit="%", unit_scale=True, desc=f"Downloading {url.split('/')[-1]}")
# Define callback function for urlretrieve
def report_progress(block_num, block_size, total_size):
progress_bar.update(block_size/total_size)
# Download file from URL to folder
try:
rq.urlretrieve(url, folder_path / url.split("/")[-1], reporthook=report_progress if callback is None else callback)
print("File downloaded successfully!")
except Exception as e:
print("Error downloading file:", e)
sys.exit(1)
def reference_model(self, path):
path = str(path).replace("\\","/")
model_name = path.split("/")[-1]+".reference"
folder_path = self.searchModelPath(model_name).replace(".reference","")
model_full_path = (folder_path / model_name)
# Check if file already exists in folder
if model_full_path.exists():
print("File already exists in folder")
return False
else:
# Create folder if it doesn't exist
folder_path.mkdir(parents=True, exist_ok=True)
with open(model_full_path,"w") as f:
f.write(path)
ASCIIColors.warning("Reference created, please make sure you don't delete or move the referenced file. This can cause the link to be broken")
return True
def sync_configuration(self, binding_config:TypedConfig, lollms_paths:LollmsPaths):
self.configuration_file_path = lollms_paths.personal_configuration_path/"bindings"/self.binding_folder_name/f"config.yaml"
self.configuration_file_path.parent.mkdir(parents=True, exist_ok=True)

View File

@ -139,71 +139,3 @@ class LOLLMSConfig(BaseConfig):
except Exception as ex:
print(f"Exception in checking model existance: {ex}")
return False
def searchModelPath(self, model_name:str):
model_path=None
for mn in self.models_folders:
if mn.name in model_name.lower():
if mn.name == "ggml":
try:
idx = model_name.index("-GGML")
models=[m for m in mn.iterdir() if model_name[:idx].lower() in m.name.lower()]
model_path = mn/models[0].name
except:
model_path = mn/model_name
elif mn.name == "gguf":
try:
idx = model_name.index("-GGUF")
models=[m for m in mn.iterdir() if model_name[:idx].lower() in m.name.lower()]
model_path = mn/models[0].name
except:
model_path = mn/model_name
else:
model_path = mn/model_name
break
if model_path is None:
model_path = self.models_folders[0]/model_name
return model_path
def download_model(self, url, binding, callback = None):
model_name = url.split("/")[-1]
folder_path = binding.searchModelPath(model_name)
model_full_path = folder_path
if binding is not None and hasattr(binding,'download_model'):
binding.download_model(url, model_full_path, callback)
else:
# Check if file already exists in folder
if model_full_path.exists():
print("File already exists in folder")
else:
# Create folder if it doesn't exist
folder_path.mkdir(parents=True, exist_ok=True)
progress_bar = tqdm(total=100, unit="%", unit_scale=True, desc=f"Downloading {url.split('/')[-1]}")
# Define callback function for urlretrieve
def report_progress(block_num, block_size, total_size):
progress_bar.update(block_size/total_size)
# Download file from URL to folder
try:
rq.urlretrieve(url, folder_path / url.split("/")[-1], reporthook=report_progress if callback is None else callback)
print("File downloaded successfully!")
except Exception as e:
print("Error downloading file:", e)
sys.exit(1)
def reference_model(self, path):
path = str(path).replace("\\","/")
model_name = path.split("/")[-1]+".reference"
folder_path = self.searchModelPath(model_name)
model_full_path = (folder_path / model_name)
# Check if file already exists in folder
if model_full_path.exists():
print("File already exists in folder")
return False
else:
# Create folder if it doesn't exist
folder_path.mkdir(parents=True, exist_ok=True)
with open(model_full_path,"w") as f:
f.write(path)
ASCIIColors.warning("Reference created, please make sure you don't delete or move the referenced file. This can cause the link to be broken")
return True

View File

@ -2,9 +2,13 @@ from fastapi import HTTPException
from ascii_colors import ASCIIColors
from pathlib import Path
def sanitize_path(path:str, error_text="Absolute database path detected", exception_text="Detected an attempt of path traversal. Are you kidding me?"):
if(".." in path or Path(path).is_absolute()):
ASCIIColors.warning()
def sanitize_path(path:str, allow_absolute_path:False, error_text="Absolute database path detected", exception_text="Detected an attempt of path traversal. Are you kidding me?"):
if(".." in path):
ASCIIColors.warning(error_text)
raise exception_text
if (not allow_absolute_path) and Path(path).is_absolute():
ASCIIColors.warning(error_text)
raise exception_text
def sanitize_path_from_endpoint(path:str, error_text="A suspected LFI attack detected. The path sent to the server has .. in it!", exception_text="Invalid path!"):
@ -12,3 +16,7 @@ def sanitize_path_from_endpoint(path:str, error_text="A suspected LFI attack det
ASCIIColors.error(error_text)
raise HTTPException(status_code=400, detail=exception_text)
def forbid_remote_access(lollmsElfServer):
if lollmsElfServer.config.host!="localhost" and lollmsElfServer.config.host!="127.0.0.1":
return {"status":False,"error":"Code execution is blocked when the server is exposed outside for very obvious reasons!"}

View File

@ -11,6 +11,7 @@ from fastapi import APIRouter, Request
from pydantic import BaseModel
import pkg_resources
from lollms.server.elf_server import LOLLMSElfServer
from lollms.security import sanitize_path, forbid_remote_access
from ascii_colors import ASCIIColors
from lollms.utilities import load_config
from pathlib import Path
@ -86,12 +87,14 @@ def get_model_status():
@router.post("/add_reference_to_local_model")
def add_reference_to_local_model(data:ModelReferenceParams):
forbid_remote_access(lollmsElfServer)
if data.path=="":
return {"status": False, "error":"Empty model path"}
path = Path(data.path)
if path.exists():
lollmsElfServer.config.reference_model(path)
lollmsElfServer.binding.reference_model(path)
return {"status": True}
else:
return {"status": False, "error":"Model not found"}