mito-ai 0.1.52__py3-none-any.whl → 0.1.53__py3-none-any.whl

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.52'
+__version__ = VERSION = '0.1.53'

mito_ai/anthropic_client.py

@@ -9,12 +9,11 @@ from anthropic.types import Message, MessageParam, TextBlockParam
 from mito_ai.completions.models import ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
 from mito_ai.constants import MESSAGE_HISTORY_TRIM_THRESHOLD
 from openai.types.chat import ChatCompletionMessageParam
-from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_server, stream_anthropic_completion_from_mito_server, get_anthropic_completion_function_params
+from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_server, select_correct_model, stream_anthropic_completion_from_mito_server, get_anthropic_completion_function_params
 
 # Max tokens is a required parameter for the Anthropic API. 
 # We set it to a high number so that we can edit large code cells
-# 8192 is the maximum allowed number of output tokens for claude-3-5-haiku-20241022
-MAX_TOKENS = 8_000
+MAX_TOKENS = 64_000
 
 def extract_and_parse_anthropic_json_response(response: Message) -> Union[object, Any]:
     """
@@ -278,6 +277,8 @@ class AnthropicClient:
                               reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]) -> str:
         try:
             anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
+            model = select_correct_model(model, message_type, anthropic_system_prompt, anthropic_messages)
+
             accumulated_response = ""
 
             if self.api_key:

mito_ai/utils/anthropic_utils.py

@@ -10,6 +10,7 @@ from mito_ai.completions.models import AgentResponse, MessageType, ResponseForma
 from mito_ai.utils.schema import UJ_STATIC_USER_ID, UJ_USER_EMAIL
 from mito_ai.utils.db import get_user_field
 from mito_ai.constants import MITO_ANTHROPIC_URL
+from mito_ai.utils.tokens import get_rough_token_estimatation_anthropic
 
 __user_email: Optional[str] = None
 __user_id: Optional[str] = None
@@ -17,7 +18,29 @@ __user_id: Optional[str] = None
 ANTHROPIC_TIMEOUT = 60
 max_retries = 1
 
-FAST_ANTHROPIC_MODEL = "claude-3-5-haiku-latest"
+FAST_ANTHROPIC_MODEL = "claude-haiku-4-5-20251001" # This should be in sync with ModelSelector.tsx
+LARGE_CONTEXT_MODEL = "claude-sonnet-4-5-20250929" # This should be in sync with ModelSelector.tsx
+
+def does_message_exceed_max_tokens(system: Union[str, List[TextBlockParam], anthropic.Omit], messages: List[MessageParam]) -> bool:
+    token_estimation = get_rough_token_estimatation_anthropic(system, messages)
+
+    if token_estimation is not None and token_estimation > 200_000:
+        return True
+    return False
+
+def select_correct_model(default_model: str, message_type: MessageType, system: Union[str, List[TextBlockParam], anthropic.Omit], messages: List[MessageParam]) -> str:
+    
+    message_exceeds_fast_model_context_limit = does_message_exceed_max_tokens(system, messages)
+    if message_exceeds_fast_model_context_limit:
+        # Anthropic lets us use beta mode to extend context window for sonnet class models
+        # but not haiku models
+        return LARGE_CONTEXT_MODEL
+    
+    message_requires_fast_model = does_message_require_fast_model(message_type)
+    if message_requires_fast_model:
+        return FAST_ANTHROPIC_MODEL
+    
+    return default_model    
 
 def _prepare_anthropic_request_data_and_headers(
     model: Union[str, None],
@@ -36,6 +59,7 @@ def _prepare_anthropic_request_data_and_headers(
         __user_email = get_user_field(UJ_USER_EMAIL)
     if __user_id is None:
         __user_id = get_user_field(UJ_STATIC_USER_ID)
+    
     # Build the inner data dict (excluding timeout, max_retries, email, user_id)
     inner_data: Dict[str, Any] = {
         "model": model,
@@ -44,6 +68,7 @@ def _prepare_anthropic_request_data_and_headers(
         "messages": messages,
         "betas": ["context-1m-2025-08-07"]
     }
+    
     # Add system to inner_data only if it is not anthropic.Omit
     if not isinstance(system, anthropic.Omit):
         inner_data["system"] = system
@@ -139,8 +164,7 @@ def get_anthropic_completion_function_params(
     Only includes fields needed for the Anthropic API.
     """
     
-    message_requires_fast_model = does_message_require_fast_model(message_type)
-    model = FAST_ANTHROPIC_MODEL if message_requires_fast_model else model
+    model = select_correct_model(model, message_type, system, messages)
     
     provider_data = {
         "model": model,
@@ -166,3 +190,4 @@ def get_anthropic_completion_function_params(
         provider_data["stream"] = stream
     # Optionally handle response_format_info if Anthropic supports it in the future
     return provider_data
+

mito_ai/utils/tokens.py

@@ -0,0 +1,29 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Union, Optional
+import anthropic
+from anthropic.types import MessageParam, TextBlockParam, ToolUnionParam
+
+
+def get_rough_token_estimatation_anthropic(system_message: Union[str, List[TextBlockParam], anthropic.Omit], messages: List[MessageParam]) -> Optional[float]: 
+    """
+    Get a very rough estimation of the number of tokens in a conversation. 
+    We bias towards overestimating to make sure we don't accidentally
+    think a conversation is safe to send to an AI without having applied an
+    optimization strategy. 
+    """
+    
+    try:
+        stringified_system_message = str(system_message)
+        stringified_messages = str(messages)
+        total_stringified_context = stringified_system_message + stringified_messages 
+        
+        # The general rule of thumb is: 1 token is about 4 characters. 
+        # To be safe we use:            1 token is about 3 characters
+        # This helps make sure we always overestimate
+        return len(total_stringified_context) / 3
+    
+    except: 
+        return None
+

{mito_ai-0.1.52.data → mito_ai-0.1.53.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

@@ -720,7 +720,7 @@
             "semver": {},
             "vscode-diff": {},
             "mito_ai": {
-              "version": "0.1.52",
+              "version": "0.1.53",
               "singleton": true,
               "import": "/home/runner/work/mito/mito/mito-ai/lib/index.js"
             }

{mito_ai-0.1.52.data → mito_ai-0.1.53.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.52",
+  "version": "0.1.53",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -140,7 +140,7 @@
     "outputDir": "mito_ai/labextension",
     "schemaDir": "schema",
     "_build": {
-      "load": "static/remoteEntry.89927e1d3b5962d57ae3.js",
+      "load": "static/remoteEntry.4395ab9342efa39fc0a2.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.52.data → mito_ai-0.1.53.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.52",
+  "version": "0.1.53",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",