Fix issue where number of tokens was too high.

Gets each model's maximum number of tokens before sending the
query. This should be consolidated into another file, as it will
be shared with Settings in a future release.

Close #14.

Author: WhatTheFuzz

src/agent.py

@@ -105,6 +105,20 @@ 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 instruction_list(self, function: Union[LowLevelILFunction,
                                          MediumLevelILFunction,
                                          HighLevelILFunction]) -> list[str]:
@@ -136,6 +150,6 @@ def send_query(self, query: str) -> str:
         response: str = openai.Completion.create(
             model=self.model,
             prompt=query,
-            max_tokens=2_048
+            max_tokens=self.max_token_count(self.model) - len(query),
         )
         return response.choices[0].text