🐛 fall back for model size

Author: Wytamma

write_the/llm.py

@@ -23,7 +23,14 @@ def __init__(self, prompt: PromptTemplate, temperature=0, model_name="gpt-3.5-tu
         self.prompt_size = self.number_of_tokens(prompt.template)
         self.temperature = temperature
         self.model_name = model_name
-        self.max_tokens = int(models[model_name]["context_window"])
+        try:
+            self.max_tokens = int(models[model_name]["context_window"])
+        except KeyError:
+            self.max_tokens = 4096
+            if model_name.startswith('gpt-4'):
+                self.max_tokens = 8192
+            elif model_name.startswith('gpt-3'):
+                self.max_tokens = 4096
 
     async def run(self, code, **kwargs):
         """