fix: do not stomp main model.info

This commit is contained in:
Paul Gauthier 2024-04-19 11:24:11 -07:00
parent 922559a15a
commit 6f1cebc4c2

View file

@ -1,7 +1,7 @@
import difflib
import sys
import json
import math
import sys
from dataclasses import dataclass, fields
import litellm
@ -170,11 +170,6 @@ class Model:
self.weak_model = self
return
try:
self.info = litellm.get_model_info(self.weak_model_name)
except Exception:
raise NoModelInfo(self.weak_model_name)
self.weak_model = Model(self.weak_model_name)
return self.weak_model
@ -242,12 +237,13 @@ def check_model_name(model):
possible_matches = fuzzy_match_models(model)
if possible_matches:
res += '\n\nDid you mean one of these:\n'
res += "\n\nDid you mean one of these:\n"
for match in possible_matches:
res += '\n- ' + match
res += "\n- " + match
return res
def fuzzy_match_models(name):
models = litellm.model_cost.keys()
@ -280,5 +276,6 @@ def main():
else:
print(f"No matching models found for '{model_name}'.")
if __name__ == "__main__":
main()