mito-ai 0.1.50__py3-none-any.whl

mito_ai/openai_client.py

@@ -0,0 +1,382 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from __future__ import annotations
+from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Union
+
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
+import openai
+from openai.types.chat import ChatCompletionMessageParam
+from traitlets import Instance, Unicode, default, validate
+from traitlets.config import LoggingConfigurable
+
+from mito_ai import constants
+from mito_ai.enterprise.utils import is_azure_openai_configured
+from mito_ai.logger import get_logger
+from mito_ai.completions.models import (
+    AICapabilities,
+    CompletionError,
+    CompletionItem,
+    CompletionItemError,
+    CompletionReply,
+    CompletionStreamChunk,
+    MessageType,
+    ResponseFormatInfo,
+)
+from mito_ai.utils.open_ai_utils import (
+    check_mito_server_quota,
+    get_ai_completion_from_mito_server,
+    get_open_ai_completion_function_params,
+    stream_ai_completion_from_mito_server,
+)
+from mito_ai.utils.server_limits import update_mito_server_quota
+from mito_ai.utils.telemetry_utils import (
+    MITO_SERVER_KEY,
+    USER_KEY,
+)
+
+OPENAI_MODEL_FALLBACK = "gpt-4.1"
+
+class OpenAIClient(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.""",
+    )
+
+    # Consider the request a failure if it takes longer than 45 seconds.
+    # We will try a total of 3 times. Once on the initial request,
+    # and then twice more if the first request fails.
+    # Note that max_retries cannot be set to None. If we want to disable it, set it to 0.
+    timeout = 30
+    max_retries = 1
+
+    def __init__(self, **kwargs: Dict[str, Any]) -> None:
+        super().__init__(log=get_logger(), **kwargs)
+        self.last_error = None
+        self._async_client: Optional[openai.AsyncOpenAI] = None
+
+    @default("api_key")
+    def _api_key_default(self) -> Optional[str]:
+        default_key = constants.OPENAI_API_KEY
+        return self._validate_api_key(default_key)
+
+    @validate("api_key")
+    def _validate_api_key(self, api_key: Optional[str]) -> Optional[str]:
+        if not api_key:
+            self.log.debug(
+                "No OpenAI API key provided; following back to Mito server API."
+            )
+            return None
+
+        client = openai.OpenAI(api_key=api_key)
+        try:
+            # Make an http request to OpenAI to make sure it works
+            client.models.list()
+        except openai.AuthenticationError as e:
+            self.log.warning(
+                "Invalid OpenAI API key provided.",
+                exc_info=e,
+            )
+            self.last_error = CompletionError.from_exception(
+                e,
+                hint="You're missing the OPENAI_API_KEY environment variable. Run the following code in your terminal to set the environment variable and then relaunch the jupyter server `export OPENAI_API_KEY=<your-api-key>`",
+            )
+            return None
+        except openai.PermissionDeniedError as e:
+            self.log.warning(
+                "Invalid OpenAI API key provided.",
+                exc_info=e,
+            )
+            self.last_error = CompletionError.from_exception(e)
+            return None
+        except openai.InternalServerError as e:
+            self.log.debug(
+                "Unable to get OpenAI models due to OpenAI error.", exc_info=e
+            )
+            return api_key
+        except openai.RateLimitError as e:
+            self.log.debug(
+                "Unable to get OpenAI models due to rate limit error.", exc_info=e
+            )
+            return api_key
+        except openai.APIConnectionError as e:
+            self.log.warning(
+                "Unable to connect to OpenAI API.",
+                exec_info=e,
+            )
+            self.last_error = CompletionError.from_exception(e)
+            return None
+        else:
+            self.log.debug("User OpenAI API key validated.")
+            return api_key
+        
+    @property
+    def capabilities(self) -> AICapabilities:
+        """Get the provider capabilities."""
+
+        if is_azure_openai_configured():
+            return AICapabilities(
+                configuration={
+                    "model": constants.AZURE_OPENAI_MODEL
+                },
+                provider="Azure OpenAI",
+            )
+
+        if constants.OLLAMA_MODEL and not self.api_key:
+            return AICapabilities(
+                configuration={
+                    "model": constants.OLLAMA_MODEL
+                },
+                provider="Ollama",
+            )
+
+        if self.api_key:
+            self._validate_api_key(self.api_key)
+
+            return AICapabilities(
+                configuration={
+                    "model": OPENAI_MODEL_FALLBACK,
+                },
+                provider="OpenAI (user key)",
+            )
+
+        try:
+            check_mito_server_quota(MessageType.CHAT)
+        except Exception as e:
+            self.log.warning("Failed to set first usage date in user.json", exc_info=e)
+            self.last_error = CompletionError.from_exception(e)
+
+        return AICapabilities(
+            configuration={
+                "model": OPENAI_MODEL_FALLBACK,
+            },
+            provider="Mito server",
+        )
+
+    @property
+    def _active_async_client(self) -> Optional[openai.AsyncOpenAI]:
+        if not self._async_client or self._async_client.is_closed():
+            self._async_client = self._build_openai_client()
+        return self._async_client
+    
+    
+    @property
+    def key_type(self) -> str:
+        """Returns the authentication key type being used."""
+
+        if self.api_key:
+            return USER_KEY
+
+        if constants.OLLAMA_MODEL:
+            return "ollama"
+
+        return MITO_SERVER_KEY
+
+    def _build_openai_client(self) -> Optional[Union[openai.AsyncOpenAI, openai.AsyncAzureOpenAI]]:
+        base_url = None
+        llm_api_key = None
+        
+        if is_azure_openai_configured():
+            self.log.debug(f"Using Azure OpenAI with model: {constants.AZURE_OPENAI_MODEL}")
+                
+            # The format for using Azure OpenAI is different than using
+            # other providers, so we have a special case for it here.
+            # Create Azure OpenAI client with explicit arguments
+            return openai.AsyncAzureOpenAI(
+                api_key=constants.AZURE_OPENAI_API_KEY,
+                api_version=constants.AZURE_OPENAI_API_VERSION,
+                azure_endpoint=constants.AZURE_OPENAI_ENDPOINT, # type: ignore
+                max_retries=self.max_retries,
+                timeout=self.timeout,
+            )
+        
+        elif constants.OLLAMA_MODEL and not self.api_key:
+            base_url = constants.OLLAMA_BASE_URL
+            llm_api_key = "ollama"
+            self.log.debug(f"Using Ollama with model: {constants.OLLAMA_MODEL}")
+        elif self.api_key:
+            llm_api_key = self.api_key
+            self.log.debug("Using OpenAI with user-provided API key")
+        else:
+            self.log.warning("No valid API key or model configuration provided")
+            return None
+
+        # Create the client with explicit arguments to satisfy type checking
+        client = openai.AsyncOpenAI(
+            api_key=llm_api_key,
+            max_retries=self.max_retries,
+            timeout=self.timeout,
+            base_url=base_url if base_url else None,
+        )
+        return client
+
+    def _adjust_model_for_azure_or_ollama(self, model: str) -> str:
+        
+        # If they have set an Azure OpenAI model, then we always use it
+        if is_azure_openai_configured() and constants.AZURE_OPENAI_MODEL is not None:
+            self.log.debug(f"Resolving to Azure OpenAI model: {constants.AZURE_OPENAI_MODEL}")
+            return constants.AZURE_OPENAI_MODEL
+        
+        # If they have set an Ollama model, then we use it
+        if constants.OLLAMA_MODEL is not None:
+            return constants.OLLAMA_MODEL
+        
+        # Otherwise, we use the model they provided
+        return model
+        
+
+    async def request_completions(
+            self,
+            message_type: MessageType,
+            messages: List[ChatCompletionMessageParam],
+            model: str,
+            response_format_info: Optional[ResponseFormatInfo] = None,
+    ) -> str:
+        """
+        Request completions from the OpenAI API.
+
+        Args:
+            message_type: The type of message to request completions for.
+            messages: The messages to request completions for.
+            model: The model to request completions for.
+        Returns:
+            The completion from the OpenAI API.
+        """
+        # Reset the last error
+        self.last_error = None
+        completion = None
+        
+        # Note: We don't catch exceptions here because we want them to bubble up 
+        # to the providers file so we can handle all client exceptions in one place.
+
+        # Handle other providers as before
+        completion_function_params = get_open_ai_completion_function_params(
+            message_type, model, messages, False, response_format_info
+        )
+        
+        # If they have set an Azure OpenAI or Ollama model, then we use it
+        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
+
+        if self._active_async_client is not None:
+            response = await self._active_async_client.chat.completions.create(**completion_function_params)
+            completion = response.choices[0].message.content or ""
+        else: 
+            last_message_content = str(messages[-1].get("content", "")) if messages else None
+            completion = await get_ai_completion_from_mito_server(
+                last_message_content,
+                completion_function_params,
+                self.timeout,
+                self.max_retries,
+                message_type,
+            )
+
+        return completion
+
+    
+    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 OpenAI API and return the accumulated response.
+        Returns: The accumulated response string.
+        """
+        # Reset the last error
+        self.last_error = None
+        accumulated_response = ""
+            
+        # Send initial acknowledgment
+        reply_fn(CompletionReply(
+            items=[
+                CompletionItem(content="", isIncomplete=True, token=message_id)
+            ],
+            parent_id=message_id,
+        ))
+
+        # Handle other providers as before
+        completion_function_params = get_open_ai_completion_function_params(
+            message_type, model, messages, True, response_format_info
+        )
+        
+        completion_function_params["model"] = self._adjust_model_for_azure_or_ollama(completion_function_params["model"])
+
+        try:
+            if self._active_async_client is not None:
+                # Stream from OpenAI
+                client = self._active_async_client
+                if client is None:
+                    raise ValueError("OpenAI client not initialized")
+                
+                stream = await client.chat.completions.create(**completion_function_params)
+
+                async for chunk in stream:
+                    if len(chunk.choices) == 0:
+                        continue
+                    
+                    is_finished = chunk.choices[0].finish_reason is not None
+                    content = chunk.choices[0].delta.content or ""
+                    accumulated_response += content
+                    
+                    reply_fn(CompletionStreamChunk(
+                        parent_id=message_id,
+                        chunk=CompletionItem(
+                            content=content,
+                            isIncomplete=True,
+                            token=message_id,
+                        ),
+                        done=is_finished,
+                    ))
+            else:
+                # Stream from Mito server
+                last_message_content = str(messages[-1].get("content", "")) if messages else ""
+                async for chunk_from_mito_server in stream_ai_completion_from_mito_server(
+                    last_message_content,
+                    completion_function_params,
+                    self.timeout,
+                    self.max_retries,
+                    message_type,
+                    reply_fn=reply_fn,
+                    message_id=message_id,
+                ):
+                    accumulated_response += str(chunk_from_mito_server)
+                
+                # Update quota after streaming is complete
+                update_mito_server_quota(message_type)
+
+            return accumulated_response
+
+        except BaseException as e:
+            self.last_error = CompletionError.from_exception(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/path_utils.py

@@ -0,0 +1,70 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import NewType
+import os
+from mito_ai.utils.error_classes import StreamlitPreviewError
+
+# Type definitions for better type safety
+AbsoluteNotebookPath = NewType('AbsoluteNotebookPath', str)
+AbsoluteNotebookDirPath = NewType('AbsoluteNotebookDirPath', str)
+AbsoluteAppPath = NewType('AbsoluteAppPath', str)
+AppFileName = NewType("AppFileName", str)
+
+def get_absolute_notebook_path(notebook_path: str) -> AbsoluteNotebookPath:
+    """
+    Convert any notebook path to an absolute path.
+    
+    Args:
+        notebook_path: Path to the notebook (can be relative or absolute)
+        
+    Returns:
+        AbsoluteNotebookPath: The absolute path to the notebook
+        
+    Raises:
+        ValueError: If the path is invalid or empty
+    """
+    if not notebook_path or not notebook_path.strip():
+        raise StreamlitPreviewError("Notebook path cannot be empty", 400)
+    
+    absolute_path = os.path.abspath(notebook_path)
+    return AbsoluteNotebookPath(absolute_path)
+    
+
+def get_absolute_notebook_dir_path(notebook_path: AbsoluteNotebookPath) -> AbsoluteNotebookDirPath:
+    """
+    Get the absolute directory containing the notebook.
+    """
+    return AbsoluteNotebookDirPath(os.path.dirname(notebook_path))
+
+def get_absolute_app_path(app_directory: AbsoluteNotebookDirPath, app_file_name: AppFileName) -> AbsoluteAppPath:
+    """
+    Get the absolute path to the app
+    """
+    return AbsoluteAppPath(os.path.join(app_directory, app_file_name))
+
+def get_app_file_name(notebook_id: str) -> AppFileName:
+    """
+    Converts the notebook id into the corresponding app id
+    """
+    mito_app_name = notebook_id.replace('mito-notebook-', 'mito-app-')
+    return AppFileName(f'{mito_app_name}.py')
+
+def does_app_path_exist(app_path: AbsoluteAppPath) -> bool:
+    """
+    Check if the app file exists
+    """
+    return os.path.exists(app_path)
+
+def does_notebook_id_have_corresponding_app(notebook_id: str, notebook_path: str) -> bool:
+    """
+    Given a notebook_id and raw notebook_path checks if the notebook has a corresponding
+    app by converting the notebook_path into an absolute path and converting the notebook_id
+    into an app name
+    """
+
+    app_file_name = get_app_file_name(notebook_id)
+    notebook_path = get_absolute_notebook_path(notebook_path)
+    app_directory = get_absolute_notebook_dir_path(notebook_path)
+    app_path = get_absolute_app_path(app_directory, app_file_name)
+    return does_app_path_exist(app_path)

mito_ai/rules/handlers.py

@@ -0,0 +1,44 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from dataclasses import dataclass
+import json
+from typing import Any, Final, Union
+import tornado
+import os
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.rules.utils import RULES_DIR_PATH, get_all_rules, get_rule, set_rules_file
+
+
+class RulesHandler(APIHandler):
+    """Handler for operations on a specific setting"""
+    
+    @tornado.web.authenticated
+    def get(self, key: Union[str, None] = None) -> None:
+        """Get a specific rule by key or all rules if no key provided"""
+        if key is None or key == '':
+            # No key provided, return all rules
+            rules = get_all_rules()
+            self.finish(json.dumps(rules))
+        else:
+            # Key provided, return specific rule
+            rule_content = get_rule(key)
+            if rule_content is None:
+                self.set_status(404)
+                self.finish(json.dumps({"error": f"Rule with key '{key}' not found"}))
+            else:
+                self.finish(json.dumps({"key": key, "content": rule_content}))
+    
+    @tornado.web.authenticated
+    def put(self, key: str) -> None:
+        """Update or create a specific setting"""
+        data = json.loads(self.request.body)
+        if 'content' not in data:
+            self.set_status(400)
+            self.finish(json.dumps({"error": "Content is required"}))
+            return
+            
+        set_rules_file(key, data['content'])
+        self.finish(json.dumps({"status": "updated", "rules file ": key}))
+
+

mito_ai/rules/urls.py

@@ -0,0 +1,22 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Any, List, Tuple
+from jupyter_server.utils import url_path_join
+from mito_ai.rules.handlers import RulesHandler
+
+def get_rules_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all rules related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    
+    return [
+        (url_path_join(BASE_URL, "rules"), RulesHandler, {}),
+        (url_path_join(BASE_URL, "rules/(.+)"), RulesHandler, {}),
+    ]

mito_ai/rules/utils.py

@@ -0,0 +1,56 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Any, Final, List, Optional
+import os
+from mito_ai.utils.schema import MITO_FOLDER
+
+RULES_DIR_PATH: Final[str] = os.path.join(MITO_FOLDER, 'rules')
+
+def set_rules_file(rule_name: str, value: Any) -> None:
+    """
+    Updates the value of a specific rule file in the rules directory
+    """
+    # Ensure the directory exists
+    if not os.path.exists(RULES_DIR_PATH):
+        os.makedirs(RULES_DIR_PATH)
+    
+    # Create the file path to the rule name as a .md file
+    file_path = os.path.join(RULES_DIR_PATH, f"{rule_name}.md")
+    
+    with open(file_path, 'w+') as f:
+        f.write(value)
+    
+
+def get_rule(rule_name: str) -> Optional[str]:
+    """
+    Retrieves the value of a specific rule file from the rules directory
+    """
+    
+    if rule_name.endswith('.md'):
+        rule_name = rule_name[:-3]
+    
+    file_path = os.path.join(RULES_DIR_PATH, f"{rule_name}.md")
+    
+    if not os.path.exists(file_path):
+        return None
+    
+    with open(file_path, 'r') as f:
+        return f.read()
+
+
+def get_all_rules() -> List[str]:
+    """
+    Retrieves all rule files from the rules directory
+    """
+    # Ensure the directory exists
+    if not os.path.exists(RULES_DIR_PATH):
+        os.makedirs(RULES_DIR_PATH)
+        return []  # Return empty list if directory didn't exist
+    
+    try:
+        return [f for f in os.listdir(RULES_DIR_PATH) if f.endswith('.md')]
+    except OSError as e:
+        # Log the error if needed and return empty list
+        print(f"Error reading rules directory: {e}")
+        return []

mito_ai/settings/handlers.py

@@ -0,0 +1,41 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import tornado
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.settings.utils import (
+    get_settings_field,
+    set_settings_field,
+    ensure_settings_file_exists,
+)
+
+
+class SettingsHandler(APIHandler):
+    """Handler for operations on a specific setting"""
+
+    @tornado.web.authenticated
+    @ensure_settings_file_exists
+    def get(self, key: str) -> None:
+        """Get a specific setting by key"""
+        setting_value = get_settings_field(key)
+        if setting_value is None:
+            self.set_status(404)
+            self.finish(json.dumps({"error": f"Setting with key '{key}' not found"}))
+        else:
+            self.finish(json.dumps({"key": key, "value": setting_value}))
+
+    @tornado.web.authenticated
+    @ensure_settings_file_exists
+    def put(self, key: str) -> None:
+        """Update or create a specific setting"""
+        data = json.loads(self.request.body)
+        if "value" not in data:
+            self.set_status(400)
+            self.finish(json.dumps({"error": "Value is required"}))
+            return
+
+        set_settings_field(key, data["value"])
+        self.finish(
+            json.dumps({"status": "updated", "key": key, "value": data["value"]})
+        )

mito_ai/settings/urls.py

@@ -0,0 +1,20 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Any, List, Tuple
+from jupyter_server.utils import url_path_join
+from mito_ai.settings.handlers import SettingsHandler
+
+def get_settings_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all settings related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai"
+    return [
+        (url_path_join(BASE_URL, "settings/(.*)"), SettingsHandler, {}),
+    ]

mito_ai/settings/utils.py

@@ -0,0 +1,42 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import os
+from typing import Any, Final, Callable
+from functools import wraps
+from mito_ai.utils.schema import MITO_FOLDER
+
+SETTINGS_PATH: Final[str] = os.path.join(MITO_FOLDER, "settings.json")
+
+
+def set_settings_field(field: str, value: Any) -> None:
+    """
+    Updates the value of a specific field in settings.json
+    """
+    with open(SETTINGS_PATH, "r") as user_file_old:
+        old_user_json = json.load(user_file_old)
+        old_user_json[field] = value
+        with open(SETTINGS_PATH, "w+") as f:
+            f.write(json.dumps(old_user_json))
+
+
+def get_settings_field(field: str) -> Any:
+    """
+    Retrieves the value of a specific field from settings.json
+    """
+    with open(SETTINGS_PATH, "r") as user_file_old:
+        return json.load(user_file_old).get(field)
+
+
+def ensure_settings_file_exists(method: Callable) -> Callable:
+    """Decorator to ensure the settings.json file exists before executing the handler method."""
+
+    @wraps(method)
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if not os.path.exists(SETTINGS_PATH):
+            with open(SETTINGS_PATH, "w") as f:
+                json.dump({}, f, indent=4)
+        return method(self, *args, **kwargs)
+
+    return wrapper