mito-ai 0.1.45__py3-none-any.whl → 0.1.46__py3-none-any.whl

mito_ai/utils/anthropic_utils.py

@@ -3,7 +3,7 @@
 
 import anthropic
 from typing import Any, Dict, List, Optional, Union, AsyncGenerator, Tuple, Callable
-from anthropic.types import MessageParam, ToolUnionParam
+from anthropic.types import MessageParam, TextBlockParam, ToolUnionParam
 from mito_ai.utils.mito_server_utils import get_response_from_mito_server, stream_response_from_mito_server
 from mito_ai.utils.provider_utils import does_message_require_fast_model
 from mito_ai.completions.models import AgentResponse, MessageType, ResponseFormatInfo, CompletionReply, CompletionStreamChunk
@@ -23,7 +23,7 @@ def _prepare_anthropic_request_data_and_headers(
     model: Union[str, None],
     max_tokens: int,
     temperature: float,
-    system: Union[str, anthropic.Omit],
+    system: Union[str, List[TextBlockParam], anthropic.Omit],
     messages: List[MessageParam],
     message_type: MessageType,
     tools: Optional[List[ToolUnionParam]],
@@ -92,7 +92,7 @@ async def stream_anthropic_completion_from_mito_server(
     model: Union[str, None],
     max_tokens: int,
     temperature: float,
-    system: Union[str, anthropic.Omit],
+    system: Union[str, List[TextBlockParam], anthropic.Omit],
     messages: List[MessageParam],
     stream: bool,
     message_type: MessageType,
@@ -102,6 +102,7 @@ async def stream_anthropic_completion_from_mito_server(
     data, headers = _prepare_anthropic_request_data_and_headers(
         model, max_tokens, temperature, system, messages, message_type, None, None, stream
     )
+    
     # Use the unified streaming function
     # If the reply_fn and message_id are empty, this function still handles those requests. This is particularly needed for the streamlit dashboard functionality
     actual_reply_fn = reply_fn if reply_fn is not None else (lambda x: None)
@@ -119,19 +120,20 @@ async def stream_anthropic_completion_from_mito_server(
         provider_name="Claude",
     ):
         yield chunk
-
+    
+    
 def get_anthropic_completion_function_params(
     message_type: MessageType,
     model: str,
     messages: List[MessageParam],
     max_tokens: int,
-    system: Union[str, anthropic.Omit],
+    system: Union[str, List[TextBlockParam], anthropic.Omit],
     temperature: float = 0.0,
     tools: Optional[List[ToolUnionParam]] = None,
     tool_choice: Optional[dict] = None,
     stream: Optional[bool] = None,
     response_format_info: Optional[ResponseFormatInfo] = None,
-) -> Dict[str, Any]:
+) -> Dict[str, Any]:    
     """
     Build the provider_data dict for Anthropic completions, mirroring the OpenAI approach.
     Only includes fields needed for the Anthropic API.

mito_ai/utils/message_history_utils.py

@@ -3,6 +3,7 @@
 
 import re
 from typing import List
+from mito_ai.constants import MESSAGE_HISTORY_TRIM_THRESHOLD
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.prompt_builders.prompt_constants import (
     ACTIVE_CELL_ID_SECTION_HEADING,
@@ -43,18 +44,18 @@ def trim_sections_from_message_content(content: str) -> str:
     return content
 
 
-def trim_old_messages(messages: List[ChatCompletionMessageParam], keep_recent: int = 3) -> List[ChatCompletionMessageParam]:
+def trim_old_messages(messages: List[ChatCompletionMessageParam]) -> List[ChatCompletionMessageParam]:
     """
     Trims metadata sections from messages that are older than the specified number of recent messages.
     We do this in order to reduce the token count of the messages, which helps us stay under the token limit for the LLM.
     """
-    if len(messages) <= keep_recent:
+    if len(messages) <= MESSAGE_HISTORY_TRIM_THRESHOLD:
         return messages
         
     # Process all messages except the keep_recent most recent ones. 
     # Only trim user messages, which is where this metadata lives. 
     # We want to not edit the system messages, as they contain important information / examples.
-    for i in range(len(messages) - keep_recent):
+    for i in range(len(messages) - MESSAGE_HISTORY_TRIM_THRESHOLD):
         content = messages[i].get("content")
         
         is_user_message = messages[i].get("role") == "user"

mito_ai/utils/telemetry_utils.py

@@ -386,10 +386,13 @@ def log_mito_server_free_tier_limit_reached(key_type: Literal['mito_server_key',
 # Streamlit Conversion
 #################################
 
-def log_streamlit_app_creation_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType) -> None:
+def log_streamlit_app_creation_success(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, edit_prompt: str) -> None:
     log(
         "mito_ai_streamlit_app_creation_success", 
-        key_type=key_type
+        key_type=key_type,
+        params={
+            "edit_prompt": edit_prompt,
+        }
     )
     
 def log_streamlit_app_creation_retry(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
@@ -401,12 +404,12 @@ def log_streamlit_app_creation_retry(key_type: Literal['mito_server_key', 'user_
         key_type=key_type
     )
     
-def log_streamlit_app_creation_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str) -> None:
+def log_streamlit_app_creation_error(key_type: Literal['mito_server_key', 'user_key'], message_type: MessageType, error: str, edit_prompt: str) -> None:
     log(
         "mito_ai_streamlit_app_creation_error", 
         params={
             "error_message": error,
+            "edit_prompt": edit_prompt,
         },
         key_type=key_type
     )
-    

{mito_ai-0.1.45.data → mito_ai-0.1.46.data}/data/share/jupyter/labextensions/mito_ai/build_log.json

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

{mito_ai-0.1.45.data → mito_ai-0.1.46.data}/data/share/jupyter/labextensions/mito_ai/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.45",
+  "version": "0.1.46",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -140,7 +140,7 @@
     "outputDir": "mito_ai/labextension",
     "schemaDir": "schema",
     "_build": {
-      "load": "static/remoteEntry.684f82575fcc2e3b350c.js",
+      "load": "static/remoteEntry.54126ab6511271265443.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.45.data → mito_ai-0.1.46.data}/data/share/jupyter/labextensions/mito_ai/schemas/mito_ai/package.json.orig

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