load litellm if model metadata json is present

This commit is contained in:
Paul Gauthier 2024-08-19 10:37:49 -07:00
parent f6023865f8
commit b8d09e1f04
2 changed files with 14 additions and 10 deletions

View file

@ -17,6 +17,8 @@ class LazyLiteLLM:
_lazy_module = None _lazy_module = None
def __getattr__(self, name): def __getattr__(self, name):
if name == "_lazy_module":
return super()
self._load_litellm() self._load_litellm()
return getattr(self._lazy_module, name) return getattr(self._lazy_module, name)

View file

@ -451,16 +451,17 @@ class Model:
self.get_weak_model(weak_model) self.get_weak_model(weak_model)
def get_model_info(self, model): def get_model_info(self, model):
# Try and do this quickly, without triggering the litellm import if not litellm._lazy_module:
spec = importlib.util.find_spec("litellm") # Try and do this quickly, without triggering the litellm import
if spec: spec = importlib.util.find_spec("litellm")
origin = Path(spec.origin) if spec:
fname = origin.parent / "model_prices_and_context_window_backup.json" origin = Path(spec.origin)
if fname.exists(): fname = origin.parent / "model_prices_and_context_window_backup.json"
data = json.loads(fname.read_text()) if fname.exists():
info = data.get(model) data = json.loads(fname.read_text())
if info: info = data.get(model)
return info if info:
return info
# Do it the slow way... # Do it the slow way...
try: try:
@ -670,6 +671,7 @@ def register_litellm_models(model_fnames):
try: try:
with open(model_fname, "r") as model_def_file: with open(model_fname, "r") as model_def_file:
model_def = json.load(model_def_file) model_def = json.load(model_def_file)
litellm._load_litellm()
litellm.register_model(model_def) litellm.register_model(model_def)
except Exception as e: except Exception as e:
raise Exception(f"Error loading model definition from {model_fname}: {e}") raise Exception(f"Error loading model definition from {model_fname}: {e}")