fix: max tokens max is 8192 now

open_notebook/graphs/chat.py

@@ -42,8 +42,7 @@ def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict
                     str(payload),
                     model_id,
                     "chat",
-                    max_tokens=10000,
+                    max_tokens=8192
-                )
             )
         finally:
             new_loop.close()
@@ -64,7 +63,7 @@ def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict
                 str(payload),
                 model_id,
                 "chat",
-                max_tokens=10000,
+                max_tokens=8192,
             )
         )
     

open_notebook/graphs/source_chat.py

@@ -26,7 +26,9 @@ class SourceChatState(TypedDict):
     context_indicators: Optional[Dict[str, List[str]]]
 
 
-def call_model_with_source_context(state: SourceChatState, config: RunnableConfig) -> dict:
+def call_model_with_source_context(
+    state: SourceChatState, config: RunnableConfig
+) -> dict:
     """
     Main function that builds source context and calls the model.
 
@@ -50,7 +52,7 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
                 source_id=source_id,
                 include_insights=True,
                 include_notes=False,  # Focus on source-specific content
-                max_tokens=50000  # Reasonable limit for source context
+                max_tokens=50000,  # Reasonable limit for source context
             )
             return new_loop.run_until_complete(context_builder.build())
         finally:
@@ -63,6 +65,7 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
         asyncio.get_running_loop()
         # If we're in an event loop, run in a thread with a new loop
         import concurrent.futures
+
         with concurrent.futures.ThreadPoolExecutor() as executor:
             future = executor.submit(build_context)
             context_data = future.result()
@@ -73,7 +76,11 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
     # Extract source and insights from context
     source = None
     insights = []
-    context_indicators: dict[str, list[str | None]] = {"sources": [], "insights": [], "notes": []}
+    context_indicators: dict[str, list[str | None]] = {
+        "sources": [],
+        "insights": [],
+        "notes": [],
+    }
 
     if context_data.get("sources"):
         source_info = context_data["sources"][0]  # First source
@@ -82,7 +89,11 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
 
     if context_data.get("insights"):
         for insight_data in context_data["insights"]:
-            insight = SourceInsight(**insight_data) if isinstance(insight_data, dict) else insight_data
+            insight = (
+                SourceInsight(**insight_data)
+                if isinstance(insight_data, dict)
+                else insight_data
+            )
             insights.append(insight)
             context_indicators["insights"].append(insight.id)
 
@@ -94,7 +105,7 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
         "source": source.model_dump() if source else None,
         "insights": [insight.model_dump() for insight in insights] if insights else [],
         "context": formatted_context,
-        "context_indicators": context_indicators
+        "context_indicators": context_indicators,
     }
 
     # Apply the source_chat prompt template
@@ -110,9 +121,10 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
             return new_loop.run_until_complete(
                 provision_langchain_model(
                     str(payload),
-                    config.get("configurable", {}).get("model_id") or state.get("model_override"),
+                    config.get("configurable", {}).get("model_id")
+                    or state.get("model_override"),
                     "chat",
-                    max_tokens=10000,
+                    max_tokens=8192,
                 )
             )
         finally:
@@ -124,6 +136,7 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
         asyncio.get_running_loop()
         # If we're in an event loop, run in a thread with a new loop
         import concurrent.futures
+
         with concurrent.futures.ThreadPoolExecutor() as executor:
             future = executor.submit(run_in_new_loop)
             model = future.result()
@@ -132,9 +145,10 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
         model = asyncio.run(
             provision_langchain_model(
                 str(payload),
-                config.get("configurable", {}).get("model_id") or state.get("model_override"),
+                config.get("configurable", {}).get("model_id")
+                or state.get("model_override"),
                 "chat",
-                max_tokens=10000,
+                max_tokens=8192,
             )
         )
 
@@ -146,7 +160,7 @@ def call_model_with_source_context(state: SourceChatState, config: RunnableConfi
         "source": source,
         "insights": insights,
         "context": formatted_context,
-        "context_indicators": context_indicators
+        "context_indicators": context_indicators,
     }
 
 
@@ -183,8 +197,12 @@ def _format_source_context(context_data: Dict) -> str:
         for insight in context_data["insights"]:
             if isinstance(insight, dict):
                 context_parts.append(f"**Insight ID:** {insight.get('id', 'Unknown')}")
-                context_parts.append(f"**Type:** {insight.get('insight_type', 'Unknown')}")
+                context_parts.append(
-                context_parts.append(f"**Content:** {insight.get('content', 'No content')}")
+                    f"**Type:** {insight.get('insight_type', 'Unknown')}"
+                )
+                context_parts.append(
+                    f"**Content:** {insight.get('content', 'No content')}"
+                )
                 context_parts.append("")  # Empty line for separation
 
     # Add metadata

open_notebook/utils/text_utils.py

@@ -12,7 +12,7 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 from .token_utils import token_count
 
 # Pattern for matching thinking content in AI responses
-THINK_PATTERN = re.compile(r'<think>(.*?)</think>', re.DOTALL)
+THINK_PATTERN = re.compile(r"<think>(.*?)</think>", re.DOTALL)
 
 
 def split_text(txt: str, chunk_size=500):
@@ -114,7 +114,7 @@ def parse_thinking_content(content: str) -> Tuple[str, str]:
     cleaned_content = THINK_PATTERN.sub("", content)
 
     # Clean up extra whitespace
-    cleaned_content = re.sub(r'\n\s*\n\s*\n', '\n\n', cleaned_content).strip()
+    cleaned_content = re.sub(r"\n\s*\n\s*\n", "\n\n", cleaned_content).strip()
 
     return thinking_content, cleaned_content