mito-ai 0.1.50__py3-none-any.whl

mito_ai/completions/prompt_builders/utils.py

@@ -0,0 +1,84 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Optional, Dict
+from mito_ai.rules.utils import get_rule
+
+
+def get_rules_str(additional_context: Optional[List[Dict[str, str]]]) -> str:
+    """
+    Extract the rules from the additional context array, and retrieve the rule content.
+    """
+    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 == "":
+            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"]
+    selected_images = [context["value"] for context in additional_context if context.get("type", "").startswith("image/")]
+
+    # 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)
+        )
+
+    if len(selected_images) > 0:
+        context_parts.append(
+            "The following images have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_images)
+        )
+
+    # STEP 3: Combine into a single string
+    return "\n\n".join(context_parts)
+
+
+def get_streamlit_app_status_str(notebook_id: str, notebook_path: str) -> str:
+    """
+    Get the streamlit app status string.
+    """
+    from mito_ai.path_utils import does_notebook_id_have_corresponding_app
+    if does_notebook_id_have_corresponding_app(notebook_id, notebook_path):
+        return "The notebook has an existing Streamlit app that you can edit"
+    return "The notebook does not have an existing Streamlit app. If you want to show an app to the user, you must create a new one."
+
+
+

mito_ai/completions/providers.py

@@ -0,0 +1,284 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from __future__ import annotations
+import asyncio
+from typing import Any, Callable, Dict, List, Optional, Union, cast
+from mito_ai import constants
+from openai.types.chat import ChatCompletionMessageParam
+from traitlets import Instance, Unicode, default, validate
+from traitlets.config import LoggingConfigurable
+from openai.types.chat import ChatCompletionMessageParam
+
+from mito_ai import constants
+from mito_ai.enterprise.utils import is_azure_openai_configured
+from mito_ai.gemini_client import GeminiClient
+from mito_ai.openai_client import OpenAIClient
+from mito_ai.anthropic_client import AnthropicClient
+from mito_ai.logger import get_logger
+from mito_ai.completions.models import (
+    AICapabilities,
+    CompletionError,
+    CompletionItem,
+    CompletionItemError,
+    CompletionReply,
+    CompletionStreamChunk,
+    MessageType,
+    ResponseFormatInfo, CompletionItemError,
+)
+from mito_ai.utils.telemetry_utils import (
+    KEY_TYPE_PARAM,
+    MITO_AI_COMPLETION_ERROR,
+    MITO_AI_COMPLETION_RETRY,
+    MITO_SERVER_KEY,
+    USER_KEY,
+    log,
+    log_ai_completion_error,
+    log_ai_completion_retry,
+    log_ai_completion_success,
+)
+from mito_ai.utils.provider_utils import get_model_provider
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
+
+__all__ = ["OpenAIProvider"]
+
+class OpenAIProvider(LoggingConfigurable):
+    """Provide AI feature through OpenAI services."""
+
+    api_key = Unicode(
+        config=True,
+        allow_none=True,
+        help="OpenAI API key. Default value is read from the OPENAI_API_KEY environment variable.",
+    )
+
+    last_error = Instance(
+        CompletionError,
+        allow_none=True,
+        help="""Last error encountered when using the OpenAI provider.
+
+This attribute is observed by the websocket provider to push the error to the client.""",
+    )
+
+    def __init__(self, **kwargs: Dict[str, Any]) -> None:
+        config = kwargs.get('config', {})
+        if 'api_key' in kwargs:
+            config['OpenAIClient'] = {'api_key': kwargs['api_key']}
+        kwargs['config'] = config
+
+        super().__init__(log=get_logger(), **kwargs)
+        self.last_error = None
+        self._openai_client: Optional[OpenAIClient] = OpenAIClient(**config)
+
+    @property
+    def capabilities(self) -> AICapabilities:
+        """
+        Returns the capabilities of the AI provider.
+        """
+        if constants.CLAUDE_API_KEY and not self.api_key:
+            return AICapabilities(
+                configuration={"model": "<dynamic>"},
+                provider="Claude",
+            )
+        if constants.GEMINI_API_KEY and not self.api_key:
+            return AICapabilities(
+                configuration={"model": "<dynamic>"},
+                provider="Gemini",
+            )
+        if self._openai_client:
+            return self._openai_client.capabilities
+        
+        return AICapabilities(
+            configuration={"model": "<dynamic>"},
+            provider="Mito server",
+        )
+
+    @property
+    def key_type(self) -> str:
+        if constants.CLAUDE_API_KEY and not self.api_key:
+            return "claude"
+        if constants.GEMINI_API_KEY and not self.api_key:
+            return "gemini"
+        if self._openai_client:
+            return self._openai_client.key_type
+        return MITO_SERVER_KEY
+
+    async def request_completions(
+        self,
+        message_type: MessageType,
+        messages: List[ChatCompletionMessageParam],
+        model: str,
+        response_format_info: Optional[ResponseFormatInfo] = None,
+        user_input: Optional[str] = None,
+        thread_id: Optional[str] = None,
+        max_retries: int = 3
+    ) -> str:
+        """
+        Request completions from the AI provider.
+        """
+        self.last_error = None
+        completion = None
+        last_message_content = str(messages[-1].get('content', '')) if messages else ""
+        model_type = get_model_provider(model)
+        
+        # Retry loop
+        for attempt in range(max_retries + 1):
+            try:
+                if model_type == "claude":
+                    api_key = constants.CLAUDE_API_KEY
+                    anthropic_client = AnthropicClient(api_key=api_key)
+                    completion = await anthropic_client.request_completions(messages, model, response_format_info, message_type)
+                elif model_type == "gemini":
+                    api_key = constants.GEMINI_API_KEY
+                    gemini_client = GeminiClient(api_key=api_key)
+                    messages_for_gemini = [dict(m) for m in messages]
+                    completion = await gemini_client.request_completions(messages_for_gemini, model, response_format_info, message_type)
+                elif model_type == "openai":
+                    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=model,
+                        response_format_info=response_format_info
+                    )
+                else:
+                    raise ValueError(f"No AI provider configured for model: {model}")
+                
+                # Success! Log and return
+                log_ai_completion_success(
+                    key_type=USER_KEY if self.key_type == "user" else MITO_SERVER_KEY,
+                    message_type=message_type,
+                    last_message_content=last_message_content,
+                    response={"completion": completion},
+                    user_input=user_input or "",
+                    thread_id=thread_id or "",
+                    model=model
+                )
+                return completion # type: ignore
+            
+            except PermissionError as e:
+                # If we hit a free tier limit, then raise an exception right away without retrying.
+                self.log.exception(f"Error during request_completions: {e}")
+                self.last_error = CompletionError.from_exception(e)
+                log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
+                raise
+
+            except BaseException as e:
+                # Check if we should retry (not on the last attempt)
+                if attempt < max_retries:
+                    # Exponential backoff: wait 2^attempt seconds
+                    wait_time = 2 ** attempt
+                    self.log.info(f"Retrying request_completions after {wait_time}s (attempt {attempt + 1}/{max_retries + 1}): {str(e)}")
+                    log_ai_completion_retry('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
+                    await asyncio.sleep(wait_time)
+                    continue
+                else:
+                    # Final failure after all retries - set error state and raise
+                    self.log.exception(f"Error during request_completions after {attempt + 1} attempts: {e}")
+                    self.last_error = CompletionError.from_exception(e)
+                    log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id or "", message_type, e)
+                    raise
+        
+        # This should never be reached due to the raise in the except block,
+        # but added to satisfy the linter
+        raise RuntimeError("Unexpected code path in request_completions")
+
+    async def stream_completions(
+        self,
+        message_type: MessageType,
+        messages: List[ChatCompletionMessageParam],
+        model: str,
+        message_id: str,
+        thread_id: str,
+        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+        user_input: Optional[str] = None,
+        response_format_info: Optional[ResponseFormatInfo] = None
+    ) -> str:
+        """
+        Stream completions from the AI provider and return the accumulated response.
+        Returns: The accumulated response string.
+        """
+        self.last_error = None
+        accumulated_response = ""
+        last_message_content = str(messages[-1].get('content', '')) if messages else ""
+        model_type = get_model_provider(model)
+        reply_fn(CompletionReply(
+            items=[
+                CompletionItem(content="", isIncomplete=True, token=message_id)
+            ],
+            parent_id=message_id,
+        ))
+
+        try:
+            if model_type == "claude":
+                api_key = constants.CLAUDE_API_KEY
+                anthropic_client = AnthropicClient(api_key=api_key)
+                accumulated_response = await anthropic_client.stream_completions(
+                    messages=messages,
+                    model=model,
+                    message_type=message_type,
+                    message_id=message_id,
+                    reply_fn=reply_fn
+                )
+            elif model_type == "gemini":
+                api_key = constants.GEMINI_API_KEY
+                gemini_client = GeminiClient(api_key=api_key)
+                # TODO: We shouldn't need to do this because the messages should already be dictionaries... 
+                # but if we do have to do some pre-processing, we should do it in the gemini_client instead.
+                messages_for_gemini = [dict(m) for m in messages]
+                accumulated_response = await gemini_client.stream_completions(
+                    messages=messages_for_gemini,
+                    model=model,
+                    message_id=message_id,
+                    reply_fn=reply_fn,
+                    message_type=message_type
+                )
+            elif model_type == "openai":
+                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=model,
+                    message_id=message_id,
+                    thread_id=thread_id,
+                    reply_fn=reply_fn,
+                    user_input=user_input,
+                    response_format_info=response_format_info
+                )
+            else:
+                raise ValueError(f"No AI provider configured for model: {model}")
+
+            # Log the successful completion
+            log_ai_completion_success(
+                key_type=USER_KEY if self.key_type == "user" else MITO_SERVER_KEY,
+                message_type=message_type,
+                last_message_content=last_message_content,
+                response={"completion": accumulated_response},
+                user_input=user_input or "",
+                thread_id=thread_id,
+                model=model
+            )
+            return accumulated_response
+
+        except BaseException as e:
+            self.log.exception(f"Error during stream_completions: {e}")
+            self.last_error = CompletionError.from_exception(e)
+            log_ai_completion_error('user_key' if self.key_type != MITO_SERVER_KEY else 'mito_server_key', thread_id, message_type, e)
+
+            # Send error message to client before raising
+            reply_fn(CompletionStreamChunk(
+                parent_id=message_id,
+                chunk=CompletionItem(
+                    content="",
+                    isIncomplete=True,
+                    error=CompletionItemError(
+                        message=f"Failed to process completion: {e!r}"
+                    ),
+                    token=message_id,
+                ),
+                done=True,
+                error=CompletionError.from_exception(e),
+            ))
+            raise
+

mito_ai/constants.py

@@ -0,0 +1,63 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+from typing import Union
+
+# Claude
+CLAUDE_API_KEY = os.environ.get("CLAUDE_API_KEY")
+
+# Gemini
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+
+# Ollama
+OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL")
+OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434/v1")
+
+# OpenAI
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+
+# Azure OpenAI Config 
+AZURE_OPENAI_API_KEY = os.environ.get("AZURE_OPENAI_API_KEY")
+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")
+
+# Mito AI Base URLs and Endpoint Paths
+MITO_PROD_BASE_URL = "https://7eax4i53f5odkshhlry4gw23by0yvnuv.lambda-url.us-east-1.on.aws/v2"
+MITO_DEV_BASE_URL = "https://g5vwmogjg7gh7aktqezyrvcq6a0hyfnr.lambda-url.us-east-1.on.aws/v2"
+MITO_LOCAL_BASE_URL = "http://127.0.0.1:8000/v2" # When you are running the mito completion server locally
+
+# Set ACTIVE_BASE_URL manually
+ACTIVE_BASE_URL = MITO_PROD_BASE_URL  # Change to MITO_DEV_BASE_URL for dev
+
+# Endpoint paths
+ANTHROPIC_PATH = "anthropic/completions"
+GEMINI_PATH = "gemini/completions"
+OPENAI_PATH = "openai/completions"
+
+# Full URLs (always use ACTIVE_BASE_URL)
+MITO_ANTHROPIC_URL = f"{ACTIVE_BASE_URL}/{ANTHROPIC_PATH}"
+MITO_GEMINI_URL = f"{ACTIVE_BASE_URL}/{GEMINI_PATH}"
+MITO_OPENAI_URL = f"{ACTIVE_BASE_URL}/{OPENAI_PATH}"
+
+# Streamlit conversion endpoints
+MITO_STREAMLIT_DEV_BASE_URL = "https://fr12uvtfy5.execute-api.us-east-1.amazonaws.com"
+MITO_STREAMLIT_TEST_BASE_URL = "https://iyual08t6d.execute-api.us-east-1.amazonaws.com"
+
+# Set ACTIVE_BASE_URL manually
+# TODO: Modify to PROD url before release
+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
+
+
+MESSAGE_HISTORY_TRIM_THRESHOLD: int = 3

mito_ai/db/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+

mito_ai/db/crawlers/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai.db.crawlers.snowflake import crawl_snowflake
+
+__all__ = ["crawl_snowflake"]

mito_ai/db/crawlers/base_crawler.py

@@ -0,0 +1,61 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Dict, Any
+from sqlalchemy import create_engine, text
+from sqlalchemy.exc import SQLAlchemyError
+from mito_ai.db.crawlers.constants import SUPPORTED_DATABASES
+from mito_ai.db.models import ColumnInfo, TableSchema
+
+
+def crawl_db(conn_str: str, db_type: str) -> Dict[str, Any]:
+    try:
+        if db_type == "mssql":
+            # For Microsoft SQL Server, we need to trust the server certificate
+            engine = create_engine(
+                conn_str, connect_args={"TrustServerCertificate": "yes"}
+            )
+        else:
+            engine = create_engine(conn_str)
+
+        tables: List[str] = []
+        schema: TableSchema = {"tables": {}}
+        tables_query = SUPPORTED_DATABASES[db_type].get("tables_query", "")
+        columns_query = SUPPORTED_DATABASES[db_type].get("columns_query", "")
+
+        # Get a list of all tables in the database
+        with engine.connect() as connection:
+            # Use parameterized query for safety
+            result = connection.execute(text(tables_query), {"schema": "public"})
+            tables = [row[0] for row in result]
+
+            # For each table, get the column names and data types
+            for table in tables:
+                if db_type == "mysql":
+                    # For MySQL we have to use string formatting
+                    # since MySQL doesn't support parameter binding
+                    query = columns_query.format(table=table)
+                    columns = connection.execute(text(query))
+                else:
+                    # For other databases, use parameter binding
+                    columns = connection.execute(text(columns_query), {"table": table})
+                # Create a list of dictionaries with column name and type
+                column_info: List[ColumnInfo] = [
+                    {"name": row[0], "type": row[1]} for row in columns
+                ]
+                schema["tables"][table] = column_info
+
+        return {
+            "schema": schema,
+            "error": None,
+        }
+    except SQLAlchemyError as e:
+        return {
+            "schema": None,
+            "error": f"Database error: {str(e)}",
+        }
+    except Exception as e:
+        return {
+            "schema": None,
+            "error": f"Unexpected error: {str(e)}",
+        }

mito_ai/db/crawlers/constants.py

@@ -0,0 +1,43 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Dict, TypedDict, List
+
+
+class DatabaseConfig(TypedDict, total=False):
+    drivers: List[str]
+    tables_query: str
+    columns_query: str
+
+
+SUPPORTED_DATABASES: Dict[str, DatabaseConfig] = {
+    "mssql": {
+        "drivers": ["pyodbc"],
+        "tables_query": "SELECT table_name FROM information_schema.tables WHERE table_schema = 'dbo'",
+        "columns_query": "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = :table",
+    },
+    "mysql": {
+        "drivers": ["PyMySQL"],
+        "tables_query": "SHOW TABLES",
+        "columns_query": "SHOW COLUMNS FROM {table}",
+    },
+    "oracle": {
+        "drivers": ["oracledb"],
+        "tables_query": "SELECT table_name FROM user_tables",
+        "columns_query": "SELECT column_name, data_type FROM user_tab_columns WHERE table_name = :table",
+    },
+    "postgres": {
+        "drivers": ["psycopg2-binary"],
+        "tables_query": "SELECT table_name FROM information_schema.tables WHERE table_schema = :schema",
+        "columns_query": "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = :table",
+    },
+    "snowflake": {
+        "drivers": ["snowflake-sqlalchemy"],
+        # Queries handled in the snowflake.py file.
+    },
+    "sqlite": {
+        "drivers": [],
+        "tables_query": "SELECT name FROM sqlite_master WHERE type='table'",
+        "columns_query": "SELECT name, type FROM pragma_table_info(:table)",
+    },
+}

mito_ai/db/crawlers/snowflake.py

@@ -0,0 +1,71 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from sqlalchemy import create_engine, text
+from mito_ai.db.models import WarehouseDetails
+
+SUPPORTED_DATABASE_KINDS = ["STANDARD", "IMPORTED DATABASE"]
+
+
+def crawl_snowflake(username: str, password: str, account: str, warehouse: str) -> dict:
+    try:
+        conn_str = (
+            f"snowflake://{username}:{password}@{account}/" f"?warehouse={warehouse}"
+        )
+        engine = create_engine(conn_str)
+
+        # Step 1: Get databases
+        db_query = text("SHOW DATABASES")
+        with engine.connect() as connection:
+            result = connection.execute(db_query)
+            databases = result.mappings().all()
+
+        # Step 2: Filter for 'STANDARD' and 'IMPORTED DATABASE' kinds
+        filtered_db_names = [
+            row["name"] for row in databases if row["kind"] in SUPPORTED_DATABASE_KINDS
+        ]
+
+        # Step 3: Escape and format for SQL IN clause
+        db_list_sql = ", ".join(f"'{name}'" for name in filtered_db_names)
+
+        # Step 4: Use filtered databases in schema query
+        schema_query = text(
+            f"""
+            SELECT 
+                c.table_catalog AS database_name,
+                c.table_schema AS schema_name,
+                c.table_name,
+                c.column_name,
+                c.data_type,
+                c.comment
+            FROM snowflake.account_usage.columns c
+            WHERE c.deleted IS NULL
+            AND c.table_catalog IN ({db_list_sql})
+        """
+        )
+
+        # Step 5: Execute the schema query and process the results
+        warehouse_details: WarehouseDetails = {"databases": {}}
+
+        with engine.connect() as connection:
+            result = connection.execute(schema_query)
+            for row in result.fetchall():
+                db = row[0]
+                schema = row[1]
+                table = row[2]
+                column_name = row[3]
+
+                db_dict = warehouse_details["databases"].setdefault(db, {"schemas": {}})
+                schema_dict = db_dict["schemas"].setdefault(schema, {"tables": {}})
+                columns_list = schema_dict["tables"].setdefault(table, [])
+                columns_list.append({"name": column_name, "type": row[4]})
+
+        return {
+            "schema": warehouse_details,
+            "error": None,
+        }
+    except Exception as e:
+        return {
+            "schema": None,
+            "error": str(e),
+        }