Added support for local models using Ollama (#36)

Major refactoring to support dynamic construction of the UI menus. This was necessary to support arbitrary model combinations installed via Ollama.
Updated translations.

gepetto/config.ini

@@ -36,3 +36,7 @@ API_KEY =
 # Base URL if you want to redirect requests to a different / local model.
 # Can also be provided via the TOGETHER_BASE_URL environment variable.
 BASE_URL =
+
+[Ollama]
+# Endpoint used to connect to the Ollama API. Default is http://localhost:11434
+HOST =

gepetto/config.py

@@ -2,7 +2,7 @@
 import gettext
 import os
 
-from gepetto.models.base import get_model
+from gepetto.models.model_manager import instantiate_model, load_available_models
 
 model = None
 parsed_ini = None
@@ -28,7 +28,8 @@ def load_config():
 
  # Select model
  requested_model = parsed_ini.get('Gepetto', 'MODEL')
- model = get_model(requested_model)
+ load_available_models()
+ model = instantiate_model(requested_model)
 
 
 def get_config(section, option, environment_variable=None, default=None):
@@ -42,10 +43,16 @@ def get_config(section, option, environment_variable=None, default=None):
  :return: The value of the requested option.
  """
  global parsed_ini
- if parsed_ini and parsed_ini.get(section, option):
- return parsed_ini.get(section, option)
- if environment_variable and os.environ.get(environment_variable):
- return os.environ.get(environment_variable)
+ try:
+ if parsed_ini and parsed_ini.get(section, option):
+ return parsed_ini.get(section, option)
+ if environment_variable and os.environ.get(environment_variable):
+ return os.environ.get(environment_variable)
+ except (configparser.NoSectionError, configparser.NoOptionError):
+ print(_("Warning: Gepetto's configuration doesn't contain option {option} in section {section}!").format(
+ option=option,
+ section=section
+ ))
  return default
 
 

gepetto/ida/handlers.py

@@ -8,7 +8,7 @@
 import idc
 
 import gepetto.config
-from gepetto.models.base import get_model
+from gepetto.models.model_manager import instantiate_model
 
 
 def comment_callback(address, view, response):
@@ -67,14 +67,13 @@ def update(self, ctx):
 
 # -----------------------------------------------------------------------------
 
-def rename_callback(address, view, response, retries=0):
+def rename_callback(address, view, response):
  """
  Callback that extracts a JSON array of old names and new names from the
  response and sets them in the pseudocode.
  :param address: The address of the function to work on
  :param view: A handle to the decompiler window
  :param response: The response from the model
- :param retries: The number of times that we received invalid JSON
  """
  names = json.loads(response)
 
@@ -148,13 +147,13 @@ def __init__(self, new_model, plugin):
 
  def activate(self, ctx):
  try:
- gepetto.config.model = get_model(self.new_model)
+ gepetto.config.model = instantiate_model(self.new_model)
  except ValueError as e: # Raised if an API key is missing. In which case, don't switch.
  print(_("Couldn't change model to {model}: {error}").format(model=self.new_model, error=str(e)))
  return
  gepetto.config.update_config("Gepetto", "MODEL", self.new_model)
  # Refresh the menus to reflect which model is currently selected.
- self.plugin.generate_plugin_select_menu()
+ self.plugin.generate_model_select_menu()
 
  def update(self, ctx):
  return idaapi.AST_ENABLE_ALWAYS

gepetto/ida/ui.py

@@ -1,12 +1,15 @@
+import functools
 import random
 import string
+import threading
 
 import idaapi
 import ida_hexrays
+import ida_kernwin
 
 import gepetto.config
 from gepetto.ida.handlers import ExplainHandler, RenameHandler, SwapModelHandler
-from gepetto.models.base import GPT4_MODEL_NAME, GPT3_MODEL_NAME, GPT4o_MODEL_NAME, GROQ_MODEL_NAME, MISTRAL_MODEL_NAME
+import gepetto.models.model_manager
 
 
 # =============================================================================
@@ -19,24 +22,12 @@ class GepettoPlugin(idaapi.plugin_t):
  explain_menu_path = "Edit/Gepetto/" + _("Explain function")
  rename_action_name = "gepetto:rename_function"
  rename_menu_path = "Edit/Gepetto/" + _("Rename variables")
-
- # Model selection menu
- select_gpt35_action_name = "gepetto:select_gpt35"
- select_gpt4_action_name = "gepetto:select_gpt4"
- select_gpt4o_action_name = "gepetto:select_gpt4o"
- select_groq_action_name = "gepetto:select_groq"
- select_mistral_action_name = "gepetto:select_mistral"
- select_gpt35_menu_path = "Edit/Gepetto/" + _("Select model") + f"/OpenAI/{GPT3_MODEL_NAME}"
- select_gpt4_menu_path = "Edit/Gepetto/" + _("Select model") + f"/OpenAI/{GPT4_MODEL_NAME}"
- select_gpt4o_menu_path = "Edit/Gepetto/" + _("Select model") + f"/OpenAI/{GPT4o_MODEL_NAME}"
- select_groq_menu_path = "Edit/Gepetto/" + _("Select model") + f"/Groq/{GROQ_MODEL_NAME}"
- select_mistral_menu_path = "Edit/Gepetto/" + _("Select model") + f"/Together/{GROQ_MODEL_NAME}"
-
  wanted_name = 'Gepetto'
  wanted_hotkey = ''
  comment = _("Uses {model} to enrich the decompiler's output").format(model=str(gepetto.config.model))
  help = _("See usage instructions on GitHub")
  menu = None
+ model_action_map = {}
 
  # -----------------------------------------------------------------------------
 
@@ -67,7 +58,7 @@ def init(self):
  idaapi.register_action(rename_action)
  idaapi.attach_action_to_menu(self.rename_menu_path, self.rename_action_name, idaapi.SETMENU_APP)
 
- self.generate_plugin_select_menu()
+ self.generate_model_select_menu()
 
  # Register context menu actions
  self.menu = ContextMenuHooks()
@@ -93,43 +84,42 @@ def bind_model_switch_action(self, menu_path, action_name, model_name):
  "",
  "",
  208 if str(gepetto.config.model) == model_name else 0) # Icon #208 == check mark.
- idaapi.register_action(action)
- idaapi.attach_action_to_menu(menu_path, action_name, idaapi.SETMENU_APP)
+ ida_kernwin.execute_sync(functools.partial(idaapi.register_action, action), ida_kernwin.MFF_FAST)
+ ida_kernwin.execute_sync(functools.partial(idaapi.attach_action_to_menu, menu_path, action_name, idaapi.SETMENU_APP),
+ ida_kernwin.MFF_FAST)
 
  # -----------------------------------------------------------------------------
 
  def detach_actions(self):
- idaapi.detach_action_from_menu(self.select_gpt35_menu_path, self.select_gpt35_action_name)
- idaapi.detach_action_from_menu(self.select_gpt4_menu_path, self.select_gpt4_action_name)
- idaapi.detach_action_from_menu(self.select_gpt4o_menu_path, self.select_gpt4o_action_name)
- idaapi.detach_action_from_menu(self.select_groq_menu_path, self.select_groq_action_name)
- idaapi.detach_action_from_menu(self.select_mistral_menu_path, self.select_mistral_action_name)
+ for provider in gepetto.models.model_manager.list_models():
+ for model in provider.supported_models():
+ if model in self.model_action_map:
+ ida_kernwin.execute_sync(functools.partial(idaapi.unregister_action, self.model_action_map[model]),
+ ida_kernwin.MFF_FAST)
+ ida_kernwin.execute_sync(functools.partial(idaapi.detach_action_from_menu,
+ "Edit/Gepetto/" + _("Select model") +
+ f"/{provider.get_menu_name()}/{model}",
+ self.model_action_map[model]),
+ ida_kernwin.MFF_FAST)
 
  # -----------------------------------------------------------------------------
 
- def generate_plugin_select_menu(self):
- # Delete any possible previous entries
- idaapi.unregister_action(self.select_gpt35_action_name)
- idaapi.unregister_action(self.select_gpt4_action_name)
- idaapi.unregister_action(self.select_gpt4o_action_name)
- idaapi.unregister_action(self.select_groq_action_name)
- idaapi.unregister_action(self.select_mistral_action_name)
- self.detach_actions()
-
- # For some reason, IDA seems to have a bug when replacing actions by new ones with identical names.
- # The old action object appears to be reused, at least partially, leading to unwanted behavior?
- # The best workaround I have found is to generate random names each time.
- self.select_gpt35_action_name = f"gepetto:{''.join(random.choices(string.ascii_lowercase, k=7))}"
- self.select_gpt4_action_name = f"gepetto:{''.join(random.choices(string.ascii_lowercase, k=7))}"
- self.select_gpt4o_action_name = f"gepetto:{''.join(random.choices(string.ascii_lowercase, k=7))}"
- self.select_groq_action_name = f"gepetto:{''.join(random.choices(string.ascii_lowercase, k=7))}"
- self.select_mistral_action_name = f"gepetto:{''.join(random.choices(string.ascii_lowercase, k=7))}"
-
- self.bind_model_switch_action(self.select_gpt35_menu_path, self.select_gpt35_action_name, GPT3_MODEL_NAME)
- self.bind_model_switch_action(self.select_gpt4_menu_path, self.select_gpt4_action_name, GPT4_MODEL_NAME)
- self.bind_model_switch_action(self.select_gpt4o_menu_path, self.select_gpt4o_action_name, GPT4o_MODEL_NAME)
- self.bind_model_switch_action(self.select_groq_menu_path, self.select_groq_action_name, GROQ_MODEL_NAME)
- self.bind_model_switch_action(self.select_mistral_menu_path, self.select_mistral_action_name, MISTRAL_MODEL_NAME)
+ def generate_model_select_menu(self):
+ def do_generate_model_select_menu():
+ # Delete any possible previous entries
+ self.detach_actions()
+
+ for provider in gepetto.models.model_manager.list_models():
+ for model in provider.supported_models():
+ # For some reason, IDA seems to have a bug when replacing actions by new ones with identical names.
+ # The old action object appears to be reused, at least partially, leading to unwanted behavior?
+ # The best workaround I have found is to generate random names each time.
+ self.model_action_map[model] = f"gepetto:{model}_{''.join(random.choices(string.ascii_lowercase, k=7))}"
+ self.bind_model_switch_action("Edit/Gepetto/" + _("Select model") + f"/{provider.get_menu_name()}/{model}",
+ self.model_action_map[model],
+ model)
+ # Building the list of available models can take a few seconds with Ollama, don't hang the UI.
+ threading.Thread(target=do_generate_model_select_menu).start()
 
  # -----------------------------------------------------------------------------
 

gepetto/locales/ca_ES/LC_MESSAGES/gepetto.po

@@ -122,3 +122,6 @@ msgstr ""
 
 msgid "Couldn't change model to {model}: {error}"
 msgstr ""
+
+msgid "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"
+msgstr "Advertència: La configuració de Gepetto no conté l'opció {option} a la secció {section}!"

gepetto/locales/es_ES/LC_MESSAGES/gepetto.po

@@ -123,3 +123,6 @@ msgstr "¡Edite este script para insertar su clave de API de {api_provider}!"
 
 msgid "Couldn't change model to {model}: {error}"
 msgstr ""
+
+msgid "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"
+msgstr "Advertencia: La configuración de Gepetto no contiene la opción {option} en la sección {section}!"

gepetto/locales/fr_FR/LC_MESSAGES/gepetto.po

@@ -122,3 +122,6 @@ msgstr "Merci d'ajouter votre clé API {api_provider} dans le fichier de configu
 
 msgid "Couldn't change model to {model}: {error}"
 msgstr "Impossible de choisir {model} comme modèle : {error}"
+
+msgid "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"
+msgstr "Attention: la configuration de Gepetto ne contient pas l'option {option} dans la section {section} !"

gepetto/locales/gepetto.pot

@@ -85,3 +85,6 @@ msgstr "Please edit the configuration file to insert your {api_provider} API key
 
 msgid "Couldn't change model to {model}: {error}"
 msgstr "Couldn't change model to {model}: {error}"
+
+msgid "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"
+msgstr "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"

gepetto/locales/it_IT/LC_MESSAGES/gepetto.po

@@ -119,4 +119,7 @@ msgid "Please edit the configuration file to insert your {api_provider} API key!
 msgstr "Per favore, modifica lo script insendo la tua {api_provider} API key!"
 
 msgid "Couldn't change model to {model}: {error}"
-msgstr ""
+msgstr ""
+
+msgid "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"
+msgstr "Avviso: La configurazione di Gepetto non contiene l'opzione {option} nella sezione {section}!"

gepetto/locales/ko_KR/LC_MESSAGES/gepetto.po

@@ -21,31 +21,31 @@ msgstr ""
 "Plural-Forms: nplurals=1; plural=0;\n"
 
 msgid "Explain function"
-msgstr ""
+msgstr "함수 설명"
 
 msgid "Rename variables"
-msgstr ""
+msgstr "변수 이름 변경"
 
 msgid "Select model"
-msgstr ""
+msgstr "모델 선택"
 
 msgid "Uses {model} to enrich the decompiler's output"
-msgstr ""
+msgstr "{model}을(를) 사용하여 디컴파일러의 출력을 향상시킵니다"
 
 msgid "See usage instructions on GitHub"
-msgstr ""
+msgstr "GitHub에서 사용법을 확인하세요"
 
 msgid "Use {model} to explain the currently selected function"
-msgstr ""
+msgstr "{model}을(를) 사용하여 선택된 함수를 설명합니다"
 
 msgid "Use {model} to rename this function's variables"
-msgstr ""
+msgstr "{model}을(를) 사용하여 이 함수의 변수 이름을 변경합니다"
 
 msgid "Comment generated by Gepetto"
-msgstr ""
+msgstr "Gepetto에서 생성된 주석"
 
 msgid "{model} query finished!"
-msgstr ""
+msgstr "{model} 쿼리가 완료되었습니다!"
 
 msgid ""
 "Can you explain what the following C function does and suggest a better name for it?\n"
@@ -57,52 +57,60 @@ msgstr ""
 msgid ""
 "Could not obtain valid data from the model, giving up. Dumping the response "
 "for manual import:"
-msgstr ""
+msgstr "모델에서 유효한 데이터를 얻지 못했습니다. 수동으로 가져오기 위해 응답을 덤프합니다:"
 
 msgid ""
 "Cannot extract valid JSON from the response. Asking the model to fix it..."
-msgstr ""
+msgstr "응답에서 유효한 JSON을 추출할 수 없습니다. 모델에 수정 요청 중..."
 
 msgid "The JSON document returned is invalid. Asking the model to fix it..."
-msgstr ""
+msgstr "반환된 JSON 문서가 유효하지 않습니다. 모델에 수정 요청 중..."
 
 msgid "Please fix the following JSON document:\n"
-msgstr ""
+msgstr "다음 JSON 문서를 수정해 주세요:\n"
 
 msgid ""
 "The JSON document provided in this response is invalid. Can you fix it?\n"
 "{response}"
 msgstr ""
+"이 응답에서 제공된 JSON 문서가 유효하지 않습니다. 수정해 주실 수 있나요?\n"
+"{response}"
 
 msgid "{model} query finished! {replaced} variable(s) renamed."
-msgstr ""
+msgstr "{model} 쿼리가 완료되었습니다! {replaced}개의 변수가 이름이 변경되었습니다."
 
 msgid ""
 "Analyze the following C function:\n"
 "{decompiler_output}\n"
 "Suggest better variable names, reply with a JSON array where keys are the original names and values are the proposed names. Do not explain anything, only print the JSON dictionary."
 msgstr ""
+"다음 C 함수에 대해 분석해 주세요:\n"
+"{decompiler_output}\n"
+"더 나은 변수 이름을 제안하고, 키는 원래 이름이고 값은 제안된 이름인 JSON 배열로 응답해 주세요. 설명하지 말고 JSON 사전만 출력해 주세요."
 
 msgid "{model} could not complete the request: {error}"
-msgstr ""
+msgstr "{model}이(가) 요청을 완료할 수 없습니다: {error}"
 
 msgid ""
 "Context length exceeded! Reducing the completion tokens to {max_tokens}..."
-msgstr ""
+msgstr "컨텍스트 길이를 초과했습니다! 완료 토큰을 {max_tokens}로 줄이는 중..."
 
 msgid ""
 "Unfortunately, this function is too big to be analyzed with the model's "
 "current API limits."
-msgstr ""
+msgstr "불행히도 이 함수는 모델의 현재 API 제한으로 분석하기에는 너무 큽니다."
 
 msgid "General exception encountered while running the query: {error}"
-msgstr ""
+msgstr "쿼리 실행 중 일반 예외가 발생했습니다: {error}"
 
 msgid "Request to {model} sent..."
-msgstr ""
+msgstr "{model}로 요청이 전송되었습니다..."
 
 msgid "Please edit the configuration file to insert your {api_provider} API key!"
-msgstr ""
+msgstr "설정 파일을 수정하여 {api_provider} API 키를 입력하세요!"
 
 msgid "Couldn't change model to {model}: {error}"
-msgstr ""
+msgstr "{model}로 모델을 변경할 수 없습니다: {error}"
+
+msgid "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"
+msgstr "경고: Gepetto의 설정에 섹션 {section}에서 옵션 {option}이(가) 포함되어 있지 않습니다!"

gepetto/locales/ru/LC_MESSAGES/gepetto.po

@@ -119,4 +119,7 @@ msgstr ""
 "Пожалуйста, отредактируйте этот скрипт, чтобы добавить свой {api_provider} API ключ!"
 
 msgid "Couldn't change model to {model}: {error}"
-msgstr ""
+msgstr ""
+
+msgid "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"
+msgstr "Предупреждение: В конфигурации Gepetto отсутствует опция {option} в разделе {section}!"

gepetto/locales/tr/LC_MESSAGES/gepetto.po

@@ -119,4 +119,7 @@ msgid "Please edit the configuration file to insert your {api_provider} API key!
 msgstr "{api_provider} API anahtarınızı eklemek için lütfen bu betiği düzenleyin!"
 
 msgid "Couldn't change model to {model}: {error}"
-msgstr ""
+msgstr ""
+
+msgid "Warning: Gepetto's configuration doesn't contain option {option} in section {section}!"
+msgstr "Uyarı: Gepetto'nun yapılandırmasında {section} bölümünde {option} seçeneği yok!"