Merge pull request #17 from WhatTheFuzz/feature/max_tokens

Allow users to set max_tokens in Binary Ninja settings

Author: WhatTheFuzz

plugin.json

@@ -27,6 +27,6 @@
          "openai"
       ]
 	},
-   "version": "1.1.0",
+   "version": "1.2.0",
    "minimumbinaryninjaversion": 3200
 }

src/agent.py

@@ -105,19 +105,16 @@ def get_model(self) -> str:
         assert self.is_valid_model('text-davinci-003')
         return 'text-davinci-003'
 
-    def max_token_count(self, model: str) -> int:
-        '''Returns the maximum number of tokens that can be generated by the
-        model. Returns a default of 2,048 if the model is not found. '''
-        # TODO: This should be somewhere else, as it's also shared by Settings.
-        models: dict[str, int] = {
-            'text-davinci-003': 4_000,
-            'text-curie-001': 2_048,
-            'text-babbage-001': 2_048,
-            'text-ada-001': 2_048,
-            'code-davinci-002': 8_000,
-            'code-cushman-001': 2_048
-        }
-        return models.get(model, 2_048)
+    def get_token_count(self) -> int:
+        '''Returns the maximum token count specified by the user. If no value is
+        set, for whatever reason, returns 1,024.'''
+        settings: Settings = Settings()
+        # Check that the key exists.
+        if settings.contains('openai.max_tokens'):
+            # Check that the value is not None.
+            if (max_tokens := settings.get_integer('openai.max_tokens')) is not None:
+                return max_tokens
+        return 1_024
 
     def instruction_list(self, function: Union[LowLevelILFunction,
                                          MediumLevelILFunction,
@@ -150,6 +147,6 @@ def send_query(self, query: str) -> str:
         response: str = openai.Completion.create(
             model=self.model,
             prompt=query,
-            max_tokens=self.max_token_count(self.model) - len(query),
+            max_tokens=self.get_token_count(),
         )
         return response.choices[0].text

src/settings.py

@@ -22,6 +22,12 @@ def __init__(self) -> None:
             raise RegisterSettingsKeyException('Failed to register OpenAI '
                                                'model settings.')
 
+        # Register the setting for the max tokens used for both the prompt and
+        # completion.
+        if not self.register_max_tokens():
+            raise RegisterSettingsKeyException('Failed to register OpenAI '
+                                               'max tokens settings.')
+
     def register_api_key_settings(self) -> bool:
         '''Register the OpenAI API key settings in Binary Ninja.'''
         # Set the attributes of the settings. Refer to:
@@ -64,3 +70,22 @@ def register_model_settings(self) -> bool:
             'default': 'text-davinci-003'
         }
         return self.register_setting('openai.model', json.dumps(properties))
+
+    def register_max_tokens(self) -> bool:
+        '''Register the OpenAI max tokens used for both the prompt and
+        completion. Defaults to 2,048. The Davinci model can use 4,000 or 8,000
+        tokens for GPT and Codex respectively. Check out the documentation here:
+        https://help.openai.com/en/articles/4936856-what-are-tokens-and-how-to-count-them
+        '''
+
+        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.',
+            'default': 1_024,
+            'minValue': 1,
+            'maxValue': 8_000,
+            'message': "Min: 1, Max: 8,000"
+        }
+        return self.register_setting('openai.max_tokens',
+                                     json.dumps(properties))