mito-ai 0.1.35__py3-none-any.whl → 0.1.37__py3-none-any.whl

mito_ai/completions/prompt_builders/utils.py

@@ -1,22 +1,65 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from typing import List, Optional
+from typing import List, Optional, Dict
 from mito_ai.rules.utils import get_rule
 
-def get_rules_str(selected_rules: Optional[List[str]]) -> str:
+
+def get_rules_str(additional_context: Optional[List[Dict[str, str]]]) -> str:
     """
-    Get a string of the rules that the user has selected.
+    Extract the rules from the additional context array, and retrieve the rule content.
     """
-    if selected_rules is None:
-        return ''
-    
-    rules_str = ''
+    if not additional_context:
+        return ""
+
+    selected_rules = [context["value"] for context in additional_context if context.get("type") == "rule"]
+    if len(selected_rules) == 0:
+        return ""
+
+    rules_str = ""
     for rule in selected_rules:
         rule_content = get_rule(rule)
-        if rule_content is None or rule_content == '':
+        if rule_content is None or rule_content == "":
             continue
-        
+
         rules_str += f"===========\n\nCustom Instructions Provided by User: {rule}\n\n{rule_content}\n\n==========="
-    
+
     return rules_str
+
+
+def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]]) -> str:
+    """
+    Get the selected context from the additional context array.
+    """
+    if not additional_context:
+        return ""
+
+    # STEP 1: Extract each context type into a separate list
+    selected_variables = [context["value"] for context in additional_context if context.get("type") == "variable"]
+    selected_files = [context["value"] for context in additional_context if context.get("type") == "file"]
+    selected_db_connections = [context["value"] for context in additional_context if context.get("type") == "db"]
+
+    # STEP 2: Create a list of strings (instructions) for each context type
+    context_parts = []
+
+    if len(selected_variables) > 0:
+        context_parts.append(
+            "The following variables have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_variables)
+        )
+
+    if len(selected_files) > 0:
+        context_parts.append(
+            "The following files have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_files)
+        )
+
+    if len(selected_db_connections) > 0:
+        context_parts.append(
+            "The following database connections have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_db_connections)
+        )
+
+    # STEP 3: Combine into a single string
+
+    return "\n\n".join(context_parts)

mito_ai/constants.py

@@ -46,4 +46,14 @@ MITO_STREAMLIT_TEST_BASE_URL = "https://iyual08t6d.execute-api.us-east-1.amazona
 
 # 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
+ACTIVE_STREAMLIT_BASE_URL = MITO_STREAMLIT_DEV_BASE_URL  # Change to MITO_STREAMLIT_DEV_BASE_URL for dev
+
+# AWS Cognito configuration
+COGNITO_CONFIG_DEV = {
+    'TOKEN_ENDPOINT': 'https://mito-app-auth.auth.us-east-1.amazoncognito.com/oauth2/token',
+    'CLIENT_ID': '6ara3u3l8sss738hrhbq1qtiqf',
+    'CLIENT_SECRET': '',
+    'REDIRECT_URI': 'http://localhost:8888/lab'
+}
+
+ACTIVE_COGNITO_CONFIG = COGNITO_CONFIG_DEV # Change to COGNITO_CONFIG_DEV for dev

mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -0,0 +1,112 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import logging
+from anthropic.types import MessageParam
+from typing import List, Tuple, cast
+
+from mito_ai.logger import get_logger
+from mito_ai.streamlit_conversion.streamlit_system_prompt import streamlit_system_prompt
+from mito_ai.streamlit_conversion.validate_and_run_streamlit_code import streamlit_code_validator
+from mito_ai.streamlit_conversion.streamlit_utils import extract_code_blocks, create_app_file, parse_jupyter_notebook_to_extract_required_content
+from mito_ai.utils.anthropic_utils import stream_anthropic_completion_from_mito_server
+from mito_ai.completions.models import MessageType
+
+STREAMLIT_AI_MODEL = "claude-3-5-haiku-latest"
+
+class StreamlitCodeGeneration:
+    def __init__(self, notebook: dict) -> None:
+
+        self.messages: List[MessageParam] = [
+            cast(MessageParam, {
+                "role": "user",
+                "content": [{
+                    "type": "text",
+                    "text": f"Here is my jupyter notebook content that I want to convert into a Streamlit dashboard - {notebook}"
+                }]
+            })
+        ]
+
+    @property
+    def log(self) -> logging.Logger:
+        """Use Mito AI logger."""
+        return get_logger()
+
+    async def get_response_from_agent(self, message_to_agent: List[MessageParam]) -> str:
+        """Gets the streaming response from the agent using the mito server"""
+        model = STREAMLIT_AI_MODEL
+        max_tokens = 8192 # 64_000
+        temperature = 0.2
+
+        self.log.info("Getting response from agent...")
+        accumulated_response = ""
+        async for stream_chunk in stream_anthropic_completion_from_mito_server(
+            model = model,
+            max_tokens = max_tokens,
+            temperature = temperature,
+            system = streamlit_system_prompt,
+            messages = message_to_agent,
+            stream=True,
+            message_type=MessageType.STREAMLIT_CONVERSION,
+            reply_fn=None,
+            message_id=""
+        ):
+            accumulated_response += stream_chunk
+        return accumulated_response
+
+    def add_agent_response_to_context(self, agent_response: str) -> None:
+        """Add the agent's response to the history"""
+        self.messages.append(
+            cast(MessageParam, {
+                "role": "assistant",
+                "content": [{
+                    "type": "text",
+                    "text": agent_response
+                }]
+            })
+        )
+
+    async def generate_streamlit_code(self) -> str:
+        """Send a query to the agent, get its response and parse the code"""
+        agent_response = await self.get_response_from_agent(self.messages)
+
+        converted_code = extract_code_blocks(agent_response)
+        self.add_agent_response_to_context(converted_code)
+        return converted_code
+
+
+    async def correct_error_in_generation(self, error: str) -> str:
+        """If errors are present, send it back to the agent to get corrections in code"""
+        self.messages.append(
+            cast(MessageParam, {
+                "role": "user",
+                "content": [{
+                    "type": "text",
+                    "text": f"When I run the streamlit app code, I get the following error: {error}\nPlease return the FULL Streamlit app code with the error corrected"
+                }]
+            })
+        )
+        agent_response = await self.get_response_from_agent(self.messages)
+        converted_code = extract_code_blocks(agent_response)
+        self.add_agent_response_to_context(converted_code)
+
+        return converted_code
+
+
+async def streamlit_handler(notebook_path: str, app_path: str) -> Tuple[bool, str]:
+    """Handler function for streamlit code generation and validation"""
+    notebook_code = parse_jupyter_notebook_to_extract_required_content(notebook_path)
+    streamlit_code_generator = StreamlitCodeGeneration(notebook_code)
+    streamlit_code = await streamlit_code_generator.generate_streamlit_code()
+    has_validation_error, error = streamlit_code_validator(streamlit_code)
+    tries = 0
+    while has_validation_error and tries < 5:
+        streamlit_code = await streamlit_code_generator.correct_error_in_generation(error)
+        has_validation_error, error = streamlit_code_validator(streamlit_code)
+        tries+=1
+
+    if has_validation_error:
+        return False, "Error generating streamlit code by agent"
+
+    success_flag, message = create_app_file(app_path, streamlit_code)
+    return success_flag, message

mito_ai/streamlit_conversion/streamlit_system_prompt.py

@@ -0,0 +1,42 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+streamlit_system_prompt = """You are a senior data scientist and Streamlit expert specializing in converting Jupyter notebooks into professional dashboard applications.
+
+ROLE AND EXPERTISE:
+- Expert in Python, Jupyter notebooks, Streamlit, and data visualization
+- Experienced in creating executive-ready dashboards for business stakeholders
+- Skilled in translating technical analysis into clear, interactive presentations
+
+TASK REQUIREMENTS:
+1. Convert Jupyter notebook content into a complete Streamlit application (app.py)
+2. Preserve ALL outputs from code cells and markdown cells as they appear in the notebook
+3. Maintain the logical flow and structure of the original analysis
+4. Create an executive-friendly dashboard suitable for company leadership
+
+STREAMLIT IMPLEMENTATION GUIDELINES:
+- Use appropriate Streamlit components (st.title, st.header, st.subheader, st.markdown, etc.)
+- Display all visualizations using st.pyplot(), st.plotly_chart(), or st.altair_chart() as appropriate
+- Show dataframes and tables using st.dataframe() or st.table()
+- Include all text explanations and insights from markdown cells
+- Add interactive elements where beneficial (filters, selectors, etc.)
+- Ensure professional styling and layout suitable for executives
+
+CODE STRUCTURE:
+- Generate a complete, runnable app.py file
+- Include all necessary imports
+- Handle data loading and processing
+- Organize content with clear sections and headers
+- Include comments explaining key sections
+
+OUTPUT FORMAT:
+- Provide the complete app.py file code
+- Ensure all notebook outputs are faithfully reproduced
+- Make the dashboard professional and presentation-ready
+- Focus on clarity and executive-level communication
+- Don't give extra explanations, just give the python code 
+- Do NOT add emojis 
+- Do NOT modify the graphs or analysis
+- Do NOT provide your own interpretations for the analysis
+
+Remember: The goal is to transform technical analysis into a polished, interactive/visually appealing dashboard that executives can easily understand and navigate."""

mito_ai/streamlit_conversion/streamlit_utils.py

@@ -0,0 +1,96 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import re
+import json
+from typing import Dict, Tuple, Any
+
+def extract_code_blocks(message_content: str) -> str:
+    """
+    Extract all code blocks from Claude's response.
+
+    Args:
+        message_content (str): The actual content from the agent's response
+
+    Returns:
+        str: Removes the ```python``` part to be able to parse the code
+    """
+    if "```python" not in message_content:
+        return message_content
+
+    # return message_content.split('```python\n')[1].split('\n```')[0]
+    # Use regex to find all Python code blocks
+    pattern = r'```python\n(.*?)```'
+    matches = re.findall(pattern, message_content, re.DOTALL)
+
+    # Concatenate with single newlines
+    return '\n'.join(matches)
+
+
+def create_app_file(file_path: str, code: str) -> Tuple[bool, str]:
+    """
+    Create app.py file and write code to it with error handling
+
+    Args:
+        file_path (str): The actual content from the agent's response
+        code (str): The actual content from the agent's response
+
+    Returns:
+        str: Removes the ```python``` part to be able to parse the code
+
+    """
+    try:
+        with open(file_path+"/app.py", 'w') as f:
+            f.write(code)
+        return True, f"Successfully created {file_path}"
+    except IOError as e:
+        return False, f"Error creating file: {str(e)}"
+    except Exception as e:
+        return False, f"Unexpected error: {str(e)}"
+
+
+def parse_jupyter_notebook_to_extract_required_content(notebook_path: str) -> Dict[str, Any]:
+    """
+    Read a Jupyter notebook and filter cells to keep only cell_type and source fields.
+
+    Args:
+        notebook_path (str): Absolute path to the .ipynb file
+
+    Returns:
+        dict: Filtered notebook dictionary with only cell_type and source in cells
+
+    Raises:
+        FileNotFoundError: If the notebook file doesn't exist
+        json.JSONDecodeError: If the file is not valid JSON
+        KeyError: If the notebook doesn't have the expected structure
+    """
+    try:
+        # Read the notebook file
+        with open(notebook_path, 'r', encoding='utf-8') as f:
+            notebook_data: Dict[str, Any] = json.load(f)
+
+        # Check if 'cells' key exists
+        if 'cells' not in notebook_data:
+            raise KeyError("Notebook does not contain 'cells' key")
+
+        # Filter each cell to keep only cell_type and source
+        filtered_cells = []
+        for cell in notebook_data['cells']:
+            filtered_cell = {
+                'cell_type': cell.get('cell_type', ''),
+                'source': cell.get('source', [])
+            }
+            filtered_cells.append(filtered_cell)
+
+        # Update the notebook data with filtered cells
+        notebook_data['cells'] = filtered_cells
+
+        return notebook_data
+
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Notebook file not found: {notebook_path}")
+    except json.JSONDecodeError as e:
+        # JSONDecodeError requires msg, doc, pos
+        raise json.JSONDecodeError(f"Invalid JSON in notebook file: {str(e)}", e.doc if hasattr(e, 'doc') else '', e.pos if hasattr(e, 'pos') else 0)
+    except Exception as e:
+        raise Exception(f"Error processing notebook: {str(e)}")

mito_ai/streamlit_conversion/validate_and_run_streamlit_code.py

@@ -0,0 +1,207 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import subprocess
+import sys
+import os
+import time
+import requests
+import tempfile
+import shutil
+import traceback
+import ast
+import importlib.util
+import warnings
+from typing import Tuple, Optional, Dict, Any
+from subprocess import Popen
+
+# warnings.filterwarnings("ignore", message=r".*missing ScriptRunContext.*")
+# warnings.filterwarnings("ignore", category=UserWarning)
+
+warnings.filterwarnings("ignore", message=".*bare mode.*")
+
+
+class StreamlitValidator:
+    def __init__(self, port: int = 8501, timeout: int = 30) -> None:
+        self.port = port
+        self.timeout = timeout
+        self.process: Optional[Popen[str]] = None
+        self.temp_dir: Optional[str] = None
+
+    def validate_syntax(self, app_code: str) -> Tuple[bool, str]:
+        """Check if the Python code has valid syntax"""
+        try:
+            ast.parse(app_code)
+            return True, "Syntax is valid"
+        except SyntaxError as e:
+            error_msg = ''.join(traceback.format_exception(type(e), e, e.__traceback__))
+            return False, f"Syntax error: {error_msg}"
+
+    def create_temp_app(self, app_code: str) -> str:
+        """Create a temporary Streamlit app file"""
+        self.temp_dir = tempfile.mkdtemp()
+        if self.temp_dir is None:
+            raise RuntimeError("Failed to create temporary directory")
+        app_path = os.path.join(self.temp_dir, "app.py")
+
+        with open(app_path, 'w') as f:
+            f.write(app_code)
+
+        return app_path
+
+    def start_streamlit_app(self, app_path: str) -> Tuple[bool, str]:
+        """Start the Streamlit app in a subprocess"""
+        try:
+            cmd = [
+                sys.executable, "-m", "streamlit", "run", app_path,
+                "--server.port", str(self.port),
+                "--server.headless", "true",
+                "--server.address", "localhost",
+                "--logger.level", "error"
+            ]
+
+            self.process = subprocess.Popen(
+                cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True
+            )
+
+            return True, "Streamlit app started"
+        except Exception as e:
+            return False, f"Failed to start Streamlit: {str(e)}"
+
+    def wait_for_app(self) -> Tuple[bool, str]:
+        """Wait for the Streamlit app to be ready"""
+        start_time = time.time()
+
+        exception_error = "Error"
+        while time.time() - start_time < self.timeout:
+            try:
+                response = requests.get(f"http://localhost:{self.port}", timeout=5)
+                if response.status_code == 200:
+                    return True, "App is running successfully"
+            except requests.exceptions.RequestException as e:
+                exception_error = str(e)
+
+            time.sleep(1)
+
+        return False, f"App failed to start within timeout - {exception_error}"
+
+    def filter_streamlit_warnings(self, text: str) -> str:
+        """Filter out known Streamlit warnings that can be safely ignored"""
+        if not text:
+            return text
+
+        filtered_lines = []
+        for line in text.split('\n'):
+            # Skip lines containing ScriptRunContext warnings
+            if any(phrase in line for phrase in [
+                'missing ScriptRunContext',
+                'bare mode',
+                'ScriptRunContext!',
+                'Thread \'MainThread\':'
+            ]):
+                continue
+            filtered_lines.append(line)
+
+        return '\n'.join(filtered_lines)
+
+    def check_for_errors(self) -> Tuple[bool, str]:
+        """Check if the Streamlit process has any errors"""
+        if self.process:
+            # Check if process is still running
+            if self.process.poll() is not None:
+                stdout, stderr = self.process.communicate()
+                # Filter out known warnings
+                filtered_stderr = self.filter_streamlit_warnings(stderr)
+                if filtered_stderr.strip():
+                    return False, f"App crashed: {filtered_stderr}"
+
+            return True, "App is running without errors"
+
+        return False, "No process found"
+
+    def cleanup(self) -> None:
+        """Clean up the temporary files and stop the process"""
+        if self.process:
+            self.process.terminate()
+            self.process.wait()
+            self.process = None
+
+        if self.temp_dir and os.path.exists(self.temp_dir):
+            shutil.rmtree(self.temp_dir)
+            self.temp_dir = None
+
+    def validate_app(self, app_code: str) -> Dict[str, Any]:
+        """Complete validation pipeline"""
+        results: Dict[str, Any] = {
+            'syntax_valid': False,
+            'app_starts': False,
+            'app_responsive': False,
+            'errors': []
+        }
+
+        try:
+            # Step 1: Check syntax
+            syntax_valid, syntax_msg = self.validate_syntax(app_code)
+            results['syntax_valid'] = syntax_valid
+            if not syntax_valid:
+                results['errors'].append(syntax_msg)
+                return results
+
+            # Step 2: Create and start app
+            app_path = self.create_temp_app(app_code)
+            app_started, start_msg = self.start_streamlit_app(app_path)
+            results['app_starts'] = app_started
+
+            if not app_started:
+                results['errors'].append(start_msg)
+                return results
+
+            # Step 3: Wait for app to be ready
+            app_ready, ready_msg = self.wait_for_app()
+            results['app_responsive'] = app_ready
+
+            if not app_ready:
+                results['errors'].append(ready_msg)
+
+            # Step 4: Check for runtime errors
+            no_errors, error_msg = self.check_for_errors()
+            if not no_errors:
+                results['errors'].append(error_msg)
+
+        except Exception as e:
+            results['errors'].append(f"Validation error: {str(e)}")
+
+        finally:
+            self.cleanup()
+
+        return results
+
+
+def streamlit_code_validator(app_code: str) -> Tuple[bool, str]:
+    """Convenience function to validate Streamlit code"""
+    has_validation_error: bool = False
+    error_message: str = ""
+
+
+    validator = StreamlitValidator()
+    results = validator.validate_app(app_code)
+
+    print("Validation Results:")
+    print(f"✓ Syntax valid: {results['syntax_valid']}")
+    print(f"✓ App starts: {results['app_starts']}")
+    print(f"✓ App responsive: {results['app_responsive']}")
+
+    if results['errors']:
+        error_message = "Errors found: "
+        print("Error detected in agent code")
+        has_validation_error = True
+        print("\nErrors found:")
+        for error in results['errors']:
+            print(f"  - {error}")
+            error_message += error + "\n"
+    if not has_validation_error:
+        print("\nAll validations passed!")
+    return has_validation_error, error_message

mito_ai/tests/providers/test_stream_mito_server_utils.py

@@ -0,0 +1,140 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import asyncio
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch, call
+
+from mito_ai.utils.mito_server_utils import stream_response_from_mito_server
+from mito_ai.completions.models import MessageType
+
+
+# Mock classes for CompletionStreamChunk, CompletionItem, etc.
+class CompletionItem:
+    def __init__(self, content: str, isIncomplete: bool, token: str):
+        self.content = content
+        self.isIncomplete = isIncomplete
+        self.token = token
+
+    def __eq__(self, other):
+        return (
+            self.content == other.content
+            and self.isIncomplete == other.isIncomplete
+            and self.token == other.token
+        )
+
+
+class CompletionStreamChunk:
+    def __init__(self, parent_id: str, chunk: CompletionItem, done: bool):
+        self.parent_id = parent_id
+        self.chunk = chunk
+        self.done = done
+
+    def __eq__(self, other):
+        return (
+            self.parent_id == other.parent_id
+            and self.chunk == other.chunk
+            and self.done == other.done
+        )
+
+
+@pytest.mark.asyncio
+async def test_stream_response_happy_path(monkeypatch):
+    # Arrange
+    url = "https://fake.mito.server/stream"
+    headers = {"Authorization": "Bearer token"}
+    data = {"prompt": "hello world"}
+    timeout = 10
+    max_retries = 2
+    message_type = MessageType.CHAT
+    message_id = "msg-123"
+
+    # Fake chunks
+    raw_chunks = [b"chunk1", b"chunk2"]
+
+    # Mock reply_fn
+    reply_fn = MagicMock()
+
+    # Mock quota check/update
+    monkeypatch.setattr(
+        "mito_ai.utils.mito_server_utils.check_mito_server_quota", lambda *_: None
+    )
+    monkeypatch.setattr(
+        "mito_ai.utils.mito_server_utils.update_mito_server_quota", lambda *_: None
+    )
+
+    # Mock HTTPClient and fetch
+    chunk_callback = MagicMock()
+
+    class FakeHTTPClient:
+        def fetch(self, *args, **kwargs):
+            nonlocal chunk_callback
+            chunk_callback = kwargs["streaming_callback"]
+
+            async def fetch_simulation():
+                # Simulate streaming data
+                for chunk in raw_chunks:
+                    await asyncio.sleep(0.01)
+                    chunk_callback(chunk)
+                return MagicMock()
+
+            return fetch_simulation()
+
+        def close(self):
+            pass
+
+    def mock_create_http_client(timeout_val, retry_val):
+        return FakeHTTPClient(), timeout_val
+
+    monkeypatch.setattr(
+        "mito_ai.utils.mito_server_utils._create_http_client", mock_create_http_client
+    )
+
+    # Act
+    gen = stream_response_from_mito_server(
+        url=url,
+        headers=headers,
+        data=data,
+        timeout=timeout,
+        max_retries=max_retries,
+        message_type=message_type,
+        reply_fn=reply_fn,
+        message_id=message_id,
+    )
+
+    results = []
+    async for chunk in gen:
+        results.append(chunk)
+
+    # Assert
+    assert results == [b"chunk1".decode(), b"chunk2".decode()]
+
+    # Check reply_fn calls
+    expected_calls = [
+        call(
+            CompletionStreamChunk(
+                parent_id=message_id,
+                chunk=CompletionItem(
+                    content="chunk1", isIncomplete=True, token=message_id
+                ),
+                done=False,
+            )
+        ),
+        call(
+            CompletionStreamChunk(
+                parent_id=message_id,
+                chunk=CompletionItem(
+                    content="chunk2", isIncomplete=True, token=message_id
+                ),
+                done=False,
+            )
+        ),
+        call(
+            CompletionStreamChunk(
+                parent_id=message_id,
+                chunk=CompletionItem(content="", isIncomplete=False, token=message_id),
+                done=True,
+            )
+        ),
+    ]
+    reply_fn.assert_has_calls(expected_calls)