mito-ai 0.1.50__py3-none-any.whl

mito_ai/app_manager/utils.py

@@ -0,0 +1,24 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from datetime import datetime, timezone
+
+def convert_utc_to_local_time(time_str: str) -> str:
+    """Convert UTC time to a user's local time"""
+    try:
+        # Remove the 'Z' suffix and parse the UTC datetime
+        utc_time_str = time_str.rstrip('Z')
+        utc_time = datetime.fromisoformat(utc_time_str)
+
+        # Set timezone to UTC
+        utc_time = utc_time.replace(tzinfo=timezone.utc)
+
+        # Convert to local timezone (system timezone)
+        local_time = utc_time.astimezone()
+
+        # Format as 'MMM DD HH:MM'
+        return local_time.strftime('%m-%d-%Y %H:%M')
+
+    except (ValueError, AttributeError) as e:
+        # Return original string if parsing fails
+        return time_str

mito_ai/auth/README.md

@@ -0,0 +1,18 @@
+# Authorization
+
+### Authorization Code Flow:
+
+1. User clicks "Sign In" → then they are redirected to Cognito's hosted UI
+2. User authenticates → then Cognito redirects back with an authorization code
+3. Our backend then exchanges the code for JWT tokens
+
+The authorization code in step2 is a short-lived, one-time use code. 
+To exchange the authorization code for tokens, we have to make a POST request to Cognito's token endpoint:
+POST https://your-domain.auth.region.amazoncognito.com/oauth2/token
+
+The JWT tokens received will provide an hour's session for the user for the next deployments, without having to re-login.
+
+### The response from this request contains 3 JWT tokens:
+1. Access Token - Used to call APIs (expires in 1 hour by default)
+2. ID Token - Contains user identity information (name, email, etc.)
+3. Refresh Token - Used to get new access/ID tokens when they expire

mito_ai/auth/__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 .handlers import AuthHandler
+
+__all__ = ['AuthHandler']

mito_ai/auth/handlers.py

@@ -0,0 +1,96 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import json
+import logging
+import requests
+import tornado
+from datetime import datetime, timezone
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.logger import get_logger
+from mito_ai.constants import ACTIVE_COGNITO_CONFIG
+from typing import Dict, Any
+
+
+class AuthHandler(APIHandler):
+    """Handler for authentication operations."""
+    
+    @property
+    def log(self) -> logging.Logger:
+        """Use Mito AI logger."""
+        return get_logger()
+    
+    @tornado.web.authenticated
+    def post(self) -> None:
+        """Exchange authorization code for JWT tokens."""
+        try:
+            data = json.loads(self.request.body)
+            code = data.get('code')
+            
+            if not code:
+                self.set_status(400)
+                self.finish(json.dumps({"error": "Authorization code is required"}))
+                return
+            
+            # Exchange authorization code for tokens
+            token_response = self._exchange_code_for_tokens(code)
+            
+            if token_response.get('error'):
+                self.set_status(400)
+                self.finish(json.dumps({"error": token_response['error']}))
+                return
+            
+            # Return the tokens to the client
+            self.finish(json.dumps({
+                "access_token": token_response.get('access_token'),
+                "id_token": token_response.get('id_token'),
+                "refresh_token": token_response.get('refresh_token'),
+                "expires_in": token_response.get('expires_in')
+            }))
+            
+        except json.JSONDecodeError:
+            self.set_status(400)
+            self.finish(json.dumps({"error": "Invalid JSON in request body"}))
+        except Exception as e:
+            self.log.error(f"Error in auth handler: {e}")
+            self.set_status(500)
+            self.finish(json.dumps({"error": "Internal server error"}))
+    
+    def _exchange_code_for_tokens(self, code: str) -> Dict[str, Any]:
+        """Exchange authorization code for JWT tokens using AWS Cognito."""
+        try:
+            # Prepare the token request
+            token_data = {
+                'grant_type': 'authorization_code',
+                'client_id': ACTIVE_COGNITO_CONFIG['CLIENT_ID'],
+                'code': code,
+                'redirect_uri': ACTIVE_COGNITO_CONFIG['REDIRECT_URI']
+            }
+            
+            # Add client secret if configured
+            if ACTIVE_COGNITO_CONFIG['CLIENT_SECRET']:
+                token_data['client_secret'] = ACTIVE_COGNITO_CONFIG['CLIENT_SECRET']
+            
+            # Make the token request
+            response = requests.post(
+                ACTIVE_COGNITO_CONFIG['TOKEN_ENDPOINT'],
+                data=token_data,
+                headers={'Content-Type': 'application/x-www-form-urlencoded'}
+            )
+            
+            if response.status_code == 200:
+                token_response: Dict[str, Any] = response.json()
+                
+                current_time = datetime.now(timezone.utc)
+                self.log.info(f"Token exchange successful at {current_time.isoformat()}")
+                return token_response
+            else:
+                self.log.error(f"Token exchange failed: {response.status_code} - {response.text}")
+                return {"error": "Failed to exchange authorization code for tokens"}
+                
+        except requests.exceptions.RequestException as e:
+            self.log.error(f"Request error during token exchange: {e}")
+            return {"error": "Network error during token exchange"}
+        except Exception as e:
+            self.log.error(f"Unexpected error during token exchange: {e}")
+            return {"error": "Unexpected error during token exchange"}

mito_ai/auth/urls.py

@@ -0,0 +1,13 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Tuple
+from jupyter_server.utils import url_path_join
+from .handlers import AuthHandler
+
+
+def get_auth_urls(base_url: str) -> List[Tuple[str, type]]:
+    """Get the auth URL patterns."""
+    return [
+        (url_path_join(base_url, "mito-ai", "auth", "token"), AuthHandler),
+    ]

mito_ai/chat_history/handlers.py

@@ -0,0 +1,63 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import tornado
+from typing import List, Any
+from jupyter_server.base.handlers import APIHandler
+from mito_ai.completions.message_history import GlobalMessageHistory
+from mito_ai.completions.models import ChatThreadMetadata
+
+
+class ChatHistoryHandler(APIHandler):
+    """
+    Endpoints for working with chat history threads.
+    """
+
+    def initialize(self, message_history: GlobalMessageHistory) -> None:
+        """Initialize the handler with the global message history instance."""
+        super().initialize()
+        self._message_history = message_history
+
+    @tornado.web.authenticated
+    def get(self, *args: Any, **kwargs: Any) -> None:
+        """Get all chat history threads or a specific thread by ID."""
+        try:
+            # Check if a specific thread ID is provided in the URL
+            thread_id = kwargs.get("thread_id")
+
+            if thread_id:
+                # Get specific thread
+                if thread_id in self._message_history._chat_threads:
+                    thread = self._message_history._chat_threads[thread_id]
+                    thread_data = {
+                        "thread_id": thread.thread_id,
+                        "name": thread.name,
+                        "creation_ts": thread.creation_ts,
+                        "last_interaction_ts": thread.last_interaction_ts,
+                        "display_history": thread.display_history,
+                        "ai_optimized_history": thread.ai_optimized_history,
+                    }
+                    self.finish(thread_data)
+                else:
+                    self.set_status(404)
+                    self.finish({"error": f"Thread with ID {thread_id} not found"})
+            else:
+                # Get all threads
+                threads: List[ChatThreadMetadata] = self._message_history.get_threads()
+
+                # Convert to dict format for JSON serialization
+                threads_data = [
+                    {
+                        "thread_id": thread.thread_id,
+                        "name": thread.name,
+                        "creation_ts": thread.creation_ts,
+                        "last_interaction_ts": thread.last_interaction_ts,
+                    }
+                    for thread in threads
+                ]
+
+                self.finish({"threads": threads_data})
+
+        except Exception as e:
+            self.set_status(500)
+            self.finish({"error": str(e)})

mito_ai/chat_history/urls.py

@@ -0,0 +1,32 @@
+# 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.chat_history.handlers import ChatHistoryHandler
+from mito_ai.completions.message_history import GlobalMessageHistory
+
+
+def get_chat_history_urls(base_url: str, message_history: GlobalMessageHistory) -> List[Tuple[str, Any, dict]]:
+    """Get all chat history related URL patterns.
+
+    Args:
+        base_url: The base URL for the Jupyter server
+        message_history: The global message history instance
+
+    Returns:
+        List of (url_pattern, handler_class, handler_kwargs) tuples
+    """
+    BASE_URL = base_url + "/mito-ai/chat-history"
+    return [
+        (
+            url_path_join(BASE_URL, "threads"),
+            ChatHistoryHandler,
+            {"message_history": message_history},
+        ),
+        (
+            url_path_join(BASE_URL, "threads", "(?P<thread_id>[^/]+)"),
+            ChatHistoryHandler,
+            {"message_history": message_history},
+        ),
+    ]

mito_ai/completions/completion_handlers/__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/completions/completion_handlers/agent_auto_error_fixup_handler.py

@@ -0,0 +1,59 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.models import AgentResponse, AgentSmartDebugMetadata, MessageType, ResponseFormatInfo
+from mito_ai.completions.prompt_builders.agent_smart_debug_prompt import create_agent_smart_debug_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
+from mito_ai.completions.completion_handlers.utils import append_agent_system_message
+
+__all__ = ["get_agent_auto_error_fixup_completion"]
+
+class AgentAutoErrorFixupHandler(CompletionHandler[AgentSmartDebugMetadata]):
+    """Handler for agent auto error fixup completions."""
+    
+    @staticmethod
+    async def get_completion(
+        metadata: AgentSmartDebugMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get an agent auto error fixup completion from the AI provider."""
+        
+        # Add the system message if it doesn't alredy exist
+        await append_agent_system_message(message_history, model, provider, metadata.threadId, metadata.isChromeBrowser)
+        
+        # Create the prompt
+        prompt = create_agent_smart_debug_prompt(metadata)
+        display_prompt = metadata.errorMessage
+        
+        # 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, metadata.threadId)
+        
+        # Get the completion
+        completion = await provider.request_completions(
+            messages=message_history.get_ai_optimized_history(metadata.threadId), 
+            model=model,
+            response_format_info=ResponseFormatInfo(
+                name='agent_response',
+                format=AgentResponse
+            ),
+            message_type=MessageType.AGENT_AUTO_ERROR_FIXUP,
+            user_input=metadata.errorMessage,
+            thread_id=metadata.threadId
+        )
+        
+        ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
+        
+        await message_history.append_message(ai_response_message, ai_response_message, model, provider, metadata.threadId)
+        
+        return completion
+
+# Use the static method directly
+get_agent_auto_error_fixup_completion = AgentAutoErrorFixupHandler.get_completion

mito_ai/completions/completion_handlers/agent_execution_handler.py

@@ -0,0 +1,66 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from typing import List, Literal, Union
+from openai.types.chat import ChatCompletionMessageParam
+from mito_ai.completions.models import AgentExecutionMetadata, MessageType, ResponseFormatInfo, AgentResponse
+from mito_ai.completions.prompt_builders.agent_execution_prompt import create_agent_execution_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
+from mito_ai.completions.completion_handlers.utils import append_agent_system_message, create_ai_optimized_message
+
+__all__ = ["get_agent_execution_completion"]
+
+class AgentExecutionHandler(CompletionHandler[AgentExecutionMetadata]):
+    """Handler for agent execution completions."""
+    
+    @staticmethod
+    async def get_completion(
+        metadata: AgentExecutionMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get an agent execution completion from the AI provider."""
+
+        if metadata.index is not None:
+            message_history.truncate_histories(
+                thread_id=metadata.threadId,
+                index=metadata.index
+            )
+
+        # Add the system message if it doesn't alredy exist
+        await append_agent_system_message(message_history, model, provider, metadata.threadId, metadata.isChromeBrowser)
+        
+        # Create the prompt
+        prompt = create_agent_execution_prompt(metadata)
+        display_prompt = metadata.input
+        
+        # Add the prompt to the message history
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
+        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, metadata.threadId)
+        
+        # Get the completion
+        completion = await provider.request_completions(
+            messages=message_history.get_ai_optimized_history(metadata.threadId), 
+            model=model,
+            response_format_info=ResponseFormatInfo(
+                name='agent_response',
+                format=AgentResponse
+            ),
+            message_type=MessageType.AGENT_EXECUTION,
+            user_input=metadata.input,
+            thread_id=metadata.threadId
+        )
+        
+        ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
+        
+        await message_history.append_message(ai_response_message, ai_response_message, model, provider, metadata.threadId)
+        
+        return completion
+
+# Use the static method directly
+get_agent_execution_completion = AgentExecutionHandler.get_completion

mito_ai/completions/completion_handlers/chat_completion_handler.py

@@ -0,0 +1,141 @@
+# 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 ChatMessageMetadata, MessageType, CompletionRequest, CompletionStreamChunk, CompletionReply
+from mito_ai.completions.prompt_builders.chat_prompt import create_chat_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
+from mito_ai.completions.completion_handlers.utils import append_chat_system_message, create_ai_optimized_message
+
+__all__ = ["get_chat_completion", "stream_chat_completion"]
+
+class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
+    """Handler for chat completions."""
+    
+    @staticmethod
+    async def get_completion(
+        metadata: ChatMessageMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get a chat completion from the AI provider."""
+
+        if metadata.index is not None:
+            message_history.truncate_histories(
+                index=metadata.index,
+                thread_id=metadata.threadId
+            )
+
+        # Add the system message if it doesn't alredy exist
+        await append_chat_system_message(message_history, model, provider, metadata.threadId)
+        
+        # Create the prompt
+        prompt = create_chat_prompt(
+            metadata.variables or [], 
+            metadata.files or [],
+            metadata.activeCellCode, 
+            metadata.activeCellId,
+            metadata.base64EncodedActiveCellOutput is not None and metadata.base64EncodedActiveCellOutput != '',
+            metadata.input,
+            metadata.additionalContext
+        )
+        display_prompt = f"```python{metadata.activeCellCode or ''}```{metadata.input}"
+        
+        # Add the prompt to the message history
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
+        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, metadata.threadId)
+        
+        # Get the completion (non-streaming)
+        completion = await provider.request_completions(
+            messages=message_history.get_ai_optimized_history(metadata.threadId), 
+            model=model,
+            message_type=MessageType.CHAT,
+            user_input=metadata.input,
+            thread_id=metadata.threadId
+        )
+        
+        ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
+        await message_history.append_message(ai_response_message, ai_response_message, model, provider, metadata.threadId)
+
+        return completion
+    
+    @staticmethod
+    async def stream_completion(
+        metadata: ChatMessageMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        message_id: str,
+        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+        model: str
+    ) -> str:
+        """Stream chat completions from the AI provider.
+        
+        Args:
+            metadata: The metadata for the chat 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.
+        """
+        index = metadata.index
+
+        if index is not None:
+            message_history.truncate_histories(
+                index=index,
+                thread_id=metadata.threadId
+            )
+        
+        # Add the system message if it doesn't already exist
+        await append_chat_system_message(message_history, model, provider, metadata.threadId)
+        
+        # Create the prompt
+        prompt = create_chat_prompt(
+            metadata.variables or [], 
+            metadata.files or [],
+            metadata.activeCellCode, 
+            metadata.activeCellId,
+            metadata.base64EncodedActiveCellOutput is not None and metadata.base64EncodedActiveCellOutput != '',
+            metadata.input,
+            metadata.additionalContext
+        )
+        display_prompt = f"```python{metadata.activeCellCode or ''}```{metadata.input}"
+        
+        # Add the prompt to the message history
+        new_ai_optimized_message = create_ai_optimized_message(prompt, metadata.base64EncodedActiveCellOutput, metadata.additionalContext)
+        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, metadata.threadId)
+        
+        # Stream the completions using the provider's stream method
+        accumulated_response = await provider.stream_completions(
+            message_type=MessageType.CHAT,
+            messages=message_history.get_ai_optimized_history(metadata.threadId),
+            model=model,
+            message_id=message_id,
+            reply_fn=reply_fn,
+            user_input=metadata.input,
+            thread_id=metadata.threadId
+        )
+
+        # Save the accumulated response to message history
+        ai_response_message: ChatCompletionMessageParam = {
+            "role": "assistant",
+            "content": accumulated_response,
+        }
+        await message_history.append_message(
+            ai_response_message, ai_response_message, model, provider, metadata.threadId
+        )
+
+        return accumulated_response
+
+# Use the static methods directly
+get_chat_completion = ChatCompletionHandler.get_completion
+stream_chat_completion = ChatCompletionHandler.stream_completion

mito_ai/completions/completion_handlers/code_explain_handler.py

@@ -0,0 +1,113 @@
+# 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 CodeExplainMetadata, MessageType, CompletionRequest, CompletionStreamChunk, CompletionReply
+from mito_ai.completions.prompt_builders.explain_code_prompt import create_explain_code_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
+from mito_ai.completions.completion_handlers.utils import append_chat_system_message
+
+__all__ = ["get_code_explain_completion", "stream_code_explain_completion"]
+
+class CodeExplainHandler(CompletionHandler[CodeExplainMetadata]):
+    """Handler for code explain completions."""
+    
+    @staticmethod
+    async def get_completion(
+        metadata: CodeExplainMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        model: str
+    ) -> str:
+        """Get a code explain completion from the AI provider."""
+        active_cell_code = metadata.activeCellCode 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_explain_code_prompt(active_cell_code)
+        display_prompt = f"```python{active_cell_code}```"
+
+        # 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.CODE_EXPLAIN,
+            thread_id=thread_id
+        )
+
+        # Add the response to message history
+        ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": completion}
+        await message_history.append_message(ai_response_message, ai_response_message, model, provider, thread_id)
+
+        return completion
+
+    @staticmethod
+    async def stream_completion(
+        metadata: CodeExplainMetadata,
+        provider: OpenAIProvider,
+        message_history: GlobalMessageHistory,
+        message_id: str,
+        reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None],
+        model: str
+    ) -> str:
+        """Stream code explain completions from the AI provider.
+
+        Args:
+            metadata: The metadata for the code explain 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.
+        """
+        active_cell_code = metadata.activeCellCode 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_explain_code_prompt(active_cell_code)
+        display_prompt = f"```python{active_cell_code}```"
+
+        # 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.CODE_EXPLAIN,
+            messages=message_history.get_ai_optimized_history(thread_id),
+            model=model,
+            message_id=message_id,
+            reply_fn=reply_fn,
+            thread_id=thread_id
+        )
+
+        # Add the response to message history
+        ai_response_message: ChatCompletionMessageParam = {"role": "assistant", "content": accumulated_response}
+        await message_history.append_message(ai_response_message, ai_response_message, model, provider, thread_id)
+
+        return accumulated_response
+
+# Use the static methods directly
+get_code_explain_completion = CodeExplainHandler.get_completion
+stream_code_explain_completion = CodeExplainHandler.stream_completion