Skip to content

Commit 4f272d6

Browse files
mahalozh114mx001
andauthored
Update config file feature. Add all settings into a single dialog (redo) (#78)
* Update config file feature. Add all settings into a single dialog. * Fix bug that doesn't load the API key on start. Some refactors on logic of load/save configs * Some fixes * Delay import * Fix setup --------- Co-authored-by: h114mx001 <[email protected]>
1 parent 02108ec commit 4f272d6

File tree

9 files changed

+365
-145
lines changed

9 files changed

+365
-145
lines changed

dailalib/__init__.py

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
__version__ = "3.15.6"
1+
__version__ = "3.16.0"
22

33
import os
44
# stop LiteLLM from querying at all to the remote server
@@ -11,13 +11,17 @@
1111

1212

1313
def create_plugin(*args, **kwargs):
14-
from libbs.api import DecompilerInterface
15-
14+
from libbs.api import DecompilerInterface
1615
#
1716
# LLM API (through LiteLLM api)
1817
#
1918

2019
litellm_api = LiteLLMAIAPI(delay_init=True)
20+
21+
# load config before creating context menus. if not, the config is only be load after "OK" button is
22+
# clicked in ask_settings :X
23+
litellm_api.load_or_create_config()
24+
2125
# create context menus for prompts
2226
gui_ctx_menu_actions = {
2327
f"DAILA/LLM/{prompt_name}": (prompt.desc, getattr(litellm_api, prompt_name))
@@ -27,12 +31,7 @@ def create_plugin(*args, **kwargs):
2731
gui_ctx_menu_actions["DAILA/LLM/chat"] = ("Open LLM Chat...", get_llm_chat_creator(litellm_api))
2832

2933
# create context menus for others
30-
gui_ctx_menu_actions["DAILA/LLM/Settings/update_api_key"] = ("Update API key...", litellm_api.ask_api_key)
31-
gui_ctx_menu_actions["DAILA/LLM/Settings/update_pmpt_style"] = ("Change prompt style...", litellm_api.ask_prompt_style)
32-
gui_ctx_menu_actions["DAILA/LLM/Settings/update_model"] = ("Change model...", litellm_api.ask_model)
33-
gui_ctx_menu_actions["DAILA/LLM/Settings/update_custom_url"] = ("Set Custom OpenAI Endpoint...", litellm_api.ask_custom_endpoint)
34-
gui_ctx_menu_actions["DAILA/LLM/Settings/update_custom_model"] = ("Set Custom OpenAI Model...", litellm_api.ask_custom_model)
35-
34+
gui_ctx_menu_actions["DAILA/LLM/Settings"] = ("Settings...", litellm_api.ask_settings)
3635
#
3736
# VarModel API (local variable renaming)
3837
#

dailalib/api/litellm/__init__.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,41 @@
1+
DEFAULT_MODEL = "gpt-4o"
2+
OPENAI_MODELS = {"gpt-4", "gpt-4o", "gpt-4-turbo", "gpt-3.5-turbo", "o1-mini", "o1-preview"}
3+
# TODO: How can I get this MODEL_TO_TOKENS in the future, without hardcopy to `configuration`
4+
MODEL_TO_TOKENS = {
5+
# TODO: update the token values for o1
6+
"o1-mini": 8_000,
7+
"o1-preview": 8_000,
8+
"gpt-4o": 8_000,
9+
"gpt-4o-mini": 16_000,
10+
"gpt-4-turbo": 128_000,
11+
"claude-3-5-sonnet-20240620": 200_000,
12+
"gemini/gemini-pro": 12_288,
13+
"vertex_ai_beta/gemini-pro": 12_288,
14+
# perplex is on legacy mode :(
15+
"perplexity/llama-3.1-sonar-small-128k-online": 127_072,
16+
"perplexity/llama-3.1-sonar-medium-128k-online": 127_072,
17+
"perplexity/llama-3.1-sonar-large-128k-online": 127_072,
18+
"sonar-pro": 127_072,
19+
"sonar": 127_072,
20+
}
21+
22+
LLM_COST = {
23+
"gpt-4o": {"prompt_price": 2.5, "completion_price": 10},
24+
"gpt-4o-mini": {"prompt_price": 0.150, "completion_price": 0.600},
25+
"gpt-4-turbo": {"prompt_price": 10, "completion_price": 30},
26+
"claude-3.5-sonnet-20240620": {"prompt_price": 3, "completion_price": 15},
27+
"gemini/gemini-pro": {"prompt_price": 0.150, "completion_price": 0.600},
28+
"vertex_ai_beta/gemini-pro": {"prompt_price": 0.150, "completion_price": 0.600},
29+
# perplex is on legacy mode not available from 02/22/25:(
30+
"perplexity/llama-3.1-sonar-small-128k-online": {"prompt_price": 0.150, "completion_price": 0.600},
31+
"perplexity/llama-3.1-sonar-large-128k-online": {"prompt_price": 0.150, "completion_price": 0.600},
32+
"perplexity/llama-3.1-sonar-huge-128k-online": {"prompt_price": 0.150, "completion_price": 0.600},
33+
# introduced the new sonar-pro/sonar
34+
"sonar": {"prompt_price": 0.150, "completion_price": 0.600},
35+
"sonar-pro": {"prompt_price": 0.150, "completion_price": 0.600},
36+
}
37+
38+
# delay import for const creation
139
from .litellm_api import LiteLLMAIAPI
40+
from .prompt_type import PromptType, ALL_STYLES, DEFAULT_STYLE
41+

dailalib/api/litellm/config_dialog.py

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
import logging
2+
from typing import Optional
3+
4+
from dailalib.configuration import DAILAConfig
5+
from .prompt_type import ALL_STYLES
6+
from . import MODEL_TO_TOKENS
7+
8+
from libbs.ui.qt_objects import (
9+
QDialog,
10+
QGridLayout,
11+
QHBoxLayout,
12+
QLabel,
13+
QLineEdit,
14+
QPushButton,
15+
QVBoxLayout,
16+
QComboBox,
17+
)
18+
19+
_l = logging.getLogger(__name__)
20+
AVAILABLE_MODELS = MODEL_TO_TOKENS.keys()
21+
22+
23+
class DAILAConfigDialog(QDialog):
24+
TITLE = "DAILA Configuration"
25+
26+
def __init__(self, config: DAILAConfig, parent=None):
27+
"""
28+
Constructor for the DAILA configuration dialog.
29+
params:
30+
+ config: config object, passed from litellm_api when calling this dialog
31+
"""
32+
33+
super().__init__(parent)
34+
self.configured = False
35+
self.DAILAConfig = config
36+
37+
self.setWindowTitle(self.TITLE)
38+
self._main_layout = QVBoxLayout()
39+
self._grid_layout = QGridLayout()
40+
self.row = 0
41+
42+
self._init_middle_widgets()
43+
self._main_layout.addLayout(self._grid_layout)
44+
45+
self._init_close_btn_widgets()
46+
47+
self.setLayout(self._main_layout)
48+
49+
def _init_middle_widgets(self):
50+
"""
51+
"""
52+
53+
# LLM Model
54+
llm_model = self.DAILAConfig.model
55+
llm_model_label = QLabel("LLM Model:")
56+
llm_model_label.setToolTip("The model to use for LiteLLM.")
57+
58+
# using dropdown for LLM model
59+
self._llm_model_edit = QComboBox(self)
60+
self._llm_model_edit.addItems(AVAILABLE_MODELS)
61+
self._llm_model_edit.setCurrentText(llm_model)
62+
self._grid_layout.addWidget(llm_model_label, self.row, 0)
63+
self._grid_layout.addWidget(self._llm_model_edit, self.row, 1)
64+
self.row += 1
65+
66+
# API Key
67+
68+
api_key = self.DAILAConfig.api_key
69+
api_key_label = QLabel("API Key:")
70+
api_key_label.setToolTip("The API key to use for LiteLLM, for the selected model.")
71+
self._api_key_edit = QLineEdit(self)
72+
self._api_key_edit.setText(api_key)
73+
self._grid_layout.addWidget(api_key_label, self.row, 0)
74+
self._grid_layout.addWidget(self._api_key_edit, self.row, 1)
75+
self.row += 1
76+
77+
# Prompt Style
78+
79+
prompt_style = self.DAILAConfig.prompt_style
80+
prompt_style_label = QLabel("Prompt Style:")
81+
prompt_style_label.setToolTip("The prompt style for DAILA to use, refer to dailalib/litellm/prompts for details.")
82+
83+
# using dropdown for prompt style
84+
self._prompt_style_edit = QComboBox(self)
85+
self._prompt_style_edit.addItems(ALL_STYLES)
86+
self._prompt_style_edit.setCurrentText(prompt_style)
87+
self._grid_layout.addWidget(prompt_style_label, self.row, 0)
88+
self._grid_layout.addWidget(self._prompt_style_edit, self.row, 1)
89+
self.row += 1
90+
91+
# Custom OpenAI Endpoint
92+
93+
custom_endpoint = self.DAILAConfig.custom_endpoint
94+
custom_endpoint_label = QLabel("Custom OpenAI Endpoint:")
95+
custom_endpoint_label.setToolTip("The custom OpenAI endpoint to use for LiteLLM.")
96+
self._custom_endpoint_edit = QLineEdit(self)
97+
self._custom_endpoint_edit.setText(custom_endpoint)
98+
self._grid_layout.addWidget(custom_endpoint_label, self.row, 0)
99+
self._grid_layout.addWidget(self._custom_endpoint_edit, self.row, 1)
100+
self.row += 1
101+
102+
# Custom OpenAI Model
103+
104+
custom_model = self.DAILAConfig.custom_model
105+
custom_model_label = QLabel("Custom OpenAI Model:")
106+
custom_model_label.setToolTip("The custom OpenAI model to use for LiteLLM.")
107+
self._custom_model_edit = QLineEdit(self)
108+
self._custom_model_edit.setText(custom_model)
109+
self._grid_layout.addWidget(custom_model_label, self.row, 0)
110+
self._grid_layout.addWidget(self._custom_model_edit, self.row, 1)
111+
self.row += 1
112+
113+
def _init_close_btn_widgets(self):
114+
# buttons
115+
self._ok_button = QPushButton(self)
116+
self._ok_button.setText("OK")
117+
self._ok_button.setDefault(True)
118+
self._ok_button.clicked.connect(self._on_ok_clicked)
119+
120+
cancel_button = QPushButton(self)
121+
cancel_button.setText("Cancel")
122+
cancel_button.clicked.connect(self._on_cancel_clicked)
123+
124+
buttons_layout = QHBoxLayout()
125+
buttons_layout.addWidget(self._ok_button)
126+
buttons_layout.addWidget(cancel_button)
127+
128+
self._main_layout.addLayout(buttons_layout)
129+
130+
def _on_cancel_clicked(self):
131+
self.close()
132+
133+
def parse_api_key(self, api_key_or_path: str) -> Optional[str]:
134+
"""
135+
Parse the API key from the input string.
136+
"""
137+
if "/" in api_key_or_path or "\\" in api_key_or_path:
138+
# treat as path
139+
with open(api_key_or_path, "r") as f:
140+
api_key = f.read().strip()
141+
else:
142+
api_key = api_key_or_path
143+
return api_key
144+
145+
def _on_ok_clicked(self):
146+
self.DAILAConfig.model = self._llm_model_edit.currentText()
147+
self.DAILAConfig.api_key = self.parse_api_key(self._api_key_edit.text())
148+
self.DAILAConfig.prompt_style = self._prompt_style_edit.currentText()
149+
self.DAILAConfig.custom_endpoint = self._custom_endpoint_edit.text()
150+
self.DAILAConfig.custom_model = self._custom_model_edit.text()
151+
self.configured = True
152+
self.close()
153+
154+
def config_dialog_exec(self):
155+
self.exec()
156+
if not self.configured:
157+
_l.warning("DAILA Configuration dialog was closed without saving changes.")
158+
else:
159+
_l.info("DAILA Configuration dialog was closed and changes were saved.")
160+
return self.DAILAConfig

0 commit comments

Comments
 (0)