mito-ai 0.1.36__py3-none-any.whl → 0.1.37__py3-none-any.whl

mito_ai/__init__.py

@@ -11,6 +11,7 @@ from mito_ai.version_check import VersionCheckHandler
 from mito_ai.db.urls import get_db_urls
 from mito_ai.settings.urls import get_settings_urls
 from mito_ai.rules.urls import get_rules_urls
+from mito_ai.auth.urls import get_auth_urls
 try:
     from _version import __version__
 except ImportError:
@@ -65,10 +66,11 @@ def _load_jupyter_server_extension(server_app) -> None: # type: ignore
     ]
     
     # REST API endpoints
-    handlers.extend(get_db_urls(base_url))
-    handlers.extend(get_settings_urls(base_url))
-    handlers.extend(get_rules_urls(base_url))
-    handlers.extend(get_log_urls(base_url))
+    handlers.extend(get_db_urls(base_url))  # type: ignore
+    handlers.extend(get_settings_urls(base_url))  # type: ignore
+    handlers.extend(get_rules_urls(base_url))  # type: ignore
+    handlers.extend(get_log_urls(base_url))  # type: ignore
+    handlers.extend(get_auth_urls(base_url))  # 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.36'
+__version__ = VERSION = '0.1.37'

mito_ai/anthropic_client.py

@@ -232,18 +232,11 @@ class AnthropicClient:
                     system=anthropic_system_prompt,
                     messages=anthropic_messages,
                     stream=True,
-                    message_type=message_type
+                    message_type=message_type,
+                    reply_fn=reply_fn,
+                    message_id=message_id
                 ):
                     accumulated_response += stram_chunk
-                    reply_fn(CompletionStreamChunk(
-                        parent_id=message_id,
-                        chunk=CompletionItem(
-                            content=stram_chunk,
-                            isIncomplete=True,
-                            token=message_id,
-                        ),
-                        done=False,
-                    ))
 
             return accumulated_response
 

mito_ai/app_builder/handlers.py

@@ -8,6 +8,7 @@ from typing import Any, Union
 import zipfile
 import tempfile
 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,
@@ -15,10 +16,12 @@ from mito_ai.app_builder.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
 
+
 class AppBuilderHandler(BaseWebSocketHandler):
     """Handler for app building requests."""
     
@@ -62,6 +65,10 @@ class AppBuilderHandler(BaseWebSocketHandler):
         self.log.debug("App builder message received: %s", message)
         
         try:
+            # Ensure message is a string before parsing
+            if not isinstance(message, str):
+                raise ValueError("Message must be a string")
+
             parsed_message = self.parse_message(message)
             message_type = parsed_message.get('type')
             
@@ -98,8 +105,10 @@ class AppBuilderHandler(BaseWebSocketHandler):
             message: The parsed message.
         """
         message_id = message.get('message_id', '')  # Default to empty string if not present
-        app_path = message.get('path')
-        
+        notebook_path = message.get('notebook_path')
+        app_path = message.get('app_path')
+        jwt_token = message.get('jwt_token', '')  # Extract JWT token from request, default to empty string
+
         if not message_id:
             self.log.error("Missing message_id in request")
             return
@@ -115,13 +124,37 @@ class AppBuilderHandler(BaseWebSocketHandler):
                 error=error
             ))
             return
-        
+
+        # Validate JWT token if provided
+        if jwt_token and jwt_token != 'placeholder-jwt-token':
+            self.log.info(f"Validating JWT token: {jwt_token[:20]}...")
+            is_valid = self._validate_jwt_token(jwt_token)
+            if not is_valid:
+                self.log.error("JWT token validation failed")
+                error = AppBuilderError(
+                    error_type="Unauthorized",
+                    title="Invalid authentication token",
+                    hint="Please sign in again to deploy your app."
+                )
+                self.reply(BuildAppReply(
+                    parent_id=message_id,
+                    url="",
+                    error=error
+                ))
+                return
+            else:
+                self.log.info("JWT token validation successful")
+        else:
+            self.log.warning("No JWT token provided or using placeholder token")
+
         try:
-            # This is a placeholder for the actual app building logic
-            # In a real implementation, this would deploy the app to a hosting service
-            # and return the URL
-            deploy_url = await self._deploy_app(app_path)
-            
+
+            success_flag, result_message = await streamlit_handler(str(notebook_path) if notebook_path else "", app_path)
+            if not success_flag:
+                raise Exception(result_message)
+
+            deploy_url = await self._deploy_app(app_path, jwt_token)
+
             # Send the response
             self.reply(BuildAppReply(
                 parent_id=message_id,
@@ -137,12 +170,47 @@ class AppBuilderHandler(BaseWebSocketHandler):
                 error=error
             ))
     
-    async def _deploy_app(self, app_path: str) -> str:
+    def _validate_jwt_token(self, token: str) -> bool:
+        """Basic JWT token validation logic.
+
+        In a production environment, you would:
+        1. Decode the JWT token
+        2. Verify the signature using AWS Cognito public keys
+        3. Check the expiration time
+        4. Validate the issuer and audience claims
+
+        For now, we'll do a basic check that the token exists and has a reasonable format.
+        """
+        try:
+            # Basic JWT format validation (header.payload.signature)
+            if not token or '.' not in token:
+                self.log.error("Token is empty or missing dots")
+                return False
+
+            parts = token.split('.')
+            if len(parts) != 3:
+                self.log.error("Token does not have 3 parts")
+                return False
+
+            # Check for placeholder token
+            if token == 'placeholder-jwt-token':
+                self.log.error("Placeholder token detected")
+                return False
+
+            return True
+
+        except Exception as e:
+            self.log.error(f"Error validating JWT token: {e}")
+            return False
+
+
+    async def _deploy_app(self, app_path: str, jwt_token: str = '') -> str:
         """Deploy the app using pre-signed URLs.
         
         Args:
             app_path: Path to the app file.
-            
+            jwt_token: JWT token for authentication (optional)
+
         Returns:
             The URL of the deployed app.
         """
@@ -153,7 +221,17 @@ class AppBuilderHandler(BaseWebSocketHandler):
         try:
             # Step 1: Get pre-signed URL from API
             self.log.info("Getting pre-signed upload URL...")
-            url_response = requests.get(f"{ACTIVE_STREAMLIT_BASE_URL}/get-upload-url?app_name={app_name}")
+
+            # Prepare headers with JWT token if provided
+            headers = {}
+            if jwt_token and jwt_token != 'placeholder-jwt-token':
+                headers['Authorization'] = f'Bearer {jwt_token}'
+            else:
+                self.log.warning("No JWT token provided for API request")
+
+            headers["Subscription-Tier"] = 'Pro' if is_pro() else 'Standard'
+
+            url_response = requests.get(f"{ACTIVE_STREAMLIT_BASE_URL}/get-upload-url?app_name={app_name}", headers=headers)
             url_response.raise_for_status()
             
             url_data = url_response.json()

mito_ai/app_builder/models.py

@@ -66,6 +66,9 @@ class BuildAppRequest:
     
     # Path to the app file.
     path: str
+    
+    # JWT token for authorization.
+    jwt_token: Optional[str] = None
 
 
 @dataclass(frozen=True)

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/completions/completion_handlers/chat_completion_handler.py

@@ -42,7 +42,7 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
             metadata.activeCellId,
             metadata.base64EncodedActiveCellOutput is not None and metadata.base64EncodedActiveCellOutput != '',
             metadata.input,
-            metadata.selectedRules
+            metadata.additionalContext
         )
         display_prompt = f"```python{metadata.activeCellCode or ''}```{metadata.input}"
         
@@ -105,7 +105,7 @@ class ChatCompletionHandler(CompletionHandler[ChatMessageMetadata]):
             metadata.activeCellId,
             metadata.base64EncodedActiveCellOutput is not None and metadata.base64EncodedActiveCellOutput != '',
             metadata.input,
-            metadata.selectedRules
+            metadata.additionalContext
         )
         display_prompt = f"```python{metadata.activeCellCode or ''}```{metadata.input}"
         

mito_ai/completions/models.py

@@ -3,11 +3,10 @@
 
 import traceback
 from dataclasses import dataclass, field
-from typing import Annotated, List, Literal, Optional, Type, Union, NewType
+from typing import List, Literal, Optional, NewType, Dict
 from openai.types.chat import ChatCompletionMessageParam
 from enum import Enum
-from pydantic import BaseModel, Field
-
+from pydantic import BaseModel
 
 # The ThreadID is the unique identifier for the chat thread.
 ThreadID = NewType('ThreadID', str)
@@ -35,6 +34,7 @@ class AgentResponse(BaseModel):
     cell_update: Optional[CellUpdate]
     get_cell_output_cell_id: Optional[str]
     next_steps: Optional[List[str]]
+    analysis_assumptions: Optional[List[str]]
     
     
 @dataclass(frozen=True)
@@ -63,6 +63,7 @@ class MessageType(Enum):
     GET_THREADS = "get_threads"
     DELETE_THREAD = "delete_thread"
     UPDATE_MODEL_CONFIG = "update_model_config"
+    STREAMLIT_CONVERSION = "streamlit_conversion"
 
     
 @dataclass(frozen=True)
@@ -84,7 +85,7 @@ class ChatMessageMetadata():
     base64EncodedActiveCellOutput: Optional[str] = None
     index: Optional[int] = None
     stream: bool = False
-    selectedRules: Optional[List[str]] = None
+    additionalContext: Optional[List[Dict[str, str]]] = None
     
     
 @dataclass(frozen=True)
@@ -98,7 +99,7 @@ class AgentExecutionMetadata():
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None
     index: Optional[int] = None
-    selectedRules: Optional[List[str]] = None
+    additionalContext: Optional[List[Dict[str, str]]] = None
     
 @dataclass(frozen=True)
 class AgentSmartDebugMetadata():
@@ -135,7 +136,7 @@ class InlineCompleterMetadata():
     suffix: str
     variables: Optional[List[str]] = None
     files: Optional[List[str]] = None
-    
+
 @dataclass(frozen=True)
 class CompletionRequest:
     """

mito_ai/completions/prompt_builders/agent_execution_prompt.py

@@ -8,14 +8,17 @@ from mito_ai.completions.prompt_builders.prompt_constants import (
     VARIABLES_SECTION_HEADING,
     cell_update_output_str
 )
-from mito_ai.completions.prompt_builders.utils import get_rules_str
+from mito_ai.completions.prompt_builders.utils import (
+    get_rules_str,
+    get_selected_context_str,
+)
 
 def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
     variables_str = '\n'.join([f"{variable}" for variable in md.variables or []])
     files_str = '\n'.join([f"{file}" for file in md.files or []])
     ai_optimized_cells_str = '\n'.join([f"{cell}" for cell in md.aiOptimizedCells or []])
-    rules_str = get_rules_str(md.selectedRules)
-    
+    rules_str = get_rules_str(md.additionalContext)
+    selected_context_str = get_selected_context_str(md.additionalContext)
     context_str = f"""Remember to choose the correct tool to respond with.
 
 {rules_str}
@@ -30,6 +33,8 @@ def create_agent_execution_prompt(md: AgentExecutionMetadata) -> str:
 {FILES_SECTION_HEADING}
 {files_str}
 
+{selected_context_str}
+
 {cell_update_output_str(md.base64EncodedActiveCellOutput is not None)}"""
 
     task_str = '' if md.input == '' else f"""Your task: 

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -53,7 +53,8 @@ Format:
         cell_type: 'code' | 'markdown'
     }}
     get_cell_output_cell_id: None,
-    next_steps: None
+    next_steps: None,
+    analysis_assumptions: None
 }}
 
 Important information:
@@ -62,6 +63,9 @@ 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"].
+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 None 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.
 
 #### Cell Addition:
 When you want to add a new cell to the notebook, respond in this format
@@ -78,7 +82,8 @@ Format:
         cell_type: 'code' | 'markdown'
     }}
     get_cell_output_cell_id: None,
-    next_steps: None
+    next_steps: None,
+    analysis_assumptions: None
 }}
 
 Important information:
@@ -87,6 +92,9 @@ 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"].
+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 None 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.
 
 <Cell Modification Example>
 Jupyter Notebook:
@@ -134,7 +142,8 @@ Output:
         cell_type: 'code'
     }},
     get_cell_output_cell_id: None,
-    next_steps: None
+    next_steps: None,
+    analysis_assumptions: None
 }}
 
 </Cell Modification Example>
@@ -184,7 +193,8 @@ Output:
         code_summary: "Plotting total_price"
     }},
     get_cell_output_cell_id: None,
-    next_steps: None
+    next_steps: None,
+    analysis_assumptions: None
 }}
 
 </Cell Addition Example>
@@ -200,7 +210,8 @@ When you want to get a base64 encoded version of a cell's output, respond with t
     message: str,
     get_cell_output_cell_id: str,
     cell_update: None,
-    next_steps: Optional[List[str]]
+    next_steps: Optional[List[str]],
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
@@ -219,7 +230,8 @@ When you have completed the user's task, respond with a message in this format:
     message: str,
     get_cell_output_cell_id: None,
     cell_update: None,
-    next_steps: Optional[List[str]]
+    next_steps: Optional[List[str]],
+    analysis_assumptions: None
 }}
 
 Important information:
@@ -230,6 +242,7 @@ Important information:
 5. If you are not sure what the user might want to do next, err on the side of suggesting next steps instead of making an assumption and using more CELL_UPDATES.
 6. If the user's task doesn't involve creating or modifying a code cell, you should respond with a FINISHED_TASK response. 
 7. If the user is just sending a friendly greeting (like "Hello", "Hi", "Hey", "How are you?", "What can you help me with?", etc.), you must respond with a FINISHED_TASK response message with a friendly message like this: "Hello! I'm Mito AI, your AI assistant for data analysis and Python programming in Jupyter notebooks. I can help you analyze datasets, create visualizations, clean data, and much more. What would you like to work on today?"
+8. Do not include any analysis_assumptions in the FINISHED_TASK response.
 
 <Finished Task Example 1>
 
@@ -253,7 +266,8 @@ Output:
     message: "Hey there! I'm Mito AI. How can I help you today?",
     get_cell_output_cell_id: None,
     cell_update: None,
-    next_steps: None
+    next_steps: None,
+    analysis_assumptions: None
 }}
 
 </Finished Task Example 2>

mito_ai/completions/prompt_builders/chat_prompt.py

@@ -1,30 +1,35 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-from typing import List, Optional
+from typing import List, Optional, Dict
 from mito_ai.completions.prompt_builders.prompt_constants import (
     ACTIVE_CELL_ID_SECTION_HEADING,
     CHAT_CODE_FORMATTING_RULES,
     FILES_SECTION_HEADING,
     VARIABLES_SECTION_HEADING,
     CODE_SECTION_HEADING,
-    get_active_cell_output_str
+    get_active_cell_output_str,
 )
-from mito_ai.completions.prompt_builders.utils import get_rules_str
+from mito_ai.completions.prompt_builders.utils import (
+    get_rules_str,
+    get_selected_context_str,
+)
+
 
 def create_chat_prompt(
     variables: List[str],
     files: List[str],
-    active_cell_code: str, 
+    active_cell_code: str,
     active_cell_id: str,
     has_active_cell_output: bool,
     input: str,
-    selected_rules: Optional[List[str]] = None
+    additional_context: Optional[List[Dict[str, str]]] = None,
 ) -> str:
-    variables_str = '\n'.join([f"{variable}" for variable in variables])
-    files_str = '\n'.join([f"{file}" for file in files])
-    rules_str = get_rules_str(selected_rules)
-    
+    variables_str = "\n".join([f"{variable}" for variable in variables])
+    files_str = "\n".join([f"{file}" for file in files])
+    selected_context_str = get_selected_context_str(additional_context)
+    rules_str = get_rules_str(additional_context)
+
     prompt = f"""{rules_str}
     
 Help me complete the following task. I will provide you with a set of variables, existing code, and a task to complete.
@@ -101,9 +106,11 @@ Hey there! I'm Mito AI. How can I help you today?
 {active_cell_code}
 ```
 
+{selected_context_str}
+
 {get_active_cell_output_str(has_active_cell_output)}
 
 Your task: {input}
 """
-    
-    return prompt
+
+    return prompt