From 2aa453f73f5f223d8126e1d4d79e900518204ed8 Mon Sep 17 00:00:00 2001
From: LUIS NOVO <lfnovo@gmail.com>
Date: Wed, 30 Oct 2024 15:39:45 -0300
Subject: [PATCH] enable support for large context fallback

---
 open_notebook/graphs/utils.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/open_notebook/graphs/utils.py b/open_notebook/graphs/utils.py
index 921afa5..f004e7f 100644
--- a/open_notebook/graphs/utils.py
+++ b/open_notebook/graphs/utils.py
@@ -1,8 +1,10 @@
 from langchain.output_parsers import OutputFixingParser
+from loguru import logger
 
 from open_notebook.config import DEFAULT_MODELS
 from open_notebook.models import get_model
 from open_notebook.prompter import Prompter
+from open_notebook.utils import token_count
 
 
 def run_pattern(
@@ -13,7 +15,22 @@ def run_pattern(
     parser=None,
     output_fixing_model_name=None,
 ) -> dict:
-    if not model_name:
+    system_prompt = Prompter(prompt_template=pattern_name, parser=parser).render(
+        data=state
+    )
+
+    tokens = token_count(str(system_prompt) + str(messages))
+    if tokens > 105_000 and DEFAULT_MODELS.large_context_model:
+        model_name = DEFAULT_MODELS.large_context_model
+        logger.debug(
+            f"Using large context model ({model_name}) because the content has {tokens} tokens"
+        )
+        logger.warning(system_prompt)
+    elif tokens > 105_000 and not DEFAULT_MODELS.large_context_model:
+        logger.critical(
+            f"Content has {tokens} tokens, but no large context model is configured"
+        )
+    elif not model_name:
         model_name = DEFAULT_MODELS.default_transformation_model
 
     chain = get_model(model_name, model_type="language")
@@ -27,10 +44,6 @@ def run_pattern(
             llm=output_fix_model,
         )
 
-    system_prompt = Prompter(prompt_template=pattern_name, parser=parser).render(
-        data=state
-    )
-
     if len(messages) > 0:
         response = chain.invoke([system_prompt] + messages)
     else: