Merge pull request #2162 from mito-ds/abacus

Abacus

Author: aarondr77

.github/workflows/test-abacus-llm-providers.yml

@@ -0,0 +1,88 @@
+name: Test - Mito AI Frontend Playwright with Abacus AI
+
+on:
+  push:
+    branches: [ dev ]
+    paths:
+      - 'mito-ai/**'
+      - 'tests/llm_providers_tests/abacus_llm_providers.spec.ts'
+      - '.github/workflows/test-abacus-llm-providers.yml'
+  pull_request:
+    paths:
+      - 'mito-ai/**'
+      - 'tests/llm_providers_tests/abacus_llm_providers.spec.ts'
+      - '.github/workflows/test-abacus-llm-providers.yml'
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test-mitoai-frontend-jupyterlab-abacus:
+    runs-on: ubuntu-24.04
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.12']
+      fail-fast: false
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: pip
+        cache-dependency-path: |
+          mito-ai/setup.py
+          tests/requirements.txt
+    - uses: actions/setup-node@v4
+      with:
+        node-version: 22
+        cache: 'npm'
+        cache-dependency-path: mito-ai/package-lock.json
+    - name: Upgrade pip
+      run: |
+        python -m pip install --upgrade pip
+    - name: Install dependencies
+      run: |
+        cd tests
+        bash mac-setup.sh
+    - name: Install mitosheet-helper-enterprise
+      run: |
+        cd tests
+        source venv/bin/activate
+        pip install mitosheet-helper-enterprise
+    - name: Install JupyterLab
+      run: |
+        python -m pip install jupyterlab
+    - name: Install Node.js dependencies
+      run: |
+        cd mito-ai
+        jlpm install
+    - name: Setup JupyterLab
+      run: |
+        cd tests
+        source venv/bin/activate
+        pip install setuptools==68.0.0
+        cd ../mito-ai
+        jupyter labextension develop . --overwrite
+        jupyter server extension enable --py mito_ai
+    - name: Start a server and run Abacus AI provider tests
+      run: |
+        cd tests
+        source venv/bin/activate
+        jupyter lab --config jupyter_server_test_config.py &
+        jlpm run test:abacus-llm-providers
+      env:
+        ABACUS_BASE_URL: ${{ secrets.ABACUS_BASE_URL }}
+        ABACUS_MODELS: ${{ secrets.ABACUS_MODELS }}
+        ABACUS_API_KEY: ${{ secrets.ABACUS_API_KEY }}
+    - name: Upload test-results
+      uses: actions/upload-artifact@v4
+      if: failure()
+      with:
+        name: mitoai-jupyterlab-playwright-abacus-report-${{ matrix.python-version }}-${{ github.run_id }}
+        path: tests/playwright-report/
+        retention-days: 14

agent.md

@@ -1,14 +1,24 @@
+
 ## Virtual Environment
 - ALWAYS activate the virtual environment before running any Python commands: `cd mito-ai && source venv/bin/activate`
 - For other components: check for venv directories and activate appropriately
 
 ## Engineering practices
 
+### Startup mentality
 We're a startup. You're probably used to writing enterprise code -- code that tries to handle every possible edge case and has fallbacks for everything. That's not how we do things around here: our number one rule is to keep things simple. We handle ONLY the most important cases.
 
 We try to only add new functionality that is small (that is, simple and few lines of code) or absolutely necessary. If a change is not small or absolutely necessary, don't make it.
 
-**Backwards-compatibility**: Since our app runs locally on a user's machine, we think about backwards compatibility very differently than enterprise code.
+### Backwards Compatibility
+Since our app runs locally on a user's machine, we think about backwards compatibility very differently than enterprise code.
 - For things like updating the environment variables structure, we don't need to worry about backwards compatibility because if the enterprise opts in to upgrading, they can update the enviornment variables at the same time. 
 - For things like adding new tool options, we don't need to worry about backwards compatibility because only users on the newest version of the tool are going to have access to the new options.
-- For things like changing how we read in old chat histories, we DO need to worry about backwards compatibility because the chat histories live on the user's machine, so we don't have a way to migrate them to the new format without adding that migration step into the tool.
+- For things like changing how we read in old chat histories, we DO need to worry about backwards compatibility because the chat histories live on the user's machine, so we don't have a way to migrate them to the new format without adding that migration step into the tool.
+
+### Use Parameterized Tests When Appropriate
+Use `@pytest.mark.parametrize` when you have multiple test cases that:
+- Test the same function/method with different inputs
+- Have the same expected output structure
+- Follow the same test logic pattern
+- Would otherwise require repetitive test code

mito-ai/docs/litellm-deployment.md

@@ -9,6 +9,7 @@
 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/mito_ai/__init__.py

@@ -1,6 +1,3 @@
-# Copyright (c) Saga Inc.
-# Distributed under the terms of the GNU Affero General Public License v3.0 License.
-
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore

mito-ai/mito_ai/_version.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_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/mito_ai/constants.py

@@ -11,6 +11,7 @@
     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 @@ def __init__(self, api_key: Optional[str], base_url: str, timeout: int = 30, max
     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 @@ async def request_completions(
         
         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 @@ async def stream_completions(
         
         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 @@ async def stream_completions(
         """
         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/mito_ai/enterprise/litellm_client.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])