mito-ai 0.1.50__py3-none-any.whl

mito_ai/completions/completion_handlers/completion_handler.py

@@ -0,0 +1,42 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import Protocol, TypeVar
+from abc import abstractmethod, ABCMeta
+from mito_ai.completions.models import ChatMessageMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, InlineCompleterMetadata, AgentSmartDebugMetadata
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.completions.message_history import GlobalMessageHistory
+
+T = TypeVar('T', ChatMessageMetadata, SmartDebugMetadata, CodeExplainMetadata, AgentExecutionMetadata, AgentSmartDebugMetadata, InlineCompleterMetadata, contravariant=True)
+
+class CompletionHandler(Protocol[T], metaclass=ABCMeta):
+    """Protocol defining the interface for completion handlers.
+    
+    All completion handler classes should implement this protocol to ensure
+    they provide a get_completion static method with the correct signature.
+    """
+    
+    @staticmethod
+    @abstractmethod
+    async def get_completion(
+        metadata: T,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get a completion from the AI provider.
+        
+        Args:
+            metadata: Metadata about the completion request
+            provider: The AI provider to use
+            message_history: The history of messages in the conversation
+            model: The model to use for the completion
+            
+        Returns:
+            The completion string from the AI
+        """
+        pass
+    
+    
+    
+    

mito_ai/completions/completion_handlers/inline_completer_handler.py

@@ -0,0 +1,48 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.models import InlineCompleterMetadata, MessageType
+from mito_ai.completions.prompt_builders.inline_completer_prompt import create_inline_prompt
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.completions.completion_handlers.completion_handler import CompletionHandler
+
+__all__ = ["get_inline_completion"]
+
+class InlineCompleterHandler(CompletionHandler[InlineCompleterMetadata]):
+    """Handler for inline completions."""
+    
+    @staticmethod
+    async def get_completion(
+        metadata: InlineCompleterMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get an inline completion from the AI provider."""
+        
+        # Create the prompt
+        prompt = create_inline_prompt(
+            metadata.prefix or '', 
+            metadata.suffix or '', 
+            metadata.variables or [],
+            metadata.files or []
+        )
+        
+        # Each inline completer is independent and ephemeral. So we do not use the message history. 
+        messages: List[ChatCompletionMessageParam] = [{"role": "user", "content": prompt}]
+        
+        # Get the completion
+        completion = await provider.request_completions(
+            messages=messages, 
+            model=model,
+            message_type=MessageType.INLINE_COMPLETION,
+            thread_id=None
+        )
+        
+        return completion
+
+# Use the static method directly
+get_inline_completion = InlineCompleterHandler.get_completion

mito_ai/completions/completion_handlers/smart_debug_handler.py

@@ -0,0 +1,160 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Union, AsyncGenerator, Callable
+
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.models import (
+    SmartDebugMetadata,
+    MessageType,
+    CompletionRequest,
+    CompletionStreamChunk,
+    CompletionReply,
+)
+from mito_ai.completions.prompt_builders.smart_debug_prompt import (
+    create_error_prompt,
+    remove_inner_thoughts_from_message,
+)
+from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.completions.completion_handlers.completion_handler import CompletionHandler
+from mito_ai.completions.completion_handlers.utils import append_chat_system_message
+
+__all__ = ["get_smart_debug_completion", "stream_smart_debug_completion"]
+
+
+class SmartDebugHandler(CompletionHandler[SmartDebugMetadata]):
+    """Handler for smart debug completions."""
+
+    @staticmethod
+    async def get_completion(
+        metadata: SmartDebugMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get a smart debug completion from the AI provider."""
+
+        error_message = metadata.errorMessage
+        active_cell_code = metadata.activeCellCode
+        active_cell_id = metadata.activeCellId
+        variables = metadata.variables or []
+        files = metadata.files or []
+        thread_id = metadata.threadId
+
+        # Add the system message if it doesn't already exist
+        await append_chat_system_message(message_history, model, provider, thread_id)
+
+        # Create the prompt
+        prompt = create_error_prompt(error_message, active_cell_code, active_cell_id, variables, files)
+        display_prompt = f"```python{active_cell_code or ''}```{error_message}"
+
+        # Add the prompt to the message history
+        new_ai_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": prompt}
+        new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": display_prompt}
+        await message_history.append_message(
+            new_ai_optimized_message, new_display_optimized_message, model, provider, thread_id
+        )
+
+        # Get the completion
+        completion = await provider.request_completions(
+            messages=message_history.get_ai_optimized_history(thread_id),
+            model=model,
+            message_type=MessageType.SMART_DEBUG,
+            user_input=error_message,
+            thread_id=thread_id
+        )
+
+        # Process the completion to remove inner thoughts
+        display_completion = remove_inner_thoughts_from_message(completion)
+
+        # Add the response to message history
+        ai_response_message: ChatCompletionMessageParam = {
+            "role": "assistant",
+            "content": completion,
+        }
+        display_response_message: ChatCompletionMessageParam = {
+            "role": "assistant",
+            "content": display_completion,
+        }
+        await message_history.append_message(
+            ai_response_message, display_response_message, model, provider, thread_id
+        )
+
+        return display_completion
+
+    @staticmethod
+    async def stream_completion(
+        metadata: SmartDebugMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        message_id: str,
+        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+        model: str
+    ) -> str:
+        """Stream smart debug completions from the AI provider.
+
+        Args:
+            metadata: The metadata for the smart debug completion request.
+            provider: The AI provider to use.
+            message_history: The message history for this conversation.
+            message_id: The ID of the message being processed.
+            reply_fn: Function to call with each chunk for streaming replies.
+
+        Returns:
+            The accumulated response string.
+        """
+        error_message = metadata.errorMessage
+        active_cell_code = metadata.activeCellCode
+        active_cell_id = metadata.activeCellId
+        variables = metadata.variables or []
+        files = metadata.files or []
+        thread_id = metadata.threadId
+
+        # Add the system message if it doesn't already exist
+        await append_chat_system_message(message_history, model, provider, thread_id)
+
+        # Create the prompt
+        prompt = create_error_prompt(error_message, active_cell_code, active_cell_id, variables, files)
+        display_prompt = f"```python{active_cell_code or ''}```{error_message}"
+
+        # Add the prompt to the message history
+        new_ai_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": prompt}
+        new_display_optimized_message: ChatCompletionMessageParam = {"role": "user", "content": display_prompt}
+        await message_history.append_message(
+            new_ai_optimized_message, new_display_optimized_message, model, provider, thread_id
+        )
+
+        # Stream the completions using the provider's stream method
+        accumulated_response = await provider.stream_completions(
+            message_type=MessageType.SMART_DEBUG,
+            messages=message_history.get_ai_optimized_history(thread_id),
+            model=model,
+            message_id=message_id,
+            reply_fn=reply_fn,
+            user_input=error_message,
+            thread_id=thread_id
+        )
+
+        # Process the completion to remove inner thoughts
+        display_completion = remove_inner_thoughts_from_message(accumulated_response)
+
+        # Save the accumulated response to message history
+        ai_response_message: ChatCompletionMessageParam = {
+            "role": "assistant",
+            "content": accumulated_response,
+        }
+        display_response_message: ChatCompletionMessageParam = {
+            "role": "assistant",
+            "content": display_completion,
+        }
+        await message_history.append_message(
+            ai_response_message, display_response_message, model, provider, thread_id
+        )
+
+        return display_completion
+
+
+# Use the static methods directly
+get_smart_debug_completion = SmartDebugHandler.get_completion
+stream_smart_debug_completion = SmartDebugHandler.stream_completion

mito_ai/completions/completion_handlers/utils.py

@@ -0,0 +1,147 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import base64
+from typing import Optional, Union, List, Dict, Any, cast
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.completions.models import ThreadID
+from mito_ai.completions.providers import OpenAIProvider
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.prompt_builders.chat_system_message import (
+    create_chat_system_message_prompt,
+)
+from mito_ai.completions.prompt_builders.agent_system_message import (
+    create_agent_system_message_prompt,
+)
+
+
+async def append_chat_system_message(
+    message_history: GlobalMessageHistory,
+    model: str,
+    provider: OpenAIProvider,
+    thread_id: ThreadID,
+) -> None:
+
+    # If the system message already exists, do nothing
+    if any(
+        msg["role"] == "system"
+        for msg in message_history.get_ai_optimized_history(thread_id)
+    ):
+        return
+
+    system_message_prompt = create_chat_system_message_prompt()
+
+    system_message: ChatCompletionMessageParam = {
+        "role": "system",
+        "content": system_message_prompt,
+    }
+
+    await message_history.append_message(
+        ai_optimized_message=system_message,
+        display_message=system_message,
+        model=model,
+        llm_provider=provider,
+        thread_id=thread_id,
+    )
+
+
+async def append_agent_system_message(
+    message_history: GlobalMessageHistory,
+    model: str,
+    provider: OpenAIProvider,
+    thread_id: ThreadID,
+    isChromeBrowser: bool,
+) -> None:
+
+    # If the system message already exists, do nothing
+    if any(
+        msg["role"] == "system"
+        for msg in message_history.get_ai_optimized_history(thread_id)
+    ):
+        return
+
+    system_message_prompt = create_agent_system_message_prompt(isChromeBrowser)
+
+    system_message: ChatCompletionMessageParam = {
+        "role": "system",
+        "content": system_message_prompt,
+    }
+
+    await message_history.append_message(
+        ai_optimized_message=system_message,
+        display_message=system_message,
+        model=model,
+        llm_provider=provider,
+        thread_id=thread_id,
+    )
+
+
+def extract_and_encode_images_from_additional_context(
+    additional_context: Optional[List[Dict[str, str]]],
+) -> List[str]:
+    encoded_images = []
+
+    for context in additional_context or []:
+        if context["type"].startswith("image/"):
+            try:
+                with open(context["value"], "rb") as image_file:
+                    image_data = image_file.read()
+                    base64_encoded = base64.b64encode(image_data).decode("utf-8")
+                    encoded_images.append(f"data:{context['type']};base64,{base64_encoded}")
+            except (FileNotFoundError, IOError) as e:
+                print(f"Error reading image file {context['value']}: {e}")
+                continue
+
+    return encoded_images
+
+
+def create_ai_optimized_message(
+    text: str,
+    base64EncodedActiveCellOutput: Optional[str] = None,
+    additional_context: Optional[List[Dict[str, str]]] = None,
+) -> ChatCompletionMessageParam:
+
+    message_content: Union[str, List[Dict[str, Any]]]
+    encoded_images = extract_and_encode_images_from_additional_context(
+        additional_context
+    )
+
+    has_uploaded_image = len(encoded_images) > 0
+    has_active_cell_output = (
+        base64EncodedActiveCellOutput is not None
+        and base64EncodedActiveCellOutput != ""
+    )
+
+    if has_uploaded_image or has_active_cell_output:
+        message_content = [
+            {
+                "type": "text",
+                "text": text,
+            }
+        ]
+
+        for img in encoded_images:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": img
+                    },
+                }
+            )
+
+        if has_active_cell_output:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{base64EncodedActiveCellOutput}"
+                    },
+                }
+            )
+    else:
+        message_content = text
+
+    return cast(
+        ChatCompletionMessageParam, {"role": "user", "content": message_content}
+    )