mito-ai 0.1.44__py3-none-any.whl → 0.1.46__py3-none-any.whl

mito_ai/__init__.py

@@ -5,6 +5,7 @@ from typing import List, Dict
 from jupyter_server.utils import url_path_join
 from mito_ai.completions.handlers import CompletionHandler
 from mito_ai.completions.providers import OpenAIProvider
+from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.app_deploy.handlers import AppDeployHandler
 from mito_ai.streamlit_preview.handlers import StreamlitPreviewHandler
 from mito_ai.log.urls import get_log_urls
@@ -16,6 +17,8 @@ from mito_ai.auth.urls import get_auth_urls
 from mito_ai.streamlit_preview.urls import get_streamlit_preview_urls
 from mito_ai.app_manager.handlers import AppManagerHandler
 from mito_ai.file_uploads.urls import get_file_uploads_urls
+from mito_ai.user.urls import get_user_urls
+from mito_ai.chat_history.urls import get_chat_history_urls
 
 # Force Matplotlib to use the Jupyter inline backend.
 # Background: importing Streamlit sets os.environ["MPLBACKEND"] = "Agg" very early.
@@ -62,13 +65,17 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     base_url = web_app.settings["base_url"]
 
     open_ai_provider = OpenAIProvider(config=server_app.config)
+    
+    # Create a single GlobalMessageHistory instance for the entire server
+    # This ensures thread-safe access to the .mito/ai-chats directory
+    global_message_history = GlobalMessageHistory()
 
     # WebSocket handlers
     handlers = [
         (
             url_path_join(base_url, "mito-ai", "completions"),
             CompletionHandler,
-            {"llm": open_ai_provider},
+            {"llm": open_ai_provider, "message_history": global_message_history},
         ),
         (
             url_path_join(base_url, "mito-ai", "app-deploy"),
@@ -100,6 +107,8 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     handlers.extend(get_auth_urls(base_url))  # type: ignore
     handlers.extend(get_streamlit_preview_urls(base_url))  # type: ignore
     handlers.extend(get_file_uploads_urls(base_url)) # type: ignore
+    handlers.extend(get_user_urls(base_url)) # type: ignore
+    handlers.extend(get_chat_history_urls(base_url, global_message_history)) # type: ignore
 
     web_app.add_handlers(host_pattern, handlers)
     server_app.log.info("Loaded the mito_ai server extension")

mito_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.44'
+__version__ = VERSION = '0.1.46'

mito_ai/anthropic_client.py

@@ -5,9 +5,9 @@ import json
 import anthropic
 from typing import Dict, Any, Optional, Tuple, Union, Callable, List, cast
 
-from anthropic.types import Message, MessageParam
-from mito_ai.completions.models import CompletionError, ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
-from mito_ai.utils.mito_server_utils import ProviderCompletionException
+from anthropic.types import Message, MessageParam, TextBlockParam
+from mito_ai.completions.models import ResponseFormatInfo, CompletionReply, CompletionStreamChunk, CompletionItem, MessageType
+from mito_ai.constants import MESSAGE_HISTORY_TRIM_THRESHOLD
 from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.utils.anthropic_utils import get_anthropic_completion_from_mito_server, stream_anthropic_completion_from_mito_server, get_anthropic_completion_function_params
 
@@ -52,12 +52,12 @@ def extract_and_parse_anthropic_json_response(response: Message) -> Union[object
 
 
 def get_anthropic_system_prompt_and_messages(messages: List[ChatCompletionMessageParam]) -> Tuple[
-    Union[str, anthropic.NotGiven], List[MessageParam]]:
+    Union[str, anthropic.Omit], List[MessageParam]]:
     """
     Convert a list of OpenAI messages to a list of Anthropic messages.
     """
     
-    system_prompt: Union[str, anthropic.NotGiven] = anthropic.NotGiven()
+    system_prompt: Union[str, anthropic.Omit] = anthropic.Omit()
     anthropic_messages: List[MessageParam] = []
 
     for message in messages:
@@ -125,6 +125,90 @@ def get_anthropic_system_prompt_and_messages(messages: List[ChatCompletionMessag
     return system_prompt, anthropic_messages
 
 
+def add_cache_control_to_message(message: MessageParam) -> MessageParam:
+    """
+    Adds cache_control to a message's content.
+    Handles both string content and list of content blocks.
+    """
+    content = message.get("content")
+    
+    if isinstance(content, str):
+        # Simple string content - convert to list format with cache_control
+        return {
+            "role": message["role"],
+            "content": [
+                {
+                    "type": "text",
+                    "text": content,
+                    "cache_control": {"type": "ephemeral"}
+                }
+            ]
+        }
+    
+    elif isinstance(content, list) and len(content) > 0:
+        # List of content blocks - add cache_control to last block
+        content_blocks = content.copy()
+        last_block = content_blocks[-1].copy()
+        last_block["cache_control"] = {"type": "ephemeral"}
+        content_blocks[-1] = last_block
+        
+        return {
+            "role": message["role"],
+            "content": content_blocks
+        }
+    
+    else:
+        # Edge case: empty or malformed content
+        return message
+
+
+def get_anthropic_system_prompt_and_messages_with_caching(messages: List[ChatCompletionMessageParam]) -> Tuple[
+    Union[str, List[TextBlockParam], anthropic.Omit], List[MessageParam]]:
+    """
+    Convert a list of OpenAI messages to a list of Anthropic messages with caching applied.
+    
+    Caching Strategy:
+    1. System prompt (static) → Always cached
+    2. Stable conversation history → Cache at keep_recent boundary
+    3. Recent messages → Never cached (always fresh)
+    
+    The keep_recent parameter determines which messages are stable and won't be trimmed.
+    We cache at the keep_recent boundary because those messages are guaranteed to be stable.
+    """
+    
+    # Get the base system prompt and messages
+    system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+    
+    # 1. Cache the system prompt always
+    # If the system prompt is something like anthropic.Omit, we don't need to cache it
+    cached_system_prompt: Union[str, List[TextBlockParam], anthropic.Omit] = system_prompt
+    if isinstance(system_prompt, str):
+        cached_system_prompt = [{
+            "type": "text",
+            "text": system_prompt,
+            "cache_control": {"type": "ephemeral"}
+        }]
+    
+    # 2. Cache conversation history at the boundary where the messages are stable.
+    # Messages are stable after they are more than MESSAGE_HISTORY_TRIM_THRESHOLD old.
+    # At this point, the messages are not edited anymore, so they will not invalidate the cache.
+    # If we included the messages before the boundary in the cache, then every time we send a new
+    # message, we would invalidate the cache and we would never get a cache hit except for the system prompt.
+    messages_with_cache = []
+    
+    if len(anthropic_messages) > 0:
+        cache_boundary = len(anthropic_messages) - MESSAGE_HISTORY_TRIM_THRESHOLD - 1
+        
+        # Add all messages, but only add cache_control to the message at the boundary
+        for i, msg in enumerate(anthropic_messages):
+            if i == cache_boundary:
+                messages_with_cache.append(add_cache_control_to_message(msg))
+            else:
+                messages_with_cache.append(msg)
+    
+    return cached_system_prompt, messages_with_cache
+
+
 class AnthropicClient:
     """
     A client for interacting with the Anthropic API or the Mito server fallback.
@@ -149,7 +233,7 @@ class AnthropicClient:
         """
         Get a response from Claude or the Mito server that adheres to the AgentResponse format.
         """
-        anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+        anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
         
         provider_data = get_anthropic_completion_function_params(
             message_type=message_type,
@@ -166,6 +250,7 @@ class AnthropicClient:
             # Unpack provider_data for direct API call
             assert self.client is not None
             response = self.client.messages.create(**provider_data)
+            
             if provider_data.get("tool_choice") is not None:
                 result = extract_and_parse_anthropic_json_response(response)
                 return json.dumps(result) if not isinstance(result, str) else result
@@ -192,7 +277,7 @@ class AnthropicClient:
     async def stream_completions(self, messages: List[ChatCompletionMessageParam], model: str, message_id: str, message_type: MessageType,
                               reply_fn: Callable[[Union[CompletionReply, CompletionStreamChunk]], None]) -> str:
         try:
-            anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages(messages)
+            anthropic_system_prompt, anthropic_messages = get_anthropic_system_prompt_and_messages_with_caching(messages)
             accumulated_response = ""
 
             if self.api_key:
@@ -206,7 +291,6 @@ class AnthropicClient:
                     stream=True
                 )
 
-
                 for chunk in stream:
                     if chunk.type == "content_block_delta" and chunk.delta.type == "text_delta":
                         content = chunk.delta.text

mito_ai/app_deploy/app_deploy_utils.py

@@ -0,0 +1,25 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import os
+import zipfile
+import logging
+from typing import List, Optional
+
+def add_files_to_zip(zip_path: str, base_path: str, files_to_add: List[str], logger: Optional[logging.Logger] = None) -> None:
+    """Create a zip file at zip_path and add the selected files/folders."""
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for rel_path in files_to_add:
+            abs_path = os.path.join(base_path, rel_path)
+
+            if os.path.isfile(abs_path):
+                zipf.write(abs_path, arcname=rel_path)
+            elif os.path.isdir(abs_path):
+                for root, _, files in os.walk(abs_path):
+                    for file in files:
+                        file_abs = os.path.join(root, file)
+                        arcname = os.path.relpath(file_abs, base_path)
+                        zipf.write(file_abs, arcname=arcname)
+            else:
+                if logger:
+                    logger.warning(f"Skipping missing file: {abs_path}")

mito_ai/app_deploy/handlers.py

@@ -4,13 +4,13 @@
 import os
 import time
 import logging
-from typing import Any, Union, Optional
-import zipfile
+from typing import Any, Union, List
 import tempfile
 from mito_ai.streamlit_conversion.streamlit_utils import get_app_path
 from mito_ai.utils.create import initialize_user
 from mito_ai.utils.version_utils import is_pro
 from mito_ai.utils.websocket_base import BaseWebSocketHandler
+from mito_ai.app_deploy.app_deploy_utils import  add_files_to_zip
 from mito_ai.app_deploy.models import (
     DeployAppReply,
     AppDeployError,
@@ -18,7 +18,6 @@ from mito_ai.app_deploy.models import (
     ErrorMessage,
     MessageType
 )
-from mito_ai.streamlit_conversion.streamlit_agent_handler import streamlit_handler
 from mito_ai.logger import get_logger
 from mito_ai.constants import ACTIVE_STREAMLIT_BASE_URL
 import requests
@@ -111,6 +110,7 @@ class AppDeployHandler(BaseWebSocketHandler):
         message_id = message.message_id
         notebook_path = message.notebook_path
         jwt_token = message.jwt_token
+        files_to_upload = message.selected_files
         
         if not message_id:
             self.log.error("Missing message_id in request")
@@ -168,7 +168,7 @@ class AppDeployHandler(BaseWebSocketHandler):
                 ))
             
             # Finally, deploy the app
-            deploy_url = await self._deploy_app(app_directory, jwt_token)
+            deploy_url = await self._deploy_app(app_directory, files_to_upload, jwt_token)
 
             # Send the response
             self.reply(DeployAppReply(
@@ -219,11 +219,12 @@ class AppDeployHandler(BaseWebSocketHandler):
             return False
 
 
-    async def _deploy_app(self, app_path: str, jwt_token: str = '') -> str:
+    async def _deploy_app(self, app_path: str, files_to_upload:List[str], jwt_token: str = '') -> str:
         """Deploy the app using pre-signed URLs.
         
         Args:
             app_path: Path to the app file.
+            files_to_upload: Files the user selected to upload for the app to run
             jwt_token: JWT token for authentication (optional)
 
         Returns:
@@ -258,16 +259,12 @@ class AppDeployHandler(BaseWebSocketHandler):
             # Step 2: Create a zip file of the app.
             temp_zip_path = None
             try:
-                # Create temp file and close it before writing to avoid file handle conflicts
-                with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as temp_zip:
+                # Create temp file
+                with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as temp_zip:
                     temp_zip_path = temp_zip.name
 
                 self.log.info("Zipping application files...")
-                with zipfile.ZipFile(temp_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-                    for root, _, files in os.walk(app_path):
-                        for file in files:
-                            file_path = os.path.join(root, file)
-                            zipf.write(file_path, arcname=os.path.relpath(file_path, app_path))
+                add_files_to_zip(temp_zip_path, app_path, files_to_upload, self.log)
 
                 upload_response = await self._upload_app_to_s3(temp_zip_path, presigned_url)
             except Exception as e:

mito_ai/app_deploy/models.py

@@ -3,7 +3,7 @@
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import Literal, Optional
+from typing import Literal, Optional, List
 
 
 class MessageType(str, Enum):
@@ -66,6 +66,9 @@ class DeployAppRequest:
     
     # Path to the app file.
     notebook_path: str
+
+    # Files to be uploaded for the app to run
+    selected_files: List[str]
     
     # JWT token for authorization.
     jwt_token: Optional[str] = None

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/handlers.py

@@ -14,6 +14,7 @@ import tornado.web
 from jupyter_core.utils import ensure_async
 from jupyter_server.base.handlers import JupyterHandler
 from tornado.websocket import WebSocketHandler
+from openai.types.chat import ChatCompletionMessageParam
 from mito_ai.completions.message_history import GlobalMessageHistory
 from mito_ai.logger import get_logger
 from mito_ai.completions.models import (
@@ -48,11 +49,8 @@ from mito_ai.utils.telemetry_utils import identify
 
 FALLBACK_MODEL = "gpt-4.1"  # Default model to use for safety
 
-# The GlobalMessageHistory is responsible for updating the message histories stored in the .mito/ai-chats directory.
-# We create one GlobalMessageHistory per backend server instance instead of one per websocket connection so that the
-# there is one manager of the locks for the .mito/ai-chats directory. This is my current understanding and it 
-# might be incorrect!
-message_history = GlobalMessageHistory()
+# The GlobalMessageHistory is now created in __init__.py and passed to handlers
+# to ensure there's only one instance managing the .mito/ai-chats directory locks
 
 # This handler is responsible for the mito_ai/completions endpoint.
 # It takes a message from the user, sends it to the OpenAI API, and returns the response.
@@ -61,10 +59,11 @@ message_history = GlobalMessageHistory()
 class CompletionHandler(JupyterHandler, WebSocketHandler):
     """Completion websocket handler."""
 
-    def initialize(self, llm: OpenAIProvider) -> None:
+    def initialize(self, llm: OpenAIProvider, message_history: GlobalMessageHistory) -> None:
         super().initialize()
         self.log.debug("Initializing websocket connection %s", self.request.path)
         self._llm = llm
+        self._message_history = message_history
         self.is_pro = is_pro()
         self._selected_model = FALLBACK_MODEL
         self.is_electron = False
@@ -149,7 +148,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
 
         # Clear history if the type is "start_new_chat"
         if type == MessageType.START_NEW_CHAT:
-            thread_id = message_history.create_new_thread()
+            thread_id = self._message_history.create_new_thread()
             
             reply = StartNewChatReply(
                 parent_id=parsed_message.get("message_id"),
@@ -160,7 +159,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
 
         # Handle get_threads: return list of chat threads
         if type == MessageType.GET_THREADS:
-            threads = message_history.get_threads()
+            threads = self._message_history.get_threads()
             reply = FetchThreadsReply(
                 parent_id=parsed_message.get("message_id"),
                 threads=threads
@@ -172,7 +171,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
         if type == MessageType.DELETE_THREAD:
             thread_id_to_delete = metadata_dict.get('thread_id')
             if thread_id_to_delete:
-                is_thread_deleted = message_history.delete_thread(thread_id_to_delete)
+                is_thread_deleted = self._message_history.delete_thread(thread_id_to_delete)
                 reply = DeleteThreadReply(
                     parent_id=parsed_message.get("message_id"),
                     success=is_thread_deleted
@@ -188,7 +187,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             
             # If a thread_id is provided, use that thread's history; otherwise, use newest.
             thread_id = metadata_dict.get('thread_id')
-            display_history = message_history.get_display_history(thread_id)
+            display_history = self._message_history.get_display_history(thread_id)
             
             reply = FetchHistoryReply(
                 parent_id=parsed_message.get('message_id'),
@@ -222,7 +221,32 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                 )
                 self.reply(reply)
             return
-        
+
+        if type == MessageType.STOP_AGENT:
+            thread_id_to_stop = metadata_dict.get('threadId')
+            if thread_id_to_stop:
+                self.log.info(f"Stopping agent, thread ID: {thread_id_to_stop}")
+                
+                ai_optimized_message: ChatCompletionMessageParam = {
+                    "role": "assistant",
+                    "content": "The user made the following request: Stop processing my last request. I want to change it. Please answer my future requests without going back and finising my previous request."
+                }
+                display_optimized_message: ChatCompletionMessageParam = {
+                    "role": "assistant",
+                    "content": "Agent interupted by user "
+                }
+                
+                await self._message_history.append_message(
+                    ai_optimized_message=ai_optimized_message,
+                    display_message=display_optimized_message,
+                    model=self._selected_model,
+                    llm_provider=self._llm,
+                    thread_id=thread_id_to_stop
+                )
+            else:
+                self.log.info("Trying to stop agent, but no thread ID available")
+            return
+
         try:
             # Get completion based on message type
             completion = None
@@ -240,7 +264,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_chat_completion(
                         chat_metadata, 
                         self._llm, 
-                        message_history, 
+                        self._message_history, 
                         message_id,
                         self.reply,
                         model
@@ -248,7 +272,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_chat_completion(chat_metadata, self._llm, message_history, model)
+                    completion = await get_chat_completion(chat_metadata, self._llm, self._message_history, model)
             elif type == MessageType.SMART_DEBUG:
                 smart_debug_metadata = SmartDebugMetadata(**metadata_dict)
                 # Handle streaming if requested and available
@@ -257,7 +281,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_smart_debug_completion(
                         smart_debug_metadata, 
                         self._llm, 
-                        message_history, 
+                        self._message_history, 
                         message_id,
                         self.reply,
                         model
@@ -265,7 +289,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, message_history, model)
+                    completion = await get_smart_debug_completion(smart_debug_metadata, self._llm, self._message_history, model)
             elif type == MessageType.CODE_EXPLAIN:
                 code_explain_metadata = CodeExplainMetadata(**metadata_dict)
 
@@ -275,7 +299,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     await stream_code_explain_completion(
                         code_explain_metadata, 
                         self._llm, 
-                        message_history,
+                        self._message_history,
                         message_id,
                         self.reply,
                         model
@@ -283,16 +307,16 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
                     return
                 else:
                     # Regular non-streaming completion
-                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, message_history, model)
+                    completion = await get_code_explain_completion(code_explain_metadata, self._llm, self._message_history, model)
             elif type == MessageType.AGENT_EXECUTION:
                 agent_execution_metadata = AgentExecutionMetadata(**metadata_dict)
-                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, message_history, model)
+                completion = await get_agent_execution_completion(agent_execution_metadata, self._llm, self._message_history, model)
             elif type == MessageType.AGENT_AUTO_ERROR_FIXUP:
                 agent_auto_error_fixup_metadata = AgentSmartDebugMetadata(**metadata_dict)
-                completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_metadata, self._llm, message_history, model)
+                completion = await get_agent_auto_error_fixup_completion(agent_auto_error_fixup_metadata, self._llm, self._message_history, model)
             elif type == MessageType.INLINE_COMPLETION:
                 inline_completer_metadata = InlineCompleterMetadata(**metadata_dict)
-                completion = await get_inline_completion(inline_completer_metadata, self._llm, message_history, model)
+                completion = await get_inline_completion(inline_completer_metadata, self._llm, self._message_history, model)
             else:
                 raise ValueError(f"Invalid message type: {type}")
             

mito_ai/completions/models.py

@@ -64,6 +64,7 @@ class MessageType(Enum):
     DELETE_THREAD = "delete_thread"
     UPDATE_MODEL_CONFIG = "update_model_config"
     STREAMLIT_CONVERSION = "streamlit_conversion"
+    STOP_AGENT = "stop_agent"
 
     
 @dataclass(frozen=True)

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -125,15 +125,33 @@ If the user has requested data that you believe is stored in the database:
 connections[connection_name]["username"]
 ```
 
+- The user may colloquially ask for a "list of x", always assume they want a pandas DataFrame. 
+- When working with dataframes created from an SQL query, ALWAYS use lowercase column names. 
+- If you think the requested data is stored in the database, but you are unsure, then ask the user for clarification.
+
+## Additional MSSQL Rules
+
+- When connecting to a Microsoft SQL Server (MSSQL) database, use the following format:
+
+```
+import urllib.parse
+
+encoded_password = urllib.parse.quote_plus(password) 
+conn_str = f"mssql+pyodbc://username:encoded_password@host:port/database?driver=ODBC+Driver+18+for+SQL+Server&TrustServerCertificate=yes"
+```
+
+- Always URL-encode passwords for MSSQL connections to handle special characters properly.
+- Include the port number in MSSQL connection strings.
+- Use "ODBC+Driver+18+for+SQL+Server" (with plus signs) in the driver parameter.
+- Always include "TrustServerCertificate=yes" for MSSQL connections to avoid SSL certificate issues.
+
+## Additional Oracle Rules
+
 - When connecting to an Oracle database, use the following format:
 ```
 conn_str = f"oracle+oracledb://username:password@host:port?service_name=service_name"
 ```
 
-- The user may colloquially ask for a "list of x", always assume they want a pandas DataFrame. 
-- When working with dataframes created from an SQL query, ALWAYS use lowercase column names. 
-- If you think the requested data is stored in the database, but you are unsure, then ask the user for clarification.
-
 Here is the schema:
 {schemas}
         """

mito_ai/constants.py

@@ -58,3 +58,6 @@ COGNITO_CONFIG_DEV = {
 }
 
 ACTIVE_COGNITO_CONFIG = COGNITO_CONFIG_DEV # Change to COGNITO_CONFIG_DEV for dev
+
+
+MESSAGE_HISTORY_TRIM_THRESHOLD: int = 3