Skip to content

Commit

Permalink
Fixes error when allowed or blocked model list is passed in config.
Browse files Browse the repository at this point in the history
  • Loading branch information
3coins committed Jun 22, 2024
1 parent 387e5b7 commit 337e4e2
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions packages/jupyter-ai/jupyter_ai/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -289,13 +289,17 @@ def filter_predicate(local_model_id: str):
else:
return model_id in self.allowed_models

def filter_model_list(model_list: Optional[List[str]]) -> List[str]:
if model_list is not None:
return list(filter(filter_predicate, provider.models))
else:
return []

# filter out every model w/ model ID according to allow/blocklist
for provider in providers:
provider.models = list(filter(filter_predicate, provider.models))
provider.chat_models = list(filter(filter_predicate, provider.chat_models))
provider.completion_models = list(
filter(filter_predicate, provider.completion_models)
)
provider.models = filter_model_list(provider.models)
provider.chat_models = filter_model_list(provider.chat_models)
provider.completion_models = filter_model_list(provider.completion_models)

# filter out every provider with no models which satisfy the allow/blocklist, then return
return filter((lambda p: len(p.models) > 0), providers)
Expand Down

0 comments on commit 337e4e2

Please sign in to comment.