Check the environment inside Model

This commit is contained in:
Paul Gauthier 2024-04-20 08:09:17 -07:00
parent e513174da4
commit 6863b47e6a
2 changed files with 24 additions and 17 deletions

View file

@ -602,28 +602,12 @@ def main(argv=None, input=None, output=None, force_git_root=None):
if args.openai_organization_id:
os.environ["OPENAI_ORGANIZATION"] = args.openai_organization_id
# Is the model known and are all needed keys/params available?
res = litellm.validate_environment(args.model)
missing_keys = res.get("missing_keys")
keys_in_environment = res.get("keys_in_environment")
if not keys_in_environment and not missing_keys:
io.tool_error(f"Unable to check environment variables for model {args.model}")
elif missing_keys:
io.tool_error(f"To use model {args.model}, please set these environment variables:")
for key in missing_keys:
io.tool_error(f"- {key}")
return 1
elif not keys_in_environment and args.require_model_info:
io.tool_error(models.check_model_name(args.model))
return 1
# Check in advance that we have model metadata
try:
main_model = models.Model(
args.model, weak_model=args.weak_model, require_model_info=args.require_model_info
)
except models.NoModelInfo as err:
except (models.NoModelInfo, models.ModelEnvironmentError) as err:
io.tool_error(str(err))
return 1

View file

@ -22,6 +22,15 @@ class NoModelInfo(Exception):
super().__init__(check_model_name(model))
class ModelEnvironmentError(Exception):
"""
Exception raised when the environment isn't setup for the model
"""
def __init__(self, message):
super().__init__(message)
@dataclass
class ModelSettings:
name: str
@ -126,6 +135,20 @@ class Model:
def __init__(self, model, weak_model=None, require_model_info=True):
self.name = model
# Are all needed keys/params available?
res = litellm.validate_environment(model)
missing_keys = res.get("missing_keys")
keys_in_environment = res.get("keys_in_environment")
if missing_keys:
res = f"To use model {model}, please set these environment variables:"
for key in missing_keys:
res += f"- {key}"
raise ModelEnvironmentError(res)
elif not keys_in_environment:
print(f"Unable to check environment variables for model {model}")
# Do we have the model_info?
try:
self.info = litellm.get_model_info(model)
except Exception: