mito-ai 0.1.58__py3-none-any.whl → 0.1.60__py3-none-any.whl

mito_ai/__init__.py

@@ -9,6 +9,7 @@ from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.app_deploy.handlers import AppDeployHandler
 from mito_ai.log.urls import get_log_urls
 from mito_ai.utils.litellm_utils import is_litellm_configured
+from mito_ai.enterprise.utils import is_abacus_configured
 from mito_ai.version_check import VersionCheckHandler
 from mito_ai.db.urls import get_db_urls
 from mito_ai.settings.urls import get_settings_urls
@@ -101,10 +102,12 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
 
     web_app.add_handlers(host_pattern, handlers)
     
-    # Log enterprise mode status and LiteLLM configuration
+    # Log enterprise mode status and router configuration
     if is_enterprise():
         server_app.log.info("Enterprise mode enabled")
-        if is_litellm_configured():
+        if is_abacus_configured():
+            server_app.log.info(f"Abacus AI configured: endpoint={constants.ABACUS_BASE_URL}, models={constants.ABACUS_MODELS}")
+        elif is_litellm_configured():
             server_app.log.info(f"LiteLLM configured: endpoint={constants.LITELLM_BASE_URL}, models={constants.LITELLM_MODELS}")
     
     server_app.log.info("Loaded the mito_ai server extension")

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.58'
+__version__ = VERSION = '0.1.60'

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -10,6 +10,7 @@ from mito_ai.completions.prompt_builders.prompt_constants import (
     get_database_rules
 )
 from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
+from mito_ai.rules.utils import get_default_rules_content
 
 def create_agent_system_message_prompt(isChromeBrowser: bool) -> str:
     
@@ -474,7 +475,12 @@ Important information:
     
     # Database rules
     sections.append(SG.Generic("Database Rules", get_database_rules()))
-    
+
+    # Default rules
+    default_rules = get_default_rules_content()
+    if default_rules:
+        sections.append(SG.Generic("Default (User Defined) Rules", default_rules))
+
     # RULES OF YOUR WORKING PROCESS
     sections.append(SG.Generic("Rules Of Working Process", f"""The user is going to ask you to guide them as through the process of completing a task. You will help them complete a task over the course of an entire conversation with them. The user will first share with you what they want to accomplish. You will then use a tool to execute the first step of the task, they will execute the tool and return to you the updated notebook state with you, and then you will give them the next step of the task. You will continue to give them the next step of the task until they have completed the task.
 

mito_ai/completions/prompt_builders/chat_system_message.py

@@ -11,6 +11,7 @@ from mito_ai.completions.prompt_builders.prompt_constants import (
     get_database_rules
 )
 from mito_ai.completions.prompt_builders.prompt_section_registry.base import PromptSection
+from mito_ai.rules.utils import get_default_rules_content
 
 def create_chat_system_message_prompt() -> str:
     sections: List[PromptSection] = []
@@ -34,6 +35,9 @@ Other useful information:
 
     sections.append(SG.Generic("Chart Config Rules", CHART_CONFIG_RULES))
     sections.append(SG.Generic("DatabaseRules", get_database_rules()))
+    default_rules = get_default_rules_content()
+    if default_rules:
+        sections.append(SG.Generic("Default (User Defined) Rules", default_rules))
     sections.append(SG.Generic("Citation Rules", CITATION_RULES))
     sections.append(SG.Generic("Cell Reference Rules", CELL_REFERENCE_RULES))
     

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -20,10 +20,23 @@ Rules:
 - All imports must appear at the top, before the chart configuration section.
 - Variables with multiple words should be underscore-separated.
 - All colors should be in hex format (e.g., "#3498db"). Use quotes around the hex string: COLOR = "#3498db" or COLOR = '#3498db'. Do NOT nest quotes.
+- Never use RGB/RGBA tuples/lists for colors (e.g. (0, 0.4, 0.8, 0.8) is forbidden).
+- If transparency is needed, store it separately as ALPHA = 0.8 and apply it in code (e.g. to_rgba(HEX_COLOR, ALPHA)).
 - Variables can only be strings, numbers, booleans, tuples, or lists. 
 - NEVER include comments on the same line as a variable assignment. Each variable assignment must be on its own line with no trailing comments.
 - For string values, use either single or double quotes (e.g., TITLE = "Sales by Product" or TITLE = 'Sales by Product'). Do not use nested quotes (e.g., do NOT use '"value"').
 
+Fixed acceptable ranges (matplotlib constraints):
+- For numeric variables that have a fixed acceptable range, add a line immediately BEFORE the variable assignment: # RANGE VARIABLE_NAME MIN MAX
+- This allows the Chart Wizard to clamp inputs and prevent invalid values. Use the following ranges when you use these variables:
+  - ALPHA (opacity): 0 1
+  - FIGURE_SIZE (tuple width, height in inches): 1 24 (each element)
+  - LINE_WIDTH, LINEWIDTH, LWD: 0 20
+  - FONT_SIZE, FONTSIZE, FONT_SIZE_TITLE, FONT_SIZE_LABEL: 0.1 72
+  - MARKER_SIZE, MARKERSIZE, S: 0 1000
+  - DPI: 1 600
+  - Any other numeric or tuple variable that you know has matplotlib constraints: add # RANGE VARIABLE_NAME MIN MAX with the appropriate min and max.
+
 Common Mistakes to Avoid:
 - WRONG: COLOR = '"#1877F2" # Meta Blue'  (nested quotes and inline comment)
 - WRONG: COLOR = "#1877F2" # Meta Blue  (inline comment)
@@ -36,6 +49,10 @@ TITLE = "Sales by Product"
 X_LABEL = "Product"
 Y_LABEL = "Sales"
 BAR_COLOR = "#000000"
+# RANGE ALPHA 0 1
+ALPHA = 0.8
+# RANGE FIGURE_SIZE 1 24
+FIGURE_SIZE = (12, 6)
 # === END CONFIG ===
 """
 

mito_ai/constants.py

@@ -2,7 +2,7 @@
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
 import os
-from typing import Union
+from typing import Union, List
 
 # Claude
 ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY")
@@ -23,12 +23,34 @@ 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 parse_comma_separated_models(models_str: str) -> List[str]:
+    """
+    Parse a comma-separated string of model names into a list.
+    Handles quoted and unquoted values, stripping whitespace and quotes.
+    
+    Args:
+        models_str: Comma-separated string of model names (e.g., "model1,model2" or '"model1","model2"')
+    
+    Returns:
+        List of model names with whitespace and quotes stripped
+    """
+    if not models_str:
+        return []
+    return [model.strip().strip('"\'') for model in models_str.split(",") if model.strip()]
+
 # LiteLLM Config (Enterprise mode only)
 LITELLM_BASE_URL = os.environ.get("LITELLM_BASE_URL")
 LITELLM_API_KEY = os.environ.get("LITELLM_API_KEY")
 LITELLM_MODELS_STR = os.environ.get("LITELLM_MODELS", "")
-# Parse comma-separated string into list, strip whitespace
-LITELLM_MODELS = [model.strip() for model in LITELLM_MODELS_STR.split(",") if model.strip()] if LITELLM_MODELS_STR else []
+# Parse comma-separated string into list, strip whitespace and quotes
+LITELLM_MODELS = parse_comma_separated_models(LITELLM_MODELS_STR)
+
+# Abacus AI Config (Enterprise mode only)
+ABACUS_BASE_URL = os.environ.get("ABACUS_BASE_URL")
+ABACUS_API_KEY = os.environ.get("ABACUS_API_KEY")
+ABACUS_MODELS_STR = os.environ.get("ABACUS_MODELS", "")
+# Parse comma-separated string into list, strip whitespace and quotes
+ABACUS_MODELS = parse_comma_separated_models(ABACUS_MODELS_STR)
 
 # Mito AI Base URLs and Endpoint Paths
 MITO_PROD_BASE_URL = "https://7eax4i53f5odkshhlry4gw23by0yvnuv.lambda-url.us-east-1.on.aws/v2"

mito_ai/enterprise/litellm_client.py

@@ -11,6 +11,7 @@ from mito_ai.completions.models import (
     CompletionItem,
 )
 from mito_ai.utils.litellm_utils import get_litellm_completion_function_params
+from mito_ai.utils.model_utils import strip_router_prefix
 import litellm
 
 class LiteLLMClient:
@@ -28,7 +29,7 @@ class LiteLLMClient:
     async def request_completions(
         self,
         messages: List[ChatCompletionMessageParam],
-        model: str,  # Should include provider prefix (e.g., "openai/gpt-4o")
+        model: str,  # Should include provider prefix (e.g., "LiteLLM/openai/gpt-4o")
         response_format_info: Optional[ResponseFormatInfo] = None,
         message_type: MessageType = MessageType.CHAT
     ) -> str:
@@ -37,16 +38,19 @@ class LiteLLMClient:
         
         Args:
             messages: List of chat messages
-            model: Model name with provider prefix (e.g., "openai/gpt-4o")
+            model: Model name with router and provider prefix (e.g., "LiteLLM/openai/gpt-4o")
             response_format_info: Optional response format specification
             message_type: Type of message (chat, agent execution, etc.)
             
         Returns:
             The completion text response
         """
+        # Strip router prefix if present (LiteLLM/ prefix)
+        model_for_litellm = strip_router_prefix(model)
+        
         # Prepare parameters for LiteLLM
         params = get_litellm_completion_function_params(
-            model=model,
+            model=model_for_litellm,
             messages=messages,
             api_key=self.api_key,
             api_base=self.base_url,
@@ -82,7 +86,7 @@ class LiteLLMClient:
         
         Args:
             messages: List of chat messages
-            model: Model name with provider prefix (e.g., "openai/gpt-4o")
+            model: Model name with router and provider prefix (e.g., "LiteLLM/openai/gpt-4o")
             message_type: Type of message (chat, agent execution, etc.)
             message_id: ID of the message being processed
             reply_fn: Function to call with each chunk for streaming replies
@@ -93,9 +97,12 @@ class LiteLLMClient:
         """
         accumulated_response = ""
         
+        # Strip router prefix if present (LiteLLM/ prefix)
+        model_for_litellm = strip_router_prefix(model)
+        
         # Prepare parameters for LiteLLM
         params = get_litellm_completion_function_params(
-            model=model,
+            model=model_for_litellm,
             messages=messages,
             api_key=self.api_key,
             api_base=self.base_url,

mito_ai/enterprise/utils.py

@@ -5,11 +5,25 @@
 # Distributed under the terms of the The Mito Enterprise license.
 
 from mito_ai.utils.version_utils import is_enterprise, is_mitosheet_private
-from mito_ai.constants import AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL
+from mito_ai.constants import (
+    AZURE_OPENAI_API_KEY, 
+    AZURE_OPENAI_ENDPOINT, 
+    AZURE_OPENAI_API_VERSION, 
+    AZURE_OPENAI_MODEL,
+    ABACUS_BASE_URL,
+    ABACUS_MODELS
+)
 
 def is_azure_openai_configured() -> bool:
     """
     Azure OpenAI is only supported for Mito Enterprise users
     """
     is_allowed_to_use_azure = is_enterprise() or is_mitosheet_private()
-    return all([is_allowed_to_use_azure, AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL])
+    return all([is_allowed_to_use_azure, AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_API_VERSION, AZURE_OPENAI_MODEL])
+
+def is_abacus_configured() -> bool:
+    """
+    Abacus AI is only supported for Mito Enterprise users.
+    Checks if Abacus AI is configured with base URL and models.
+    """
+    return all([is_enterprise(), ABACUS_BASE_URL, ABACUS_MODELS])

mito_ai/openai_client.py

@@ -11,8 +11,9 @@ from traitlets import Instance, default, validate
 from traitlets.config import LoggingConfigurable
 
 from mito_ai import constants
-from mito_ai.enterprise.utils import is_azure_openai_configured
+from mito_ai.enterprise.utils import is_azure_openai_configured, is_abacus_configured
 from mito_ai.logger import get_logger
+from mito_ai.utils.model_utils import strip_router_prefix
 from mito_ai.completions.models import (
     AICapabilities,
     CompletionError,
@@ -24,12 +25,11 @@ from mito_ai.completions.models import (
     ResponseFormatInfo,
 )
 from mito_ai.utils.open_ai_utils import (
-    check_mito_server_quota,
     get_ai_completion_from_mito_server,
     get_open_ai_completion_function_params,
     stream_ai_completion_from_mito_server,
 )
-from mito_ai.utils.server_limits import update_mito_server_quota
+from mito_ai.utils.server_limits import update_mito_server_quota, check_mito_server_quota
 
 OPENAI_MODEL_FALLBACK = "gpt-4.1"
 
@@ -68,6 +68,14 @@ This attribute is observed by the websocket provider to push the error to the cl
                 provider="Azure OpenAI",
             )
 
+        if is_abacus_configured():
+            return AICapabilities(
+                configuration={
+                    "model": "<dynamic>"
+                },
+                provider="Abacus AI",
+            )
+
         if constants.OLLAMA_MODEL:
             return AICapabilities(
                 configuration={
@@ -121,6 +129,10 @@ This attribute is observed by the websocket provider to push the error to the cl
                 timeout=self.timeout,
             )
         
+        elif is_abacus_configured():
+            base_url = constants.ABACUS_BASE_URL
+            llm_api_key = constants.ABACUS_API_KEY
+            self.log.debug(f"Using Abacus AI with base URL: {constants.ABACUS_BASE_URL}")
         elif constants.OLLAMA_MODEL:
             base_url = constants.OLLAMA_BASE_URL
             llm_api_key = "ollama"
@@ -141,17 +153,25 @@ This attribute is observed by the websocket provider to push the error to the cl
         )
         return client
 
-    def _adjust_model_for_azure_or_ollama(self, model: str) -> str:
+    def _adjust_model_for_provider(self, model: str) -> str:
         
         # If they have set an Azure OpenAI model, then we always use it
         if is_azure_openai_configured() and constants.AZURE_OPENAI_MODEL is not None:
             self.log.debug(f"Resolving to Azure OpenAI model: {constants.AZURE_OPENAI_MODEL}")
+            # TODO: We should update Azure so it works the way LiteLLM and Abacus do: 
+            # when configured, we only show models from Azure in the UI. 
             return constants.AZURE_OPENAI_MODEL
         
         # If they have set an Ollama model, then we use it
         if constants.OLLAMA_MODEL is not None:
             return constants.OLLAMA_MODEL
         
+        # If using Abacus, strip the "Abacus/" prefix from the model name
+        if is_abacus_configured() and model.lower().startswith('abacus/'):
+            stripped_model = strip_router_prefix(model)
+            self.log.debug(f"Stripping Abacus prefix: {model} -> {stripped_model}")
+            return stripped_model
+        
         # Otherwise, we use the model they provided
         return model
         
@@ -186,7 +206,7 @@ This attribute is observed by the websocket provider to push the error to the cl
         )
         
         # If they have set an Azure OpenAI or Ollama model, then we use it
-        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
+        completion_function_params["model"] = self._adjust_model_for_provider(completion_function_params["model"])
 
         if self._active_async_client is not None:
             response = await self._active_async_client.chat.completions.create(**completion_function_params)
@@ -236,7 +256,7 @@ This attribute is observed by the websocket provider to push the error to the cl
             model, messages, True, response_format_info
         )
         
-        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
+        completion_function_params["model"] = self._adjust_model_for_provider(completion_function_params["model"])
 
         try:
             if self._active_async_client is not None:

mito_ai/provider_manager.py

@@ -27,6 +27,7 @@ from mito_ai.completions.models import (
     ResponseFormatInfo,
 )
 from mito_ai.utils.litellm_utils import is_litellm_configured
+from mito_ai.enterprise.utils import is_abacus_configured
 from mito_ai.utils.telemetry_utils import (
     MITO_SERVER_KEY,
     USER_KEY,
@@ -79,6 +80,12 @@ This attribute is observed by the websocket provider to push the error to the cl
         # TODO: We should validate that these keys are actually valid for the provider
         # otherwise it will look like we are using the user_key when actually falling back 
         # to the mito server because the key is invalid. 
+        if is_abacus_configured():
+            return AICapabilities(
+                configuration={"model": "<dynamic>"},
+                provider="Abacus AI",
+            )
+            
         if is_litellm_configured():
             return AICapabilities(
                 configuration={"model": "<dynamic>"},
@@ -116,6 +123,9 @@ This attribute is observed by the websocket provider to push the error to the cl
         # TODO: We should validate that these keys are actually valid for the provider
         # otherwise it will look like we are using the user_key when actually falling back 
         # to the mito server because the key is invalid. 
+        if is_abacus_configured():
+            return USER_KEY
+            
         if is_litellm_configured():
             return USER_KEY
         
@@ -172,7 +182,16 @@ This attribute is observed by the websocket provider to push the error to the cl
         # Retry loop
         for attempt in range(max_retries + 1):
             try:
-                if model_type == "litellm":
+                if model_type == "abacus":
+                    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=resolved_model,
+                        response_format_info=response_format_info
+                    )
+                elif model_type == "litellm":
                     from mito_ai.enterprise.litellm_client import LiteLLMClient
                     if not constants.LITELLM_BASE_URL:
                         raise ValueError("LITELLM_BASE_URL is required for LiteLLM models")
@@ -299,7 +318,20 @@ This attribute is observed by the websocket provider to push the error to the cl
         ))
 
         try:
-            if model_type == "litellm":
+            if model_type == "abacus":
+                if not self._openai_client:
+                    raise RuntimeError("OpenAI client is not initialized.")
+                accumulated_response = await self._openai_client.stream_completions(
+                    message_type=message_type,
+                    messages=messages,
+                    model=resolved_model,
+                    message_id=message_id,
+                    thread_id=thread_id,
+                    reply_fn=reply_fn,
+                    user_input=user_input,
+                    response_format_info=response_format_info
+                )
+            elif model_type == "litellm":
                 from mito_ai.enterprise.litellm_client import LiteLLMClient
                 if not constants.LITELLM_BASE_URL:
                     raise ValueError("LITELLM_BASE_URL is required for LiteLLM models")

mito_ai/rules/handlers.py

@@ -7,7 +7,16 @@ from typing import Any, Final, Union
 import tornado
 import os
 from jupyter_server.base.handlers import APIHandler
-from mito_ai.rules.utils import RULES_DIR_PATH, get_all_rules, get_rule, set_rules_file
+from mito_ai.rules.utils import (
+    RULES_DIR_PATH,
+    cleanup_rules_metadata,
+    delete_rule,
+    get_all_rules,
+    get_rule,
+    get_rule_default,
+    set_rule_default,
+    set_rules_file,
+)
 
 
 class RulesHandler(APIHandler):
@@ -17,17 +26,26 @@ class RulesHandler(APIHandler):
     def get(self, key: Union[str, None] = None) -> None:
         """Get a specific rule by key or all rules if no key provided"""
         if key is None or key == '':
-            # No key provided, return all rules
-            rules = get_all_rules()
+            # No key provided, return all rules with is_default flag
+            rule_files = get_all_rules()
+            rules = [
+                {"name": name, "is_default": get_rule_default(name)}
+                for name in rule_files
+            ]
             self.finish(json.dumps(rules))
         else:
             # Key provided, return specific rule
-            rule_content = get_rule(key)
-            if rule_content is None:
-                self.set_status(404)
-                self.finish(json.dumps({"error": f"Rule with key '{key}' not found"}))
-            else:
-                self.finish(json.dumps({"key": key, "content": rule_content}))
+            try:
+                rule_content = get_rule(key)
+                if rule_content is None:
+                    self.set_status(404)
+                    self.finish(json.dumps({"error": f"Rule with key '{key}' not found"}))
+                else:
+                    is_default = get_rule_default(key)
+                    self.finish(json.dumps({"key": key, "content": rule_content, "is_default": is_default}))
+            except ValueError as e:
+                self.set_status(400)
+                self.finish(json.dumps({"error": str(e)}))
     
     @tornado.web.authenticated
     def put(self, key: str) -> None:
@@ -37,8 +55,24 @@ class RulesHandler(APIHandler):
             self.set_status(400)
             self.finish(json.dumps({"error": "Content is required"}))
             return
-            
-        set_rules_file(key, data['content'])
-        self.finish(json.dumps({"status": "updated", "rules file ": key}))
 
+        try:
+            set_rules_file(key, data['content'])
+            if 'is_default' in data:
+                set_rule_default(key, bool(data['is_default']))
+            cleanup_rules_metadata()
+            self.finish(json.dumps({"status": "updated", "rules_file": key}))
+        except ValueError as e:
+            self.set_status(400)
+            self.finish(json.dumps({"error": str(e)}))
 
+    @tornado.web.authenticated
+    def delete(self, key: str) -> None:
+        """Delete a rule by key (rule name)."""
+        try:
+            delete_rule(key)
+            cleanup_rules_metadata()
+            self.finish(json.dumps({"status": "deleted", "key": key}))
+        except ValueError as e:
+            self.set_status(400)
+            self.finish(json.dumps({"error": str(e)}))