Fix truncation logic

llm_toolkit/models.py

@@ -678,22 +678,22 @@ def truncate_prompt(self,
 
     extra_text_token_count = self.estimate_token_num(extra_text)
     # Reserve 10000 tokens for raw prompt wrappers.
-    max_raw_prompt_token_size = (self.MAX_INPUT_TOKEN - extra_text_token_count -
-                                 10000)
-
-    while token_count > max_raw_prompt_token_size // 4:
+    max_raw_prompt_token_size = (self.MAX_INPUT_TOKEN * 0.9 -
+                                 extra_text_token_count) // 4
+    while token_count > max_raw_prompt_token_size:
       estimate_truncate_size = int(
           (1 - max_raw_prompt_token_size / token_count) * len(raw_prompt_text))
 
-      logger.warning('estimate_truncate_size: %s', estimate_truncate_size)
-      raw_prompt_init = raw_prompt_text[:100] + (
+      num_init_tokens = min(100, int(max_raw_prompt_token_size * 0.1))
+      raw_prompt_init = raw_prompt_text[:num_init_tokens] + (
           '\n...(truncated due to exceeding input token limit)...\n')
       raw_prompt_text = raw_prompt_init + raw_prompt_text[
-          100 + estimate_truncate_size + 1:]
+          min(num_init_tokens + estimate_truncate_size + 1,
+              len(raw_prompt_text) - 100):]
 
       token_count = self.estimate_token_num(raw_prompt_text)
-      logger.warning('Truncated raw prompt from %d to %d tokens:',
-                     original_token_count, token_count)
+      logger.warning('Truncated %d raw prompt chars from %d to %d tokens.',
+                     estimate_truncate_size, original_token_count, token_count)
 
     return raw_prompt_text