add support for configuring GPT model

README.md

@@ -29,8 +29,11 @@ To do that, run the command
 
 ```chatgdb -k <API KEY> ```
 
-Without the API key, you won't be able to make requests to OpenAI. The API key is stored in
-text in the same directory as the installed script, which is currently in your python site packages
+You also need to set the model to use. There are two possible options, ```gpt-3.5-turbo``` and ```gpt-4```:
+
+```chatgdb -m <MODEL>```
+
+This information is stored in text in the same directory as the installed script, which is currently in your python site packages
 folder along with the main script. You can easily find this location by running the following in your terminal:
 
 ``` python -m site --user-site```

chatgdb/cli.py

@@ -4,14 +4,21 @@
 from urllib.request import Request, urlopen
 import json
 
+PATH = dirname(abspath(getfile(currentframe())))
+
 
 def set_key(key):
  """Set the api key for ChatGDB"""
- path = dirname(abspath(getfile(currentframe()))) + "/.secret.txt"
- with open(path, "w") as f:
+ with open(PATH + "/.secret.txt", "w") as f:
  f.write("OPENAI_KEY=\"" + key + "\"")
 
 
+def set_model(model):
+ """Set the model for ChatGDB"""
+ with open(PATH + "/.model.txt", "w") as f:
+ f.write("MODEL=\"" + model + "\"")
+
+
 def version():
  """Return version information"""
  with urlopen(Request("https://pypi.org/pypi/chatgdb/json"), timeout=10) as f:
@@ -27,6 +34,12 @@ def main():
  "--key",
  type=str,
  help="Provide an api key for ChatGDB")
+ parser.add_argument(
+ '-m',
+ "--model",
+ type=str,
+ choices=["gpt-3.5-turbo", "gpt-4"],
+ help="Provide a model for ChatGDB (gpt-3.5-turbo or gpt-4)")
  parser.add_argument(
  '-v',
  "--version",
@@ -37,6 +50,8 @@ def main():
  args = parser.parse_args()
  if args.key:
  set_key(args.key)
+ elif args.model:
+ set_model(args.model)
  else:
  parser.print_help()
 

chatgdb/utils.py

@@ -7,7 +7,7 @@
 
 
 def get_key():
- """Gets api key from .env file
+ """Gets api key from secret file
 
  Returns: (str) api key
  """
@@ -24,12 +24,36 @@ def get_key():
  secret = k.split('"')[1::2]
  except FileNotFoundError:
  print("Could not find api key. Please make sure you've run the CLI "
- "tool and set up your api key")
+ "tool and set up your model")
  quit("Exiting...")
 
  return secret[0]
 
 
+def get_model():
+ """Gets model from model file
+
+ Returns: (str) model
+ """
+ model = []
+ model_name = ""
+ # gets path of this script - OS independent
+ path = dirname(abspath(getfile(currentframe()))) + "/.model.txt"
+ try:
+ # get appropriate api key
+ with open(path) as f:
+ model = [line.strip() for line in f]
+ for m in model:
+ if m.startswith("MODEL"):
+ model_name = m.split('"')[1::2]
+ except FileNotFoundError:
+ print("Could not find model. Please make sure you've run the CLI "
+ "tool and set up your model")
+ quit("Exiting...")
+
+ return model_name[0]
+
+
 def make_request(url, headers=None, data=None):
  """Makes API request
 
@@ -74,7 +98,6 @@ def chat_help():
  "Content-Type": "application/json"
 }
 URL = "https://api.openai.com/v1/chat/completions"
-MODEL = "gpt-3.5-turbo"
 
 
 def explain_helper(prev_command, command, prompt):
@@ -86,7 +109,7 @@ def explain_helper(prev_command, command, prompt):
  prompt (str): prompt to use for explanation
  """
  question = prompt + prev_command if command == "" else command
- data = {"model": MODEL,
+ data = {"model": get_model(),
  "messages": [{"role": "user",
  "content": question}]}
  body, response = make_request(URL, HEADERS, data=bytes(json.dumps(data),
@@ -98,12 +121,12 @@ def explain_helper(prev_command, command, prompt):
 
 def chat_helper(command, prompt):
  """Generates GDB/LLDB command based on user input
- 
+
  Params:
  command (str): user input
  prompt (str): prompt to use for command generation
  """
- data = {"model": MODEL,
+ data = {"model": get_model(),
  "messages": [{"role": "user",
  "content": prompt + command}]}
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "ChatGDB"
-version = "1.1.0"
+version = "1.2.0"
 authors = ["Pranay Gosar <gosarpranay@gmail.com>"]
 description = "Harness the power of ChatGPT directly inside the GDB debugger!"
 readme = "README.md"