mito-ai 0.1.50__py3-none-any.whl

mito_ai/file_uploads/handlers.py

@@ -0,0 +1,248 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import tempfile
+import tornado
+from typing import Dict, Any
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.utils.telemetry_utils import (
+    log_file_upload_attempt,
+    log_file_upload_failure,
+)
+
+MAX_IMAGE_SIZE_MB = 3
+
+
+def _is_image_file(filename: str) -> bool:
+    image_extensions = {
+        ".jpg",
+        ".jpeg",
+        ".png",
+        ".gif",
+        ".bmp",
+        ".tiff",
+        ".tif",
+        ".webp",
+        ".svg",
+    }
+    file_extension = os.path.splitext(filename)[1].lower()
+    return file_extension in image_extensions
+
+
+def _check_image_size_limit(file_data: bytes, filename: str) -> None:
+    if not _is_image_file(filename):
+        return
+
+    file_size_mb = len(file_data) / (1024 * 1024)  # Convert bytes to MB
+
+    if file_size_mb > MAX_IMAGE_SIZE_MB:
+        raise ValueError(f"Image exceeded {MAX_IMAGE_SIZE_MB}MB limit.")
+
+
+class FileUploadHandler(APIHandler):
+    # Class-level dictionary to store temporary directories for each file upload
+    # This persists across handler instances since Tornado recreates handlers per request
+    # Key: filename, Value: dict with temp_dir, total_chunks, received_chunks, logged_upload
+    _temp_dirs: Dict[str, Dict[str, Any]] = {}
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+    @tornado.web.authenticated
+    def post(self) -> None:
+        """Handle file upload with multipart form data."""
+        try:
+            # Validate request has file
+            if not self._validate_file_upload():
+                return
+
+            uploaded_file = self.request.files["file"][0]
+            filename = uploaded_file["filename"]
+            file_data = uploaded_file["body"]
+
+            # Get notebook directory from request
+            notebook_dir = self.get_argument("notebook_dir", ".")
+
+            # Check if this is a chunked upload
+            chunk_number = self.get_argument("chunk_number", None)
+            total_chunks = self.get_argument("total_chunks", None)
+
+            if chunk_number and total_chunks:
+                self._handle_chunked_upload(
+                    filename, file_data, chunk_number, total_chunks, notebook_dir
+                )
+            else:
+                # Log the file upload attempt for regular (non-chunked) uploads
+                file_extension = filename.split(".")[-1].lower()
+                log_file_upload_attempt(filename, file_extension, False, 0)
+                self._handle_regular_upload(filename, file_data, notebook_dir)
+
+            self.finish()
+
+        except Exception as e:
+            self._handle_error(str(e))
+
+    def _validate_file_upload(self) -> bool:
+        """Validate that a file was uploaded in the request."""
+        if "file" not in self.request.files:
+            self._handle_error("No file uploaded", status_code=400)
+            return False
+        return True
+
+    def _handle_chunked_upload(
+        self,
+        filename: str,
+        file_data: bytes,
+        chunk_number: str,
+        total_chunks: str,
+        notebook_dir: str,
+    ) -> None:
+        """Handle chunked file upload."""
+        chunk_num = int(chunk_number)
+        total_chunks_num = int(total_chunks)
+
+        # Log the file upload attempt only for the first chunk
+        if chunk_num == 1:
+            file_extension = filename.split(".")[-1].lower()
+            log_file_upload_attempt(filename, file_extension, True, total_chunks_num)
+
+        # Save chunk to temporary file
+        self._save_chunk(filename, file_data, chunk_num, total_chunks_num)
+
+        # Check if all chunks are received and reconstruct if complete
+        if self._are_all_chunks_received(filename, total_chunks_num):
+            self._reconstruct_file(filename, total_chunks_num, notebook_dir)
+            self._send_chunk_complete_response(filename, notebook_dir)
+        else:
+            self._send_chunk_received_response(chunk_num, total_chunks_num)
+
+    def _handle_regular_upload(
+        self, filename: str, file_data: bytes, notebook_dir: str
+    ) -> None:
+        """Handle regular (non-chunked) file upload."""
+        # Check image file size limit before saving
+        _check_image_size_limit(file_data, filename)
+
+        file_path = os.path.join(notebook_dir, filename)
+        with open(file_path, "wb") as f:
+            f.write(file_data)
+
+        self.write({"success": True, "filename": filename, "path": file_path})
+
+    def _save_chunk(
+        self, filename: str, file_data: bytes, chunk_number: int, total_chunks: int
+    ) -> None:
+        """Save a chunk to a temporary file."""
+        # Initialize temporary directory for this file if it doesn't exist
+        if filename not in self._temp_dirs:
+            temp_dir = tempfile.mkdtemp(prefix=f"mito_upload_{filename}_")
+            self._temp_dirs[filename] = {
+                "temp_dir": temp_dir,
+                "total_chunks": total_chunks,
+                "received_chunks": set(),
+            }
+
+        # Save the chunk to the temporary directory
+        chunk_filename = os.path.join(
+            self._temp_dirs[filename]["temp_dir"], f"chunk_{chunk_number}"
+        )
+        with open(chunk_filename, "wb") as f:
+            f.write(file_data)
+
+        # Mark this chunk as received
+        self._temp_dirs[filename]["received_chunks"].add(chunk_number)
+
+    def _are_all_chunks_received(self, filename: str, total_chunks: int) -> bool:
+        """Check if all chunks for a file have been received."""
+        if filename not in self._temp_dirs:
+            return False
+
+        received_chunks = self._temp_dirs[filename]["received_chunks"]
+        is_complete = len(received_chunks) == total_chunks
+        return is_complete
+
+    def _reconstruct_file(
+        self, filename: str, total_chunks: int, notebook_dir: str
+    ) -> None:
+        """Reconstruct the final file from all chunks and clean up temporary directory."""
+
+        if filename not in self._temp_dirs:
+            raise ValueError(f"No temporary directory found for file: {filename}")
+
+        temp_dir = self._temp_dirs[filename]["temp_dir"]
+        file_path = os.path.join(notebook_dir, filename)
+
+        try:
+            # First, read all chunks to check total file size for images
+            all_file_data = b""
+            for i in range(1, total_chunks + 1):
+                chunk_filename = os.path.join(temp_dir, f"chunk_{i}")
+                with open(chunk_filename, "rb") as chunk_file:
+                    chunk_data = chunk_file.read()
+                    all_file_data += chunk_data
+
+            # Check image file size limit before saving
+            _check_image_size_limit(all_file_data, filename)
+
+            # Write the complete file
+            with open(file_path, "wb") as final_file:
+                final_file.write(all_file_data)
+        finally:
+            # Clean up the temporary directory
+            self._cleanup_temp_dir(filename)
+
+    def _cleanup_temp_dir(self, filename: str) -> None:
+        """Clean up the temporary directory for a file."""
+        if filename in self._temp_dirs:
+            temp_dir = self._temp_dirs[filename]["temp_dir"]
+            try:
+                import shutil
+
+                shutil.rmtree(temp_dir)
+            except Exception as e:
+                # Log the error but don't fail the upload
+                print(
+                    f"Warning: Failed to clean up temporary directory {temp_dir}: {e}"
+                )
+            finally:
+                # Remove from tracking dictionary
+                del self._temp_dirs[filename]
+
+    def _send_chunk_complete_response(self, filename: str, notebook_dir: str) -> None:
+        """Send response indicating all chunks have been processed and file is complete."""
+        file_path = os.path.join(notebook_dir, filename)
+        self.write(
+            {
+                "success": True,
+                "filename": filename,
+                "path": file_path,
+                "chunk_complete": True,
+            }
+        )
+
+    def _send_chunk_received_response(
+        self, chunk_number: int, total_chunks: int
+    ) -> None:
+        """Send response indicating a chunk was received but file is not yet complete."""
+        self.write(
+            {
+                "success": True,
+                "chunk_received": True,
+                "chunk_number": chunk_number,
+                "total_chunks": total_chunks,
+            }
+        )
+
+    def _handle_error(self, error_message: str, status_code: int = 500) -> None:
+        """Handle errors and send appropriate error response."""
+        log_file_upload_failure(error_message)
+        self.set_status(status_code)
+        self.write({"error": error_message})
+        self.finish()
+
+    def on_finish(self) -> None:
+        """Clean up any remaining temporary directories when the handler is finished."""
+        super().on_finish()
+        # Note: We don't clean up here anymore since we want to preserve state across requests
+        # The cleanup happens when the file is fully reconstructed

mito_ai/file_uploads/urls.py

@@ -0,0 +1,21 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple, Any
+from jupyter_server.utils import url_path_join
+from mito_ai.file_uploads.handlers import FileUploadHandler
+
+
+def get_file_uploads_urls(base_url: str) -> List[Tuple[str, Any, dict]]:
+    """Get all file uploads 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, "upload"), FileUploadHandler, {}),
+    ]

mito_ai/gemini_client.py

@@ -0,0 +1,232 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+import base64
+from typing import Any, Callable, Dict, List, Optional, Union, Tuple
+from google import genai
+from google.genai import types
+from google.genai.types import GenerateContentConfig, Part, Content, GenerateContentResponse
+from mito_ai.completions.models import CompletionError, CompletionItem, CompletionReply, CompletionStreamChunk, MessageType, ResponseFormatInfo
+from mito_ai.utils.gemini_utils import get_gemini_completion_from_mito_server, stream_gemini_completion_from_mito_server, get_gemini_completion_function_params
+from mito_ai.utils.mito_server_utils import ProviderCompletionException
+
+def extract_and_parse_gemini_json_response(response: GenerateContentResponse) -> Optional[str]:
+    """
+    Extracts and parses the JSON response from the Gemini API.
+    """
+    if hasattr(response, 'text') and response.text:
+        return response.text
+
+    if hasattr(response, 'candidates') and response.candidates:
+        candidate = response.candidates[0]
+        if hasattr(candidate, 'content') and candidate.content:
+            content = candidate.content
+            if hasattr(content, 'parts') and content.parts:
+                return " ".join(str(part) for part in content.parts)
+            return str(content)
+        return str(candidate)
+    
+    return None
+
+def get_gemini_system_prompt_and_messages(messages: List[Dict[str, Any]]) -> Tuple[str, List[Dict[str, Any]]]:
+    """
+    Converts a list of OpenAI messages to a list of Gemini messages.
+    
+    IMPORTANT: THIS FUNCTION IS ALSO USED IN THE LAMDBA FUNCTION. IF YOU UPDATE IT HERE,
+    YOU PROBABLY NEED TO UPDATE THE LAMBDA FUNCTION AS WELL.
+    """
+
+    system_prompt = ""
+    gemini_messages: List[Dict[str, Any]] = []
+
+    for msg in messages:
+        role = msg.get("role")
+        content = msg.get("content", "")
+
+        if role == "system":
+            if content:
+                # We assume that that there is only one system message 
+                system_prompt = content
+        elif role in ("user", "assistant"):
+            parts: List[Union[Dict[str, str], Part]] = []
+
+            # Handle different content types
+            if isinstance(content, str):
+                # Simple text content
+                if content:
+                    parts.append({"text": content})
+            elif isinstance(content, list):
+                # Mixed content (text + images)
+                for item in content:
+                    if item.get("type") == "text":
+                        text_content = item.get("text", "")
+                        if text_content:
+                            parts.append({"text": text_content})
+                    elif item.get("type") == "image_url":
+                        image_url_data = item.get("image_url", {})
+                        image_url = image_url_data.get("url", "")
+
+                        if image_url and image_url.startswith("data:"):
+                            # Handle base64 data URLs
+                            try:
+                                # Extract the base64 data and mime type
+                                header, base64_data = image_url.split(",", 1)
+                                mime_type = header.split(";")[0].split(":")[1]
+
+                                # Decode base64 to bytes
+                                image_bytes = base64.b64decode(base64_data)
+
+                                # Create Gemini image part
+                                image_part = types.Part.from_bytes(
+                                    data=image_bytes,
+                                    mime_type=mime_type
+                                )
+                                parts.append(image_part)
+                            except Exception as e:
+                                print(f"Error processing image: {e}")
+                                # Skip this image if there's an error
+                                continue
+                            
+            # Only add to contents if we have parts
+            if parts:
+                # Map assistant role to model role for Gemini API
+                gemini_role = "model" if role == "assistant" else "user"
+                gemini_messages.append({
+                    "role": gemini_role,
+                    "parts": parts
+                })
+
+    return system_prompt, gemini_messages
+
+
+class GeminiClient:
+    def __init__(self, api_key: Optional[str]):
+        self.api_key = api_key
+        if api_key:
+            self.client = genai.Client(api_key=api_key)
+
+    async def request_completions(
+        self,
+        messages: List[Dict[str, Any]],
+        model: str,
+        response_format_info: Optional[ResponseFormatInfo] = None,
+        message_type: MessageType = MessageType.CHAT
+    ) -> str:
+        # Extract system instructions and contents
+        system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
+
+        # Get provider data for Gemini completion
+        provider_data = get_gemini_completion_function_params(
+            model=model,
+            contents=contents,
+            message_type=message_type,
+            response_format_info=response_format_info
+        )
+
+        if self.api_key:
+            # Generate content using the Gemini client
+            response_config = GenerateContentConfig(
+                system_instruction=system_instructions,
+                response_mime_type=provider_data.get("config", {}).get("response_mime_type"),
+                response_schema=provider_data.get("config", {}).get("response_schema")
+            )
+            response = self.client.models.generate_content(
+                model=provider_data["model"],
+                contents=contents,  # type: ignore
+                config=response_config
+            )
+            
+            result = extract_and_parse_gemini_json_response(response)
+            
+            if not result:
+                return "No response received from Gemini API"
+
+            return result
+        else:
+            # Fallback to Mito server for completion
+            return await get_gemini_completion_from_mito_server(
+                model=provider_data["model"],
+                contents=messages, # Use the extracted contents instead of converted messages to avoid serialization issues
+                message_type=message_type,
+                config=provider_data.get("config", None),
+                response_format_info=response_format_info,
+            )
+
+    async def stream_completions(
+            self,
+            messages: List[Dict[str, Any]],
+            model: str,
+            message_id: str,
+            reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+            message_type: MessageType = MessageType.CHAT
+    ) -> str:
+        accumulated_response = ""
+        try:
+            # Extract system instructions and Gemini-compatible contents
+            system_instructions, contents = get_gemini_system_prompt_and_messages(messages)
+            if self.api_key:
+                for chunk in self.client.models.generate_content_stream(
+                        model=model,
+                        contents=contents,  # type: ignore
+                        config=GenerateContentConfig(
+                            system_instruction=system_instructions
+                        )
+                ):
+
+                    next_chunk = ""
+                    if hasattr(chunk, 'text'):
+                        next_chunk = chunk.text or ''
+                    else:
+                        next_chunk = str(chunk)
+
+                    accumulated_response += next_chunk
+
+                    # Return the chunk to the frontend
+                    reply_fn(CompletionStreamChunk(
+                        parent_id=message_id,
+                        chunk=CompletionItem(
+                            content=next_chunk or '',
+                            isIncomplete=True,
+                            token=message_id,
+                        ),
+                        done=False,
+                    ))
+
+                # Send final chunk
+                reply_fn(CompletionStreamChunk(
+                    parent_id=message_id,
+                    chunk=CompletionItem(
+                        content="",
+                        isIncomplete=False,
+                        token=message_id,
+                    ),
+                    done=True,
+                ))
+                return accumulated_response
+            else:
+                async for chunk_text in stream_gemini_completion_from_mito_server(
+                        model=model,
+                        contents=messages,  # Use the extracted contents instead of converted messages to avoid serialization issues
+                        message_type=message_type,
+                        message_id=message_id,
+                        reply_fn=reply_fn
+                ):
+                    # Clean and decode the chunk text
+                    clean_chunk = chunk_text.strip('"')
+                    decoded_chunk = clean_chunk.encode().decode('unicode_escape')
+                    accumulated_response += decoded_chunk or ''
+
+                # Send final chunk with the complete response
+                reply_fn(CompletionStreamChunk(
+                    parent_id=message_id,
+                    chunk=CompletionItem(
+                        content=accumulated_response,
+                        isIncomplete=False,
+                        token=message_id,
+                    ),
+                    done=True,
+                ))
+
+                return accumulated_response
+
+        except Exception as e:
+            return f"Error streaming content: {str(e)}"

mito_ai/log/handlers.py

@@ -0,0 +1,38 @@
+# 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, Literal
+import tornado
+import os
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.utils.telemetry_utils import MITO_SERVER_KEY, USER_KEY, log
+
+
+class LogHandler(APIHandler):
+    """Handler for logging"""
+    
+    def initialize(self, key_type: Literal['mito_server_key', 'user_key']) -> None:
+        """Initialize the log handler"""
+        
+        # The key_type is required so that we know if we can log pro users
+        self.key_type = key_type
+        
+    @tornado.web.authenticated
+    def put(self) -> None:
+        """Log an event"""
+        data = json.loads(self.request.body)
+        
+        if 'log_event' not in data:
+            self.set_status(400)
+            self.finish(json.dumps({"error": "Log event is required"}))
+            return
+        
+        log_event = data['log_event']
+        params = data.get('params', {})
+        
+        key_type = MITO_SERVER_KEY if self.key_type == "mito_server_key" else USER_KEY
+        log(log_event, params, key_type=key_type)
+
+

mito_ai/log/urls.py

@@ -0,0 +1,21 @@
+# 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.log.handlers import LogHandler
+
+def get_log_urls(base_url: str, key_type: str) -> List[Tuple[str, Any, dict]]:
+    """Get all log 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, "log"), LogHandler, {"key_type": key_type}),
+    ]

mito_ai/logger.py

@@ -0,0 +1,37 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import logging
+
+from traitlets.config import Application
+
+
+_LOGGER = None  # type: logging.Logger | None
+
+
+def get_logger() -> logging.Logger:
+    """Create a logger for the Mito AI module.
+    
+    The logger will be attached as a child of the Jupyter server logger.
+    This allows for easier filtering/flagging of log messages in the server logs.
+    
+    Example:
+
+    The following snippet shows a log message produced by JupyterLab followed by
+    two log messages produced by the Mito AI module.
+
+        [D 2024-12-16 15:49:07.333 LabApp] 204 PUT /lab/api/workspaces/default?1734360547329 (ea7486428da24ff3921aebd4422611d9@::1) 1.85ms
+        [D 2024-12-16 15:49:08.293 ServerApp.mito_ai] Message received: {...}
+        [D 2024-12-16 15:49:08.293 ServerApp.mito_ai] Requesting completion from Mito server.
+
+    You can filter the server logs to only show messages produced by the Mito AI module:
+
+        jupyter lab --debug 2>&1 | egrep mito_ai
+    """
+    global _LOGGER
+    if _LOGGER is None:
+        app = Application.instance()
+        _LOGGER = logging.getLogger("{!s}.mito_ai".format(app.log.name))
+        Application.clear_instance()
+
+    return _LOGGER