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 difflib
import sys
import json import json
import math import math
import sys
from dataclasses import dataclass, fields from dataclasses import dataclass, fields
import litellm import litellm
@ -170,11 +170,6 @@ class Model:
self.weak_model = self self.weak_model = self
return 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) self.weak_model = Model(self.weak_model_name)
return self.weak_model return self.weak_model
@ -242,12 +237,13 @@ def check_model_name(model):
possible_matches = fuzzy_match_models(model) possible_matches = fuzzy_match_models(model)
if possible_matches: 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: for match in possible_matches:
res += '\n- ' + match res += "\n- " + match
return res return res
def fuzzy_match_models(name): def fuzzy_match_models(name):
models = litellm.model_cost.keys() models = litellm.model_cost.keys()
@ -280,5 +276,6 @@ def main():
else: else:
print(f"No matching models found for '{model_name}'.") print(f"No matching models found for '{model_name}'.")
if __name__ == "__main__": if __name__ == "__main__":
main() main()