all models to be registered with litellm

This commit is contained in:
Krazer 2024-06-11 13:07:19 -05:00
parent d05f8e77b4
commit d91707f73e
3 changed files with 39 additions and 0 deletions

View file

@ -135,6 +135,23 @@ def get_parser(default_config_files, git_root):
env_var="OPENAI_ORGANIZATION_ID",
help="Specify the OpenAI organization ID",
)
group.add_argument(
"--model_file",
metavar="MODEL_FILE",
default=None,
help={
"File with model definitions to be registered for info/cost, json formated",
" {"
" \"gpt-4\": {",
" \"max_tokens\": 8192,",
" \"input_cost_per_token\": 0.00003,",
" \"output_cost_per_token\": 0.00006,",
" \"litellm_provider\": \"openai\",",
" \"mode\": \"chat\"",
" },",
" }"
}
)
group.add_argument(
"--edit-format",
metavar="EDIT_FORMAT",

View file

@ -332,6 +332,16 @@ def main(argv=None, input=None, output=None, force_git_root=None, return_coder=F
if args.openai_organization_id:
os.environ["OPENAI_ORGANIZATION"] = args.openai_organization_id
model_def_files = []
model_def_fname = Path(".aider.models.json")
model_def_files.append(Path.home() / model_def_fname) # homedir
if git_root:
model_def_files.append(Path(git_root) / model_def_fname) # git root
if args.models:
model_def_files.append(args.models)
model_def_files = list(map(str, model_def_files))
models.register_models(model_def_files)
main_model = models.Model(args.model, weak_model=args.weak_model)
lint_cmds = parse_lint_cmds(args.lint_cmd, io)

View file

@ -426,6 +426,18 @@ class Model:
return res
def register_models(model_def_fnames):
for model_def_fname in model_def_fnames:
if not os.path.exists(model_def_fname):
continue
print(f"Registering model definition from {model_def_fname}")
try:
with open(model_def_fname, "r") as model_def_file:
model_def = json.load(model_def_file)
except json.JSONDecodeError as e:
print(f"Error opening/decoding model definition: {e}")
litellm.register_model(model_def)
def validate_variables(vars):
missing = []