mito-ai 0.1.58__tar.gz → 0.1.59__tar.gz

{mito_ai-0.1.58 → mito_ai-0.1.59}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mito_ai
-Version: 0.1.58
+Version: 0.1.59
 Summary: AI chat for JupyterLab
 Project-URL: Homepage, https://trymito.io
 Project-URL: Bug Tracker, https://github.com/mito-ds/monorepo/issues

{mito_ai-0.1.58 → mito_ai-0.1.59}/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-0.1.58 → mito_ai-0.1.59}/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.59'

{mito_ai-0.1.58 → mito_ai-0.1.59}/mito_ai/completions/prompt_builders/prompt_constants.py

@@ -20,6 +20,8 @@ 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"').

{mito_ai-0.1.58 → mito_ai-0.1.59}/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-0.1.58 → mito_ai-0.1.59}/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-0.1.58 → mito_ai-0.1.59}/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-0.1.58 → mito_ai-0.1.59}/mito_ai/labextension/build_log.json

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

{mito_ai-0.1.58/venv/share/jupyter/labextensions/mito_ai → mito_ai-0.1.59/mito_ai/labextension}/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mito_ai",
-  "version": "0.1.58",
+  "version": "0.1.59",
   "description": "AI chat for JupyterLab",
   "keywords": [
     "jupyter",
@@ -141,7 +141,7 @@
     "schemaDir": "schema",
     "themePath": "style/theme/theme.css",
     "_build": {
-      "load": "static/remoteEntry.570df809a692f53a7ab7.js",
+      "load": "static/remoteEntry.f7decebaf69618541e0f.js",
       "extension": "./extension",
       "style": "./style"
     }

{mito_ai-0.1.58/venv/share/jupyter/labextensions/mito_ai → mito_ai-0.1.59/mito_ai/labextension}/schemas/mito_ai/package.json.orig

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