mito-ai 0.1.50__py3-none-any.whl

mito_ai/tests/chat_history/test_chat_history.py

@@ -0,0 +1,211 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import requests
+import time
+from unittest.mock import patch, MagicMock
+from mito_ai.tests.conftest import TOKEN
+from mito_ai.completions.message_history import GlobalMessageHistory, ChatThread
+from mito_ai.completions.models import ThreadID
+
+
+@pytest.fixture
+def mock_chat_threads():
+    """Fixture that creates mock chat threads for testing"""
+    thread_id_1 = ThreadID("test-thread-1")
+    thread_id_2 = ThreadID("test-thread-2")
+
+    # Create mock threads with different timestamps
+    thread_1 = ChatThread(
+        thread_id=thread_id_1,
+        creation_ts=time.time() - 3600,  # 1 hour ago
+        last_interaction_ts=time.time() - 1800,  # 30 minutes ago
+        name="Test Chat 1",
+        ai_optimized_history=[
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there!"},
+        ],
+        display_history=[
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there!"},
+        ],
+    )
+
+    thread_2 = ChatThread(
+        thread_id=thread_id_2,
+        creation_ts=time.time() - 7200,  # 2 hours ago
+        last_interaction_ts=time.time() - 900,  # 15 minutes ago (more recent)
+        name="Test Chat 2",
+        ai_optimized_history=[
+            {"role": "user", "content": "How are you?"},
+            {"role": "assistant", "content": "I'm doing well, thanks!"},
+        ],
+        display_history=[
+            {"role": "user", "content": "How are you?"},
+            {"role": "assistant", "content": "I'm doing well, thanks!"},
+        ],
+    )
+
+    return {thread_id_1: thread_1, thread_id_2: thread_2}
+
+
+@pytest.fixture
+def mock_message_history(mock_chat_threads):
+    """Fixture that mocks the GlobalMessageHistory with test data"""
+    mock_history = MagicMock(spec=GlobalMessageHistory)
+    mock_history._chat_threads = mock_chat_threads
+
+    # Mock the get_threads method to return sorted threads
+    def mock_get_threads():
+        from mito_ai.completions.models import ChatThreadMetadata
+
+        threads = []
+        for thread in mock_chat_threads.values():
+            threads.append(
+                ChatThreadMetadata(
+                    thread_id=thread.thread_id,
+                    name=thread.name,
+                    creation_ts=thread.creation_ts,
+                    last_interaction_ts=thread.last_interaction_ts,
+                )
+            )
+        # Sort by last_interaction_ts (newest first)
+        threads.sort(key=lambda x: x.last_interaction_ts, reverse=True)
+        return threads
+
+    mock_history.get_threads = mock_get_threads
+    return mock_history
+
+
+# --- GET ALL THREADS ---
+
+
+def test_get_all_threads_success(jp_base_url: str, mock_message_history):
+    """Test successful GET all threads endpoint"""
+    # Since the server extension is already loaded, we need to work with the actual instance
+    # Let's just test that the endpoint works and returns the expected structure
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert "threads" in response_json
+    # The actual number of threads will depend on what's in the .mito/ai-chats directory
+    # So we'll just check that it's a list
+    assert isinstance(response_json["threads"], list)
+
+    # Check thread structure for any threads that exist
+    for thread in response_json["threads"]:
+        assert "thread_id" in thread
+        assert "name" in thread
+        assert "creation_ts" in thread
+        assert "last_interaction_ts" in thread
+
+
+def test_get_all_threads_empty(jp_base_url: str):
+    """Test GET all threads endpoint when no threads exist"""
+    # This test will work with whatever threads exist in the actual .mito/ai-chats directory
+    # We'll just verify the endpoint works and returns the expected structure
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert "threads" in response_json
+    assert isinstance(response_json["threads"], list)
+
+
+def test_get_all_threads_with_no_auth(jp_base_url: str):
+    """Test GET all threads endpoint without authentication"""
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_all_threads_with_incorrect_auth(jp_base_url: str):
+    """Test GET all threads endpoint with incorrect authentication"""
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+# --- GET SPECIFIC THREAD ---
+
+
+def test_get_specific_thread_success(jp_base_url: str, mock_message_history):
+    """Test successful GET specific thread endpoint"""
+    # First, get all threads to see what's available
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    threads = response.json()["threads"]
+    if not threads:
+        # If no threads exist, skip this test
+        pytest.skip("No threads available for testing")
+
+    # Use the first available thread
+    thread_id = threads[0]["thread_id"]
+
+    response = requests.get(
+        jp_base_url + f"/mito-ai/chat-history/threads/{thread_id}",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert response_json["thread_id"] == thread_id
+    assert "name" in response_json
+    assert "creation_ts" in response_json
+    assert "last_interaction_ts" in response_json
+    assert "display_history" in response_json
+    assert "ai_optimized_history" in response_json
+
+    # Check message history structure
+    display_history = response_json["display_history"]
+    assert isinstance(display_history, list)
+    ai_optimized_history = response_json["ai_optimized_history"]
+    assert isinstance(ai_optimized_history, list)
+
+
+def test_get_specific_thread_not_found(jp_base_url: str, mock_message_history):
+    """Test GET specific thread endpoint with non-existent thread ID"""
+    # Use a clearly non-existent thread ID
+    fake_thread_id = "non-existent-thread-12345"
+
+    response = requests.get(
+        jp_base_url + f"/mito-ai/chat-history/threads/{fake_thread_id}",
+        headers={"Authorization": f"token {TOKEN}"},
+    )
+    assert response.status_code == 404
+
+    response_json = response.json()
+    assert "error" in response_json
+    assert fake_thread_id in response_json["error"]
+
+
+def test_get_specific_thread_with_no_auth(jp_base_url: str):
+    """Test GET specific thread endpoint without authentication"""
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads/test-thread-1",
+    )
+    assert response.status_code == 403  # Forbidden
+
+
+def test_get_specific_thread_with_incorrect_auth(jp_base_url: str):
+    """Test GET specific thread endpoint with incorrect authentication"""
+    response = requests.get(
+        jp_base_url + "/mito-ai/chat-history/threads/test-thread-1",
+        headers={"Authorization": f"token incorrect-token"},
+    )
+    assert response.status_code == 403  # Forbidden

mito_ai/tests/completions/completion_handlers_utils_test.py

@@ -0,0 +1,190 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import base64
+import os
+import tempfile
+from contextlib import contextmanager
+from mito_ai.completions.completion_handlers.utils import (
+    create_ai_optimized_message,
+    extract_and_encode_images_from_additional_context,
+)
+
+
+@contextmanager
+def temporary_image_file(suffix=".png", content=b"fake_image_data"):
+    """Context manager that creates a temporary image file for testing."""
+    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as temp_file:
+        temp_file.write(content)
+        temp_file_path = temp_file.name
+
+    try:
+        yield temp_file_path
+    finally:
+        # Clean up the temporary file
+        os.unlink(temp_file_path)
+
+
+def test_text_only_message():
+    """Test scenario where the user only inputs text"""
+    result = create_ai_optimized_message("Hello world")
+
+    assert result["role"] == "user"
+    assert result["content"] == "Hello world"
+
+
+def test_message_with_uploaded_image():
+    """Test scenario where the user uploads an image"""
+    with temporary_image_file() as temp_file_path:
+        result = create_ai_optimized_message(
+            text="Analyze this",
+            additional_context=[{"type": "image/png", "value": temp_file_path}],
+        )
+
+        assert result["role"] == "user"
+        assert isinstance(result["content"], list)
+        assert result["content"][0]["type"] == "text"
+        assert result["content"][1]["type"] == "image_url"
+
+
+def test_message_with_multiple_uploaded_images():
+    """Test scenario where the user uploads multiple images"""
+    with temporary_image_file(suffix=".png", content=b"image1_data") as temp_file1:
+        with temporary_image_file(suffix=".jpg", content=b"image2_data") as temp_file2:
+            result = create_ai_optimized_message(
+                text="Analyze these images",
+                additional_context=[
+                    {"type": "image/png", "value": temp_file1},
+                    {"type": "image/jpeg", "value": temp_file2},
+                ],
+            )
+
+            assert result["role"] == "user"
+            assert isinstance(result["content"], list)
+            assert len(result["content"]) == 3  # text + 2 images
+            assert result["content"][0]["type"] == "text"
+            assert result["content"][0]["text"] == "Analyze these images"
+            assert result["content"][1]["type"] == "image_url"
+            assert result["content"][2]["type"] == "image_url"
+            
+            # Verify the image URLs are properly formatted
+            assert result["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
+            assert result["content"][2]["image_url"]["url"].startswith("data:image/jpeg;base64,")
+
+
+def test_message_with_active_cell_output():
+    """Test scenario where the active cell has an output"""
+    result = create_ai_optimized_message(
+        text="Analyze this", base64EncodedActiveCellOutput="cell_output_data"
+    )
+
+    assert result["role"] == "user"
+    assert isinstance(result["content"], list)
+    assert result["content"][0]["type"] == "text"
+    assert result["content"][1]["type"] == "image_url"
+
+
+def test_message_with_uploaded_image_and_active_cell_output():
+    """Test scenario where the user uploads an image and the active cell has an output"""
+    with temporary_image_file() as temp_file_path:
+        result = create_ai_optimized_message(
+            text="Analyze this",
+            additional_context=[{"type": "image/png", "value": temp_file_path}],
+            base64EncodedActiveCellOutput="cell_output_data",
+        )
+
+        assert result["role"] == "user"
+        assert isinstance(result["content"], list)
+        assert result["content"][0]["type"] == "text"
+        assert result["content"][1]["type"] == "image_url"
+        assert result["content"][2]["type"] == "image_url"
+
+
+def test_extract_and_encode_images_from_additional_context_valid_image():
+    """Test extracting and encoding a valid image file"""
+    with temporary_image_file() as temp_file_path:
+        additional_context = [{"type": "image/png", "value": temp_file_path}]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        assert len(encoded_images) == 1
+        assert encoded_images[0].startswith("data:image/png;base64,")
+        # Verify it's valid base64 by checking it can be decoded
+        base64_data = encoded_images[0].split(",")[1]
+        decoded_data = base64.b64decode(base64_data)
+        assert decoded_data == b"fake_image_data"
+
+
+def test_extract_and_encode_images_from_additional_context_multiple_images():
+    """Test extracting and encoding multiple image files"""
+    with temporary_image_file(suffix=".png", content=b"image1_data") as temp_file1:
+        with temporary_image_file(suffix=".jpg", content=b"image2_data") as temp_file2:
+            additional_context = [
+                {"type": "image/png", "value": temp_file1},
+                {"type": "image/jpeg", "value": temp_file2},
+            ]
+
+            encoded_images = extract_and_encode_images_from_additional_context(
+                additional_context
+            )
+
+            assert len(encoded_images) == 2
+            assert encoded_images[0].startswith("data:image/png;base64,")
+            assert encoded_images[1].startswith("data:image/jpeg;base64,")
+
+
+def test_extract_and_encode_images_from_additional_context_invalid_file():
+    """Test handling of invalid/non-existent image files"""
+    additional_context = [{"type": "image/png", "value": "non_existent_file.png"}]
+
+    encoded_images = extract_and_encode_images_from_additional_context(
+        additional_context
+    )
+
+    assert len(encoded_images) == 0
+
+
+def test_extract_and_encode_images_from_additional_context_non_image_types():
+    """Test that non-image types are ignored"""
+    with temporary_image_file(suffix=".txt", content=b"text_data") as temp_file:
+        additional_context = [
+            {"type": "text/plain", "value": temp_file},
+            {"type": "application/pdf", "value": "document.pdf"},
+        ]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        assert len(encoded_images) == 0
+
+
+def test_extract_and_encode_images_from_additional_context_mixed_types():
+    """Test handling of mixed image and non-image types"""
+    with temporary_image_file() as temp_image_file:
+        additional_context = [
+            {"type": "image/png", "value": temp_image_file},
+            {"type": "text/plain", "value": "document.txt"},
+            {"type": "image/jpeg", "value": "non_existent.jpg"},
+        ]
+
+        encoded_images = extract_and_encode_images_from_additional_context(
+            additional_context
+        )
+
+        # Should only have the valid PNG image
+        assert len(encoded_images) == 1
+        assert encoded_images[0].startswith("data:image/png;base64,")
+
+
+def test_extract_and_encode_images_from_additional_context_empty():
+    """Test handling of empty or None additional_context"""
+    # Test with None
+    encoded_images = extract_and_encode_images_from_additional_context(None)
+    assert len(encoded_images) == 0
+
+    # Test with empty list
+    encoded_images = extract_and_encode_images_from_additional_context([])
+    assert len(encoded_images) == 0

mito_ai/tests/conftest.py

@@ -0,0 +1,53 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+import threading
+import time
+from jupyter_server.serverapp import ServerApp
+from traitlets.config import Config
+
+TOKEN = "test-token"
+
+@pytest.fixture
+def jp_server_config(token_fixture):
+    """Configure the Jupyter server for testing."""
+    config = Config()
+    config.ServerApp.jpserver_extensions = {"mito_ai": True}
+    # Disable password requirement for testing
+    config.ServerApp.password = ""
+    # Set the token for testing
+    config.ServerApp.token = token_fixture
+    # Enable authentication
+    config.ServerApp.allow_unauthenticated_access = False
+    return config
+
+@pytest.fixture
+def jp_serverapp(jp_server_config, tmp_path):
+    """Create a Jupyter server instance for testing."""
+    app = ServerApp(config=jp_server_config)
+    app.root_dir = str(tmp_path)
+    app.initialize(argv=[])
+
+    # Start the server in a non-blocking way
+    def start_server():
+        app.start()
+
+    server_thread = threading.Thread(target=start_server)
+    server_thread.daemon = True  # This ensures the thread will be killed when the main program exits
+    server_thread.start()
+
+    # Give the server a moment to start
+    time.sleep(1)
+
+    yield app
+    app.stop()
+
+@pytest.fixture
+def jp_base_url(jp_serverapp):
+    """Get the base URL of the Jupyter server."""
+    return jp_serverapp.connection_url 
+
+@pytest.fixture
+def token_fixture():
+    return TOKEN

mito_ai/tests/create_agent_system_message_prompt_test.py

@@ -0,0 +1,22 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import pytest
+from mito_ai.completions.prompt_builders.agent_system_message import create_agent_system_message_prompt
+
+def test_create_agent_system_message_prompt_browser_conditional() -> None:
+    """
+    Test that verifies the create_agent_system_message_prompt function produces different output
+    based on the browser type.
+    
+    Specifically:
+    - When isChromeBrowser=False, the prompt should NOT contain GET_CELL_OUTPUT
+    - When isChromeBrowser=True, the prompt SHOULD contain GET_CELL_OUTPUT
+    """
+    # Test non-Chrome browser case
+    non_chrome_prompt = create_agent_system_message_prompt(isChromeBrowser=False)
+    assert "GET_CELL_OUTPUT" not in non_chrome_prompt, "Non-Chrome prompt should not contain GET_CELL_OUTPUT"
+    
+    # Test Chrome browser case
+    chrome_prompt = create_agent_system_message_prompt(isChromeBrowser=True)
+    assert "GET_CELL_OUTPUT" in chrome_prompt, "Chrome prompt should contain GET_CELL_OUTPUT"

mito_ai/tests/data/prompt_lg.py

@@ -0,0 +1,69 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+prompt_lg = [
+    {
+        "role": "user",
+        "content": """You are an expert python programmer writing a script in a Jupyter notebook. You are given a set of variables, existing code, and a task.
+
+There are two possible types of responses you might give:
+1. Code Update: If the task requires modifying or extending the existing code, respond with the updated active code cell and a short explanation of the changes made. 
+2. Explanation/Information: If the task does not require a code update, provide an explanation, additional information about a package, method, or general programming question, without writing any code. Keep your response concise and to the point.
+
+When responding:
+- Do not use the word "I"
+- Do not recreate variables that already exist
+- Keep as much of the original code as possible
+
+<Example>
+
+Files in the current directory:
+file_name: sales.csv
+
+Defined Variables:
+{{
+    'loan_multiplier': 1.5,
+    'sales_df': pd.DataFrame({{
+        'transaction_date': ['2024-01-02', '2024-01-02', '2024-01-02', '2024-01-02', '2024-01-03'],
+        'price_per_unit': [10, 9.99, 13.99, 21.00, 100],
+        'units_sold': [1, 2, 1, 4, 5],
+        'total_price': [10, 19.98, 13.99, 84.00, 500]
+    }})
+}}
+
+Code in the active code cell:
+```python
+import pandas as pd
+sales_df = pd.read_csv('./sales.csv')
+```
+
+Your task: convert the transaction_date column to datetime and then multiply the total_price column by the sales_multiplier.
+
+Output:
+```python
+import pandas as pd
+sales_df = pd.read_csv('./sales.csv')
+sales_df['transaction_date'] = pd.to_datetime(sales_df['transaction_date'])
+sales_df['total_price'] = sales_df['total_price'] * sales_multiplier
+```
+
+Converted the `transaction_date` column to datetime using the built-in pd.to_datetime function and multiplied the `total_price` column by the `sales_multiplier` variable.
+
+</Example>
+
+Files in the current directory:
+
+Defined Variables:
+
+Code in the active code cell:
+```python
+```
+
+Your task: 
+Create a dataset about the stock market. It should have fundamental data 
+going back 10 years for 100 companies. There should be some correlation in 
+the dataset, I don't care what columns. Once the dataset is created, create 
+three charts to help me understand the data. Please select the topic for 
+these charts based on details you think are pertinent.""",
+    },
+]

mito_ai/tests/data/prompt_sm.py

@@ -0,0 +1,6 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+prompt_sm = """
+Can you create a function that perfoms fizz buzz. It should take an integer n and return a list of strings.
+"""

mito_ai/tests/data/prompt_xl.py

@@ -0,0 +1,13 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+prompt_xl = [
+    {
+        "role": "system",
+        "content": "You are Mito Data Copilot, an AI assistant for Jupyter. You\u2019re a great python programmer, a seasoned data scientist and a subject matter expert.\n\nThe user is going to ask you to guide them as they complete 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 give them the first step of the task, they will apply that first step, share 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.\n\nTo communicate with the user, you will send them a CELL_UPDATE or FINISHED_TASK response in each message. Each time you respond with a CELL_UPDATE, the user will apply the CELL_UPDATE to the notebook, run the new code cell, and provide you with updated information about the notebook and variables defined in the kernel to help you decide what to do next.\n\n====\n\nCELL_UPDATES\n\nCELL_UPDATES are how you communicate to the user about the changes you want to make to the notebook. Each CELL_UPDATE can either modify an existing cell or create a new cell. \n\nThere are two types of CELL_UPDATES:\n\n1. CellModification\n2. CellAddition\n\nEach time you want to make a change to the notebook, you will respond with a CellModification or CellAddition.\n\n#### Cell Modification\nWhen you want to modify an existing cell in the notebook, respond in this format.\n\nFormat:\n{{\n    is_finished: false, \n    cell_update: {{\n        type: modification\n        id: str,\n        code: str\n    }}\n}}\n\nImportant information:\n1. 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.\n2. 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.\n\n#### Cell Addition:\nWhen you want to add a new cell to the notebook, respond in this format\n\nFormat: \n{{\n    is_finished: false, \n    cell_update: {{\n        type: 'new'\n        index: int\n        code: str   \n    }}\n}}\n\nImportant information:\n1. The index should be the 0-index position of where you want the new code cell to be added in the notebook.\n2. 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.\n\n<Cell Modification Example>\nJupyter Notebook:\n[\n    {{\n        cell_type: 'markdown'\n        id: '9e38c62b-38f8-457d-bb8d-28bfc52edf2c'\n        code: \"\"\" # Used Car Sales Analysis \"\"\"\n    }},\n    {{\n        cell_type: 'code'\n        id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc'\n        code: \"\"\"import pandas as pd\nsales_df = pd.read_csv('./sales.csv') \nloan_multiplier = 1.5\"\"\"\n    }},\n]\n\nDefined Variables:\n{{\n    'loan_multiplier': 1.5,\n    'sales_df': pd.DataFrame({{\n        'transaction_date': ['2024-01-02', '2024-01-02', '2024-01-02', '2024-01-02', '2024-01-03'],\n        'price_per_unit': [10, 9.99, 13.99, 21.00, 100],\n        'units_sold': [1, 2, 1, 4, 5],\n        'total_price': [10, 19.98, 13.99, 84.00, 500]\n    }})\n}}\n\nFiles in the current directory:\nfile_name: sales.csv\n\nYour task: \nConvert the transaction_date column to datetime and then multiply the total_price column by the sales_multiplier.\n\nOutput:\n{{\n    is_finished: false, \n    cell_update: {{\n        type: 'modification'\n        id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc',\n        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\"\n    }}\n}}\n\n</Cell Modification Example>\n\n<Cell Addition Example>\nJupyter Notebook:\n[\n    {{\n        cell_type: 'markdown'\n        id: '9e38c62b-38f8-457d-bb8d-28bfc52edf2c'\n        code: \"\"\"# Used Car Sales Analysis \"\"\"\n    }},\n    {{\n        cell_type: 'code'\n        id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc'\n        code: \"\"\"import pandas as pd\nsales_df = pd.read_csv('./sales.csv')\nsales_df['transaction_date'] = pd.to_datetime(sales_df['transaction_date'])\"\"\"\n    }},\n]\n\nDefined Variables:\n{{\n    'sales_df': pd.DataFrame({{\n        'transaction_date': ['2024-01-02', '2024-01-02', '2024-01-02', '2024-01-02', '2024-01-03'],\n        'price_per_unit': [10, 9.99, 13.99, 21.00, 100],\n        'units_sold': [1, 2, 1, 4, 5],\n        'total_price': [10, 19.98, 13.99, 84.00, 500]\n    }})\n}}\n\nFiles in the current directory:\nfile_name: sales.csv\n\nYour task: \nGraph the total_price for each sale\n\nOutput:\n{{\n    is_finished: false, \n    cell_update: {{\n        type: 'add'\n        index: 2\n        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()\"\n    }}\n}}\n\n</Cell Addition Example>\n\n====\n\nFINISHED_TASK\n\nWhen you have completed the user's task, respond with a message in this format:\n\n{{\n    is_finished: true\n}}\n\n====\n\nRULES\n\n- You are working in a Jupyter Lab environment in a .ipynb file. \n- You can only respond with CELL_UPDATES or FINISHED_TASK responses.\n- In each message you send to the user, you can send one CellModification, one CellAddition, or one FINISHED_TASK response. BUT YOU WILL GET TO SEND MULTIPLE MESSAGES TO THE USER TO ACCOMPLISH YOUR TASK SO DO NOT TRY TO ACCOMPLISH YOUR TASK IN A SINGLE MESSAGE.\n- After you send a CELL_UPDATE, the user will send you a message with the updated variables, code, and files in the current directory. You will use this information to decide what to do next, so it is critical that you wait for the user's response after each CELL_UPDATE before deciding your next action.\n- When updating code, keep as much of the original code as possible and do not recreate variables that already exist.\n- When you want to display a dataframe to the user, just write the dataframe on the last line of the code cell instead of writing print(<dataframe name>). Jupyter will automatically display the dataframe in the notebook.\n\n====\n\nRULES OF YOUR WORKING PROCESS\n\n\nThe 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 give them the first step of the task, they will apply that first step, share 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.\n\nAs you are guiding the user through the process of completing the task, send them a CELL_UPDATE message to give them the next step of the task. When you have finished the task, send a FINISHED_TASK message. \n\nThe user is a beginning Python user, so you will need to be careful to send them only small steps to complete. Don't try to complete the task in a single response to the user. Instead, each message you send to the user should only contain a single, small step towards the end goal. When the user has completed the step, they will let you know that they are ready for the next step. \n\nYou will keep working in the following iterative format until you have decided that you have finished the user's request. When you decide that you have finished the user's request, respond with a FINISHED_TASK message. Otherwise, if you have not finished the user's request, respond with a CELL_UPDATE. When you respond with a CELL_UPDATE, the user will apply the CELL_UPDATE to the notebook and run the new code cell. The user will then send you a message with an updated version of the variables defined in the kernel, code in the notebook, and files in the current directory. In addition, the user will check if the code you provided produced an errored when executed. If it did produce an error, the user will share the error message with you.\n\nWhenever you get a message back from the user, you should:\n1. Ask yourself if the previous CELL_UDPATE is correct. You can answer this question by reviewing the updated variable and code. \n2. Ask yourself if you can improve the code or results you got from the previous CELL_UPDATE. If you can, send a new CELL_UPDATE to modify the code you just wrote. \n3. Decide if you have finished the user's request to you. If you have, respond with a FINISHED_TASK message.\n4. If you have not finished the user's request, create the next CELL_UPDATE. \n\nREMEMBER, YOU ARE GOING TO COMPLETE THE USER'S TASK OVER THE COURSE OF THE ENTIRE CONVERSATION -- YOU WILL GET TO SEND MULTIPLE MESSAGES TO THE USER TO ACCOMPLISH YOUR TASK SO DO NOT TRY TO ACCOMPLISH YOUR TASK IN A SINGLE MESSAGE. IT IS CRUCIAL TO PROCEED STEP-BY-STEP WITH THE SMALLEST POSSIBLE CELL_UPDATES. For example, if asked to build a new dataframe, then analyze it, and then graph the results, you should proceed as follows. \n1. Send a CellAddition to add a new code cell to the notebook that creates the dataframe.\n2. Wait for the user to send you back the updated variables and notebook state so you can decide how to analyze the dataframe.\n3. Use the data that the user sent you to decide how to analyze the dataframe. Send a CellAddition to add the dataframe analysis code to the notebook.\n4. Wait for the user to send you back the updated variables and notebook state so you can decide how to proceed. \n5. If after reviewing the updates provided by the user, you decide that you want to update the analysis code, send a CellModification to modify the code you just wrote.\n6. Wait for the user to send you back the updated variables and notebook state so you can decide how to proceed.\n7. If you are happy with the analysis, refer back to the original task provided by the user to decide your next steps. In this example, it is to graph the results, so you will send a CellAddition to construct the graph. \n8. Wait for the user to send you back the updated variables and notebook state so you can decide how to proceed.\n9. If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK message.",
+    },
+    {
+        "role": "user",
+        "content": "Jupyter Notebook:\n{'id': '6c622de6-1eaa-48df-97c9-c6a52b095ac1', 'cell_type': 'code', 'code': \"import pandas as pd\\n\\n# Create a dataframe of NBA players with some sample data\\ndata = {\\n    'Name': ['LeBron James', 'Kevin Durant', 'Stephen Curry', 'Giannis Antetokounmpo', 'Luka Doncic', 'Nikola Jokic'],\\n    'Position': ['SF', 'PF', 'PG', 'PF', 'PG', 'C'],\\n    'Team': ['Lakers', 'Nets', 'Warriors', 'Bucks', 'Mavericks', 'Nuggets'],\\n    'Points': [27.0, 26.9, 28.1, 29.5, 27.7, 27.0],\\n    'Rebounds': [7.4, 7.1, 5.3, 13.6, 8.0, 10.8],\\n    'Assists': [7.4, 5.6, 6.3, 5.6, 8.6, 8.5]\\n}\\n\\nnba_players = pd.DataFrame(data)\\nnba_players \"}\n{'id': 'a173a08a-5afa-4449-bb4b-499ae8861b01', 'cell_type': 'code', 'code': \"## Build an All-Star Team\\n# We choose one best player per position based on the highest Points\\nall_star_team = nba_players.sort_values('Points', ascending=False).drop_duplicates('Position')\\n\\n# Display the All-Star team\\nall_star_team\\n\\n## Graphs for the All-Star Team\\nimport matplotlib.pyplot as plt\\n\\n# Graph 1: Points per All-Star Team member\\nplt.figure(figsize=(8, 4))\\nplt.bar(all_star_team['Name'], all_star_team['Points'], color='skyblue')\\nplt.title('All-Star Team: Points')\\nplt.xlabel('Player Name')\\nplt.ylabel('Points')\\nplt.xticks(rotation=45)\\nplt.show()\\n\\n# Graph 2: Rebounds per All-Star Team member\\nplt.figure(figsize=(8, 4))\\nplt.bar(all_star_team['Name'], all_star_team['Rebounds'], color='lightgreen')\\nplt.title('All-Star Team: Rebounds')\\nplt.xlabel('Player Name')\\nplt.ylabel('Rebounds')\\nplt.xticks(rotation=45)\\nplt.show()\\n\\n# Graph 3: Assists per All-Star Team member\\nplt.figure(figsize=(8, 4))\\nplt.bar(all_star_team['Name'], all_star_team['Assists'], color='salmon')\\nplt.title('All-Star Team: Assists')\\nplt.xlabel('Player Name')\\nplt.ylabel('Assists')\\nplt.xticks(rotation=45)\\nplt.show() \"}\n{'id': 'e25f1dd6-66b8-4769-97f8-bd956502ee4d', 'cell_type': 'code', 'code': '## Additional Graph: Scatter Plot of Points vs Rebounds (Bubble size indicates Assists)\\nimport matplotlib.pyplot as plt\\n\\nplt.figure(figsize=(8, 6))\\nplt.scatter(all_star_team[\\'Points\\'], all_star_team[\\'Rebounds\\'], s=all_star_team[\\'Assists\\']*20, color=\\'purple\\', alpha=0.6)\\n\\n# Annotate each point with the player\\'s name\\nfor idx, row in all_star_team.iterrows():\\n    plt.annotate(row[\\'Name\\'], (row[\\'Points\\'], row[\\'Rebounds\\']), textcoords=\"offset points\", xytext=(5,5))\\n\\nplt.title(\\'All-Star Team: Points vs Rebounds (Bubble size by Assists)\\')\\nplt.xlabel(\\'Points\\')\\nplt.ylabel(\\'Rebounds\\')\\nplt.show()\\n '}\n{'id': 'e2447e90-6645-4f59-a31c-3d160292ed42', 'cell_type': 'code', 'code': ''}\n{'id': '41f6c119-e21d-4154-a299-211287bb3c37', 'cell_type': 'code', 'code': ''}\n{'id': 'd0c4c623-7052-40b2-b864-400260c2bdce', 'cell_type': 'code', 'code': ''}\n{'id': 'ff0ff671-e054-4ca4-b337-4b237a4f602b', 'cell_type': 'code', 'code': ''}\n\nDefined Variables:\n{'variable_name': 'pd', 'type': \"<class 'module'>\", 'value': \"<module 'pandas' from '/Users/aarondiamond-reivich/Mito/mito/mito-ai/venv/lib/python3.11/site-packages/pandas/__init__.py'>\"}\n{'variable_name': 'data', 'type': \"<class 'dict'>\", 'value': \"{'Name': ['LeBron James', 'Kevin Durant', 'Stephen Curry', 'Giannis Antetokounmpo', 'Luka Doncic', 'Nikola Jokic'], 'Position': ['SF', 'PF', 'PG', 'PF', 'PG', 'C'], 'Team': ['Lakers', 'Nets', 'Warriors', 'Bucks', 'Mavericks', 'Nuggets'], 'Points': [27.0, 26.9, 28.1, 29.5, 27.7, 27.0], 'Rebounds': [7.4, 7.1, 5.3, 13.6, 8.0, 10.8], 'Assists': [7.4, 5.6, 6.3, 5.6, 8.6, 8.5]}\"}\n{'variable_name': 'nba_players', 'type': 'pd.DataFrame', 'value': {'Name': {'dtype': 'object', 'samples': ['LeBron James', 'Kevin Durant', 'Stephen Curry', 'Giannis Antetokounmpo', 'Luka Doncic']}, 'Position': {'dtype': 'object', 'samples': ['SF', 'PF', 'PG', 'PF', 'PG']}, 'Team': {'dtype': 'object', 'samples': ['Lakers', 'Nets', 'Warriors', 'Bucks', 'Mavericks']}, 'Points': {'dtype': 'float64', 'samples': [27, 26.9, 28.1, 29.5, 27.7]}, 'Rebounds': {'dtype': 'float64', 'samples': [7.4, 7.1, 5.3, 13.6, 8]}, 'Assists': {'dtype': 'float64', 'samples': [7.4, 5.6, 6.3, 5.6, 8.6]}}}\n{'variable_name': 'all_star_team', 'type': 'pd.DataFrame', 'value': {'Name': {'dtype': 'object', 'samples': ['Giannis Antetokounmpo', 'Stephen Curry', 'LeBron James', 'Nikola Jokic']}, 'Position': {'dtype': 'object', 'samples': ['PF', 'PG', 'SF', 'C']}, 'Team': {'dtype': 'object', 'samples': ['Bucks', 'Warriors', 'Lakers', 'Nuggets']}, 'Points': {'dtype': 'float64', 'samples': [29.5, 28.1, 27, 27]}, 'Rebounds': {'dtype': 'float64', 'samples': [13.6, 5.3, 7.4, 10.8]}, 'Assists': {'dtype': 'float64', 'samples': [5.6, 6.3, 7.4, 8.5]}}}\n{'variable_name': 'plt', 'type': \"<class 'module'>\", 'value': \"<module 'matplotlib.pyplot' from '/Users/aarondiamond-reivich/Mito/mito/mito-ai/venv/lib/python3.11/site-packages/matplotlib/pyplot.py'>\"}\n{'variable_name': 'idx', 'type': \"<class 'int'>\", 'value': '5'}\n{'variable_name': 'row', 'type': \"<class 'pandas.core.series.Series'>\", 'value': 'Name        Nikola Jokic\\nPosition               C\\nTeam             Nuggets\\nPoints              27.0\\nRebounds            10.8\\nAssists              8.5\\nName: 5, dtype: object'}\n\nFiles in the current directory:\n{'file_name': 'used_cars.csv'}\n\nYour task: \nUsing the analysis that you started, add another 20 players to the NBA dataset that you created and then use a more advanced strategy for figuring out who should be on the east and west coast all star team",
+    },
+]

mito_ai/tests/db/conftest.py

@@ -0,0 +1,39 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import json
+import pytest
+import shutil
+from typing import Generator
+from mito_ai.tests.db.test_db_constants import (
+    DB_DIR_PATH,
+    CONNECTIONS_PATH,
+    BACKUP_DB_DIR_PATH,
+)
+
+
+@pytest.fixture(scope="module", autouse=True)
+def backup_db_folder() -> Generator[None, None, None]:
+    """Backup the DB folder before tests and restore it after."""
+    if os.path.exists(DB_DIR_PATH):
+        # Create backup, so we can test with a clean db dir
+        shutil.move(DB_DIR_PATH, BACKUP_DB_DIR_PATH)
+
+    yield
+
+    # Cleanup after tests
+    if os.path.exists(DB_DIR_PATH):
+        shutil.rmtree(DB_DIR_PATH)
+    if os.path.exists(BACKUP_DB_DIR_PATH):
+        shutil.move(BACKUP_DB_DIR_PATH, DB_DIR_PATH)
+
+
+@pytest.fixture
+def first_connection_id() -> str:
+    # Manually open the connections.json file and get the first connection ID
+    with open(CONNECTIONS_PATH, "r") as f:
+        connections = json.load(f)
+    # Get the first connection ID from the object keys
+    connection_id = next(iter(connections.keys()))
+    return str(connection_id)