Skip to content

Commit

Permalink
Implement Favorite Models Feature
Browse files Browse the repository at this point in the history
Fixes #90

Favorite AI models by focusing on the list and using `Shift+Space` or the context menu. Favorites appear at the top.
  • Loading branch information
AAClause committed Mar 30, 2024
1 parent 3f6866f commit 76c5b7d
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 31 deletions.
6 changes: 3 additions & 3 deletions addon/globalPlugins/openai/apikeymanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@

AVAILABLE_PROVIDERS = [
"OpenAI",
"OpenRouter",
"MistralAI"
"MistralAI",
"OpenRouter"
]

_managers = {}
Expand All @@ -24,7 +24,7 @@ def __init__(
self.data_dir = data_dir
self.provider = provider
self.api_key_path = os.path.join(
data_dir,
data_dir,
f"{provider}.key"
)
self.api_key_org_path = os.path.join(
Expand Down
90 changes: 62 additions & 28 deletions addon/globalPlugins/openai/maindialog.py
Original file line number Diff line number Diff line change
Expand Up @@ -546,7 +546,6 @@ def __init__(
)
self.previousPrompt = None
self._lastSystem = None
self._model_ids = [model.id for model in self._models]
if self.conf["saveSystem"]:
# If the user has chosen to save the system prompt, use the last system prompt used by the user as the default value, otherwise use the default system prompt.
if "system" in self.data:
Expand Down Expand Up @@ -727,25 +726,7 @@ def __init__(
self.modelsListCtrl.Bind(wx.EVT_CONTEXT_MENU, self.onModelContextMenu)
self.modelsListCtrl.Bind(wx.EVT_LIST_ITEM_RIGHT_CLICK, self.onModelContextMenu)
self.modelsListCtrl.Bind(wx.EVT_RIGHT_UP, self.onModelContextMenu)

for i, model in enumerate(self._models):
self.modelsListCtrl.InsertItem(i, model.name)
self.modelsListCtrl.SetItem(i, 1, model.provider)
self.modelsListCtrl.SetItem(i, 2, model.id)
self.modelsListCtrl.SetItem(i, 3, str(model.contextWindow))
self.modelsListCtrl.SetItem(
i,
4,
str(model.maxOutputToken) if model.maxOutputToken > 1 else ""
)
model_id = conf["modelVision" if self.pathList else "model"]
model_index = self._getModelIndex(model_id)
self.modelsListCtrl.SetItemState(
model_index,
wx.LIST_STATE_SELECTED | wx.LIST_STATE_FOCUSED,
wx.LIST_STATE_SELECTED | wx.LIST_STATE_FOCUSED
)
self.modelsListCtrl.EnsureVisible(model_index)
self._refreshModelsList()
mainSizer.Add(modelsLabel, 0, wx.ALL, 5)
mainSizer.Add(self.modelsListCtrl, 0, wx.ALL, 5)

Expand Down Expand Up @@ -896,12 +877,6 @@ def __init__(
self.Bind(wx.EVT_CHAR_HOOK, self.onCharHook)
self.Bind(wx.EVT_CLOSE, self.onCancel)


def _getModelIndex(self, model_id):
return list(self._model_ids).index(model_id) if model_id in self._model_ids else (
list(self._model_ids).index(DEFAULT_MODEL_VISION) if self.pathList else 0
)

def addImageToList(
self,
path,
Expand Down Expand Up @@ -1035,11 +1010,65 @@ def showModelDetails(self, evt=None):
True
)

def _getModelIndex(self, model_id):
for i, model in enumerate(self._models):
if model.id == model_id:
return i
return -1

def _refreshModelsList(self, model_to_select=None):
self.modelsListCtrl.DeleteAllItems()
favorite_models = self.data.get("favorite_models", [])
self._models = sorted(
self._models,
key=lambda model: (
not model.id in favorite_models,
apikeymanager.AVAILABLE_PROVIDERS.index(model.provider),
model.id.lower()
),
reverse=False
)
for i, model in enumerate(self._models):
self.modelsListCtrl.InsertItem(i, model.name)
self.modelsListCtrl.SetItem(i, 1, model.provider)
self.modelsListCtrl.SetItem(i, 2, model.id)
self.modelsListCtrl.SetItem(i, 3, str(model.contextWindow))
self.modelsListCtrl.SetItem(
i,
4,
str(model.maxOutputToken) if model.maxOutputToken > 1 else ""
)
model_id = model_to_select or self.conf["modelVision" if self.pathList else "model"]
model_index = self._getModelIndex(model_id)
if model_index == -1:
model_index = 0
self.modelsListCtrl.SetItemState(
model_index,
wx.LIST_STATE_SELECTED | wx.LIST_STATE_FOCUSED,
wx.LIST_STATE_SELECTED | wx.LIST_STATE_FOCUSED
)
self.modelsListCtrl.EnsureVisible(model_index)

def onFavoriteModel(self, evt):
model = self.getCurrentModel()
if not "favorite_models" in self.data:
self.data["favorite_models"] = []
if model.id in self.data.get("favorite_models", []):
self.data["favorite_models"].remove(model.id)
else:
self.data["favorite_models"].append(model.id)
self.saveData(True)
self._refreshModelsList(model.id)

def onModelKeyDown(self, evt):
if evt.GetKeyCode() == wx.WXK_SPACE:
self.showModelDetails()
if evt.GetModifiers() == wx.MOD_SHIFT:
self.onFavoriteModel(evt)
elif evt.GetModifiers() == wx.MOD_NONE:
self.showModelDetails()
else:
evt.Skip()

def onSubmit(self, evt):
if not self.promptTextCtrl.GetValue().strip() and not self.pathList:
self.promptTextCtrl.SetFocus()
Expand Down Expand Up @@ -1660,8 +1689,13 @@ def message(
def onModelContextMenu(self, evt):
menu = wx.Menu()
item_id = wx.NewIdRef()
menu.Append(item_id, _("Show model details") + " (Space)")
menu.Append(item_id, _("Show model &details") + " (Space)")
self.Bind(wx.EVT_MENU, self.showModelDetails, id=item_id)
isFavorite = self.getCurrentModel().id in self.data.get("favorite_models", [])
item_id = wx.NewIdRef()
label = _("Add to &favorites") if not isFavorite else _("Remove from &favorites")
menu.Append(item_id, f"{label} (Shift+Space)")
self.Bind(wx.EVT_MENU, self.onFavoriteModel, id=item_id)
menu.AppendSeparator()
self.modelsListCtrl.PopupMenu(menu)
menu.Destroy()
Expand Down

0 comments on commit 76c5b7d

Please sign in to comment.