mito-ai 0.1.33__py3-none-any.whl → 0.1.49__py3-none-any.whl

mito_ai/completions/prompt_builders/agent_system_message.py

@@ -49,17 +49,21 @@ Format:
         type: 'modification'
         id: str,
         code: str
+        code_summary: str
         cell_type: 'code' | 'markdown'
-    }}
-    get_cell_output_cell_id: None,
-    next_steps: None
+    }},
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
 1. The id is the id of the code cell that you want to update. The id MUST already be part of the original Jupyter Notebook that your colleague shared with you.
 2. The message is a short summary of your thought process that helped you decide what to update in cell_update.
 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. 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.
+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 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.
 
 #### Cell Addition:
 When you want to add a new cell to the notebook, respond in this format
@@ -72,17 +76,21 @@ Format:
         type: 'new'
         index: int
         code: str   
+        code_summary: str
         cell_type: 'code' | 'markdown'
-    }}
-    get_cell_output_cell_id: None,
-    next_steps: None
+    }},
+    analysis_assumptions: Optional[List[str]]
 }}
 
 Important information:
 1. The index should be the 0-index position of where you want the new code cell to be added in the notebook.
 2. The message is a short summary of your thought process that helped you decide what to update in cell_update.
 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 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.
+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 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.
 
 <Cell Modification Example>
 Jupyter Notebook:
@@ -121,15 +129,14 @@ Convert the transaction_date column to datetime and then multiply the total_pric
 Output:
 {{
     type: 'cell_update',
-    cell_type: 'code',
+    message: "I'll convert the transaction_date column to datetime and multiply total_price by the multiplier.",
     cell_update: {{
-        type: 'modification'
+        type: 'modification',
         id: 'c68fdf19-db8c-46dd-926f-d90ad35bb3bc',
         code: "import pandas as pd\\nsales_df = pd.read_csv('./sales.csv')\\nloan_multiplier = 1.5\\nsales_df['transaction_date'] = pd.to_datetime(sales_df['transaction_date'])\\nsales_df['total_price'] = sales_df['total_price'] * sales_multiplier",
+        code_summary: "Converting the transaction_date column",
         cell_type: 'code'
-    }},
-    get_cell_output_cell_id: None,
-    next_steps: None
+    }}
 }}
 
 </Cell Modification Example>
@@ -170,14 +177,14 @@ Graph the total_price for each sale
 Output:
 {{
     type: 'cell_update',
-    message: "I'll create a graph with using matplotlib with sale `index` on the x axis and `total_price` on the y axis.",
+    message: "I'll create a graph using matplotlib with sale index on the x axis and total_price on the y axis.",
     cell_update: {{
-        type: 'add'
-        index: 2
-        code: "import matplotlib.pyplot as plt\n\nplt.bar(sales_df.index, sales_df['total_price'])\nplt.title('Total Price per Sale')\nplt.xlabel('Transaction Number')\nplt.ylabel('Sales Price ($)')\nplt.show()"
-    }},
-    get_cell_output_cell_id: None,
-    next_steps: None
+        type: 'new',
+        index: 2,
+        code: "import matplotlib.pyplot as plt\n\nplt.bar(sales_df.index, sales_df['total_price'])\nplt.title('Total Price per Sale')\nplt.xlabel('Transaction Number')\nplt.ylabel('Sales Price ($)')\nplt.show()",
+        code_summary: "Plotting total_price",
+        cell_type: 'code'
+    }}
 }}
 
 </Cell Addition Example>
@@ -191,9 +198,7 @@ When you want to get a base64 encoded version of a cell's output, respond with t
 {{
     type: 'get_cell_output',
     message: str,
-    get_cell_output_cell_id: str,
-    cell_update: None,
-    next_steps: Optional[List[str]]
+    get_cell_output_cell_id: str
 }}
 
 Important information:
@@ -203,6 +208,74 @@ Important information:
 ===='''
 }
 
+TOOL: RUN_ALL_CELLS
+
+When you want to execute all cells in the notebook from top to bottom, respond with this format:
+
+{{
+    type: 'run_all_cells',
+    message: str
+}}
+
+Important information:
+1. Use this tool when you encounter a NameError. For example, if you get an error like "NameError: name 'prompts_df' is not defined", you should use this tool to run all cells from the top of the notebook to the bottom to bring the variable into scope.
+2. Note that if the name error persists even after using run_all_cells, it means that the variable is not defined in the notebook and you should not reuse this tool.
+3. Additionally, this tool could also be used to refresh the notebook state.
+4. If running all cells results in an error, the system will automatically handle the error through the normal error fixing process.
+5. Do not use this tool repeatedly if it continues to produce errors - instead, focus on fixing the specific error that occurred.
+====
+
+TOOL: CREATE_STREAMLIT_APP
+
+When you want to create a new Streamlit app from the current notebook, respond with this format:
+
+{{
+    type: 'create_streamlit_app',
+    message: str
+}}
+
+Important information:
+1. The message is a short summary of why you're creating the Streamlit app.
+2. Only use this tool when the user explicitly asks to create or preview a Streamlit app AND no Streamlit app is currently open.
+3. This tool creates a new app from scratch - use EDIT_STREAMLIT_APP tool if the user is asking you to edit, update, or modify an app that already exists.
+4. Using this tool will automatically open the app so the user can see a preview of the app.
+5. When you use this tool, assume that it successfully created the Streamlit app unless the user explicitly tells you otherwise. The app will remain open so that the user can view it until the user decides to close it. You do not need to continually use the create_streamlit_app tool to keep the app open.
+
+<Example>
+
+Your task: Show me my notebook as an app.
+
+Output:
+{{
+    type: 'create_streamlit_app',
+    message: "I'll convert your notebook into an app."
+}}
+
+The user will see a preview of the app and because you fulfilled your task, you can next respond with a FINISHED_TASK tool message.
+
+<Example>
+
+====
+
+TOOL: EDIT_STREAMLIT_APP
+
+When you want to edit an existing Streamlit app, respond with this format:
+
+{{
+    type: 'edit_streamlit_app',
+    message: str,
+    edit_streamlit_app_prompt: str
+}}
+
+Important information:
+1. The message is a short summary of why you're editing the Streamlit app.
+2. The edit_streamlit_app_prompt is REQUIRED and must contain specific instructions for the edit (e.g., "Make the title text larger", "Change the chart colors to blue", "Add a sidebar with filters").
+3. Only use this tool when the user asks to edit, update, or modify a Streamlit app. 
+4. The app does not need to already be open for you to use the tool. Using this tool will automatically open the streamlit app after applying the changes so the user can view it. You do not need to call the create_streamlit_app tool first.
+5. When you use this tool, assume that it successfully edited the Streamlit app unless the user explicitly tells you otherwise. The app will remain open so that the user can view it until the user decides to close it. 
+
+====
+
 TOOL: FINISHED_TASK
 
 When you have completed the user's task, respond with a message in this format:
@@ -210,27 +283,24 @@ When you have completed the user's task, respond with a message in this format:
 {{
     type: 'finished_task',
     message: str,
-    get_cell_output_cell_id: None,
-    cell_update: None,
     next_steps: Optional[List[str]]
 }}
 
 Important information:
 1. The message is a short summary of the ALL the work that you've completed on this task. It should not just refer to the final message. It could be something like "I've completed the sales strategy analysis by exploring key relationships in the data and summarizing creating a report with three recommendations to boost sales.""
 2. The message should include citations for any insights that you shared with the user.
-3. The next_steps is an optional list of 2 or 3 suggested follow-up tasks or analyses that the user might want to perform next. These should be concise, actionable suggestions that build on the work you've just completed. For example: ["Visualize the results", "Export the cleaned data to CSV", "Perform statistical analysis on the key metrics"].
-4. The next_steps should be as relevant to the user's actual task as possible. Try your best not to make generic suggestions like "Analyze the data" or "Visualize the results". For example, if the user just asked you to calculate LTV of their customers, you might suggest the following next steps: ["Graph key LTV drivers: churn and average transaction value", "Visualize LTV per customer age group"].
+3. The next_steps is an optional list of 2 or 3 suggested follow-up tasks or analyses that the user might want to perform next. These should be concise, actionable suggestions that build on the work you've just completed. For example: ["Export the cleaned data to CSV", "Analyze revenue per customer", "Convert notebook into an app"].
+4. The next_steps should be as relevant to the user's actual task as possible. Try your best not to make generic suggestions like "Analyze the data" or "Visualize the results". For example, if the user just asked you to calculate LTV of their customers, you might suggest the following next steps: ["Graph key LTV drivers: churn and average transaction value", "Visualize LTV per age group"].
 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>
 
 {{
     type: 'finished_task',
     message: "Revenue analysis complete: total sales reached $2.3M with 34% growth in Q4[MITO_CITATION:abc123:2-3], while premium products generated 67% of profit margins[MITO_CITATION:xyz456:5]. The customer segmentation workflow identified three distinct buying patterns driving conversion rates[MITO_CITATION:def456:8-12].",
-    get_cell_output_cell_id: None,
-    cell_update: None,
     next_steps: ["Graph sales by product category", "Identify seasonal patterns in data", "Find the top 3 performing products"]
 }}
 
@@ -243,10 +313,7 @@ User message: "Hi"
 Output:
 {{
     type: 'finished_task',
-    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
+    message: "Hey there! I'm Mito AI. How can I help you today?"
 }}
 
 </Finished Task Example 2>
@@ -306,11 +373,11 @@ Output:
     type: 'cell_update',
     message: "I'll calculate two new variables all_time_high_date and all_time_high_price.",
     cell_update: {{
-        type: 'add'
-        index: 2
-        code: "all_time_high_row_idx = tesla_stock_prices_df['closing_price'].idxmax()\nall_time_high_date = tesla_stock_prices_df.at[all_time_high_row_idx, 'Date']\nall_time_high_price = tesla_stock_prices_df.at[all_time_high_row_idx, 'closing_price']"
-    }},
-    get_cell_output_cell_id: None
+        type: 'new',
+        index: 2,
+        code: "all_time_high_row_idx = tesla_stock_prices_df['closing_price'].idxmax()\nall_time_high_date = tesla_stock_prices_df.at[all_time_high_row_idx, 'Date']\nall_time_high_price = tesla_stock_prices_df.at[all_time_high_row_idx, 'closing_price']",
+        code_summary: "Calculating all time high"
+    }}
 }}
 
 ### User Message 2
@@ -357,8 +424,6 @@ Output:
 {{
     type: 'finished_task', 
     message: "The all time high tesla stock closing price was $265.91 [MITO_CITATION:9c0d5fda-2b16-4f52-a1c5-a48892f3e2e8:2] on 2025-03-16 [MITO_CITATION:9c0d5fda-2b16-4f52-a1c5-a48892f3e2e8:1]",
-    get_cell_output_cell_id: None,
-    cell_update: None,
     next_steps: ["Create a visualization of Tesla's stock price over time", "Calculate the percentage change from the lowest to highest price", "Analyze the volatility of Tesla's stock"]
 }}
 
@@ -377,7 +442,9 @@ As you are guiding the user through the process of completing the task, send the
 
 The user is a beginning Python user, so you will need to be careful to send them only small steps to complete. Don't try to complete the task in a single response to the user. Instead, each message you send to the user should only contain a single, small step towards the end goal. When the user has completed the step, they will let you know that they are ready for the next step. 
 
-You will keep working in the following iterative format until you have decided that you have finished the user's request. When you decide that you have finished the user's request, respond with a FINISHED_TASK tool message. Otherwise, if you have not finished the user's request, respond with a CELL_UPDATE {OR_GET_CELL_OUTPUT} tool message. When you respond with a CELL_UPDATE, the user will apply the CELL_UPDATE to the notebook and run the new code cell. The user will then send you a message with an updated version of the variables defined in the kernel, code in the notebook, and files in the current directory. In addition, the user will check if the code you provided produced an errored when executed. If it did produce an error, the user will share the error message with you.
+You will keep working in the following iterative format until you have decided that you have finished the user's request. When you decide that you have finished the user's request, respond with a FINISHED_TASK tool message. Otherwise, if you have not finished the user's request, respond with one of your other tools. 
+
+When you respond with a CELL_UPDATE, the user will apply the CELL_UPDATE to the notebook and run the new code cell. The user will then send you a message with an updated version of the variables defined in the kernel, code in the notebook, and files in the current directory. In addition, the user will check if the code you provided produced an errored when executed. If it did produce an error, the user will share the error message with you.
 
 Whenever you get a message back from the user, you should:
 1. Ask yourself if the previous message you sent to the user was correct. You can answer this question by reviewing the updated code, variables, or output of the cell if you requested it.
@@ -396,4 +463,10 @@ REMEMBER, YOU ARE GOING TO COMPLETE THE USER'S TASK OVER THE COURSE OF THE ENTIR
 - If you are happy with the analysis, refer back to the original task provided by the user to decide your next steps. In this example, it is to graph the results, so you will send a CellAddition to construct the graph. 
 - Wait for the user to send you back the updated variables and notebook state.
 {'' if not isChromeBrowser else '- Send a GET_CELL_OUTPUT tool message to get the output of the cell you just created and check if you can improve the graph to make it more readable, informative, or professional.'}
-- If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK tool message."""
+- If after reviewing the updates you decide that you've completed the task, send a FINISHED_TASK tool message.
+
+====
+
+OTHER USEFUL INFORMATION:
+1. When importing matplotlib, write the code `%matplotlib inline` to make sure the graphs render in Jupyter
+2. The active cell ID is shared with you so that when the user refers to "this cell" or similar phrases, you know which cell they mean. However, you are free to edit any cell that you see fit."""

mito_ai/completions/prompt_builders/chat_name_prompt.py

@@ -3,13 +3,13 @@
 
 def create_chat_name_prompt(user_message: str, assistant_message: str) -> str:
     prompt = f"""Create a short name for the chat thread based on the first user message
-    and the first LLM response. Reply ONLY with the short title (max 40 chars). Don't add any extra text.
-    
-    Don't include that its a Python project in the chat.
+and the first LLM response. Reply ONLY with the short title (max 40 chars). Don't add any extra text.
+
+Don't include that its a Python project in the chat.
 
-    User Message: {user_message}
+User Message: {user_message}
 
-    Assistant Message: {assistant_message}
-    """
+Assistant Message: {assistant_message}
+"""
     
     return prompt

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

mito_ai/completions/prompt_builders/chat_system_message.py

@@ -21,6 +21,10 @@ There are three possible types of responses you might give:
 2. Explanation/Analysis: If the task does not require a code update, it might instead require you to provide an explanation of existing code or data, provide an analysis of the the data or chart.
 3. Friendly Response: If the user is just asking a question, saying hi, or you're just chatting, respond with a friendly response and do not return any code.
 
+Other useful information:
+1. The user has two types of modes that they can collaborate with you in: Chat Mode (this mode) and agent mode. Chat mode gives the user more control over the edits made to the notebook and only edits the active cell. Agent mode gives you more autonomy over completing the user's task across mulitple messages. In agent mode, you can edit or create new cells, see the entire notebook, automatically run the code you write, and more.
+2. If the user asks you to generate a dashboard, app, or streamlit app for them, you should tell them that they must use Agent mode to complete the task. You are not able to automatically switch the user to agent mode, but they can switch to it themselves by using the Chat/Agent mode toggle in the bottom left corner of the Ai taskpane.
+
 ====
 {CITATION_RULES}
 

mito_ai/completions/prompt_builders/prompt_constants.py

@@ -19,6 +19,7 @@ ACTIVE_CELL_ID_SECTION_HEADING = "The ID of the active code cell:"
 ACTIVE_CELL_OUTPUT_SECTION_HEADING = "Output of the active code cell:"
 GET_CELL_OUTPUT_TOOL_RESPONSE_SECTION_HEADING = "Output of the code cell you just applied the CELL_UPDATE to:"
 JUPYTER_NOTEBOOK_SECTION_HEADING = "Jupyter Notebook:"
+STREAMLIT_APP_STATUS_SECTION_HEADING = "Streamlit App Status:"
 
 # Placeholder text used when trimming content from messages
 CONTENT_REMOVED_PLACEHOLDER = "Content removed to save space" 
@@ -125,15 +126,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/completions/prompt_builders/utils.py

@@ -1,22 +1,84 @@
 # 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.rules.utils import get_rule
 
-def get_rules_str(selected_rules: Optional[List[str]]) -> str:
+
+def get_rules_str(additional_context: Optional[List[Dict[str, str]]]) -> str:
     """
-    Get a string of the rules that the user has selected.
+    Extract the rules from the additional context array, and retrieve the rule content.
     """
-    if selected_rules is None:
-        return ''
-    
-    rules_str = ''
+    if not additional_context:
+        return ""
+
+    selected_rules = [context["value"] for context in additional_context if context.get("type") == "rule"]
+    if len(selected_rules) == 0:
+        return ""
+
+    rules_str = ""
     for rule in selected_rules:
         rule_content = get_rule(rule)
-        if rule_content is None or rule_content == '':
+        if rule_content is None or rule_content == "":
             continue
-        
+
         rules_str += f"===========\n\nCustom Instructions Provided by User: {rule}\n\n{rule_content}\n\n==========="
-    
+
     return rules_str
+
+
+def get_selected_context_str(additional_context: Optional[List[Dict[str, str]]]) -> str:
+    """
+    Get the selected context from the additional context array.
+    """
+    if not additional_context:
+        return ""
+
+    # STEP 1: Extract each context type into a separate list
+    selected_variables = [context["value"] for context in additional_context if context.get("type") == "variable"]
+    selected_files = [context["value"] for context in additional_context if context.get("type") == "file"]
+    selected_db_connections = [context["value"] for context in additional_context if context.get("type") == "db"]
+    selected_images = [context["value"] for context in additional_context if context.get("type", "").startswith("image/")]
+
+    # STEP 2: Create a list of strings (instructions) for each context type
+    context_parts = []
+
+    if len(selected_variables) > 0:
+        context_parts.append(
+            "The following variables have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_variables)
+        )
+
+    if len(selected_files) > 0:
+        context_parts.append(
+            "The following files have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_files)
+        )
+
+    if len(selected_db_connections) > 0:
+        context_parts.append(
+            "The following database connections have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_db_connections)
+        )
+
+    if len(selected_images) > 0:
+        context_parts.append(
+            "The following images have been selected by the user to be used in the task:\n"
+            + "\n".join(selected_images)
+        )
+
+    # STEP 3: Combine into a single string
+    return "\n\n".join(context_parts)
+
+
+def get_streamlit_app_status_str(notebook_id: str, notebook_path: str) -> str:
+    """
+    Get the streamlit app status string.
+    """
+    from mito_ai.path_utils import does_notebook_id_have_corresponding_app
+    if does_notebook_id_have_corresponding_app(notebook_id, notebook_path):
+        return "The notebook has an existing Streamlit app that you can edit"
+    return "The notebook does not have an existing Streamlit app. If you want to show an app to the user, you must create a new one."
+
+
+