Merge pull request #28 from WhatTheFuzz/feature/new-models

Feature/new models

Author: WhatTheFuzz

src/query.py

@@ -22,16 +22,25 @@ def run(self) -> None:
 
         log_debug(f'Sending query: {self.query_string}')
 
-        response = openai.Completion.create(
-            model=self.model,
-            prompt=self.query_string,
-            max_tokens=self.max_token_count,
-        )
-        # Get the response text.
-        result: str = response.choices[0].text
+        if self.model in ["gpt-3.5-turbo","gpt-4","gpt-4-32k"]:
+            response = openai.ChatCompletion.create(
+                model=self.model,
+                messages=[{"role":"user","content":self.query_string}],
+                max_tokens=self.max_token_count,
+            )
+            # Get the response text.
+            result: str = response.choices[0].message.content
+        else:
+            response = openai.Completion.create(
+                model=self.model,
+                prompt=self.query_string,
+                max_tokens=self.max_token_count,
+            )
+            # Get the response text.
+            result: str = response.choices[0].text
         # If there is a callback, do something with it.
         if self.callback:
             self.callback(result)
         # Otherwise, assume we just want to log it.
         else:
-            log_info(result)
+            log_info(result)

src/settings.py

@@ -42,7 +42,7 @@ def register_api_key_settings(self) -> bool:
 
     def register_model_settings(self) -> bool:
         '''Register the OpenAI model settings in Binary Ninja.
-        Defaults to text-davinci-003.
+        Defaults to gpt-3.5-turbo.
         '''
         # Set the attributes of the settings. Refer to:
         # https://api.binary.ninja/binaryninja.settings-module.html
@@ -52,22 +52,30 @@ def register_model_settings(self) -> bool:
             'description': 'The OpenAI model used to generate the response.',
             # https://beta.openai.com/docs/models
             'enum': [
+                'gpt-4',
+                'gpt-4-32k',
+                'gpt-3.5-turbo',
                 'text-davinci-003',
+                'text-davinci-002',
                 'text-curie-001',
                 'text-babbage-001',
                 'text-babbage-002',
                 'code-davinci-002',
                 'code-cushman-001'
             ],
             'enumDescriptions': [
-                'Most capable GPT-3 model. Can do any task the other models can do, often with higher quality, longer output and better instruction-following. Also supports inserting completions within text.',
+                'More capable than any GPT-3.5 model, able to do more complex tasks, and optimized for chat. Will be updated with our latest model iteration.',
+                'Same capabilities as the base gpt-4 mode but with 4x the context length. Will be updated with our latest model iteration.',
+                'Most capable GPT-3.5 model and optimized for chat at 1/10th the cost of text-davinci-003. Will be updated with our latest model iteration.',
+                'Can do any language task with better quality, longer output, and consistent instruction-following than the curie, babbage, or ada models.',
+                'Similar capabilities to text-davinci-003 but trained with supervised fine-tuning instead of reinforcement learning'
                 'Very capable, but faster and lower cost than Davinci.',
                 'Capable of straightforward tasks, very fast, and lower cost.',
                 'Capable of very simple tasks, usually the fastest model in the GPT-3 series, and lowest cost.',
                 'Most capable Codex model. Particularly good at translating natural language to code. In addition to completing code, also supports inserting completions within code.',
                 'Almost as capable as Davinci Codex, but slightly faster. This speed advantage may make it preferable for real-time applications.'
             ],
-            'default': 'text-davinci-003'
+            'default': 'gpt-3.5-turbo'
         }
         return self.register_setting('openai.model', json.dumps(properties))
 
@@ -81,11 +89,11 @@ def register_max_tokens(self) -> bool:
         properties = {
             'title': 'OpenAI Max Completion Tokens',
             'type': 'number',
-            'description': 'The maximum number of tokens used for completion. Tokens do not necessarily align with word or instruction count. Typically, each token is four characters. If your function is very large, you may need to decrease this value, as the number of tokens in your prompt counts against the total number of tokens supported by the model. Not all models support the same number of maximum tokens; most support 2,048 tokens. For larger functions, check out text-davinci-003 and code-davinci-002 which support 4,000 and 8,000 respectively.',
+            'description': 'The maximum number of tokens used for completion. Tokens do not necessarily align with word or instruction count. Typically, each token is four characters. If your function is very large, you may need to decrease this value, as the number of tokens in your prompt counts against the total number of tokens supported by the model. Not all models support the same number of maximum tokens; most support 2,048 tokens. For larger functions, check out text-davinci-003 and code-davinci-002 which support 4,000 and 8,000 respectively. For the maximum amount of tokens, check out gpt-4-32k which supports up to 32,768 tokens. This may be costly, however.',
             'default': 1_024,
             'minValue': 1,
-            'maxValue': 8_000,
-            'message': "Min: 1, Max: 8,000"
+            'maxValue': 32_768,
+            'message': "Min: 1, Max: 32,768"
         }
         return self.register_setting('openai.max_tokens',
                                      json.dumps(properties))