mito-ai 0.1.43__py3-none-any.whl → 0.1.45__py3-none-any.whl

mito_ai/__init__.py

@@ -5,7 +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.app_builder.handlers import AppBuilderHandler
+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
 from mito_ai.version_check import VersionCheckHandler
@@ -71,8 +71,8 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
             {"llm": open_ai_provider},
         ),
         (
-            url_path_join(base_url, "mito-ai", "app-builder"),
-            AppBuilderHandler,
+            url_path_join(base_url, "mito-ai", "app-deploy"),
+            AppDeployHandler,
             {}
         ),
         (

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.43'
+__version__ = VERSION = '0.1.45'

mito_ai/anthropic_client.py

@@ -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:
@@ -206,7 +206,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_builder → app_deploy}/__init__.py

@@ -3,4 +3,4 @@
 
 """App builder module for Mito AI."""
 
-from .handlers import AppBuilderHandler 
+from .handlers import AppDeployHandler 

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_builder → app_deploy}/handlers.py

@@ -4,27 +4,27 @@
 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_builder.models import (
-    BuildAppReply,
-    AppBuilderError,
-    BuildAppRequest,
+from mito_ai.app_deploy.app_deploy_utils import  add_files_to_zip
+from mito_ai.app_deploy.models import (
+    DeployAppReply,
+    AppDeployError,
+    DeployAppRequest,
     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
 
 
-class AppBuilderHandler(BaseWebSocketHandler):
-    """Handler for app building requests."""
+class AppDeployHandler(BaseWebSocketHandler):
+    """Handler for app deploy requests."""
     
     def initialize(self) -> None:
         """Initialize the WebSocket handler."""
@@ -57,6 +57,7 @@ class AppBuilderHandler(BaseWebSocketHandler):
         Args:
             message: The message received on the WebSocket.
         """
+        
         start = time.time()
         
         # Convert bytes to string if needed
@@ -73,13 +74,13 @@ class AppBuilderHandler(BaseWebSocketHandler):
             parsed_message = self.parse_message(message)
             message_type = parsed_message.get('type')
             
-            if message_type == MessageType.BUILD_APP.value:
+            if message_type == MessageType.DEPLOY_APP.value:
                 # Handle build app request
-                build_app_request = BuildAppRequest(**parsed_message)
-                await self._handle_build_app(build_app_request)
+                deploy_app_request = DeployAppRequest(**parsed_message)
+                await self._handle_deploy_app(deploy_app_request)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
-                error = AppBuilderError(
+                error = AppDeployError(
                     error_type="InvalidRequest",
                     title=f"Unknown message type: {message_type}"
                 )
@@ -87,11 +88,11 @@ class AppBuilderHandler(BaseWebSocketHandler):
                 
         except ValueError as e:
             self.log.error("Invalid app builder request", exc_info=e)
-            error = AppBuilderError.from_exception(e)
+            error = AppDeployError.from_exception(e)
             self.reply(ErrorMessage(**error.__dict__))
         except Exception as e:
             self.log.error("Error handling app builder message", exc_info=e)
-            error = AppBuilderError.from_exception(
+            error = AppDeployError.from_exception(
                 e, 
                 hint="An error occurred while building the app. Please check the logs for details."
             )
@@ -100,7 +101,7 @@ class AppBuilderHandler(BaseWebSocketHandler):
         latency_ms = round((time.time() - start) * 1000)
         self.log.info(f"App builder handler processed in {latency_ms} ms.")
     
-    async def _handle_build_app(self, message: BuildAppRequest) -> None:
+    async def _handle_deploy_app(self, message: DeployAppRequest) -> None:
         """Handle a build app request.
         
         Args:
@@ -109,17 +110,18 @@ class AppBuilderHandler(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")
             return
         
         if not notebook_path:
-            error = AppBuilderError(
+            error = AppDeployError(
                 error_type="InvalidRequest",
                 title="Missing 'notebook_path' parameter"
             )
-            self.reply(BuildAppReply(
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error
@@ -132,12 +134,12 @@ class AppBuilderHandler(BaseWebSocketHandler):
         is_valid = self._validate_jwt_token(jwt_token) if jwt_token else False
         if not is_valid or not jwt_token:
             self.log.error("JWT token validation failed")
-            error = AppBuilderError(
+            error = AppDeployError(
                 error_type="Unauthorized",
                 title="Invalid authentication token",
                 hint="Please sign in again to deploy your app."
             )
-            self.reply(BuildAppReply(
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error
@@ -150,25 +152,34 @@ class AppBuilderHandler(BaseWebSocketHandler):
             notebook_path = str(notebook_path) if notebook_path else ""
 
             app_directory = os.path.dirname(notebook_path)
-            app_path = os.path.join(app_directory, "app.py")
-
-            if not os.path.exists(app_path):
-                success_flag, app_path_result, result_message = await streamlit_handler(notebook_path)
-                if not success_flag or app_path_result is None:
-                    raise Exception(result_message)
-
-            deploy_url = await self._deploy_app(app_directory, jwt_token)
+            
+            # Check if the app.py file exists
+            app_path = get_app_path(app_directory)
+            if app_path is None:
+                error = AppDeployError(
+                    error_type="AppNotFound",
+                    title="App not found",
+                    hint="Please make sure the app.py file exists in the same directory as the notebook."
+                )
+                self.reply(DeployAppReply(
+                    parent_id=message_id,
+                    url="",
+                    error=error
+                ))
+            
+            # Finally, deploy the app
+            deploy_url = await self._deploy_app(app_directory, files_to_upload, jwt_token)
 
             # Send the response
-            self.reply(BuildAppReply(
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url=deploy_url
             ))
             
         except Exception as e:
             self.log.error(f"Error building app: {e}", exc_info=e)
-            error = AppBuilderError.from_exception(e)
-            self.reply(BuildAppReply(
+            error = AppDeployError.from_exception(e)
+            self.reply(DeployAppReply(
                 parent_id=message_id,
                 url="",
                 error=error
@@ -208,11 +219,12 @@ class AppBuilderHandler(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:
@@ -247,16 +259,12 @@ class AppBuilderHandler(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_builder → app_deploy}/models.py

@@ -3,17 +3,17 @@
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import Literal, Optional
+from typing import Literal, Optional, List
 
 
 class MessageType(str, Enum):
-    """Types of app builder messages."""
-    BUILD_APP = "build-app"
+    """Types of app deploy messages."""
+    DEPLOY_APP = "deploy-app"
 
 
 @dataclass(frozen=True)
-class AppBuilderError:
-    """Error information for app builder operations."""
+class AppDeployError:
+    """Error information for app deploy operations."""
     
     # Error type.
     error_type: str
@@ -28,7 +28,7 @@ class AppBuilderError:
     hint: Optional[str] = None
     
     @classmethod
-    def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppBuilderError":
+    def from_exception(cls, e: Exception, hint: Optional[str] = None) -> "AppDeployError":
         """Create an error from an exception.
 
         Args:
@@ -47,7 +47,7 @@ class AppBuilderError:
 
 
 @dataclass(frozen=True)
-class ErrorMessage(AppBuilderError):
+class ErrorMessage(AppDeployError):
     """Error message."""
     
     # Message type.
@@ -55,25 +55,28 @@ class ErrorMessage(AppBuilderError):
 
 
 @dataclass(frozen=True)
-class BuildAppRequest:
-    """Request to build an app."""
+class DeployAppRequest:
+    """Request to deploy an app."""
     
     # Request type.
-    type: Literal["build-app"]
+    type: Literal["deploy-app"]
     
     # Message ID.
     message_id: str
     
     # 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
 
 
 @dataclass(frozen=True)
-class BuildAppReply:
-    """Reply to a build app request."""
+class DeployAppReply:
+    """Reply to a deplpy app request."""
     
     # ID of the request message this is replying to.
     parent_id: str
@@ -82,7 +85,7 @@ class BuildAppReply:
     url: str
     
     # Optional error information.
-    error: Optional[AppBuilderError] = None
+    error: Optional[AppDeployError] = None
     
     # Type of reply.
-    type: Literal["build-app"] = "build-app" 
+    type: Literal["deploy-app"] = "deploy-app" 

mito_ai/app_manager/handlers.py

@@ -12,6 +12,8 @@ from mito_ai.app_manager.models import (
     AppManagerError,
     ManageAppRequest,
     ManageAppReply,
+    CheckAppStatusRequest,
+    CheckAppStatusReply,
     ErrorMessage,
     MessageType
 )
@@ -57,6 +59,10 @@ class AppManagerHandler(BaseWebSocketHandler):
                 # Handle manage app request
                 manage_app_request = ManageAppRequest(**parsed_message)
                 await self._handle_manage_app(manage_app_request)
+            elif message_type == MessageType.CHECK_APP_STATUS.value:
+                # Handle check app status request
+                check_status_request = CheckAppStatusRequest(**parsed_message)
+                await self._handle_check_app_status(check_status_request)
             else:
                 self.log.error(f"Unknown message type: {message_type}")
                 error_response = ErrorMessage(
@@ -131,4 +137,31 @@ class AppManagerHandler(BaseWebSocketHandler):
                 error=error,
                 message_id=request.message_id
             )
+            self.reply(error_reply)
+
+    async def _handle_check_app_status(self, request: CheckAppStatusRequest) -> None:
+        """Handle a check app status request."""
+        self.log.info("In check app status")
+        try:
+            # Make a HEAD request to check if the app URL is accessible
+            response = requests.head(request.app_url, timeout=10, verify=False)
+            self.log.debug(f"Is app accessible: {response.status_code}")
+            is_accessible = response.status_code==200
+
+            # Create successful response
+            reply = CheckAppStatusReply(
+                is_accessible=is_accessible
+            )
+
+            self.reply(reply)
+
+        except Exception as e:
+            self.log.error(f"Error checking app status: {e}", exc_info=e)
+            error = AppManagerError.from_exception(e)
+
+            # Return error response
+            error_reply = CheckAppStatusReply(
+                is_accessible=False,
+                error=error
+            )
             self.reply(error_reply)

mito_ai/app_manager/models.py

@@ -8,6 +8,7 @@ from typing import List, Optional
 class MessageType(str, Enum):
     """Types of app manager messages."""
     MANAGE_APP = "manage-app"
+    CHECK_APP_STATUS = "check-app-status"
 
 
 @dataclass(frozen=True)
@@ -15,7 +16,7 @@ class ManageAppRequest:
     """Request to manage apps."""
     type: str = "manage-app"
     jwt_token: Optional[str] = None
-    message_id: Optional[str] = None  
+    message_id: Optional[str] = None
 
 @dataclass(frozen=True)
 class App:
@@ -48,6 +49,19 @@ class ManageAppReply:
     error: Optional[AppManagerError] = None
     message_id: Optional[str] = None  
 
+@dataclass(frozen=True)
+class CheckAppStatusRequest:
+    """Request to check app status."""
+    app_url: str
+    type: str = "check-app-status"
+
+@dataclass(frozen=True)
+class CheckAppStatusReply:
+    """Reply to a check app status request."""
+    is_accessible: bool
+    type: str = "check-app-status"
+    error: Optional[AppManagerError] = None
+
 @dataclass(frozen=True)
 class ErrorMessage:
     """Error message."""

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 (
@@ -67,6 +68,7 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
         self._llm = llm
         self.is_pro = is_pro()
         self._selected_model = FALLBACK_MODEL
+        self.is_electron = False
         identify(llm.key_type)
         
     @property
@@ -128,6 +130,18 @@ class CompletionHandler(JupyterHandler, WebSocketHandler):
             parsed_message = json.loads(message)
             metadata_dict = parsed_message.get('metadata', {})
             type: MessageType = MessageType(parsed_message.get('type'))
+            
+            # Extract environment information from the message
+            environment = parsed_message.get('environment', {})
+            if environment:
+                is_electron = environment.get('isElectron', None)
+                if is_electron is not None:
+                    if is_electron != self.is_electron:
+                        # If the is_electron status is different, log it
+                        identify(key_type=self._llm.key_type, is_electron=is_electron)
+                    
+                    self.is_electron = is_electron
+
         except ValueError as e:
             self.log.error("Invalid completion request.", exc_info=e)
             return
@@ -209,7 +223,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 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

mito_ai/completions/models.py

@@ -3,7 +3,7 @@
 
 import traceback
 from dataclasses import dataclass, field
-from typing import List, Literal, Optional, NewType, Dict
+from typing import List, Literal, Optional, NewType, Dict, Any
 from openai.types.chat import ChatCompletionMessageParam
 from enum import Enum
 from pydantic import BaseModel
@@ -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)
@@ -155,6 +156,9 @@ class CompletionRequest:
     # Whether to stream the response (if supported by the model).
     stream: bool = False
     
+    # Environment information from the client
+    environment: Optional[Dict[str, Any]] = None
+    
     
 @dataclass(frozen=True)
 class AICapabilities:

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -51,7 +51,8 @@ Format:
         code: str
         code_summary: str
         cell_type: 'code' | 'markdown'
-    }}
+    }},
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
@@ -60,7 +61,7 @@ Important information:
 3. The code should be the full contents of that updated code cell. The code that you return will overwrite the existing contents of the code cell so it must contain all necessary code.
 4. The code_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Loading data", "Filtering rows", "Calculating average", "Plotting revenue"). Avoid full sentences or explanations—this should read like a quick commit message or code label, not a description.
 5. Important: Only use the CELL_UPDATE tool if you want to add/modify a notebook cell in response to the user's request. If the user is just sending you a friendly greeting or asking you a question about yourself, you SHOULD NOT USE A CELL_UPDATE tool because it does not require modifying the notebook. Instead, just use the FINISHED_TASK response.
-6. The assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
+6. The analysis_assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
 7. Only include important data and analytical assumptions that if incorrect would fundamentally change your analysis conclusions. These should be data handling decisions, methodological choices, and definitional boundaries. Do not include: obvious statements ("Each record is counted once"), result interpretation guidance ("Gaps in the plot represent zero values"), display choices ("Data is sorted for clarity"), internal reasoning ("Bar chart is better than line plot"), or environment assumptions ("Library X is installed"). Prioritize quality over quantity - include only the most critical assumptions or omit the field entirely if there are no critical assumptions made in this step that have not already be shared with the user. If you ever doubt whether an assumption is critical enough to be shared with the user as an assumption, don't include it. Most messages should not include an assumption. 
 8. Do not include the same assumption or variations of the same assumption multiple times in the same conversation. Once you have presented the assumption to the user, they will already have the opportunity to confirm or correct it so do not include it again.
 
@@ -77,7 +78,8 @@ Format:
         code: str   
         code_summary: str
         cell_type: 'code' | 'markdown'
-    }}
+    }},
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
@@ -86,7 +88,7 @@ Important information:
 3. The code should be the full contents of that updated code cell. The code that you return will overwrite the existing contents of the code cell so it must contain all necessary code.
 4. code_summary must be a very short phrase (1–5 words maximum) that begins with a verb ending in "-ing" (e.g., "Loading data", "Filtering rows", "Calculating average", "Plotting revenue"). Avoid full sentences or explanations—this should read like a quick commit message or code label, not a description.
 5. The cell_type should only be 'markdown' if there is no code to add. There may be times where the code has comments. These are still code cells and should have the cell_type 'code'. Any cells that are labeled 'markdown' will be converted to markdown cells by the user.
-6. The assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
+6. The analysis_assumptions is an optional list of critical assumptions that you made about the data or analysis approach. The assumptions you list here will be displayed to the user so that they can confirm or correct the assumptions. For example: ["NaN values in the impressions column represent 0 impressions", "Only crashes with pedestrian or cyclist fatalities are considered fatal crashes", "Intervention priority combines both volume and severity to identify maximum impact opportunities"].
 7. Only include important data and analytical assumptions that if incorrect would fundamentally change your analysis conclusions. These should be data handling decisions, methodological choices, and definitional boundaries. Do not include: obvious statements ("Each record is counted once"), result interpretation guidance ("Gaps in the plot represent zero values"), display choices ("Data is sorted for clarity"), internal reasoning ("Bar chart is better than line plot"), or environment assumptions ("Library X is installed"). Prioritize quality over quantity - include only the most critical assumptions or omit the field entirely if there are no critical assumptions made in this step that have not already be shared with the user. If you ever doubt whether an assumption is critical enough to be shared with the user as an assumption, don't include it. Most messages should not include an assumption. 
 8. Do not include the same assumption or variations of the same assumption multiple times in the same conversation. Once you have presented the assumption to the user, they will already have the opportunity to confirm or correct it so do not include it again.
 

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}
         """