mito-ai 0.1.32__py3-none-any.whl → 0.1.34__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.32'
+__version__ = VERSION = '0.1.34'

mito_ai/anthropic_client.py

@@ -6,7 +6,8 @@ import anthropic
 from typing import Dict, Any, Optional, Tuple, Union, Callable, List, cast
 
 from anthropic.types import Message, MessageParam
-from mito_ai.completions.models import ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
+from mito_ai.completions.models import CompletionError, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
 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
 
@@ -15,8 +16,6 @@ from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_ser
 # 8192 is the maximum allowed number of output tokens for claude-3-5-haiku-20241022
 MAX_TOKENS = 8_000
 
-ANTHROPIC_FAST_MODEL = "claude-3-5-haiku-latest"
-
 def extract_and_parse_anthropic_json_response(response: Message) -> Union[object, Any]:
     """
     Extracts and parses the JSON response from the Claude API.
@@ -131,9 +130,8 @@ class AnthropicClient:
     A client for interacting with the Anthropic API or the Mito server fallback.
     """
 
-    def __init__(self, api_key: Optional[str], model: str, timeout: int = 30, max_retries: int = 1):
+    def __init__(self, api_key: Optional[str], timeout: int = 30, max_retries: int = 1):
         self.api_key = api_key
-        self.model = model
         self.timeout = timeout
         self.max_retries = max_retries
         self.client: Optional[anthropic.Anthropic]
@@ -142,57 +140,56 @@ class AnthropicClient:
         else:
             self.client = None
 
-    async def request_completions(self, messages: List[ChatCompletionMessageParam],
-                                  response_format_info: Optional[ResponseFormatInfo] = None,
-                                  message_type: MessageType = MessageType.CHAT) -> Any:
+    async def request_completions(
+        self, messages: List[ChatCompletionMessageParam],
+        model: str,
+        response_format_info: Optional[ResponseFormatInfo] = None,
+        message_type: MessageType = MessageType.CHAT
+    ) -> Any:
         """
         Get a response from Claude or the Mito server that adheres to the AgentResponse format.
         """
-        try:
-            anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
-            
-            provider_data = get_anthropic_completion_function_params(
-                model=self.model if response_format_info else ANTHROPIC_FAST_MODEL,
-                messages=anthropic_messages,
-                max_tokens=MAX_TOKENS,
-                temperature=0,
-                system=anthropic_system_prompt,
-                stream=None,
-                response_format_info=response_format_info
-            )
-            
-            if self.api_key:
-                # Unpack provider_data for direct API call
-                assert self.client is not None
-                response = self.client.messages.create(**provider_data)
-                if provider_data.get("tool_choice") is not None:
-                    result = extract_and_parse_anthropic_json_response(response)
-                    return json.dumps(result) if not isinstance(result, str) else result
-                else:
-                    content = response.content
-                    if content[0].type == "text":
-                        return content[0].text
-                    else:
-                        return ""
+        anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+        
+        provider_data = get_anthropic_completion_function_params(
+            message_type=message_type,
+            model=model,
+            messages=anthropic_messages,
+            max_tokens=MAX_TOKENS,
+            temperature=0,
+            system=anthropic_system_prompt,
+            stream=None,
+            response_format_info=response_format_info
+        )
+        
+        if self.api_key:
+            # Unpack provider_data for direct API call
+            assert self.client is not None
+            response = self.client.messages.create(**provider_data)
+            if provider_data.get("tool_choice") is not None:
+                result = extract_and_parse_anthropic_json_response(response)
+                return json.dumps(result) if not isinstance(result, str) else result
             else:
-                # Only pass provider_data to the server
-                response = await get_anthropic_completion_from_mito_server(
-                    model=provider_data["model"],
-                    max_tokens=provider_data["max_tokens"],
-                    temperature=provider_data["temperature"],
-                    system=provider_data["system"],
-                    messages=provider_data["messages"],
-                    tools=provider_data.get("tools"),
-                    tool_choice=provider_data.get("tool_choice"),
-                    message_type=message_type
-                )
-                return response
-        except anthropic.RateLimitError:
-            raise Exception("Rate limit exceeded. Please try again later or reduce your request frequency.")
-        except Exception as e:
-            return f"Error streaming content: {str(e)}"
+                content = response.content
+                if content[0].type == "text":
+                    return content[0].text
+                else:
+                    return ""
+        else:
+            # Only pass provider_data to the server
+            response = await get_anthropic_completion_from_mito_server(
+                model=provider_data["model"],
+                max_tokens=provider_data["max_tokens"],
+                temperature=provider_data["temperature"],
+                system=provider_data["system"],
+                messages=provider_data["messages"],
+                tools=provider_data.get("tools"),
+                tool_choice=provider_data.get("tool_choice"),
+                message_type=message_type
+            )
+            return response
 
-    async def stream_response(self, messages: List[ChatCompletionMessageParam], message_id: str, message_type: MessageType,
+    async def stream_completions(self, messages: List[ChatCompletionMessageParam], model: str, message_id: str, message_type: MessageType,
                               reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]) -> str:
         try:
             anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
@@ -201,7 +198,7 @@ class AnthropicClient:
             if self.api_key:
                 assert self.client is not None
                 stream = self.client.messages.create(
-                    model=self.model,
+                    model=model,
                     max_tokens=MAX_TOKENS,
                     temperature=0,
                     system=anthropic_system_prompt,
@@ -229,7 +226,7 @@ class AnthropicClient:
 
             else:
                 async for stram_chunk in stream_anthropic_completion_from_mito_server(
-                    model=self.model,
+                    model=model,
                     max_tokens=MAX_TOKENS,
                     temperature=0,
                     system=anthropic_system_prompt,
@@ -254,6 +251,7 @@ class AnthropicClient:
             raise Exception("Rate limit exceeded. Please try again later or reduce your request frequency.")
 
         except Exception as e:
-            return f"Error streaming content: {str(e)}"
+            print(f"Error streaming content: {str(e)}")
+            raise e
 
 

mito_ai/app_builder/handlers.py

@@ -16,11 +16,9 @@ from mito_ai.app_builder.models import (
     MessageType
 )
 from mito_ai.logger import get_logger
+from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
 import requests
 
-# API endpoint for getting pre-signed URL
-API_BASE_URL = "https://fr12uvtfy5.execute-api.us-east-1.amazonaws.com"
-
 class AppBuilderHandler(BaseWebSocketHandler):
     """Handler for app building requests."""
     
@@ -155,7 +153,7 @@ class AppBuilderHandler(BaseWebSocketHandler):
         try:
             # Step 1: Get pre-signed URL from API
             self.log.info("Getting pre-signed upload URL...")
-            url_response = requests.get(f"{API_BASE_URL}/get-upload-url?app_name={app_name}")
+            url_response = requests.get(f"{ACTIVE_STREAMLIT_BASE_URL}/get-upload-url?app_name={app_name}")
             url_response.raise_for_status()
             
             url_data = url_response.json()

mito_ai/completions/models.py

@@ -8,6 +8,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from enum import Enum
 from pydantic import BaseModel, Field
 
+
 # The ThreadID is the unique identifier for the chat thread.
 ThreadID = NewType('ThreadID', str)
   
@@ -20,6 +21,7 @@ class CellUpdate(BaseModel):
     index: Optional[int]
     id: Optional[str]
     code: str
+    code_summary: str
     cell_type: Optional[Literal['code', 'markdown']]
 
 
@@ -225,6 +227,19 @@ class CompletionError:
         While mypy doesn't know about this attribute on BaseException, we need to handle it
         to properly extract error messages from OpenAI API responses.
         """
+        from mito_ai.utils.mito_server_utils import ProviderCompletionException
+
+        
+        # Handle ProviderCompletionException specially
+        if isinstance(exception, ProviderCompletionException):
+            return CompletionError(
+                error_type="LLM Provider Error", 
+                title=exception.user_friendly_title, 
+                traceback=traceback.format_exc(),
+                hint=exception.user_friendly_hint
+            )
+        
+        # Handle all other exceptions as before
         error_type = type(exception)
         error_module = getattr(error_type, "__module__", "")
         
@@ -249,7 +264,6 @@ class CompletionError:
             hint=hint,
         )
 
-
 @dataclass(frozen=True)
 class ErrorMessage(CompletionError):
     """

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -49,6 +49,7 @@ Format:
         type: 'modification'
         id: str,
         code: str
+        code_summary: str
         cell_type: 'code' | 'markdown'
     }}
     get_cell_output_cell_id: None,
@@ -59,7 +60,8 @@ Important information:
 1. The id is the id of the code cell that you want to update. The id MUST already be part of the original Jupyter Notebook that your colleague shared with you.
 2. The message is a short summary of your thought process that helped you decide what to update in cell_update.
 3. The code should be the full contents of that updated code cell. The code that you return will overwrite the existing contents of the code cell so it must contain all necessary code.
-4. Important: Only use the CELL_UPDATE tool if you want to add/modify a notebook cell in response to the user's request. If the user is just sending you a friendly greeting or asking you a question about yourself, you SHOULD NOT USE A CELL_UPDATE tool because it does not require modifying the notebook. Instead, just use the FINISHED_TASK response.
+4. The code_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Loading data", "Filtering rows", "Calculating average", "Plotting revenue"). Avoid full sentences or explanations—this should read like a quick commit message or code label, not a description.
+5. Important: Only use the CELL_UPDATE tool if you want to add/modify a notebook cell in response to the user's request. If the user is just sending you a friendly greeting or asking you a question about yourself, you SHOULD NOT USE A CELL_UPDATE tool because it does not require modifying the notebook. Instead, just use the FINISHED_TASK response.
 
 #### Cell Addition:
 When you want to add a new cell to the notebook, respond in this format
@@ -72,6 +74,7 @@ Format:
         type: 'new'
         index: int
         code: str   
+        code_summary: str
         cell_type: 'code' | 'markdown'
     }}
     get_cell_output_cell_id: None,
@@ -82,7 +85,8 @@ Important information:
 1. The index should be the 0-index position of where you want the new code cell to be added in the notebook.
 2. The message is a short summary of your thought process that helped you decide what to update in cell_update.
 3. The code should be the full contents of that updated code cell. The code that you return will overwrite the existing contents of the code cell so it must contain all necessary code.
-4. The cell_type should only be 'markdown' if there is no code to add. There may be times where the code has comments. These are still code cells and should have the cell_type 'code'. Any cells that are labeled 'markdown' will be converted to markdown cells by the user.
+4. code_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Loading data", "Filtering rows", "Calculating average", "Plotting revenue"). Avoid full sentences or explanations—this should read like a quick commit message or code label, not a description.
+5. The cell_type should only be 'markdown' if there is no code to add. There may be times where the code has comments. These are still code cells and should have the cell_type 'code'. Any cells that are labeled 'markdown' will be converted to markdown cells by the user.
 
 <Cell Modification Example>
 Jupyter Notebook:
@@ -126,6 +130,7 @@ Output:
         type: 'modification'
         id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc',
         code: "import pandas as pd\\nsales_df = pd.read_csv('./sales.csv')\\nloan_multiplier = 1.5\\nsales_df['transaction_date'] = pd.to_datetime(sales_df['transaction_date'])\\nsales_df['total_price'] = sales_df['total_price'] * sales_multiplier",
+        code_summary: "Converting the transaction_date column",
         cell_type: 'code'
     }},
     get_cell_output_cell_id: None,
@@ -175,6 +180,8 @@ Output:
         type: 'add'
         index: 2
         code: "import matplotlib.pyplot as plt\n\nplt.bar(sales_df.index, sales_df['total_price'])\nplt.title('Total Price per Sale')\nplt.xlabel('Transaction Number')\nplt.ylabel('Sales Price ($)')\nplt.show()"
+        code_summary: "Plotting total_price",
+        code_summary: "Plotting total_price"
     }},
     get_cell_output_cell_id: None,
     next_steps: None
@@ -309,6 +316,7 @@ Output:
         type: 'add'
         index: 2
         code: "all_time_high_row_idx = tesla_stock_prices_df['closing_price'].idxmax()\nall_time_high_date = tesla_stock_prices_df.at[all_time_high_row_idx, 'Date']\nall_time_high_price = tesla_stock_prices_df.at[all_time_high_row_idx, 'closing_price']"
+        code_summary: "Calculating all time high"
     }},
     get_cell_output_cell_id: None
 }}

mito_ai/completions/providers.py

@@ -2,6 +2,7 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 from __future__ import annotations
+import asyncio
 from typing import Any, Callable, Dict, List, Optional, Union, cast
 from mito_ai import constants
 from openai.types.chat import ChatCompletionMessageParam
@@ -28,12 +29,16 @@ from mito_ai.completions.models import (
 from mito_ai.utils.telemetry_utils import (
     KEY_TYPE_PARAM,
     MITO_AI_COMPLETION_ERROR,
+    MITO_AI_COMPLETION_RETRY,
     MITO_SERVER_KEY,
     USER_KEY,
     log,
+    log_ai_completion_error,
+    log_ai_completion_retry,
     log_ai_completion_success,
 )
-from mito_ai.constants import get_model_provider
+from mito_ai.utils.provider_utils import get_model_provider
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
 
 __all__ = ["OpenAIProvider"]
 
@@ -66,6 +71,9 @@ This attribute is observed by the websocket provider to push the error to the cl
 
     @property
     def capabilities(self) -> AICapabilities:
+        """
+        Returns the capabilities of the AI provider.
+        """
         if constants.CLAUDE_API_KEY and not self.api_key:
             return AICapabilities(
                 configuration={"model": "<dynamic>"},
@@ -78,6 +86,7 @@ This attribute is observed by the websocket provider to push the error to the cl
             )
         if self._openai_client:
             return self._openai_client.capabilities
+        
         return AICapabilities(
             configuration={"model": "<dynamic>"},
             provider="Mito server",
@@ -100,7 +109,8 @@ This attribute is observed by the websocket provider to push the error to the cl
         model: str,
         response_format_info: Optional[ResponseFormatInfo] = None,
         user_input: Optional[str] = None,
-        thread_id: Optional[str] = None
+        thread_id: Optional[str] = None,
+        max_retries: int = 3
     ) -> str:
         """
         Request completions from the AI provider.
@@ -109,43 +119,69 @@ This attribute is observed by the websocket provider to push the error to the cl
         completion = None
         last_message_content = str(messages[-1].get('content', '')) if messages else ""
         model_type = get_model_provider(model)
-        try:
-            if model_type == "claude":
-                api_key = constants.CLAUDE_API_KEY
-                anthropic_client = AnthropicClient(api_key=api_key, model=model)
-                completion = await anthropic_client.request_completions(messages, response_format_info, message_type)
-            elif model_type == "gemini":
-                api_key = constants.GEMINI_API_KEY
-                gemini_client = GeminiClient(api_key=api_key, model=model)
-                messages_for_gemini = [dict(m) for m in messages]
-                completion = await gemini_client.request_completions(messages_for_gemini, response_format_info, message_type)
-            elif model_type == "openai":
-                if not self._openai_client:
-                    raise RuntimeError("OpenAI client is not initialized.")
-                completion = await self._openai_client.request_completions(
+        
+        # Retry loop
+        for attempt in range(max_retries + 1):
+            try:
+                if model_type == "claude":
+                    api_key = constants.CLAUDE_API_KEY
+                    anthropic_client = AnthropicClient(api_key=api_key)
+                    completion = await anthropic_client.request_completions(messages, model, response_format_info, message_type)
+                elif model_type == "gemini":
+                    api_key = constants.GEMINI_API_KEY
+                    gemini_client = GeminiClient(api_key=api_key)
+                    messages_for_gemini = [dict(m) for m in messages]
+                    completion = await gemini_client.request_completions(messages_for_gemini, model, response_format_info, message_type)
+                elif model_type == "openai":
+                    if not self._openai_client:
+                        raise RuntimeError("OpenAI client is not initialized.")
+                    completion = await self._openai_client.request_completions(
+                        message_type=message_type,
+                        messages=messages,
+                        model=model,
+                        response_format_info=response_format_info
+                    )
+                else:
+                    raise ValueError(f"No AI provider configured for model: {model}")
+                
+                # Success! Log and return
+                log_ai_completion_success(
+                    key_type=USER_KEY if self.key_type == "user" else MITO_SERVER_KEY,
                     message_type=message_type,
-                    messages=messages,
-                    model=model,
-                    response_format_info=response_format_info
+                    last_message_content=last_message_content,
+                    response={"completion": completion},
+                    user_input=user_input or "",
+                    thread_id=thread_id or "",
+                    model=model
                 )
-            else:
-                raise ValueError(f"No AI provider configured for model: {model}")
-            log_ai_completion_success(
-                key_type=USER_KEY if self.key_type == "user" else MITO_SERVER_KEY,
-                message_type=message_type,
-                last_message_content=last_message_content,
-                response={"completion": completion},
-                user_input=user_input or "",
-                thread_id=thread_id or "",
-                model=model
-            )
-            return completion
+                return completion # type: ignore
+            
+            except PermissionError as e:
+                # If we hit a free tier limit, then raise an exception right away without retrying.
+                self.log.exception(f"Error during request_completions: {e}")
+                self.last_error = CompletionError.from_exception(e)
+                log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                raise
 
-        except BaseException as e:
-            self.log.exception(f"Error during request_completions: {e}")
-            self.last_error = CompletionError.from_exception(e)
-            log(MITO_AI_COMPLETION_ERROR, params={KEY_TYPE_PARAM: self.key_type}, error=e)
-            raise
+            except BaseException as e:
+                # Check if we should retry (not on the last attempt)
+                if attempt < max_retries:
+                    # Exponential backoff: wait 2^attempt seconds
+                    wait_time = 2 ** attempt
+                    self.log.info(f"Retrying request_completions after {wait_time}s (attempt {attempt + 1}/{max_retries + 1}): {str(e)}")
+                    log_ai_completion_retry('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                    await asyncio.sleep(wait_time)
+                    continue
+                else:
+                    # Final failure after all retries - set error state and raise
+                    self.log.exception(f"Error during request_completions after {attempt + 1} attempts: {e}")
+                    self.last_error = CompletionError.from_exception(e)
+                    log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', message_type, e)
+                    raise
+        
+        # This should never be reached due to the raise in the except block,
+        # but added to satisfy the linter
+        raise RuntimeError("Unexpected code path in request_completions")
 
     async def stream_completions(
         self,
@@ -176,19 +212,23 @@ This attribute is observed by the websocket provider to push the error to the cl
         try:
             if model_type == "claude":
                 api_key = constants.CLAUDE_API_KEY
-                anthropic_client = AnthropicClient(api_key=api_key, model=model)
-                accumulated_response = await anthropic_client.stream_response(
+                anthropic_client = AnthropicClient(api_key=api_key)
+                accumulated_response = await anthropic_client.stream_completions(
                     messages=messages,
+                    model=model,
                     message_type=message_type,
                     message_id=message_id,
                     reply_fn=reply_fn
                 )
             elif model_type == "gemini":
                 api_key = constants.GEMINI_API_KEY
-                gemini_client = GeminiClient(api_key=api_key, model=model)
+                gemini_client = GeminiClient(api_key=api_key)
+                # TODO: We shouldn't need to do this because the messages should already be dictionaries... 
+                # but if we do have to do some pre-processing, we should do it in the gemini_client instead.
                 messages_for_gemini = [dict(m) for m in messages]
                 accumulated_response = await gemini_client.stream_completions(
                     messages=messages_for_gemini,
+                    model=model,
                     message_id=message_id,
                     reply_fn=reply_fn,
                     message_type=message_type

mito_ai/constants.py

@@ -23,30 +23,9 @@ AZURE_OPENAI_API_VERSION = os.environ.get("AZURE_OPENAI_API_VERSION")
 AZURE_OPENAI_ENDPOINT = os.environ.get("AZURE_OPENAI_ENDPOINT")
 AZURE_OPENAI_MODEL = os.environ.get("AZURE_OPENAI_MODEL")
 
-def get_model_provider(model: str) -> Union[str, None]:
-    """
-    Determine the model type based on the model name prefix
-    """
-    if not model:
-        return None
-
-    model_lower = model.lower()
-
-    if model_lower.startswith('claude'):
-        return 'claude'
-    elif model_lower.startswith('gemini'):
-        return 'gemini'
-    elif model_lower.startswith('ollama'):
-        return 'ollama'
-    elif model_lower.startswith('gpt'):
-        return 'openai'
-
-    return None
-
-
 # Mito AI Base URLs and Endpoint Paths
-MITO_PROD_BASE_URL = "https://yxwyadgaznhavqvgnbfuo2k6ca0jboku.lambda-url.us-east-1.on.aws"
-MITO_DEV_BASE_URL = "https://x3rafympznv4abp7phos44gzgu0clbui.lambda-url.us-east-1.on.aws"
+MITO_PROD_BASE_URL = "https://7eax4i53f5odkshhlry4gw23by0yvnuv.lambda-url.us-east-1.on.aws/v1"
+MITO_DEV_BASE_URL = "https://g5vwmogjg7gh7aktqezyrvcq6a0hyfnr.lambda-url.us-east-1.on.aws/v1"
 
 # Set ACTIVE_BASE_URL manually
 ACTIVE_BASE_URL = MITO_PROD_BASE_URL  # Change to MITO_DEV_BASE_URL for dev
@@ -59,4 +38,12 @@ OPENAI_PATH = "openai/completions"
 # Full URLs (always use ACTIVE_BASE_URL)
 MITO_ANTHROPIC_URL = f"{ACTIVE_BASE_URL}/{ANTHROPIC_PATH}"
 MITO_GEMINI_URL = f"{ACTIVE_BASE_URL}/{GEMINI_PATH}"
-MITO_OPENAI_URL = f"{ACTIVE_BASE_URL}/{OPENAI_PATH}"
+MITO_OPENAI_URL = f"{ACTIVE_BASE_URL}/{OPENAI_PATH}"
+
+# Streamlit conversion endpoints
+MITO_STREAMLIT_DEV_BASE_URL = "https://fr12uvtfy5.execute-api.us-east-1.amazonaws.com"
+MITO_STREAMLIT_TEST_BASE_URL = "https://iyual08t6d.execute-api.us-east-1.amazonaws.com"
+
+# Set ACTIVE_BASE_URL manually
+# TODO: Modify to PROD url before release
+ACTIVE_STREAMLIT_BASE_URL = MITO_STREAMLIT_TEST_BASE_URL  # Change to MITO_STREAMLIT_DEV_BASE_URL for dev