return fq and model names mixed together from fuzzy match

This commit is contained in:
Paul Gauthier 2024-07-14 18:26:05 +01:00
parent 29afc85957
commit 50353acb1a

View file

@ -654,11 +654,7 @@ def sanity_check_model(io, model):
if possible_matches: if possible_matches:
io.tool_output("Did you mean one of these?") io.tool_output("Did you mean one of these?")
for match in possible_matches: for match in possible_matches:
fq, m = match io.tool_output(f"- {model}")
if fq == m:
io.tool_output(f"- {m}")
else:
io.tool_output(f"- {m} ({fq})")
if show: if show:
io.tool_output(f"For more info, see: {urls.model_warnings}\n") io.tool_output(f"For more info, see: {urls.model_warnings}\n")
@ -667,7 +663,7 @@ def sanity_check_model(io, model):
def fuzzy_match_models(name): def fuzzy_match_models(name):
name = name.lower() name = name.lower()
chat_models = [] chat_models = set()
for model, attrs in litellm.model_cost.items(): for model, attrs in litellm.model_cost.items():
model = model.lower() model = model.lower()
if attrs.get("mode") != "chat": if attrs.get("mode") != "chat":
@ -679,8 +675,10 @@ def fuzzy_match_models(name):
else: else:
fq_model = provider + model fq_model = provider + model
chat_models.append((fq_model, model)) chat_models.add(fq_model)
chat_models.add(model)
chat_models = sorted(chat_models)
# exactly matching model # exactly matching model
# matching_models = [ # matching_models = [
# (fq,m) for fq,m in chat_models # (fq,m) for fq,m in chat_models
@ -690,19 +688,14 @@ def fuzzy_match_models(name):
# return matching_models # return matching_models
# Check for model names containing the name # Check for model names containing the name
matching_models = [(fq, m) for fq, m in chat_models if name in fq] matching_models = [m for m in chat_models if name in m]
if matching_models: if matching_models:
return matching_models return matching_models
# Check for slight misspellings # Check for slight misspellings
models = [m for fq, m in chat_models] models = list(chat_models)
matching_models = difflib.get_close_matches(name, models, n=3, cutoff=0.8) matching_models = difflib.get_close_matches(name, models, n=3, cutoff=0.8)
if matching_models: return sorted(matching_models)
return list(zip(matching_models, matching_models))
fq_models = [fq for fq, m in chat_models]
matching_models = difflib.get_close_matches(name, fq_models, n=3, cutoff=0.8)
return list(zip(matching_models, matching_models))
def print_matching_models(io, search): def print_matching_models(io, search):
@ -710,8 +703,7 @@ def print_matching_models(io, search):
if matches: if matches:
io.tool_output(f'Models which match "{search}":') io.tool_output(f'Models which match "{search}":')
for model in matches: for model in matches:
fq, m = model io.tool_output(f"- {model}")
io.tool_output(f"- {fq}")
else: else:
io.tool_output(f'No models match "{search}".') io.tool_output(f'No models match "{search}".')