Proper TokenTrimming for new models

OpenInterpreter · Sep 27, 2023 · 2324241 · 2324241
1 parent 3671d9f
commit 2324241
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/interpreter/llm/setup_openai_coding_llm.py b/interpreter/llm/setup_openai_coding_llm.py
@@ -2,6 +2,7 @@
 from ..utils.merge_deltas import merge_deltas
 from ..utils.parse_partial_json import parse_partial_json
 from ..utils.convert_to_openai_messages import convert_to_openai_messages
+from ..utils.display_markdown_message import display_markdown_message
 import tokentrim as tt
 
 
@@ -47,7 +48,17 @@ def coding_llm(messages):
         messages = messages[1:]
 
         # Trim messages, preserving the system_message
-        messages = tt.trim(messages=messages, system_message=system_message, model=interpreter.model)
+        try:
+            messages = tt.trim(messages=messages, system_message=system_message, model=interpreter.model)
+        except:
+            if interpreter.context_window:
+                messages = tt.trim(messages=messages, system_message=system_message, max_tokens=interpreter.context_window)
+            else:
+                display_markdown_message("""
+                **We were unable to determine the context window of this model.** Defaulting to 3000.
+                If your model can handle more, run `interpreter --context_window {token limit}` or `interpreter.context_window = {token limit}`.
+                """)
+                messages = tt.trim(messages=messages, system_message=system_message, max_tokens=3000)
 
         if interpreter.debug_mode:
             print("Sending this to the OpenAI LLM:", messages)

diff --git a/interpreter/llm/setup_text_llm.py b/interpreter/llm/setup_text_llm.py
@@ -79,7 +79,11 @@ def base_llm(messages):
             try:
                 messages = tt.trim(messages, system_message=system_message, model=interpreter.model)
             except:
-                # If we don't know the model, just do 3000.
+                display_markdown_message("""
+                **We were unable to determine the context window of this model.** Defaulting to 3000.
+                If your model can handle more, run `interpreter --context_window {token limit}` or `interpreter.context_window = {token limit}`.
+                Also, please set max_tokens: `interpreter --max_tokens {max tokens per response}` or `interpreter.max_tokens = {max tokens per response}`
+                """)
                 messages = tt.trim(messages, system_message=system_message, max_tokens=3000)
 
         if interpreter.debug_mode: